Purpose

With sleep mode on, CUDA graphs keep their private tensor pools resident after sleep() has unmapped weights and KV cache. On large MoE deployments this is GiBs per GPU that sleep cannot release. This PR adds an opt-in, sleep_mode_offload_cudagraph (off by default), that captures CUDA graphs into a dedicated cuMem-backed pool, so sleep releases their physical memory and wake remaps the same virtual addresses. Graphs are not recaptured. With the option off, behaviour is identical to main: plain graph pools, custom allreduce registers graph buffers, NCCL_GRAPH_REGISTER is untouched.

The design follows SGLang's SGLANG_MEMORY_SAVER_CUDA_GRAPH (off by default; sgl-project/sglang#7873, #32718, #40648); the table below compares the two item by item.

It is based on the graph-pool design from aoshen02#50 (same author), ported to current main, plus two fixes needed on multi-GPU deployments.

Based on main, which now includes #59156 and #59158 (both merged): per-tag pool lists in CuMemAllocator.allocator_and_pools, and the rule that any wake restores every tag other than weights and kv_cache. With #59158's runtime tag, CuMemBackend.suspend() offloads ("weights", "runtime", "cudagraph") at level 1 and ("runtime", "cudagraph") at level 2.

Must merge after #59228 (free-while-asleep fix in csrc/cumem_allocator.cpp). Without it, tearing down a cuMem graph pool while asleep (for example deleting the engine during sleep) double-unmaps it and leaves a sticky CUDA error. The weights and KV pools already have this failure on main; this PR extends it to the graph pool.

Enabling it

llm = LLM(model, enable_sleep_mode=True, sleep_mode_offload_cudagraph=True)

vllm serve <model> --enable-sleep-mode --sleep-mode-offload-cudagraph

sleep_mode_offload_cudagraph is a ModelConfig field next to enable_sleep_mode and sleep_mode_backend (the name keeps their sleep_mode_ prefix and says what it does to graphs). It reaches ModelConfig through EngineArgs like its neighbours, with the CLI flag --sleep-mode-offload-cudagraph next to --enable-sleep-mode.

The pool is active when the option is on, enable_sleep_mode is on, the sleep backend is cumem, CUDA graphs are on and the platform is CUDA (VllmConfig.use_cumem_cudagraph_pool, the single predicate that the capture routing, CudaCommunicator and the NCCL default all read). If the option is on but another requirement is not met, the pool is simply not used and startup logs one warning saying which requirement is missing.

Benefit when on: graph memory is released while asleep at both levels (DeepSeek-V4-Flash TP4 + EP: 2.3 GiB/GPU; DeepSeek-V4.1 TP8: 7.43 GiB/GPU; Qwen3-8B TP1: 332 MiB). Cost when on: pinned host memory for the pool's backup, also at level 2; about +3 % decode TPOT at concurrency 1 from the custom-allreduce copy (within noise at concurrency 16 and above); 0.05–0.1 s more per sleep.

SGLang parity

Scope

