[GG] fix/perf: isolate B12X graph channels and capture DSpark context KV - #251
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📥 CommitsReviewing files that changed from the base of the PR and between 45849684bea6266acec672a968db7e33b79daf0c and c66ce73. 📒 Files selected for processing (9)
🚧 Files skipped from review as they are similar to previous changes (6)
📝 WalkthroughWalkthroughThe PR propagates stable semantic channel IDs through B12X communication, distributed graph capture, DCP all-to-all, DFlash context graphs, and model-runner CUDA graph capture. It also changes communicator cleanup ordering and adds lifecycle, validation, rollback, and padding tests. ChangesB12X communication and lifecycle
CUDA graph capture
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related issues
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
left a comment
There was a problem hiding this comment.
🧹 Nitpick comments (6)
tests/distributed/test_dcp_a2a.py (1)
749-814: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the new
channel_idvalidation branches.This test passes a context that already carries a channel ID, so it exercises only the pass-through branch of
graph_capture. Two new branches invllm/distributed/parallel_state.pystay untested: theValueErrorraised whenchannel_idconflicts withcontext.channel_id, and the clone performed when the context has no ID. Add two short cases for them.🤖 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_dcp_a2a.py` around lines 749 - 814, Extend test_global_graph_capture_enters_b12x_dcp_pool with two focused cases for graph_capture: assert it raises ValueError when an explicit channel_id conflicts with context.channel_id, and verify it clones a context lacking an ID with the requested channel_id while preserving the original context unchanged. Reuse the existing fake groups/context setup and validate the resulting context’s channel_id.vllm/distributed/parallel_state.py (1)
1682-1690: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePreserve the original
GraphCaptureContextidentity when injectingchannel_id.
context = GraphCaptureContext(context.stream, channel_id=channel_id)at line 1690 returns a different object for the caller ifcontext.channel_idisNone. A caller that readschannel_idfrom its own reference after the block will still seeNone, even though the yielded context carries the ID. IfGraphCaptureContextis mutable, update the existing instance instead of cloning it. If it is frozen, document the object substitution in the docstring.🤖 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 `@vllm/distributed/parallel_state.py` around lines 1682 - 1690, The channel_id injection in the graph capture context path must preserve the caller’s original GraphCaptureContext identity. Update the existing context’s channel_id when it is None, while retaining the conflict validation for mismatched IDs; only use object replacement if GraphCaptureContext is immutable, and document that substitution in the surrounding API docstring.vllm/v1/worker/gpu_model_runner.py (2)
6835-6854: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueCleanup after a
clear_all_graphsfailure is skipped.The inner
tryguarantees pool restore, and the outertryguarantees channel rollback. It does not guarantee the remaining cleanup. IfCUDAGraphWrapper.clear_all_graphs()orencoder_cudagraph_manager.clear()raises, the code skips dispatcher key reset, LoRA removal,_cleanup_profiling_kv_cache(), and the capture-counter restore. The profiling KV cache then stays allocated. The parameterized test attests/v1/worker/test_gpu_model_runner.pylines 2143-2146 encodes this behavior, so the gap is intentional today. Consider moving_cleanup_profiling_kv_cache()into its ownfinallyso GPU memory is released even when graph teardown fails.🤖 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 `@vllm/v1/worker/gpu_model_runner.py` around lines 6835 - 6854, Ensure _cleanup_profiling_kv_cache() runs even when CUDAGraphWrapper.clear_all_graphs(), BreakableCUDAGraphWrapper.clear_all_graphs(), or encoder_cudagraph_manager.clear() raises. Move the profiling KV-cache cleanup into an appropriate finally block while preserving graph-pool restoration, dispatcher reset, LoRA removal, counter restoration, and channel rollback behavior.
6725-6739: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the duplicated component filter.
profile_cudagraph_memoryandcapture_modelbuildcomponent_descswith identical logic. The two copies must stay in sync, because a change to the drafter-separation rule in one place silently diverges from the other.♻️ Proposed helper
def _component_capture_descs( self, component: str, capture_descs: list[tuple[CUDAGraphMode, list[BatchDescriptor]]], ) -> list[tuple[CUDAGraphMode, list[BatchDescriptor]]]: """Select the capture descriptors owned by one graph component. Args: component: Either ``"target"`` or ``"draft"``. capture_descs: All capture descriptors grouped by CUDA graph mode. Returns: The descriptor groups this component must capture. """ return [ (mode, descs) for mode, descs in capture_descs if descs and ( component == "target" or self._captures_independent_drafter_graphs(mode) ) ]Also applies to: 6899-6913
🤖 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 `@vllm/v1/worker/gpu_model_runner.py` around lines 6725 - 6739, Extract the duplicated component-desc filtering into a shared helper, such as _component_capture_descs, near the existing capture logic. Update both profile_cudagraph_memory and capture_model to call it for their target and draft components, preserving the current drafter-separation condition and empty-result handling.tests/v1/worker/test_gpu_model_runner.py (1)
1786-1863: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for capture-mode reset when
capture_modelfails.
capture_modelnow disables capture mode in afinallyblock (vllm/v1/worker/gpu_model_runner.pylines 6940-6942). No test in this file exercises that path.set_cudagraph_capturing_enabledis global state, so a regression leaves capture mode enabled for the rest of the process. The existing fixture already patchesset_cudagraph_capturing_enabled; recording its argument and raising from_capture_cudagraphscovers the case.Line 1801 also assigns
runner.encoder_cudagraph_manager = Noneand line 1825 overwrites it. Remove the first assignment.🤖 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/v1/worker/test_gpu_model_runner.py` around lines 1786 - 1863, The test test_v1_capture_separates_target_and_draft_semantic_channels should also cover capture_model failure and verify set_cudagraph_capturing_enabled is called with False in the finally path; record the patched setter’s arguments and make _capture_cudagraphs raise, while preserving the existing setup as appropriate. Remove the redundant initial runner.encoder_cudagraph_manager = None assignment since it is overwritten later.vllm/v1/worker/gpu/spec_decode/dflash/speculator.py (1)
197-237: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReset
context_cudagraph_manageron the ineligible branch.
init_cudagraph_managerassignsquery_cudagraph_managerunconditionally, but it only assignscontext_cudagraph_managerinside thewants_full and supports_full and self._speculator_name == "DSpark"branch. Ifinit_cudagraph_managerruns a second time with a mode that no longer qualifies, the stale context manager survives andproposekeeps dispatching context graphs that were captured for the previous configuration.♻️ Proposed change
if wants_full and supports_full and self._speculator_name == "DSpark": self.context_cudagraph_manager = DFlashContextCudaGraphManager( self.vllm_config, self.device, max_num_context_tokens=self.max_num_tokens, ) + else: + self.context_cudagraph_manager = None🤖 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 `@vllm/v1/worker/gpu/spec_decode/dflash/speculator.py` around lines 197 - 237, Update init_cudagraph_manager to explicitly set context_cudagraph_manager to None before the wants_full/supports_full/DSpark eligibility branch, so repeated initialization cannot retain a manager from a previous configuration; leave creation through DFlashContextCudaGraphManager unchanged for eligible DSpark modes.
🤖 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.
Nitpick comments:
In `@tests/distributed/test_dcp_a2a.py`:
- Around line 749-814: Extend test_global_graph_capture_enters_b12x_dcp_pool
with two focused cases for graph_capture: assert it raises ValueError when an
explicit channel_id conflicts with context.channel_id, and verify it clones a
context lacking an ID with the requested channel_id while preserving the
original context unchanged. Reuse the existing fake groups/context setup and
validate the resulting context’s channel_id.
In `@tests/v1/worker/test_gpu_model_runner.py`:
- Around line 1786-1863: The test
test_v1_capture_separates_target_and_draft_semantic_channels should also cover
capture_model failure and verify set_cudagraph_capturing_enabled is called with
False in the finally path; record the patched setter’s arguments and make
_capture_cudagraphs raise, while preserving the existing setup as appropriate.
Remove the redundant initial runner.encoder_cudagraph_manager = None assignment
since it is overwritten later.
In `@vllm/distributed/parallel_state.py`:
- Around line 1682-1690: The channel_id injection in the graph capture context
path must preserve the caller’s original GraphCaptureContext identity. Update
the existing context’s channel_id when it is None, while retaining the conflict
validation for mismatched IDs; only use object replacement if
GraphCaptureContext is immutable, and document that substitution in the
surrounding API docstring.
In `@vllm/v1/worker/gpu_model_runner.py`:
- Around line 6835-6854: Ensure _cleanup_profiling_kv_cache() runs even when
CUDAGraphWrapper.clear_all_graphs(),
BreakableCUDAGraphWrapper.clear_all_graphs(), or
encoder_cudagraph_manager.clear() raises. Move the profiling KV-cache cleanup
into an appropriate finally block while preserving graph-pool restoration,
dispatcher reset, LoRA removal, counter restoration, and channel rollback
behavior.
- Around line 6725-6739: Extract the duplicated component-desc filtering into a
shared helper, such as _component_capture_descs, near the existing capture
logic. Update both profile_cudagraph_memory and capture_model to call it for
their target and draft components, preserving the current drafter-separation
condition and empty-result handling.
In `@vllm/v1/worker/gpu/spec_decode/dflash/speculator.py`:
- Around line 197-237: Update init_cudagraph_manager to explicitly set
context_cudagraph_manager to None before the wants_full/supports_full/DSpark
eligibility branch, so repeated initialization cannot retain a manager from a
previous configuration; leave creation through DFlashContextCudaGraphManager
unchanged for eligible DSpark modes.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 8ae7897d-143d-4646-85fc-5d922ecc1977
📥 Commits
Reviewing files that changed from the base of the PR and between 3003860 and 45849684bea6266acec672a968db7e33b79daf0c.
📒 Files selected for processing (20)
tests/distributed/test_b12x_fused_all_reduce.pytests/distributed/test_custom_allreduce_lifecycle.pytests/distributed/test_dcp_a2a.pytests/v1/cudagraph/test_breakable_cudagraph.pytests/v1/spec_decode/test_dflash_context_cudagraph.pytests/v1/spec_decode/test_dflash_cudagraph_lifetime.pytests/v1/worker/test_gpu_autoregressive_speculator.pytests/v1/worker/test_gpu_model_runner.pyvllm/distributed/device_communicators/cuda_communicator.pyvllm/distributed/device_communicators/custom_all_reduce.pyvllm/distributed/parallel_state.pyvllm/v1/attention/ops/dcp_alltoall.pyvllm/v1/worker/gpu/cudagraph_utils.pyvllm/v1/worker/gpu/model_runner.pyvllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.pyvllm/v1/worker/gpu/spec_decode/autoregressive/speculator.pyvllm/v1/worker/gpu/spec_decode/dflash/cudagraph.pyvllm/v1/worker/gpu/spec_decode/dflash/speculator.pyvllm/v1/worker/gpu/spec_decode/speculator.pyvllm/v1/worker/gpu_model_runner.py
Assisted-by: OpenAI Codex Signed-off-by: Michel Belleau <michel.belleau@malaiwah.com>
Signed-off-by: Michel Belleau <michel.belleau@malaiwah.com>
Signed-off-by: Michel Belleau <michel.belleau@malaiwah.com>
Signed-off-by: Michel Belleau <michel.belleau@malaiwah.com>
Address semantic-channel review feedback by closing custom all-reduce before its process groups, keeping finalizers non-collective, and hardening CUDA graph profiling cleanup. Assisted-by: OpenAI Codex Signed-off-by: Michel Belleau <michel.belleau@malaiwah.com>
4584968 to
f8fb9ad
Compare
voipmonitor
changed the title
commented
Aug 7, 2026
|
Final r30 composition exposed and fixed one stale call-site after the current GG graph-manager API change: DFlash and autoregressive speculators still passed Commit The clean r30 image then completed live TP2 qualification on physical GPUs 4-5:
This confirms the semantic capture-channel integration against the current GG API rather than only against the earlier stacked source state. |
3 participants
Purpose
Make B12X CUDA-graph transport ownership deterministic and move the remaining DSpark context-KV projection/write stage into a dedicated FULL CUDA graph family.
These changes are intentionally reviewed together. The context-KV graph creates an additional distributed graph owner, so it requires the semantic channel lifecycle in the same source state. Keeping the two changes in independent PRs against
dev/gilded-gnosisproduced an order-dependent conflict and allowed the context graph to compile only through hidden stacked ancestry.This PR supersedes #247 and the previous stacked form of #251. Its history now starts at the current
dev/gilded-gnosishead; there is no private integration commit.Distributed graph lifecycle
DSpark context FULL graphs
0andPAD_SLOT_ID.Measured result
deepseek-ai/DeepSeek-V4-Flash-0731, TP2/DCP1, B12X W4A8, FP8 DS-MLA KV, InstantTensor, fixed probabilistic K5, same physical GPUs 4-5:The rank-0 four-step trace reduced DSpark propose mean from
2.408 msto1.510 ms, eager CUDA launches from296to188, and traced GPU span from68.44 msto63.74 ms. CC32 remained within run variance.At MNS64 / graph-row envelope 384, the graph pool grew by about
0.01 GiB, estimated graph reservation by about0.03 GiB, and engine KV capacity decreased by 74 tokens.Validation
e2666d9a65).git diff --check, and clean release composition: pass.Prefill remains PIECEWISE. Host metadata preparation, rejection sampling, and output bookkeeping remain outside CUDA graphs; this PR removes the remaining eager DSpark model-compute stage rather than claiming a literally zero-eager API step.
Summary by CodeRabbit
New Features
Bug Fixes