[GG] Prewarm CuTe PCIe one-shot all-reduce before graph capture

[voipmonitor]

Summary

Prewarm the B12X CuTe one-shot all-reduce specialization during vLLM's CUDA graph warmup, before the nested Inductor capture begins.

B12X commit 680d819 migrated the PCIe communication kernels to CuTe DSL and correctly rejects a cold one-shot launch during CUDA graph capture. vLLM's existing graph-warmup branch intentionally skipped the collective and returned empty_like; that preserved communication semantics, but it also meant the new CuTe specialization was never prepared. A cold GLM TP4/DCP4 boot then failed in the first PIECEWISE graph capture with:

cold PCIe oneshot CUDA graph capture is not allowed;
call prepare_graph_all_reduce() before capture

Change

When vLLM is in the non-capturing warmup branch:

• prepare the CuTe graph specialization only for inputs eligible for the one-shot path;
• keep returning the same out-of-place placeholder;
• do not launch communication;
• leave DMA-only inputs and the runtime hot path unchanged;
• retain B12X's cold-capture guard.

Validation

• Focused unit tests: 2 passed
  • one-shot warmup prepares once without calling all_reduce;
  • DMA-only input does not prepare the one-shot path.
• Ruff check and format: pass.
• Before fix: clean GLM-5.2 EXL3 TP4/DCP4 cold boot failed deterministically during the first PIECEWISE graph capture.
• After fix: the same image/config completed profiling and production CUDA graph capture and became healthy.
• API coherence: 42/42 exact responses passed.
• Prefill remained within run variance against the pre-CuTe candidate:

Context	Pre-CuTe candidate	Fixed CuTe path
8k	3,352 tok/s	3,342 tok/s
64k	3,393 tok/s	3,438 tok/s

Runtime gate used current GG plus current B12X master and the GLM TP4/DCP4 one-shard query-split profile.

Summary by CodeRabbit

• Performance Improvements
  • Improved CUDA graph warmup for eligible custom all-reduce operations.
  • Prepares supported PCIe graph specializations during warmup to improve subsequent execution.
  • Skips preparation for unsupported or DMA-only inputs.

[coderabbitai]

Warning

Review limit reached

@voipmonitor, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 52 minutes

You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository.

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 5e5eac73-01b0-4c79-b231-e03a5671c873

📥 Commits

Reviewing files that changed from the base of the PR and between ad6d782 and 85bc977.

📒 Files selected for processing (1)

• tests/distributed/test_b12x_custom_all_reduce_warmup.py

📝 Walkthrough

Walkthrough

CUDA graph warmup now prepares PCIe graph all-reduce for eligible plain inputs. Tests verify preparation on the capture stream and confirm that DMA-only inputs skip graph preparation.

Changes

Custom all-reduce warmup

Layer / File(s)	Summary
Graph warmup preparation vllm/distributed/device_communicators/custom_all_reduce.py	Eligible inputs within the configured oneshot limit call prepare_graph_all_reduce on the current PCIe capture stream. Unsupported runtimes and oversized inputs keep the existing behavior.
Warmup validation tests/distributed/test_b12x_custom_all_reduce_warmup.py	Mocked capture warmup verifies graph preparation for plain inputs and no preparation for DMA-only inputs.

Estimated code review effort: 3 (Moderate) | ~15–30 minutes

Possibly related PRs

• local-inference-lab/vllm#131: Changes related PCIe channel behavior in CustomAllreduce.

Suggested reviewers: lukealonso

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: prewarming CuTe PCIe one-shot all-reduce before CUDA graph capture.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

tests/distributed/test_b12x_custom_all_reduce_warmup.py (1)

77-80: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Verify that DMA-only warmup does not communicate.

Assert that custom_allreduce._pcie_dma.all_reduce was not called. Also assert that output.shape matches inp.shape. The current assertions pass if a future change launches DMA communication or returns a tensor with the wrong shape.

Proposed test additions

     assert output is not None
+    assert output.shape == inp.shape
     runtime.prepare_graph_all_reduce.assert_not_called()
+    runtime.all_reduce.assert_not_called()
+    custom_allreduce._pcie_dma.all_reduce.assert_not_called()

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/distributed/test_b12x_custom_all_reduce_warmup.py` around lines 77 -
80, Strengthen the test around custom_allreduce.custom_all_reduce by asserting
that custom_allreduce._pcie_dma.all_reduce was not called and that output.shape
equals inp.shape, while retaining the existing output and
runtime.prepare_graph_all_reduce assertions.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/distributed/test_b12x_custom_all_reduce_warmup.py`:
- Around line 12-14: Add Google-style docstrings to _make_custom_allreduce, the
capture-warmup mock helper, and the plain-input and DMA-only test functions in
tests/distributed/test_b12x_custom_all_reduce_warmup.py at lines 12-14, 31-34,
50-52, and 68-70; document parameters with Args:, return values with Returns:
where applicable, and exceptions with Raises: where applicable.

---

Nitpick comments:
In `@tests/distributed/test_b12x_custom_all_reduce_warmup.py`:
- Around line 77-80: Strengthen the test around
custom_allreduce.custom_all_reduce by asserting that
custom_allreduce._pcie_dma.all_reduce was not called and that output.shape
equals inp.shape, while retaining the existing output and
runtime.prepare_graph_all_reduce assertions.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: fa8f3094-b9f6-4b78-8c53-e08e338980a6

📥 Commits

Reviewing files that changed from the base of the PR and between 3003860 and ad6d782.

📒 Files selected for processing (2)

• tests/distributed/test_b12x_custom_all_reduce_warmup.py
• vllm/distributed/device_communicators/custom_all_reduce.py