• Runner-agnostic predicate; the capture sites and tests are those of Model Runner V2, and V1 code is unchanged. Covered capture paths all go through one helper, use_cudagraph_pool() in vllm/compilation/cudagraph_pool.py (next to the capture sites, which keeps the NCCL set_graph_pool_id coupling at the compilation layer). It gets the allocator through get_mem_allocator_instance(); the predicate already requires CUDA, so that is CuMemAllocator (a narrow cast at the call site rather than new MemAllocator protocol methods, which XPU's allocator would not implement):
  • V2 FULL graphs (CudaGraphManager.capture, including DBO microbatch capture on the ubatch capture stream);
  • compiled PIECEWISE graphs (CUDAGraphWrapper);
  • breakable CUDA graph segments (BreakableCUDAGraphWrapper);
  • speculator graphs (autoregressive/EAGLE/MTP, DFlash, DSpark), which subclass CudaGraphManager and reuse its capture().
• Memory profiling is deliberately not routed through cuMem: V2's profile_cudagraph_memory() captures inside plain_cudagraph_capture() (a ContextVar), so destroying the profiling graphs frees their pool as before.
• V1 limitation: with the option on under Model Runner V1, V1's profiling graphs stay in the cuMem pool (awake +326 MiB on Qwen3-8B, +1.6 GiB on DeepSeek-V4-Flash TP4, not counted in the KV-cache estimate) and V1 DBO microbatch graphs stay in the plain pool; outputs are unaffected.
• Out of scope: encoder CUDA graphs (the encoder runner keeps its own graph_pool_handle()), and the Gemma4 centroids graphs in the V1 spec-decode path.

Design

• One graph pool, tagged cudagraph. All routed captures share the process's single graph pool, as they share the global graph pool without sleep mode. CuMemAllocator.use_cudagraph_pool() creates one pluggable-allocator MemPool on the first capture, keeps it as a one-entry cudagraph list in allocator_and_pools, and yields its id for torch.cuda.graph(pool=...) while the allocator tags allocations cudagraph. There is no use_mem_pool (capture_begin routes allocations itself) and no trim loop, so reusing the pool is safe. The module-level use_cudagraph_pool() also points NCCL's graph allocator (set_graph_pool_id) at whichever pool is used, so the capture sites need no separate call.
• CPU-backed at level 1 and level 2. A captured pool is not only scratch activations. It can hold values written once at capture time (constants, persistent kernel or attention state) that replay reads but never rewrites. Discarding the pool corrupts those graphs. SGLang hit this and moved to backing up the graph pool across pause/resume ([sglang-miles] Back up the CUDA graph pool across TMS pause/resume sgl-project/sglang#32718) and to keeping DSA graph state and the pool intact ([RL] Keep DSA cuda-graph state and the graph pool intact across TMS pause/resume sgl-project/sglang#40804). The pool is therefore always backed up to host memory, at level 2 as well. This costs pinned host RAM at level 2 as well: the pool is 2.3 GiB/GPU on DeepSeek-V4-Flash TP4 and 7.43 GiB/GPU on DeepSeek-V4.1 TP8 (about 60 GiB per 8-GPU node), in addition to the level-1 weight backup. Discarding the pool was also bit-exact on DeepSeek-V4-Flash (TP2 and TP4, full and selective wakes), but it is not safe in general, so the default stays offload.
• Sleep policy lives in CuMemBackend; the allocator is mechanism only. suspend() offloads cudagraph at both levels; like every tag other than weights and kv_cache, it is restored automatically by any wake, selective or not. The RL flow wake_up(["weights"]) → update → wake_up(["kv_cache"]) has graphs replayable once both are awake.
• Pool lifetime. The graph MemPool is kept together with its pluggable allocator wrapper and released by release_pools() like every other pool. Freeing an allocation while asleep (e.g. a graph pool at teardown) needs no spare VRAM and no Python remap: this PR depends on [Bugfix] Don't unmap cuMem allocations freed while asleep #59228, which makes the C++ my_free skip unmapping an allocation that sleep() already released and only free its address reservation (ROCm already skipped via an empty chunk list).
• Recapture (elastic EP). Elastic EP is the only runtime path that drops and recaptures graphs (on every scale). Each capture runs on a new stream (graph_capture()), and the caching allocator reuses a freed block of a live pool only on the stream that freed it, so recapturing into the kept pool grew it: DeepSeek-V2-Lite DP2↔4 with sleep mode went from 586 to 5098 MiB of graph pool on rank 0 over five scale cycles. Elastic EP now calls release_cudagraph_pool() (a no-op unless the pool is active) right after release_graphs() on the V2 path; the V1 path is as on main. Once no graph uses the pool, it is dropped, its segments are unmapped through the free callback, and the next capture starts a fresh pool. It refuses while the pool is asleep (freeing would unmap a second time) or while something still references it, and the message counts the other references and names the possible holders (live graphs, or a capture-time pool reference taken without a release, as NCCL symmetric memory did). That check found that graph release left the PIECEWISE graphs alive in the compiled model's wrappers until the recompile replaced them; CudaGraphManager.release_graphs() now clears them too, as profiling teardown already does. No other runtime path releases or recaptures graphs: LoRA and weight updates replay the existing graphs, and profiling captures into a throwaway plain pool.

Custom allreduce fix

Custom allreduce IPC-registers the buffers that captured graphs read (register_graph_buffers). cudaIpcGetMemHandle rejects cuMem memory, so capture fails with csrc/custom_all_reduce.cuh:164 'invalid argument' and the worker exits. This was reproduced on GB300 with Qwen3-8B TP4, VLLM_ALLREDUCE_USE_FLASHINFER=0 VLLM_ALLREDUCE_USE_SYMM_MEM=0 and --compilation-config '{"pass_config":{"fuse_allreduce_rms":false}}'.

When the cuMem graph pool is active, custom allreduce captures with registered=False: it copies the input into the pre-registered IPC buffer, which is cudaMalloc memory. SGLang does the same (sgl-project/sglang#7873), keyed on its env alone; here it follows the same predicate as the pool. CudaCommunicator reads use_cumem_cudagraph_pool from the current config once at construction (workers build communicators under set_current_vllm_config; without a config there is no sleep-mode policy, so graph buffers stay registered) and passes register_graph_buffers=False to CustomAllreduce, which itself reads no config and does not look at allocator state during capture. Custom all-gather and reduce-scatter always copy into the pre-registered buffer already, so they need no change. The cost is one device-to-device copy per captured allreduce, the same path eager mode uses, and only when the pool is active. Measured on GB300 (DeepSeek-V4-Flash TP4 + EP, MTP, sleep mode on and awake): over five server restarts per build (four of them with a warmup), median decode TPOT at concurrency 1 is about +3 % for the whole PR (5.35 vs 5.21 ms) and +3.6 % for a control that only forces the copy (no cuMem pool), so the cost is the copy; the cuMem pool itself adds nothing measurable. Main's own restart-to-restart spread is 2.4 %. NCCL_GRAPH_REGISTER=0 alone costs nothing measurable. At concurrency 16 and above the difference is within run-to-run noise. On MI355X (#59523) the same copy costs +3.4 % (vLLM custom allreduce) and +2.7 % (AITER).

NCCL graph registration

With NCCL_GRAPH_REGISTER enabled, NCCL registers graph-captured user buffers and retains their cuMem handles. That pins about 400–510 MiB/GPU through sleep. After wake remaps new physical memory at the same address, NCCL also keeps stale registrations; in SGLang this caused hangs and wrong results (fzyzcjy/torch_memory_saver#88, sgl-project/sglang#40648). As in #40648, startup (VllmConfig.__post_init__, _verify_cumem_cudagraph_pool) sets NCCL_GRAPH_REGISTER=0 when the pool is active and it is unset, before any worker or communicator exists (workers inherit the environment; Ray copies NCCL_*). An explicit value is kept, with a warning giving the reason. An earlier revision also refused TORCH_NCCL_USE_TENSOR_REGISTER_ALLOCATOR_HOOK. On GB300 (DeepSeek-V4-Flash TP4 + EP and DP4 + EP, hook=1 with that check removed, NCCL_DEBUG_SUBSYS=REG), NCCL logged no buffer registration, memory while asleep and after each wake matched hook=0, and outputs matched hook=0, so the check was removed.

What is not released

Graph executables (cudaGraphExec_t) and the driver memory they own are not released. They live outside any PyTorch pool. On DeepSeek-V4-Flash DP4, about 1944 MiB of graph-exec memory plus 1142 MiB charged during capture stay resident per GPU after sleep. Releasing them would require destroying and recapturing graphs, or checkpointing the CUDA process.

Related work

• [Bugfix][Worker] Track CUDA graph memory in CuMemAllocator to fix sle… #47700 captures CUDA graphs into the weights cuMem pool. Graph memory then follows the weights policy: it is discarded at level 2 and restored only with weights. This PR gives graphs their own tag with an always-backup policy, so level-2 sleep and selective wake cannot corrupt captured state. It also covers FULL, compiled PIECEWISE, breakable and speculator capture in both runners.
• [V2] Offload CUDA process state during sleep without graph recapture #56373 (closed) used CUDA checkpoint of the whole process. That releases everything, including graph executables and the context, but it supports a single GPU only. This PR works with TP/DP and keeps communicators alive, at the cost of the resident graph executables described above.

Test Plan

pytest -v tests/basic_correctness/memory/cumem/test_cumem.py -k cudagraph_pool
pytest -v tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py -k cudagraph_pool  # TP=2 case needs 2 GPUs
pytest -v tests/v1/worker/test_sleep_mode_backend.py -k cumem_cudagraph_pool

• test_cudagraph_pool_sleep_and_release[sleep-1|sleep-2] (allocator level): captures are not routed with the policy off or under plain_cudagraph_capture(), and are routed with NCCL's graph pool id following the pool. A captured graph that reads a pool-resident constant survives suspend(level) → resume(["kv_cache"]) (which already restores the pool at the same addresses) → resume(["weights"]) and replays correctly. Release is refused while asleep and while the graph is alive; afterwards the pool is empty and a recapture on a new stream refills it to the same size.
• test_sleep_cudagraph_pool[...] (engine level, tiny Llama): for each Model Runner V2 capture site (FULL, PIECEWISE, breakable) graph memory carries the cudagraph tag, NCCL_GRAPH_REGISTER defaults to 0, and greedy output is exact after level-1 and level-2 selective wakes; on V2, elastic EP's sequence (release_graphs(), release_cudagraph_pool(), capture_model()) empties the pool and refills it to the same size, twice. The TP=2 case forces the legacy custom allreduce. The off-by-default case checks that without the option there is no cuMem graph memory and NCCL_GRAPH_REGISTER is untouched.
• test_cumem_cudagraph_pool_startup_checks[...]: active pool → NCCL_GRAPH_REGISTER=0; an explicit value wins with a warning; not opted in → untouched, no log; opted in but inactive → one warning, env untouched.

Test Result

Opt-in and runner-agnostic revision (h200, Qwen3-8B TP1, sleep mode on, this branch on the CI image):

• V1 (code as on main): KV cache size is the same with the option on and off, but awake memory is +326 MiB and peak during GSM8K +266 MiB (124913 vs 124647 MiB), because V1's profiling graphs stay in the pool; outputs after full, selective and level-2 wakes are 16/16.
• Probe outputs (16 prompts, greedy) after three rounds of full, selective and level-2 wakes are 14–16/16 in all four arms, the same as three awake repeats without sleep (13–15/16 against the first run), so the differences are run-to-run batching variation, not the option.
• With the option off: no cudagraph tag, NCCL_GRAPH_REGISTER untouched (engine-level off-by-default test).

Current head (GB300 rack 4, nightly 0cbac6cd1 + #59158 + this PR per file; DeepSeek-V4-Flash TP4 + EP, MTP, V2):

• GSM8K awake / full wake / selective wake: 0.944 / 0.945 / 0.946 vs 0.947 / 0.951 / 0.941 without this PR (level 2 + reload is broken on V4 either way). Asleep 8945 → 6379 MiB/GPU; cudagraph tag 2342 MiB.

• 21-cycle soak (full, weights-first and kv-first selective wakes): memory flat, every probe exact, sleep about 1.0 s and wake 0.6–1.7 s per cycle.

• KV cache size matches the unpatched build to 0.01 %; the profiling graphs are released.

• TP2: exact after full and selective wakes. Qwen3-VL-2B (encoder graphs on) and a Qwen3-0.6B LoRA: exact after every wake order (h200).

• Elastic EP (h200, DeepSeek-V2-Lite, DP2↔4 ×5 with a level-1 sleep/wake after each cycle, Ray backend): the cudagraph tag stays at 586 MiB on every scale (it grew to 5098 MiB before the fix). Awake memory on rank 0 after each scale-down is 117341–117355 MiB, versus 117321–117325 MiB for the same build without sleep mode and 118135–118143 MiB on main with sleep mode. The 0.8 GiB below main comes from clearing the previous PIECEWISE graphs in release_graphs(), which the build without sleep shows too. Asleep: 6845 MiB vs 8219 MiB on main. Probe outputs match main in every arm (4/4 at DP2; 2/4 at DP4 in every arm, including main, from DP batching nondeterminism).

• Dual batch overlap could not be run on the current GB300 image: Model Runner V2 rejects --enable-dbo with speculative decoding, microbatching without MTP needs the DeepEP or NIXL EP all2all backend, and DeepEP low-latency fails at startup on that image (cudaErrorSymbolNotFound in internode_ll.cu) on main as well. The DBO capture path is covered by code review only.

• Unit tests on h200: test_cudagraph_pool_policy.py, test_cumem.py, test_sleep_mode_backend.py, test_workspace.py, the capture-site and TP2 custom-allreduce sleep tests: all pass.

• pre-commit run --files <changed files>: all hooks pass (ruff check/format, mypy, typos, SPDX, forbidden imports, torch.cuda API check, test tethering). python -m py_compile passes.
  All runs are on GB300 with the nightly image, torch 2.13.0+cu130. The branch also needs the C++ fix from [Bugfix] Don't unmap cuMem allocations freed while asleep #59228.

Re-validation after the restructure (head d5f3bd94fb, with #59156 underneath), nightly image, GB300.

• Unit tests: test_cumem.py + test_sleep_mode_backend.py: 26 passed (single GPU). The graph-pool tests also pass with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, so the capture path needs no expandable-segments toggle.

• DeepSeek-V4-Flash TP4 + EP4 + MTP, gated bit-exact harness (553 requests per round): asleep 11517 MiB/GPU with [Feature] Release WorkspaceManager scratch on sleep #59156 underneath (this PR alone measured 12839 MiB before the rebase; the difference is [Feature] Release WorkspaceManager scratch on sleep #59156 workspace scratch). Every round 553/553 bit-identical to the awake reference after a level-1 full wake and after a level-1 selective wake (weights, then kv_cache). Sleep 1.1-1.2 s, wake 0.6 s.

• Qwen3-8B with custom allreduce forced (FlashInfer and symm-mem allreduce off, allreduce+RMSNorm fusion off): TP4 and TP2 both start (the engine died at capture before the custom-allreduce fix), and single-request text and logprobs are identical to before sleep after two sleep/wake cycles; asleep 3497 MiB/GPU at TP4. The TP=2 pytest regression needs hmellor/tiny-random-LlamaForCausalLM, which is not in the offline cache there, so the TP2 check used Qwen3-8B.

• All sleep PRs merged ([Feature] Release WorkspaceManager scratch on sleep #59156 + [Feature] Offload KV-init runtime state on sleep #59158 + this PR): asleep 6313 MiB/GPU, 553/553 bit-identical after full and selective wakes; merged unit tests 29 passed.

• Unit tests: test_cumem.py 14/14, including both test_cudagraph_pool_survives_sleep cases; test_sleep_mode_backend.py 9/9.

• TP=2 custom allreduce regression test: passed, not skipped. Custom allreduce was active and graph memory carried the cudagraph tag.

• Qwen3-8B TP4 with custom allreduce forced (FlashInfer/symm-mem allreduce and allreduce+RMSNorm fusion off):

  • Before this change the engine died at capture with custom_all_reduce.cuh:164 'invalid argument'. It now starts, and still registers 11972 graph addresses, which are the pre-registered IPC buffer.
  • After a full wake and a selective wake, single-request text and logprobs are identical to before sleep.

• DeepSeek-V4-Flash TP4 + EP4, MTP, bit-exact: 553 deterministic requests (batch sizes 1–256) per round.

  • 13 rounds are identical to the awake baseline: 3× full wake, 3× selective wake weights-first, 3× selective wake kv_cache-first, plus awake rounds.
  • Level 2 + reload_weights is broken for DeepSeek-V4 on unpatched vLLM as well.

• Memory while asleep, level 1, per GPU:

  Model / setup	Unpatched	This PR	cudagraph tag (backed up)
  DeepSeek-V4-Flash, TP4 + EP4, MTP	15178 MiB	12839 MiB (−2339)	2.29 GiB
  GLM-5.3, TP8 over 2 nodes, MTP	9449 MiB	8199 MiB (−1250)	1.22 GiB
  DeepSeek-V4.1, TP8 over 2 nodes	21035 MiB	13425 MiB (−7610)	7.43 GiB

• Other model families (all sleep PRs combined, TP4): level-1 full wake, level-1 selective wake and level 2 + reload_weights are all bit-exact against the first awake run. The engine used Model Runner V2 on every model.

  Model	Asleep MiB/GPU, without → with this PR	cudagraph tag
  Qwen3-8B	3605 → 3561	0.04 GiB
  Llama-3.1-8B NVFP4	3535 → 3491	0.04 GiB
  Qwen3-VL-8B (text + image)	3629 → 3573	0.05 GiB
  Kimi-Linear-48B (MLA + KDA), EP	4733 → 4505	0.23 GiB
  Qwen3-235B-A22B FP8, EP	6829 → 6585	0.24 GiB

  GSM8K on Qwen3-8B across awake / full wake / selective wake / level 2 + reload: 0.885 / 0.882 / 0.892 / 0.885.

• Sleep/wake latency:

  • V4-Flash TP4: first sleep 3.09 s, later sleeps 1.2 s, wake 0.8 s. That is 0.05–0.1 s more than discarding the pool.
  • V4.1: the first level-1 sleep is about 1.7 s slower, because it backs up 7.4 GiB/GPU. Wake time is unchanged.

• GSM8K on V4.1: 0.897 / 0.901 / 0.899 for awake / full wake / selective wake, against 0.892 / 0.899 / 0.902 unpatched. Level 2 + reload is broken upstream for DeepSeek-V4.

• GSM8K on GLM-5.3 (1319 questions, 5-shot, temperature 0):

  Build	Awake	Level-1 full wake	Level-1 selective wake	Level 2 + reload_weights
  Unpatched	0.914	0.914	0.918	0.911
  This PR	0.910	0.916	0.915	0.917

• NCCL guards:

  • NCCL_GRAPH_REGISTER=1 fails at startup with a clear error (the register-hook check was removed, see "NCCL graph registration").
  • With sleep mode off, NCCL_GRAPH_REGISTER is left unset and graph pools stay plain PyTorch pools.

• Teardown while asleep with the freed VRAM taken by another process (this PR + [Bugfix] Don't unmap cuMem allocations freed while asleep #59228): no error, no OOM.

AI assistance (Codex, Claude Code) was used for the implementation, porting, review and drafting of this PR. The human submitter is responsible for reviewing the changes and owns them end to end.

AMD MI355X (ROCm)

This PR is gated on current_platform.is_cuda(), so it is inert on ROCm (verified: memory matches #59156 within 6 MiB). Enabling it on ROCm was tried as a local patch, not included here:

• is_cuda() → is_cuda_alike() alone fails at the first capture: a blocking hipMemset in csrc/cumem_allocator.cpp makes the legacy stream depend on the capturing stream. Running it as hipMemsetAsync on a private non-blocking stream fixes it.
• vLLM custom allreduce needs the same unregistered capture path as on CUDA (hipIpcGetMemHandle fails on cuMem memory). QuickReduce needs no change. RCCL honours NCCL_GRAPH_REGISTER, so the existing guard applies.
• With that patch, DeepSeek-V2-Lite TP2 stays bit-exact over 10 cycles plus a level-2 reload, but the graph pool is only 72–244 MiB per GPU: most graph memory (0.95–1.7 GiB) is held by the HIP runtime outside torch pools. Unlike CUDA, ROCm returns it when graphs are destroyed, so destroy-and-recapture would be the better option there (9–50 s per wake).

ROCm support is left for a follow-up.

🤖 Generated with Claude Code