[GG] fix(speculation): isolate graph resources and profile MRV2 memory

[voipmonitor]

Summary

Establish one ownership contract for target and draft CUDA-graph resources, then use it to make MRV2 memory profiling accurate and disposable.

This is intentionally one atomic PR: MRV2 profiling captures both target and speculator graphs. That capture is only safe when their workspace views and SparkInfer communication channels have independent owners. MoE scheduling is unrelated and remains in #150.

What changes

• workspace.py gives each (ubatch, execution owner) an independent reusable buffer. Target uses lane 0; speculative execution uses lane 1.
• The model runner applies lane 1 to every speculator phase, including the temporary MRV2 profiling capture.
• SparkInfer PCIe all-reduce and DCP A2A bind channels to the enclosing graph capture instead of sharing target/draft channel lifetime.
• Oneshot buffers are sized to their dispatch limits; the larger DMA buffer remains sized for scheduled prefill traffic.
• MRV2 profiling captures target and draft graphs in a private pool, includes sparse-DCP transient memory, then destroys graphs and rolls back channels before restoring production pools.

Root causes and lifecycle

Previously, graph-retained workspace views could alias, communication channels could outlive the graph that created them, and MRV2 reserved no production graph footprint. Sparse MLA transient buffers were also absent from the dummy profile.

A disposable profile now follows this order:

1. checkpoint production communication channels
2. capture target and draft graphs in a private pool
3. measure graph and sparse-DCP transient memory
4. synchronize and destroy every disposable graph while its pool is still owned
5. roll back transient channels and restore production pool references

Tests assert lane separation and restoration, independent graph channels, cleanup ordering, retained-pool accounting, and sparse-DCP workspace accounting.

Scope and dependencies

• Directly based on dev/gilded-gnosis; no Docker-only overlay.
• Replaces the separately opened [GG] fix(workspace): isolate target and speculative execution buffers #151 and [GG] fix(pcie): isolate target and draft graph channels #152, removing their merge-order dependency.
• Runtime channel checkpoint/capture support is provided by SparkInfer #59.
• Shared-expert/MoE scheduling remains isolated in [GG] fix(moe): safely overlap shared experts before resident grids #150 with SparkInfer #60.

Validation

Release CUDA image with four visible GPUs:

• tests/v1/worker/test_workspace.py
• tests/distributed/test_b12x_fused_all_reduce.py
• tests/distributed/test_dcp_a2a.py
• tests/v1/attention/test_b12x_mla_dcp_workspace.py
• tests/v1/cudagraph/test_breakable_cudagraph.py
• 103 passed
• Ruff lint and format checks passed
• git diff --check passed

[coderabbitai]

Warning

Review limit reached

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

Next review available in: 42 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

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: 7bded19a-792d-4e01-8fde-9632634cae1d

📥 Commits

Reviewing files that changed from the base of the PR and between 91788aa and c1b446a.

📒 Files selected for processing (18)

• tests/distributed/test_b12x_fused_all_reduce.py
• tests/distributed/test_dcp_a2a.py
• tests/v1/attention/test_b12x_mla_dcp_workspace.py
• tests/v1/cudagraph/test_breakable_cudagraph.py
• tests/v1/worker/test_workspace.py
• vllm/distributed/device_communicators/custom_all_reduce.py
• vllm/distributed/parallel_state.py
• vllm/envs.py
• vllm/model_executor/layers/attention/mla_attention.py
• vllm/v1/attention/ops/dcp_alltoall.py
• vllm/v1/worker/gpu/cudagraph_utils.py
• vllm/v1/worker/gpu/model_runner.py
• vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
• vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
• vllm/v1/worker/gpu/spec_decode/speculator.py
• vllm/v1/worker/gpu/warmup.py
• vllm/v1/worker/gpu_worker.py
• vllm/v1/worker/workspace.py

📝 Walkthrough

Walkthrough

Adds B12X PCIe/DCP channel checkpointing, complete disposable CUDA graph memory profiling with cleanup and pool restoration, and sparse MLA DCP transient-memory profiling controlled by a new environment flag.

Changes

B12X profiling and graph lifecycle

