[GG] fix(distributed): isolate semantic PCIe graph channels - #216
[GG] fix(distributed): isolate semantic PCIe graph channels#216malaiwah wants to merge 5 commits into
Conversation
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>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
📝 WalkthroughWalkthroughThe change adds stable semantic channel IDs across B12X, DCP, CUDA graph, and speculative-decoding capture paths. It separates target, draft, and encoder captures and adds lifecycle, propagation, failure, and rollback coverage. ChangesSemantic channel isolation
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 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}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (5)
tests/v1/worker/test_gpu_autoregressive_speculator.py (1)
171-172: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe test name claims rank symmetry that the fixture cannot verify.
capture_for_ranktakes no rank input and no rank-dependent state changes between the two calls. The assertion proves that capture is deterministic in one process. Rename the test to state determinism, or add a rank-dependent attribute to the fake speculator so a rank-derived channel id would fail.♻️ Proposed rename
-def test_autoregressive_capture_ids_are_cross_rank_symmetric_and_separated(): - def capture_for_rank() -> list[str]: +def test_autoregressive_capture_ids_are_deterministic_and_phase_separated(): + def capture_channel_ids() -> list[str]:🤖 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_autoregressive_speculator.py` around lines 171 - 172, Rename test_autoregressive_capture_ids_are_cross_rank_symmetric_and_separated to describe single-process capture determinism, since capture_for_rank has no rank input or rank-dependent state. Keep the existing assertions and test behavior unchanged.vllm/distributed/parallel_state.py (1)
1650-1667: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd an
Args:section for the newchannel_idparameter.The docstring describes
channel_idin prose only. Add Google-style sections so the new public parameter and the conflict behavior are documented.As per coding guidelines: "Use Google-style docstrings in Python code, with
Args:/Returns:/Raises:sections".📝 Proposed docstring sections
A caller may pass an explicit ``graph_capture_context`` to control the stream used (e.g. to capture on the default stream). + + Args: + device: Device that owns the capture stream. + graph_capture_context: Existing capture context to reuse. + channel_id: Stable distributed identity for this graph owner. + + Raises: + ValueError: If ``channel_id`` conflicts with the channel identity + already stored on ``graph_capture_context``. """🤖 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 1650 - 1667, Add a Google-style Args section to the graph_capture docstring documenting the channel_id parameter, including its optional semantic channel identity and the behavior when it conflicts with an explicitly provided graph_capture_context; preserve the existing prose and document other parameters only if already required by the surrounding docstring conventions.Source: Coding guidelines
vllm/distributed/device_communicators/custom_all_reduce.py (1)
672-699: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the new
channel_idparameter in the docstring.
capturenow takes a keyword-onlychannel_idand raisesRuntimeErrorwhen it is missing for a PCIe runtime. AddArgs:andRaises:sections that describe this contract.As per coding guidelines: "Use Google-style docstrings in Python code, with
Args:/Returns:/Raises:sections".📝 Proposed docstring update
"""Bind communicator resources to the enclosing CUDA graph capture. Legacy custom all-reduce registers graph buffers on exit. SparkInfer PCIe channels are instead created on the graph's owning stream and remain valid for the graph lifetime. + + Args: + stream: CUDA stream owned by the enclosing graph capture. + channel_id: Stable identity shared by every rank for this graph + owner. Required when the SparkInfer PCIe runtime is active. + + Raises: + RuntimeError: If the PCIe runtime is active and ``channel_id`` is + ``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/distributed/device_communicators/custom_all_reduce.py` around lines 672 - 699, Update the `capture` method docstring to use Google-style `Args:` and `Raises:` sections: document the optional `channel_id` keyword parameter and state that `RuntimeError` is raised when a PCIe runtime is active without an explicit channel ID. Preserve the existing description and document other parameters as needed for a complete signature contract.Source: Coding guidelines
vllm/v1/worker/gpu_model_runner.py (2)
6844-6865: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winConsider isolating the graph-pool restoration from earlier cleanup failures.
The cleanup block runs several independent steps in sequence. If an earlier step raises, for example
CUDAGraphWrapper.clear_all_graphs()on ROCm, the loop at lines 6853-6855 never restoresinstance.graph_pool. The persistent wrappers then keep the disposableprofiling_poolhandle for later production capture.The outer
finallyalready protects the channel rollback. Apply the same protection to the pool restoration, so the persistent wrapper state is always returned to its original value.♻️ Proposed restructuring
finally: try: - set_cudagraph_capturing_enabled(False) - CUDAGraphWrapper.clear_all_graphs() - BreakableCUDAGraphWrapper.clear_all_graphs() - if encoder_cudagraph_manager is not None: - encoder_cudagraph_manager.clear() - all_wrappers = list(CUDAGraphWrapper._all_instances) + list( - BreakableCUDAGraphWrapper._all_instances - ) - for instance in all_wrappers: - if id(instance) in original_pools: - instance.graph_pool = original_pools[id(instance)] + try: + set_cudagraph_capturing_enabled(False) + CUDAGraphWrapper.clear_all_graphs() + BreakableCUDAGraphWrapper.clear_all_graphs() + if encoder_cudagraph_manager is not None: + encoder_cudagraph_manager.clear() + finally: + all_wrappers = list(CUDAGraphWrapper._all_instances) + list( + BreakableCUDAGraphWrapper._all_instances + ) + for instance in all_wrappers: + if id(instance) in original_pools: + instance.graph_pool = original_pools[id(instance)] for key_set in self.cudagraph_dispatcher.cudagraph_keys.values(): key_set.clear()🤖 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 6844 - 6865, Ensure graph-pool restoration always runs even when earlier cleanup in the try block fails. In the cleanup flow around CUDAGraphWrapper.clear_all_graphs and the all_wrappers restoration loop, isolate the graph_pool restoration in its own finally-protected section, while preserving the existing outer rollback via rollback_b12x_graph_channels.
6970-6984: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winShare the drafter-eligibility predicate with
_dummy_run.The method repeats the same drafter-method test used in
_dummy_runat lines 6186-6191:use_eagle() or uses_draft_model() or uses_extract_hidden_states(). The two sites must stay in agreement, becauserun_drafterrouting depends on it. If a future change adds a method to one site only, target-only passes and draft-only passes disagree and a drafter graph is captured twice or never.Extract the method test into a single helper and call it from both sites. A short Google-style docstring on
_captures_independent_drafter_graphswould also record whyPIECEWISEis the only qualifying mode.🤖 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 6970 - 6984, Extract the shared drafter-method predicate currently duplicated in `_dummy_run` and `_captures_independent_drafter_graphs` into a helper, then call that helper from both sites so their eligibility decisions remain identical. Preserve the existing `use_eagle()`, `uses_draft_model()`, and `uses_extract_hidden_states()` behavior, and add a brief Google-style docstring to `_captures_independent_drafter_graphs` explaining why only `PIECEWISE` qualifies.Source: Coding guidelines
🤖 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/distributed/device_communicators/custom_all_reduce.py`:
- Around line 982-995: The __del__ path currently retains PCIe-backed
CustomAllreduce instances indefinitely in
_ABANDONED_B12X_PCIE_ALLREDUCE_QUARANTINE. Replace this unbounded
object-retention approach with a flag that prevents unsafe finalizer cleanup, or
add a coordinated drain/cleanup mechanism that releases every quarantined
instance and its PCIe runtime, DMA resources, and shared buffers; ensure
explicit close() remains the coordinated cleanup path.
In `@vllm/v1/worker/gpu_model_runner.py`:
- Around line 6807-6816: Update the CUDA graph memory debug message in the
logger.debug call to replace the Unicode multiplication sign (×) with the ASCII
character x, leaving the formatting arguments and message meaning unchanged.
- Around line 6765-6805: Filter out entries with empty descs when constructing
component_descs in the current capture flow, so the loop never reaches
mem_samples[0] without profiling data. Apply the same descs-nonempty filter to
the matching comprehension in capture_model, preserving existing behavior for
populated descriptor lists.
---
Nitpick comments:
In `@tests/v1/worker/test_gpu_autoregressive_speculator.py`:
- Around line 171-172: Rename
test_autoregressive_capture_ids_are_cross_rank_symmetric_and_separated to
describe single-process capture determinism, since capture_for_rank has no rank
input or rank-dependent state. Keep the existing assertions and test behavior
unchanged.
In `@vllm/distributed/device_communicators/custom_all_reduce.py`:
- Around line 672-699: Update the `capture` method docstring to use Google-style
`Args:` and `Raises:` sections: document the optional `channel_id` keyword
parameter and state that `RuntimeError` is raised when a PCIe runtime is active
without an explicit channel ID. Preserve the existing description and document
other parameters as needed for a complete signature contract.
In `@vllm/distributed/parallel_state.py`:
- Around line 1650-1667: Add a Google-style Args section to the graph_capture
docstring documenting the channel_id parameter, including its optional semantic
channel identity and the behavior when it conflicts with an explicitly provided
graph_capture_context; preserve the existing prose and document other parameters
only if already required by the surrounding docstring conventions.
In `@vllm/v1/worker/gpu_model_runner.py`:
- Around line 6844-6865: Ensure graph-pool restoration always runs even when
earlier cleanup in the try block fails. In the cleanup flow around
CUDAGraphWrapper.clear_all_graphs and the all_wrappers restoration loop, isolate
the graph_pool restoration in its own finally-protected section, while
preserving the existing outer rollback via rollback_b12x_graph_channels.
- Around line 6970-6984: Extract the shared drafter-method predicate currently
duplicated in `_dummy_run` and `_captures_independent_drafter_graphs` into a
helper, then call that helper from both sites so their eligibility decisions
remain identical. Preserve the existing `use_eagle()`, `uses_draft_model()`, and
`uses_extract_hidden_states()` behavior, and add a brief Google-style docstring
to `_captures_independent_drafter_graphs` explaining why only `PIECEWISE`
qualifies.
🪄 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: 22931018-78fe-47e9-819d-adc23809eca9
📒 Files selected for processing (17)
tests/distributed/test_b12x_fused_all_reduce.pytests/distributed/test_dcp_a2a.pytests/v1/cudagraph/test_breakable_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/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
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>
Review closeout at
|
1 participant
Purpose
Pair vLLM CUDA-graph callers with the semantic multi-channel PCIe ownership contract in SparkInfer #105, and ensure normal teardown closes those collective owners while their process groups are still live.
A process-local CUDA stream handle is not a sufficient identity for transport resources owned by independently captured/replayed target, draft, encoder, profiling, production, and eager paths. This series:
channel_idvalues through V1/V2 graph capture, custom all-reduce, and DCP A2A;__del__non-collective without an unbounded vLLM wrapper quarantine.Distributed PCIe graph capture fails closed when a caller omits its semantic identity. Related defects remain tracked by #208 and #215 until this draft is reviewed and merged.
Status and exact provenance
The first four clean commits are patch-identical to the field lineage; stable patch IDs and
git range-diffreport=:a89816a3fc186b0bbd09382752cfe40dbe1e289a7dbe5a27f99e1e7bd344da3652201611Do not combine this with rejected SparkInfer #101/#103 heads or intermediate
31fa6a4. That pair passed focused tests but failed full GLM startup when descriptor warmup insidevllm:target:profilecollided with the static eager identity. SparkInferbc629805repaired that owner-stream handoff without weakening side-stream collision checks.Review follow-up at
52201611All eight CodeRabbit findings were valid and addressed:
CustomAllreduceclose before NCCL/process-group destruction;Validation:
git diff --check: pass;parallel_state.pylines 2115/2133/2134/2159; those are outside this diff;Exact GPU qualification at
f99e1e7/d344da36Paired with SparkInfer
bc629805on 4x RTX PRO 6000 Blackwell:This closes the 131K text-only/offload-disabled repair-attribution profile. It does not by itself qualify 512K, vision, or live LMCache. Healthy serve/shutdown proves normal-path compatibility; the abnormal finalizer and new explicit-close ordering are covered by focused tests, not by claiming that ordinary shutdown exercised Python GC at exactly that boundary.
Claim boundary and evidence
This is correctness and lifecycle work. No direct VRAM, PP, TG, or acceptance-length gain is claimed. The approximately 1.1 GiB/rank full-model recovery belongs to vLLM #210.
Contribution disclosure
This series was developed with OpenAI Codex assistance; assisted commits carry the trailer. Duplicate checks covered #208, #215, #210, and open semantic-channel/finalizer PRs; no duplicate caller implementation was found. SparkInfer #105 is the required backend half, not a duplicate.
Per repository policy, this PR remains a draft until the human submitter has reviewed every changed line and can defend the current
52201611successor end-to-end. Human maintainer review and independent counter-validation are requested.