[GG] Fix MRV2 profiled CUDA graph pool lifecycle - #180
Conversation
|
Warning Review limit reached
Next review available in: 39 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 (2)
📝 WalkthroughWalkthroughChangesThe CUDA graph memory profiling path now reuses the production graph pool through a temporary pool anchor. Production capture releases the anchor on skipped, successful, and failed capture paths, with tests covering pool bindings, lifecycle events, and exception cleanup. CUDA graph pool anchor lifecycle
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant GPUModelRunner
participant GraphManagers
participant CUDA
GPUModelRunner->>CUDA: obtain global production pool
GPUModelRunner->>GraphManagers: profile and capture using production pool
GPUModelRunner->>CUDA: create pool anchor
GPUModelRunner->>GraphManagers: destroy profiling graphs
GPUModelRunner->>GraphManagers: run production capture
GraphManagers-->>GPUModelRunner: success or failure
GPUModelRunner->>CUDA: release pool anchor
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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: 1
🧹 Nitpick comments (2)
vllm/v1/worker/gpu/model_runner.py (1)
154-158: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse a Google-style docstring.
Add
Args:andReturns:sections forpool,device, and the retained graph/token pair.Proposed fix
- """Keep a graph-private pool live between profiling and real capture. + """Keep a graph-private pool live between profiling and real capture. PyTorch cannot reopen a pool whose last graph was reset while allocations remain. This tiny graph holds the pool reference until production capture. + + Args: + pool: CUDA graph pool to retain. + device: CUDA device on which to create the anchor. + + Returns: + The anchor CUDA graph and its retained token tensor. """As per coding guidelines, Python code must use Google-style docstrings with
Args:/Returns:/Raises:sections.🤖 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 154 - 158, Update the docstring for the helper that retains the graph-private pool to Google style, adding an Args: section documenting pool and device and a Returns: section documenting the retained graph/token pair; preserve the existing lifecycle explanation.Source: Coding guidelines
tests/v1/cudagraph/test_breakable_cudagraph.py (1)
220-263: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCover the skipped-capture cleanup path.
The suite covers failed capture but not
needs_capture() == False, so the new release at Line 1014 is unverified. Add a fake manager returningFalseand assert the anchor is reset and cleared.🤖 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 220 - 263, Add a test for the capture-skipped path in GPUModelRunner.capture_model using a fake cudagraph manager whose needs_capture() returns False; initialize _cudagraph_pool_anchor with a reset tracker, invoke capture_model(), and assert the anchor is reset, cleared, and no capture event occurs.
🤖 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/v1/worker/gpu/model_runner.py`:
- Around line 966-969: Update shutdown() to release and clear
_cudagraph_pool_anchor after shutdown synchronization completes, including the
path where startup exits before capture_model(). Preserve the existing assertion
and anchor creation flow in the initialization code.
---
Nitpick comments:
In `@tests/v1/cudagraph/test_breakable_cudagraph.py`:
- Around line 220-263: Add a test for the capture-skipped path in
GPUModelRunner.capture_model using a fake cudagraph manager whose
needs_capture() returns False; initialize _cudagraph_pool_anchor with a reset
tracker, invoke capture_model(), and assert the anchor is reset, cleared, and no
capture event occurs.
In `@vllm/v1/worker/gpu/model_runner.py`:
- Around line 154-158: Update the docstring for the helper that retains the
graph-private pool to Google style, adding an Args: section documenting pool and
device and a Returns: section documenting the retained graph/token pair;
preserve the existing lifecycle explanation.
🪄 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: 804e2e9d-5ca5-4115-a51c-452357555469
📒 Files selected for processing (2)
tests/v1/cudagraph/test_breakable_cudagraph.pyvllm/v1/worker/gpu/model_runner.py
7538b65 to
7fce7ed
Compare
|
Review follow-up finalized in |
Reuse the production pool during graph-memory profiling, keep it alive through production recapture, and account the full graph high-water mark. Assisted-by: OpenAI Codex Signed-off-by: Martin Vit <martin@voipmonitor.org>
7fce7ed to
0b25fdb
Compare
1391c9c
into
local-inference-lab:dev/gilded-gnosis
2 participants
Summary
Preserve the CUDA graph private-pool lifecycle across MRV2 graph-memory
profiling and production capture.
The previous implementation profiled into a graph pool, reset every graph,
kept allocations from that pool alive, and then attempted to capture new
graphs into the same pool handle. PyTorch's caching allocator rejects that
state: once the last graph reference is gone, a private pool with live
allocations has
use_count == 0and cannot be reopened.This change:
PyTorch-reserved and the outer profiler intentionally restores the pre-graph
Torch peak.
There is no tail padding, allocator over-read allowance, disabled graph mode,
or backend-specific workaround in this PR.
Why a new PR
This supersedes #168. That PR is hosted on an external fork that the current
maintainer account cannot update, and its one-line pool reuse is incomplete:
it reproduces PyTorch's
use_count > 0allocator assertion when retainedallocations outlive the last profiling graph. This PR retains the valid pool
reuse idea and adds the missing lifecycle and accounting contract.
Upstream PR vllm-project#47925 was also checked. It pre-sizes model output
staging buffers and does not address private-pool teardown/reopen semantics, so
this is not duplicate work.
Root-cause proof
A minimal CUDA repro holds an output allocated by a private graph pool, resets
the last graph, empties the cache, and recaptures into the same pool handle.
Without a live graph it hits the same allocator assertion. A one-kernel anchor
kept alive across teardown makes recapture succeed without adding pool storage.
The unit tests cover pool identity, anchor ordering, full high-water accounting,
and anchor release when production capture throws.
Validation
Exact PR-only image:
git diff --checktests/v1/cudagraph/test_breakable_cudagraph.pyandtests/v1/worker/test_gpu_worker.py: 27 passedThe complete
rtx6kpro#34
stack was then tested in an isolated image containing only current GG, #172,
this runtime change, and SparkInfer #76:
The DCP2 profile measured 1.00 GiB total graph high-water with 0.72 GiB already
retained in the reusable pool and produced 1,102,720 KV tokens. The older
gross - retainedcalculation reported more KV because it omitted those livePyTorch-reserved pages; that was overcommit, not free capacity. Operators can
recover equivalent KV capacity by raising GMU by the logged graph-memory delta
(0.91 to approximately 0.9205 in this test).
No new GPU fault or server error appeared in either long-context gate.
AI assistance disclosure
Root-cause analysis, implementation, tests, hardware validation, and this
write-up were prepared with OpenAI Codex assistance and reviewed by the human
submitter.