[GG] fix DCP packed A2A buffer lifetime during graph prewarm

[voipmonitor]

Summary

• Allocate retained packed NCCL A2A staging buffers during the outer vLLM CUDA graph prewarm scope.
• Reuse the same exact-shape owners when the corresponding FULL graph is captured.
• Keep ordinary eager calls outside graph setup temporary, as before.

Root cause

The existing graph-buffer fix retains exact-shape tensors only when torch.cuda.is_current_stream_capturing() is true. vLLM performs an eager descriptor prewarm before entering the inner torch.cuda.graph context, so the first retained allocation still comes from the shared graph pool during capture.

For DCP sizes unsupported by the B12X PCIe pool, notably virtual TP6 with DCP3 or DCP6, NCCL graph nodes retain these staging addresses. The shared pool can recycle an address while later FULL descriptors are captured, producing intermittent Xid 31 illegal reads. Two preserved failures occurred at FULL graph 12/16 for MTP0 on rank 0 and 15/16 for MTP3 on rank 4; the Python frame where the asynchronous fault surfaced differed, but both NCCL watchdogs reported the same illegal-access class.

The outer vllm_cudagraph_capture_scope already spans descriptor prewarm and capture. Treating that scope as graph setup allocates each exact-shape pair before the inner CUDA graph pool is active.

Scope

This does not change transport selection, tensor contents, CUDA graph sizes, or steady-state execution. It only advances allocation and extends the lifetime of buffers that were already intended to remain graph-owned.

Measured overhead on TP6/DCP3 was about 10-20 MiB per GPU of graph memory. KV capacity was unchanged.

Validation

• 3 passed, 38 deselected for the packed A2A capture/prewarm/eager ownership tests.
• Fresh-cache E2E startup passed with the patched file for TP6/DCP3/MTP0 and TP6/DCP3/MTP3.
• Three dual baseline/patch startup rounds completed; the baseline also passed those retries, consistent with the original fault being nondeterministic.
• Exact v19 production image, InstantTensor BUFFERED, 16 FULL graph descriptors, no backend or KV dtype changes.

Summary by CodeRabbit

• Bug Fixes
  • Improved CUDA graph capture handling for distributed attention operations.
  • Ensured prewarmed buffers are retained and reused during capture, helping prevent buffer address conflicts and improving reliability.

[coderabbitai]

📝 Walkthrough

Walkthrough

The DCP A2A buffer path now treats vLLM cudagraph capture as capture-active alongside CUDA stream capture. Tests verify prewarmed buffer reuse before stream capture and confirm eager buffers remain unretained.

Changes

A2A capture buffer retention

Layer / File(s)	Summary
Capture detection and retention validation vllm/v1/attention/ops/dcp_alltoall.py, tests/distributed/test_dcp_a2a.py	Buffer retention now recognizes vLLM cudagraph capture, with tests covering prewarm reuse and eager-path non-retention.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

• local-inference-lab/vllm#117: Updates the same DCP A2A CUDA graph buffer-retention logic and related tests.

Suggested reviewers: lukealonso

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
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.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the main change: fixing packed A2A buffer lifetime during graph prewarm.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/gg-dcp-a2a-prewarm-20260719

---

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.}