[II] Bind B12X TP pools to CUDA graph owners - #324
voipmonitor wants to merge 1 commit into
Conversation
Include B12X pools registered for tensor-parallel process groups in graph capture and disposable-profile rollback. Retain separate TP and DCP capture scopes even when coordinators share a process group, while checkpointing shared pool storage once. The helper remains a no-op when a process group has no registered B12X pool. Models and deployments without B12X tensor-parallel collectives retain their existing execution path. Validation: tests/distributed/test_dcp_a2a.py passes 33 tests with 20 hardware-dependent skips in the PyTorch 2.13/CUDA 13.3 Kimi-K3 image. Ruff, formatting, Python compilation, and whitespace validation pass.
|
Warning Review limit reached
Next review available in: 56 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 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 configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
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. Comment |
Record each active B12X DCP graph channel by process-group identity. A pool initialized after the graph context opens now prepares and enters that channel before its first operation, and collective calls select the active channel instead of the eager scheduler channel. Eager execution is unchanged. The change depends on the graph-owner integration in vLLM #324. Validation: Ruff format/check and git diff --check pass. Thirty-seven relevant DCP tests pass in the CUDA 13.3/PyTorch 2.13 Kimi-K3 image; four focused eager/graph channel tests pass.
Record each active B12X DCP graph channel by process-group identity. A pool initialized after the graph context opens now prepares and enters that channel before its first operation, and collective calls select the active channel instead of the eager scheduler channel. Eager execution is unchanged. The change depends on the graph-owner integration in vLLM #324. Validation: Ruff format/check and git diff --check pass. Thirty-seven relevant DCP tests pass in the CUDA 13.3/PyTorch 2.13 Kimi-K3 image; four focused eager/graph channel tests pass.
Record each active B12X DCP graph channel by process-group identity. A pool initialized after the graph context opens now prepares and enters that channel before its first operation, and collective calls select the active channel instead of the eager scheduler channel. Eager execution is unchanged. The change depends on the graph-owner integration in vLLM #324. Validation: Ruff format/check and git diff --check pass. Thirty-seven relevant DCP tests pass in the CUDA 13.3/PyTorch 2.13 Kimi-K3 image; four focused eager/graph channel tests pass.
1 participant
Behavior
Registers B12X PCIe pools associated with tensor-parallel process groups with the owning CUDA graph capture. Disposable profiling captures also checkpoint and roll back both tensor-parallel and decode-context-parallel B12X pool state.
Status: implemented and qualified.
Technical reason
B12X tensor-parallel projection collectives allocate graph-owned communication channels through the same pool implementation used by DCP attention collectives. Registering only the DCP coordinator leaves TP projection channels outside the CUDA graph lifecycle and allows disposable profiling captures to leak or reuse channels with the wrong owner.
Compatibility
The change does not enable a B12X kernel or collective. It only extends lifecycle handling to a B12X pool after another runtime path registers that pool for a tensor-parallel process group. Models without B12X TP collectives retain their existing execution path.
The patch changes only distributed graph lifecycle code and its tests. It contains no Kimi-K3 model logic, DCP algorithm changes, speculative decoding, loader changes, or launch scripts.
Validation
Validation used
voipmonitor/vllm:kimi-k3-qsrt-ii-vllm735952b-b12x180ccab-cu133-torch213-20260815-r3with PyTorch 2.13 and CUDA 13.3:tests/distributed/test_dcp_a2a.py: 33 passed, 20 hardware-dependent skipped.git diff --checkpass.Review scope
This pull request replaces only the B12X graph-pool lifecycle responsibility contained in #317. The TP projection operation that registers the pool is reviewed separately.