[Core] Release CUDA graphs before sleep() unmap, re-capture on wake

[terafin]

What

With CUDA-graph capture pools resident, per-allocation cumem VMM unmap/remap operations slow dramatically during sleep()/wake_up() — field-measured ~6-9s for a 32B model versus ~2s without graphs — because every block the graph capture pool touches forces an individual VMM op.

This PR releases the captured graphs before the cumem unmap so the unmap takes the fast path, then re-captures them on the next full wake once weights and the KV cache have been remapped.

• GPUModelRunner.release_cudagraphs() — drops all captured graph entries (CUDAGraphWrapper / BreakableCUDAGraphWrapper instances + the encoder cudagraph manager), forces finalization (gc.collect() + empty_cache()) so the capture pool returns to the driver before the unmap, and reports whether re-capture is needed.
• GPUModelRunner.recapture_cudagraphs() — rebuilds graphs via capture_model().
• Worker.sleep() — releases graphs before allocator.sleep() (the unmap).
• Worker.wake_up() — re-captures after allocator.wake_up() (the remap), only on a full wake (no tags / kv_cache); a weights-only wake defers re-capture until the KV cache is back.

All paths are no-ops when cudagraph_mode is NONE.

Why this is complementary to #45623 (not a duplicate)

#45623 routes the graph pool through cumem and restores it in place on wake (a capacity optimization — frees graph memory while asleep, no re-capture). This PR takes the opposite, complementary tack for wake latency: it releases the graphs so the slow VMM unmap/remap path is avoided entirely, then re-captures. The release/recapture strategy is the workaround field-confirmed by @AlanFokCo on #45398 (no clean allocator-level fix was found for the in-place case).

Tests

tests/v1/cudagraph/test_cudagraph_release_on_sleep.py asserts the ordering contract:

• graphs are released before the allocator unmap in sleep() (the core assertion — fails pre-fix, passes post-fix),
• graphs are re-captured after the allocator remap in wake_up(),
• a weights-only wake defers re-capture until a kv_cache wake,
• release_cudagraphs() clears all wrapper instances + drops the encoder manager,
• cudagraph_mode == NONE is a no-op on both paths.

The ordering/no-op tests run on CPU (allocator + model interactions mocked); the real captured-graph release test is GPU-gated via current_platform.is_cuda_alike().

🤖 Generated with Claude Code

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

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 ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: 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.

🚀

[AlanFokCo]

Hey @terafin — great to see this formalized into a proper PR. This is exactly the workaround we've been running in production for multi-model hot-swap on A10 and H20 clusters, so I can confirm the approach is sound end-to-end.

A few notes from running this pattern at scale:

The ordering contract is critical and your implementation gets it right. We originally tried releasing graphs after the unmap (thinking it was a cleanup step) and hit intermittent segfaults — the capture pool holds VA references that the unmap invalidates. Your release → unmap → remap → recapture sequence is the only safe order, and the test asserting it is valuable.

Re-capture cost is worth calling out. For a 32B model the capture_model() call takes ~2-3s, which is acceptable for a full sleep/wake cycle but adds up in multi-model switching scenarios where you might switch several times per minute. Two optimizations we've validated:

• Same-architecture switches (e.g., Qwen3-32B → another Qwen3-32B variant): the captured graphs are structurally identical, so re-capture can be skipped entirely — the existing graph entries remain valid after a weights-only reload.
• Reduced batch-size set: capturing with fewer decode batch sizes (e.g., 8 instead of the full 52) cuts re-capture from ~40s to ~2.3s on 32B TP=8 with piecewise mode, at minimal throughput cost.

These are follow-up optimizations, not blockers for this PR — just sharing in case they inform the design.

On interaction with #45623: the in-place restore approach (routing the graph pool through cumem) is strictly better when it works — zero re-capture cost. But as @aoshen02 noted, it doesn't catch all graph memory on MoE/EP configs. The release/recapture path in this PR serves as a reliable fallback. If both land, it might be worth a flag or auto-detection to prefer in-place when coverage is sufficient and fall back to release/recapture otherwise.

One minor suggestion: in release_cudagraphs(), the gc.collect() + synchronize() + empty_cache() triple is necessary (we tried each subset — all three are needed for the capture pool to actually return to the driver), but it might be worth a brief inline comment explaining why all three are required, since it looks like defensive over-flushing at first glance.

Overall LGTM — this is a meaningful latency improvement for anyone using sleep mode with CUDA graphs enabled.

[aoshen02]

Hi Alan, thanks for the help, will take a look at it right now.

[terafin]

Thanks @AlanFokCo / @aoshen02 — much appreciated. The inline comment explaining why gc.collect() + synchronize() + empty_cache() are all required is in. On the #45623 interplay: I'd keep this PR as the reliable release/recapture path and add the "prefer in-place when coverage is sufficient, fall back to release/recapture otherwise" auto-detect as a follow-up once #45623 lands — the MoE/EP coverage gap @aoshen02 noted means the recapture fallback needs to exist regardless. The same-architecture skip-recapture and reduced-batch-capture-set optimizations I'll bank as separate follow-ups to keep this PR focused.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @terafin.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork