[Core] Offload CUDA graph + NCCL communicator memory in sleep mode#45623
Closed
aoshen02 wants to merge 1 commit into
Closed
[Core] Offload CUDA graph + NCCL communicator memory in sleep mode#45623aoshen02 wants to merge 1 commit into
aoshen02 wants to merge 1 commit into
Conversation
18 tasks
Extend sleep/wake (CuMemAllocator) to release two GPU memory classes that
previously stayed resident while an engine sleeps, restoring them on wake:
CUDA graphs: route the CUDA-graph capture pool through the cumem allocator under
a new 'graphs' tag, so the graph memory is offloaded to CPU and restored in
place on wake without re-capture (cumem preserves virtual addresses).
- cumem.py: graphs_tag, get_graph_pool_handle(); keep graphs resident when not
selected for offload; always attempt torch.cuda.empty_cache() on sleep
(guarded) so freed blocks return to the driver.
- platforms/cuda.py: graph_pool_handle() routes through the cumem graph pool
when the allocator singleton exists (sleep mode).
- gpu_worker.py: pin the cumem graph pool BEFORE capture
(_pin_sleep_mode_graph_pool overwrites CudaPlatform._global_graph_pool and
re-points every live CUDAGraphWrapper/BreakableCUDAGraphWrapper) and hold the
'graphs' tag across capture_model(); offload 'graphs' at both sleep levels.
The pre-capture pin is required on the v1 runner, where the global graph pool
is otherwise cached before the cumem instance exists and capture lands in the
native allocator.
- executor: 'graphs' added to sleeping_tags.
NCCL communicators: release idle PyNccl communicator memory via native
ncclCommSuspend(NCCL_SUSPEND_MEM)/ncclCommResume (NCCL >= 2.29.7), a ctypes shim
over torch's libnccl (nccl_suspend.py). Graceful no-op on older NCCL / 1 GPU.
Hooked into gpu_worker sleep()/wake_up(), collective across ranks.
custom all-reduce guard: routing the graph pool through cumem makes captured
graph buffers VMM-backed; CustomAllreduce registers them for IPC via
cudaIpcGetMemHandle, which fails on VMM memory (needs cuMemExportToShareableHandle)
and aborts FULL-cudagraph capture. Disable custom all-reduce while graph offload
is active (falls back to symm-mem / pynccl).
Scope: graph offload covers allocations that route through the cumem graph
MemPool. On dense models nearly all captured graph memory routes in; on large
MoE+EP models a large share of capture-time memory (attention/expert/EP
workspace) is allocated outside torch's pool routing and is not yet captured, so
NCCL communicator release is the dominant saving there. Full coverage via global
allocation interception (torch_memory_saver style) is a follow-up.
AI-assisted (Claude). Validated on H200: Qwen3-4B/8B/30B-A3B 1-GPU bit-exact
sleep/wake under VLLM_BATCH_INVARIANT (levels 1 and 2); TP=2/8 and MoE+EP NCCL
release ~0.5-1 GiB/rank; 2-node cross-node ncclCommSuspend/resume. Qwen3-235B-A22B
TP=8+EP: post-fix sleep frees an extra ~1 GiB/GPU vs baseline (NCCL 958 MiB +
graphs); all-reduce microbench confirms disabling custom all-reduce is ~free.
Multi-GPU sleep/wake VMM-mutation cross-rank safety (the discard/remap race) is
out of scope here and handled separately by vllm-project#45519/vllm-project#45554.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>
aoshen02
force-pushed
the
offload-cudagraph-nccl-main
branch
from
199b5cc to
f9f1cb9
Compare
intarweb-sync-bot Bot
pushed a commit
to intarweb/vllm
that referenced
this pull request
Jun 15, 2026
With CUDA-graph capture pools resident, per-allocation cumem VMM unmap/remap operations slow dramatically -- 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 during sleep()/wake_up(). Release the captured graphs (CUDAGraphWrapper / BreakableCUDAGraphWrapper entries + the encoder cudagraph manager) before the cumem unmap so the unmap takes the fast path, then re-capture them on the next full wake once weights and the KV cache have been remapped. - GPUModelRunner.release_cudagraphs(): drops all captured graph entries, forces finalization (gc + 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(): release graphs before allocator.sleep() (the unmap). - Worker.wake_up(): re-capture after allocator.wake_up() (the remap), only on a full wake (no tags / kv_cache); a weights-only wake defers until the KV cache is back. All paths are no-ops when cudagraph_mode is NONE. This is complementary to (not a duplicate of) the offload-in-place approach in vllm-project#45623: that routes the graph pool through cumem and restores it in place (a capacity optimization), whereas this releases-and-recaptures to avoid the slow VMM path entirely (a wake-latency optimization). The release/recapture approach matches the workaround field-confirmed by AlanFokCo on vllm-project#45398. Adds tests asserting the ordering contract: graphs are released before the allocator unmap and re-captured after the remap, the weights-only-wake defer, and the cudagraph_mode==NONE no-op. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
terafin
added a commit
to intarweb/vllm
that referenced
this pull request
Jun 15, 2026
With CUDA-graph capture pools resident, per-allocation cumem VMM unmap/remap operations slow dramatically -- 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 during sleep()/wake_up(). Release the captured graphs (CUDAGraphWrapper / BreakableCUDAGraphWrapper entries + the encoder cudagraph manager) before the cumem unmap so the unmap takes the fast path, then re-capture them on the next full wake once weights and the KV cache have been remapped. - GPUModelRunner.release_cudagraphs(): drops all captured graph entries, forces finalization (gc + 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(): release graphs before allocator.sleep() (the unmap). - Worker.wake_up(): re-capture after allocator.wake_up() (the remap), only on a full wake (no tags / kv_cache); a weights-only wake defers until the KV cache is back. All paths are no-ops when cudagraph_mode is NONE. This is complementary to (not a duplicate of) the offload-in-place approach in vllm-project#45623: that routes the graph pool through cumem and restores it in place (a capacity optimization), whereas this releases-and-recaptures to avoid the slow VMM path entirely (a wake-latency optimization). The release/recapture approach matches the workaround field-confirmed by AlanFokCo on vllm-project#45398. Adds tests asserting the ordering contract: graphs are released before the allocator unmap and re-captured after the remap, the weights-only-wake defer, and the cudagraph_mode==NONE no-op. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: terafin <terafin@users.noreply.github.com>
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Extends vLLM sleep/wake (
CuMemAllocator) to release two GPU-memory classesthat currently stay resident while an engine sleeps, restoring them on wake —
beyond the existing weights/KV-cache offload:
allocator under a new
graphstag. The graph memory is offloaded to CPU andrestored in place on wake (no re-capture; cumem preserves virtual
addresses).
ncclCommSuspend(NCCL_SUSPEND_MEM)/ncclCommResume(NCCL ≥ 2.29.7), actypes shim over torch's loaded
libnccl. Graceful no-op on older NCCL orsingle-GPU.
To make graph capture actually land in the cumem pool, the pool is pinned
before capture:
gpu_worker._pin_sleep_mode_graph_pool()overwritesCudaPlatform._global_graph_pooland re-points every liveCUDAGraphWrapper/BreakableCUDAGraphWrapper, and thegraphstag is held acrosscapture_model(). (Without this, on the v1 runner the global graph pool iscached before the cumem instance exists and capture lands in the native
allocator.)
sleep()also always attempts a guardedtorch.cuda.empty_cache()so freed blocks return to the driver.
It also disables custom all-reduce while graph offload is active: captured
graph buffers become VMM-backed, and custom all-reduce registers them for IPC
via
cudaIpcGetMemHandle, which fails on VMM memory (needscuMemExportToShareableHandle) and aborts FULL-cudagraph capture. The enginefalls back to symm-mem / pynccl (a microbenchmark shows these are
equal-or-faster at 8-GPU NVLink anyway).
Scope / honesty about coverage
Graph offload only catches allocations that route through the cumem graph
MemPool. Measured on Qwen3-235B-A22B, TP=8 + EP, H200:dominant, reliable saving.
capture-time memory (attention / expert / EP workspace) is allocated outside
torch's pool routing, so
pool=does not capture it. On dense models nearlyall captured graph memory routes in (Qwen3-4B: ~99%).
(NCCL 958 MiB + graphs). The pre-fix version was actually net-negative on the
v1 runner (graph capture bypassed the pool and
empty_cachewas skipped); bothare fixed here.
Full graph-memory coverage on MoE+EP needs global allocation interception
(torch_memory_saver style, as SGLang does) rather than torch
MemPoolrouting —left as a follow-up.
Why this is not a duplicate
(auto-closed for inactivity). This revives the graph-tag approach and ports it
to current code.
the comm alive and only releases its buffers (
ncclCommSuspend).graphs/nccltags and implement no graph/NCCL offload.Tests (H200; container
vllm/vllm-openai:v0.23.0, NCCL upgraded to 2.30.7)Correctness uses a variant-set /
VLLM_BATCH_INVARIANT=1bit-exact criterion(greedy decode is otherwise nondeterministic on near-tie tokens, independent of
sleep).
VLLM_BATCH_INVARIANT=1:sleep→wake bit-exact across cycles at level 1 and level 2; graph pool routed
through cumem (graphs tag isolated).
offload; NCCL releases ~0.5–1 GiB/rank; no hang; outputs within no-sleep
variance.
ncclCommSuspend/resumeover inter-node transport,correct all-reduce before and after.
above. Sleep wall ~2.1 s, wake ~3.3 s (NCCL suspend/resume dominates).
~free.
pre-commit runruff/ruff-format/typos/SPDX/torch.cuda policy pass.AI assistance
Developed with AI assistance (Claude). The human submitter has reviewed every
changed line and run the tests above.