Layer / File(s)	Summary
B12X channel checkpointing vllm/distributed/device_communicators/custom_all_reduce.py, vllm/distributed/parallel_state.py, vllm/v1/attention/ops/dcp_alltoall.py, tests/distributed/*	PCIe and DCP A2A channel state can be checkpointed and rolled back across tensor, data, and pipeline parallel transports, with ordering and delegation tests.
CUDA graph manager lifecycle vllm/v1/worker/gpu/cudagraph_utils.py, vllm/v1/worker/gpu/spec_decode/*, tests/v1/cudagraph/*	CUDA graph managers and speculator managers expose cleanup hooks that clear captured graphs and model-specific cached state.
Disposable CUDA graph memory profiling vllm/v1/worker/gpu/model_runner.py, tests/v1/cudagraph/test_breakable_cudagraph.py	profile_cudagraph_memory() initializes minimal KV cache state, captures disposable graphs with temporary pools, restores channels and pools, cleans up state, and measures retained memory.
Sparse MLA attention profiling vllm/envs.py, vllm/model_executor/layers/attention/mla_attention.py, tests/v1/attention/*	Adds an environment-controlled estimate and profiling allocation for sparse DCP MLA transient memory across workspace, projection, dispatch, and pure_a2a configurations.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant GPUModelRunner
  participant CudaGraphManagers
  participant parallel_state
  participant MLAAttention
  GPUModelRunner->>CudaGraphManagers: assign profiling pools
  GPUModelRunner->>parallel_state: checkpoint B12X channels
  GPUModelRunner->>MLAAttention: profile sparse transient memory
  GPUModelRunner->>CudaGraphManagers: capture and clear disposable graphs
  GPUModelRunner->>parallel_state: rollback B12X channels
  GPUModelRunner->>CudaGraphManagers: restore production pools
  GPUModelRunner->>GPUModelRunner: measure retained capture memory

Loading

Possibly related PRs

• local-inference-lab/vllm#131: Adds the same B12X CustomAllreduce channel checkpoint/rollback delegation regression test.

Suggested reviewers: lukealonso, woosukkwon

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 20.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
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 is concise and accurately reflects the main changes around graph resource isolation and MRV2 memory profiling.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/gg-concurrency-lifecycle-consolidated-20260721

---

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

[voipmonitor]

Split completed. This PR head was reduced from the former aggregate commit a6cb402f to the MRV2-only commit 91788aac.

Extracted direct-to-GG PRs:

• [GG] fix(moe): safely overlap shared experts before resident grids #150: MoE shared-expert ordering
• [GG] fix(workspace): isolate target and speculative execution buffers #151: target/speculator workspace lanes
• [GG] fix(pcie): isolate target and draft graph channels #152: PCIe/DCP graph-channel isolation

SparkInfer companion split:

• fix(pcie): isolate CUDA graph channel lifecycle sparkinfer#59: PCIe lifecycle
• fix(moe): expose safe shared-expert overlap plans sparkinfer#60: MoE plan capability

Nothing has been silently discarded. The two SparkInfer replacement branches reconstruct #58 exactly. On the vLLM side, #151 and #152 merge first; this PR is then refreshed against GG to apply workspace lane 1 to the profiling-only speculator capture and resolve the overlapping DCP test/module hunks. After that mechanical refresh, the intended functional difference from the former aggregate #149 is only the profiler-guided MoE ordering correction documented in #150.

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (1)

tests/v1/cudagraph/test_breakable_cudagraph.py (1)

42-135: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Test only covers pre-existing wrappers, not ones created mid-capture.

FakeWrapper is instantiated before profile_cudagraph_memory() runs, so this test only exercises wrappers that are already in _all_instances prior to capture. It doesn't cover a wrapper being lazily created during capture() (e.g. a BreakableCUDAGraphWrapper created by init_breakable_cg_runner), which is the exact gap flagged in model_runner.py's profile_cudagraph_memory review comment — such a wrapper's graph_pool is never redirected to the disposable profiling_pool.

🤖 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/cudagraph/test_breakable_cudagraph.py` around lines 42 - 135, Extend
test_memory_profile_destroys_graphs_before_restoring_pools to create a new
wrapper during FakeManager.capture, register it in the relevant wrapper instance
collection, and assert its graph_pool uses profile_pool before cleanup and
production_pool afterward. Preserve the existing assertions for pre-existing
wrappers and event ordering while covering lazy wrapper creation during capture.

🤖 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 `@vllm/model_executor/layers/attention/mla_attention.py`:
- Around line 784-793: In the candidate construction logic, update the projected
AG_RS branch in the dcp_project_before_merge path to add (last_ag_rs_row, True)
whenever last_ag_rs_row exceeds project_threshold, regardless of
workspace_eligible. Keep the existing row bounds and unprojected candidate logic
unchanged.

In `@vllm/v1/worker/gpu/model_runner.py`:
- Around line 876-915: Update ModelCudaGraphManager.init_breakable_cg_runner so
newly created BreakableCUDAGraphWrapper instances inherit the manager’s
currently active pool when self.pool is set, rather than always retaining the
global graph pool. Preserve the existing wrapper creation behavior and allow
normal captures to continue using the global pool when no manager pool is
configured.

---

Nitpick comments:
In `@tests/v1/cudagraph/test_breakable_cudagraph.py`:
- Around line 42-135: Extend
test_memory_profile_destroys_graphs_before_restoring_pools to create a new
wrapper during FakeManager.capture, register it in the relevant wrapper instance
collection, and assert its graph_pool uses profile_pool before cleanup and
production_pool afterward. Preserve the existing assertions for pre-existing
wrappers and event ordering while covering lazy wrapper creation during capture.

🪄 Autofix (Beta)

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: 8e2ea35f-a3e5-49b3-881a-039e6e063db0

📥 Commits

Reviewing files that changed from the base of the PR and between 73e4a8c and 91788aa.

📒 Files selected for processing (14)

• tests/distributed/test_b12x_fused_all_reduce.py
• tests/distributed/test_dcp_a2a.py
• tests/v1/attention/test_b12x_mla_dcp_workspace.py
• tests/v1/cudagraph/test_breakable_cudagraph.py
• vllm/distributed/device_communicators/custom_all_reduce.py
• vllm/distributed/parallel_state.py
• vllm/envs.py
• vllm/model_executor/layers/attention/mla_attention.py
• vllm/v1/attention/ops/dcp_alltoall.py
• vllm/v1/worker/gpu/cudagraph_utils.py
• vllm/v1/worker/gpu/model_runner.py
• vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
• vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
• vllm/v1/worker/gpu/spec_decode/speculator.py