[GG] Fix MRV2 CUDA graph and sparse DCP memory profiling - #131
Conversation
|
Warning Review limit reached
Next review available in: 23 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. 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 (1)
📝 WalkthroughWalkthroughThe changes add B12X sparse MLA transient-memory profiling and implement CUDA graph memory profiling with minimal KV-cache setup, capture cleanup, memory measurement, and restoration of PCIe and DCP A2A channel state. ChangesSparse MLA memory profiling
CUDA graph memory profiling
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant GPUModelRunner
participant CudaGraphManager
participant Speculator
participant parallel_state
participant CUDA
GPUModelRunner->>GPUModelRunner: Initialize minimal KV cache
GPUModelRunner->>parallel_state: Checkpoint B12X channels
GPUModelRunner->>CudaGraphManager: Capture CUDA graphs
GPUModelRunner->>Speculator: Get and capture speculator managers
GPUModelRunner->>CUDA: Measure memory deltas
GPUModelRunner->>CudaGraphManager: Clear captured graphs
GPUModelRunner->>Speculator: Clear cudagraphs
GPUModelRunner->>parallel_state: Roll back B12X channels
CUDA-->>GPUModelRunner: Return estimated graph memory
🚥 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.
🧹 Nitpick comments (1)
vllm/distributed/device_communicators/custom_all_reduce.py (1)
653-668: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider adding Google-style
Args:/Returns:/Raises:sections.As per coding guidelines, Python code must use Google-style docstrings with
Args:/Returns:/Raises:sections instead of reStructuredText/Sphinx fields. While these single-line docstrings don't use Sphinx fields, expanding them to include explicitArgs:,Returns:, andRaises:sections would fully align with the guideline. Consider applying this to the newly added docstrings across the PR.📝 Proposed docstring updates
def checkpoint_pcie_channels(self) -> Any | None: - """Snapshot B12X channels before a throwaway graph capture.""" + """Snapshot B12X channels before a throwaway graph capture. + + Returns: + The checkpoint object, or None if the runtime does not support it. + """ runtime = self._pcie_runtime checkpoint = getattr(runtime, "checkpoint_channels", None) if checkpoint is None: return None return checkpoint() def rollback_pcie_channels(self, checkpoint: Any) -> None: - """Release B12X channels created after ``checkpoint``.""" + """Release B12X channels created after ``checkpoint``. + + Args: + checkpoint: The checkpoint object returned by checkpoint_pcie_channels. + + Raises: + RuntimeError: If the B12X PCIe all-reduce runtime is unavailable. + """ runtime = self._pcie_runtime rollback = getattr(runtime, "rollback_channels", None) if rollback 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 653 - 668, Expand the docstrings for checkpoint_pcie_channels and rollback_pcie_channels with Google-style sections: document the checkpoint returned by checkpoint_pcie_channels, the checkpoint argument accepted by rollback_pcie_channels, and the optional return value and RuntimeError raised when the PCIe runtime is unavailable. Preserve the existing behavior and descriptions.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.
Nitpick comments:
In `@vllm/distributed/device_communicators/custom_all_reduce.py`:
- Around line 653-668: Expand the docstrings for checkpoint_pcie_channels and
rollback_pcie_channels with Google-style sections: document the checkpoint
returned by checkpoint_pcie_channels, the checkpoint argument accepted by
rollback_pcie_channels, and the optional return value and RuntimeError raised
when the PCIe runtime is unavailable. Preserve the existing behavior and
descriptions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: e1d6c703-50dc-4c73-b37f-5474cb0abdce
📒 Files selected for processing (7)
tests/distributed/test_b12x_fused_all_reduce.pytests/distributed/test_dcp_a2a.pytests/v1/cudagraph/test_breakable_cudagraph.pyvllm/distributed/device_communicators/custom_all_reduce.pyvllm/distributed/parallel_state.pyvllm/v1/attention/ops/dcp_alltoall.pyvllm/v1/worker/gpu/model_runner.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/v1/cudagraph/test_breakable_cudagraph.py
- vllm/v1/worker/gpu/model_runner.py
|
Addressed the lifecycle API docstring review note in 454b853. Full focused MRV2/DCP/CUDA-graph validation in the release environment: 94 tests passed. |
1 participant
Summary
VLLM_MEMORY_PROFILE_INCLUDE_ATTN=1Root causes
MRV2 previously returned
0fromprofile_cudagraph_memory(). On large TP6/DCP6 configurations this filled remaining VRAM with KV cache without reserving the production graph footprint. Sparse MLA also skips attention during the dummy profile, so its eager AG/RS transient was omitted.The first real profiling implementation exposed a separate lifetime bug. It originally checkpointed only immediately before draft capture, leaving target-created B12X graph channels outside rollback. It also restored the production graph pool before destroying disposable graphs. Cleanup must instead cover both target and draft captures and happen while every owner still points at the private pool.
The fixed order is:
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:Trueis not a fix. Validation usesbackend:native; expandable segments can change address layout enough to hide a lifetime failure and are incompatible with native KV-offload address registration.Validation
git diff --checkpassedThis PR is based directly on
dev/gilded-gnosis; it contains no Docker-only source overlay. Pair it with local-inference-lab/sparkinfer#47 for coordinated IPC teardown.