[Feature] Offload KV-init runtime state on sleep

[aoshen02]

Based on main, which now includes #59156 (merged). The runner-construction part (former #59162) has moved to a separate follow-up; this PR covers only the state built in initialize_kv_cache().

Purpose

/sleep releases the weights and the KV cache, but not the other persistent GPU state built during KV-cache initialization. On DeepSeek-V4-Flash with TP4 on GB300 that state is about 4.6 GiB per GPU:

Buffer	Size
V2 block tables (staged + input)	3.3 GiB
DSV4 sparse-MLA metadata	1.0 GiB
Indexer metadata	0.25 GiB
Sampler state and similar	small

These buffers are allocated in Worker.initialize_from_config, but only the KV cache inside that call uses a sleep-mode pool.

This PR:

• Runs model_runner.initialize_kv_cache(...) inside a new cumem tag, runtime. That is where block tables and attention metadata are built. The KV cache keeps its own nested kv_cache pool.
• Needs no KV-connector exception. What in-tree connectors allocate on the GPU in register_kv_caches is scratch or pointer tables (e.g. the HF3FS staging buffers), which the offload restores at the same address with the same contents. The memory connectors register with a transport (NIXL, Mooncake, MoRIIO, LMCache-MP IPC) is the KV cache, which keeps its own kv_cache (or connector) pool. The aux-output connector's RoutedExpertsCapturer.device_buffer is likewise a plain tensor. XPU is the one exception, below.
• Has CuMemBackend.suspend offload runtime to host memory at every sleep level. These buffers hold live state, so they are copied rather than discarded. Wake copies them back to the same addresses, so captured CUDA graphs stay valid.
• Needs no wake-side change: since [Feature] Release WorkspaceManager scratch on sleep #59156 any wake, selective or not, restores every tag other than weights and kv_cache, so runtime comes back automatically, for example wake_up(["weights"]) → reload_weights → wake_up(["kv_cache"]).
• Model Runner V1: with CUDA graphs, the memory profiler runs initialize_kv_cache() outside any pool, and for MLA, hybrid and multi-group models it rebuilds the InputBatch there. The real initialize_kv_cache() then saw matching sizes and kept that InputBatch, so the live block tables stayed in the default pool. _cleanup_profiling_kv_cache(), which already undoes the rest of the profiling KV init, now also resets _init_block_sizes, so the real call rebuilds the InputBatch under runtime, as it already does without CUDA graphs. Model Runner V2 needs no change: its initialize_kv_cache() always builds new BlockTables, inside the runtime pool for the real call.
• Leaves behavior unchanged when the sleep-mode allocator is off. On XPU the KV-init state stays outside runtime (see the XPU note below); CPU (no sleep mode) is unchanged.

Cost: sleep and wake copy the runtime bytes over the host link, about 4.6 GiB in the case above.

#59156 (merged) covers the workspace scratch.

Test Plan

• New GPU unit test test_runtime_state_survives_sleep[sleep-1|sleep-2]:
  • runtime with a nested kv_cache pool keeps the tags separate;
  • a sleep at either level releases everything;
  • wake_up(["weights"]) restores the runtime contents;
  • wake_up(["kv_cache"]) restores the KV cache.
• End-to-end on 4× GB300 with DeepSeek-V4-Flash (TP4 + EP4, MTP, CUDA graphs):
  • memory after /sleep;
  • prompt logprobs before sleep and after each wake (level 1 twice, then level 2);
  • sleep and wake time.

Test Result

Current head 658bdd4b52 (on main with #59156 merged), H200, CI image deb148e730: test_cumem.py 19, test_sleep_mode_backend.py 9, test_gpu_worker.py 20 passed. Model Runner V1, DeepSeek-V2-Lite with FlashMLA (block size 64, so the profiling pass rebuilds the InputBatch): with the _cleanup_profiling_kv_cache() reset the live block table is in runtime, without it the block table, slot mapping and sampling tensors of the InputBatch are not. Model Runner V2: the block tables are built in initialize_kv_cache() and are in runtime. opt-125m + ExampleConnector (kv_both), V1 and V2: outputs equal the awake run after level-1 and level-2 (+ reload_weights) wakes. An allocation audit of everything initialize_kv_cache() allocates under runtime (V1 and V2, FlashMLA, Triton MLA, FA3, FlashInfer, draft-model and ngram speculation, LoRA, a GDN hybrid) found no memory that is registered with a transport or exported, and no allocation or free in runtime after startup.

Earlier measurements. The re-validation sections from head b0a91fd8f8 on were measured with heads that also built the model runner under runtime, which this PR no longer does: with a speculator, those asleep numbers are lower than this PR's by the MTP hidden-states buffer (512 MiB on DeepSeek-V4-Flash, 192 MiB on GLM-5.3); without one the difference is negligible (DeepSeek-V4.1: 18583 MiB either way).

TP4 + EP4, auto MoE. Memory per GPU after /sleep:

	Unpatched	Patched
Residual	15178 MiB	10487 MiB (−4691 MiB)
Block tables / sparse-MLA / indexer metadata	3361 / 1024 / 257 MiB	0
runtime tag (offloaded at level 1 and level 2)	—	4704 MiB

• Correctness: prompt logprobs are bit-identical to before sleep after both level-1 wakes. The level-2 round completes without errors; its output is garbage only because weights are discarded, same as unpatched.
• Timing: sleep 9.2–11.3 s (unpatched 9–11 s). Wake 1.0 s (unpatched 0.7 s); the extra ~0.3 s is copying the runtime bytes back.
• Unit tests pass. Pre-commit passes.

Re-validation of the combined PR (head b0a91fd8f8), DeepSeek-V4-Flash TP4 + EP4 + MTP, 4× GB300. Gated bit-exact harness (553 requests per round, batch composition fixed by /pause gating; tokens, output logprobs and prompt logprobs compared with the first awake round):

• Asleep: 9977 MiB on two GPUs / 10103 MiB on the other two, same after a second sleep (the PR2 + [Feature] Offload model runner buffers on sleep #59162 measurement was 9975 MiB).
• Every round is 553/553 bit-identical: after a level-1 full wake and after a level-1 selective wake (weights, then kv_cache).
• Sleep 1.2–1.3 s, wake 0.95 s.
• Unit tests test_cumem.py and test_sleep_mode_backend.py: 24 passed.
• With [Feature] Release WorkspaceManager scratch on sleep #59156 and [Feature] Release the CUDA graph pool on sleep #59160 merged on top: 6313 MiB asleep (6439 on one GPU), every round 553/553 bit-identical after full and selective wakes; merged unit tests 29 passed.

Re-validation stacked on #59156 (head 0354f3d8a9, #59156 at 21c7a0be4c), DeepSeek-V4-Flash TP4 + EP4 + MTP, 4× GB300, same image, another rack. Gated bit-exact harness, 553 requests per round:

• Model Runner V2 (default): asleep 8775–8781 MiB per GPU (15181 unpatched on these nodes; 9165–9293 with the previous heads). Every round 553/553 bit-identical after two level-1 full wakes and two level-1 selective wakes, plus 4 more sleep/wake cycles exact. KV cache 202.31 GiB / 13,036,348 tokens (unpatched 202.44 GiB / 13,044,408).
• Model Runner V1 (VLLM_USE_V2_MODEL_RUNNER=0): asleep 7507 MiB with the V1 InputBatch fix vs 9187 MiB without it (−1680 MiB: the block tables), same KV cache (12.26 M tokens). Both arms give the same 552/553 in every round, including the awake-only rounds, so the one differing request is V1 run-to-run variation, not sleep.
• V1 probe (Qwen3.5-0.8B, hybrid, 1 GPU, in-process): the live block tables move from the default pool to runtime; KV blocks equal unpatched (34471); full and selective wakes exact. Qwen3-8B (geometry matches the placeholder): identical with and without the fix, block tables already in runtime.
• Unit tests: test_cumem.py 15, test_sleep_mode_backend.py 11, test_workspace.py 11 passed.

Large model accuracy (GLM-5.3, TP8 over 2 GB300 nodes, mp backend, MTP on). GSM8K, 1319 questions, 5-shot, temperature 0. Each cell is accuracy, and every run had 0 invalid answers. The unpatched numbers come from the same nodes and image (see #59156).

State	Unpatched	Patched
Awake	0.914	0.914
Level-1 sleep, full wake	0.914	0.913
Level-1 sleep, selective wake (weights, then kv_cache)	0.918	0.917
Level-2 sleep, reload_weights	0.911	0.911

• Memory while asleep, per GPU: 9449 MiB unpatched, 8847 MiB patched, so the runtime state saves 602 MiB. The level-2 sleep saves the same amount, and the state comes back correctly after reload_weights.
• Timing: sleep and wake timings are unchanged.
• Errors: no errors on either node.

Large model accuracy (DeepSeek-V4.1, TP8 over 2 GB300 nodes, mp backend). GSM8K, 1319 questions, 5-shot, temperature 0. Each cell is accuracy / invalid rate.

State	Unpatched	Patched
Awake	0.892 / 0.002	0.906 / 0.001
Level-1 sleep, full wake	0.899 / 0.002	0.905 / 0.002
Level-1 sleep, selective wake	0.902 / 0.003	0.896 / 0.002
Memory while asleep, per GPU	21035 MiB	18583 MiB (−2452 MiB)

• reload_weights after a level-2 sleep produces garbage for this model with or without this PR: 172 "Failed to load weights" warnings from layerwise.py for DeepseekV4MegaAttnAttention, DeepseekV4ScalingRotaryEmbedding and NgramHashState. It is a separate, existing layerwise-reload gap, so level-2 accuracy is not compared here.

Other model families (with #59156, TP4, 4× GB300). In every phase, every request matches the first awake run exactly, text and all logprobs. The phases are level-1 full wake, level-1 selective wake, and level 2 + reload_weights. Level 2 + reload being exact confirms the runtime state comes back.

Model	Asleep MiB/GPU, unpatched → patched
Qwen3-8B	4019 → 3605
Llama-3.1-8B NVFP4	3949 → 3535
Qwen3-VL-8B (text + image)	4045 → 3629
Kimi-Linear-48B (MLA + KDA)	4809–4937 → 4733
Qwen3-235B-A22B FP8, EP	7363–7491 → 6829

• GSM8K, Qwen3-8B (awake / full wake / selective / level 2 + reload): 0.888 / 0.886 / 0.887 / 0.885.

Lifecycle edge cases (with #59156, DeepSeek-V4-Flash TP4): all pass, bit-exact after each.

• 300 consecutive sleep/wake cycles with all sleep-mode PRs (alternating full and selective wake): every cycle bit-exact (34,140/34,140 requests), memory after wake equal to before sleep on every GPU, device-memory slope 0.000 MiB/cycle. Host RSS grows by about +5 MB/cycle over the first ~50 cycles of every launch, unpatched too; later slopes change sign between launches of the same configuration in both arms (patched −1.8 to +3.6, unpatched −0.5 to +0.6 MB/cycle) and the shared-memory backup stays byte-for-byte constant.
• Sleep right after startup.
• A request sent while asleep, or after a weights-only wake: it is queued and completes after the full wake.
• Double sleep and double wake: no-ops.
• Level 1 followed by level 2 without a wake: ignored.
• 256 concurrent mixed-length requests right after wake: no errors, no NaN.

Configuration variants (with #59156, DeepSeek-V4-Flash TP4). Every post-wake round is bit-identical to the awake reference: 553 requests, tokens plus output and prompt logprobs.

Variant	Asleep MiB/GPU, unpatched → patched	runtime tag
Default (MTP)	15181 → 9293	4.59 GiB
MTP off	13593 → 7589	4.57 GiB
Model Runner V1	12641 → 10083	1.31 GiB
--enforce-eager	9851 → 3827	4.59 GiB
Chunked prefill, --max-num-batched-tokens 2048	13561 → 8841	3.45 GiB
--max-num-seqs 8	6973 → 4391	1.33 GiB
Qwen3-8B TP1 / TP2 (incl. level 2 + reload_weights)	3123 → 2069 / 4133 → 3725	1.03 / 0.53 GiB

• KV connector (earlier heads, which skipped the KV-init runtime pool with a connector): Qwen3-8B TP2 + ExampleConnector, asleep memory equal to unpatched. The current head drops that exception; see the current-head results above. DeepSeek-V4-Flash + ExampleConnector fails to start on unpatched vLLM too (MLAAttentionSpec ... 'bounded_replay'), so it is not covered.
• Scaling: the runtime state scales with max_num_seqs × max_model_len, as expected for block tables.

Review: three rounds of adversarial review (Codex, Kimi K3). The final round found no blocker or major issue.

---

This PR contains AI-assisted code (Claude Code); commits carry a Co-authored-by trailer.

AMD MI355X (ROCm)

Works on ROCm unchanged, stacked on #59156 (vLLM ROCm nightly e975732152). Nested pools behave as on CUDA: the innermost pool wins, so the KV cache and the runtime state land in separate tags. Additional saving over #59156, memory per GPU while asleep (level 1): DeepSeek-V2-Lite TP1 −104, TP2 −110, DP2 −92 MiB; DeepSeek-V4-Flash TP4 −143 MiB. Outputs are bit-identical after full and selective wake.

AMD 50-cycle leak test (DeepSeek-V2-Lite TP1 / TP2 / DP2, DeepSeek-V4-Flash TP4; with #59156): device-memory slope 0.00 MiB/cycle, 50/50 cycles bit-exact, no host-RSS growth over level-1 cycles, in-process memory after wake equal to before sleep. GSM8K on V4-Flash TP4 (1319 questions, 5-shot, temperature 0), awake R1 / R2 → full wake / selective wake / +10 cycles: unpatched 0.949 / 0.945 → 0.948 / 0.944 / 0.945; with #59156 0.947 / 0.948 → 0.949 / 0.948 / 0.947.

XPU: the XPU caching allocator allocates from the outermost active pool (get_pool() walks the active pools FIFO; CUDA uses LIFO since pytorch/pytorch#168382), so the nested runtime/kv_cache pools would share one pool there. On XPU the KV-init state therefore stays outside the runtime pool, as on main. A hardware check is requested in the comments.

🤖 Generated with Claude Code

[mergify]

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

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

[aoshen02]

Request for an XPU check (nested memory pools)

Could someone with an Intel XPU run the snippet below? This PR allocates the KV-init state under the runtime pool, and the KV cache itself goes into a nested kv_cache pool:

with allocator.use_memory_pool("runtime"):
    with allocator.use_memory_pool("kv_cache"):   # inside initialize_kv_cache
        kv = torch.empty(...)
    block_table = torch.empty(...)

On CUDA, torch allocates from the innermost active pool, so the KV cache and the runtime state stay in separate pools. The XPU caching allocator picks the outermost one instead. get_pool() walks the active pools in opposite orders:

• CUDA: captures_underway.rbegin(), i.e. LIFO, the innermost pool (c10/cuda/CUDACachingAllocator.cpp, v2.11.0 L3080-3084);
• XPU: for (auto& entry : captures_underway), i.e. FIFO, the outermost pool (c10/xpu/XPUCachingAllocator.cpp, v2.11.0 L670-672). This is unchanged in v2.12, v2.13 and main (renamed allocation_scopes_).

As a result:

• the KV cache lands in the runtime pool (still tagged kv_cache);
• a later small runtime tensor (for example a block table) can be carved from the free tail of a kv_cache-tagged segment;
• that tensor is then discarded on sleep instead of offloaded, and reads garbage after wake.

Repro, run on this branch (stacked on #59156) on a single XPU:

import torch
from vllm.device_allocator.xpumem import XpuMemAllocator

a = XpuMemAllocator.get_instance()
with a.use_memory_pool(tag="runtime"):
    with a.use_memory_pool(tag="kv_cache"):
        kv = torch.empty(3 << 20, dtype=torch.uint8, device="xpu")
    rt = torch.empty(3 << 19, dtype=torch.uint8, device="xpu")

for name, t in (("kv", kv), ("rt", rt)):
    for tag, entries in a.allocator_and_pools.items():
        for pool, _ in entries:
            if any(s["address"] <= t.data_ptr() < s["address"] + s["total_size"]
                   for s in pool.snapshot()):
                print(f"{name} is in the {tag!r} pool")
    for d in a.pointer_to_data.values():
        base, size = d.handle[2], d.handle[1]
        if base <= t.data_ptr() < base + size:
            print(f"{name} is in a segment tagged {d.tag!r}")

Expected (what CUDA prints): kv is in the kv_cache pool and a kv_cache segment; rt is in the runtime pool and a runtime segment. The problem case is kv in the runtime pool, or rt in a segment tagged kv_cache.

Please post the output, so we can confirm this on hardware. Unless it shows otherwise, I'll skip the runtime pool around initialize_kv_cache on XPU (one line), so XPU keeps main's behaviour for that state. The runner-construction buffers stay under runtime either way.

🤖 Generated with Claude Code

[jikunshang]

Request for an XPU check (nested memory pools)

Could someone with an Intel XPU run the snippet below? This PR allocates the KV-init state under the runtime pool, and the KV cache itself goes into a nested kv_cache pool:

with allocator.use_memory_pool("runtime"):
    with allocator.use_memory_pool("kv_cache"):   # inside initialize_kv_cache
        kv = torch.empty(...)
    block_table = torch.empty(...)

On CUDA, torch allocates from the innermost active pool, so the KV cache and the runtime state stay in separate pools. The XPU caching allocator picks the outermost one instead. get_pool() walks the active pools in opposite orders:

• CUDA: captures_underway.rbegin(), i.e. LIFO, the innermost pool (c10/cuda/CUDACachingAllocator.cpp, v2.11.0 L3080-3084);
• XPU: for (auto& entry : captures_underway), i.e. FIFO, the outermost pool (c10/xpu/XPUCachingAllocator.cpp, v2.11.0 L670-672). This is unchanged in v2.12, v2.13 and main (renamed allocation_scopes_).

As a result:

• the KV cache lands in the runtime pool (still tagged kv_cache);
• a later small runtime tensor (for example a block table) can be carved from the free tail of a kv_cache-tagged segment;
• that tensor is then discarded on sleep instead of offloaded, and reads garbage after wake.

Repro, run on this branch (stacked on #59156) on a single XPU:

import torch
from vllm.device_allocator.xpumem import XpuMemAllocator

a = XpuMemAllocator.get_instance()
with a.use_memory_pool(tag="runtime"):
    with a.use_memory_pool(tag="kv_cache"):
        kv = torch.empty(3 << 20, dtype=torch.uint8, device="xpu")
    rt = torch.empty(3 << 19, dtype=torch.uint8, device="xpu")

for name, t in (("kv", kv), ("rt", rt)):
    for tag, entries in a.allocator_and_pools.items():
        for pool, _ in entries:
            if any(s["address"] <= t.data_ptr() < s["address"] + s["total_size"]
                   for s in pool.snapshot()):
                print(f"{name} is in the {tag!r} pool")
    for d in a.pointer_to_data.values():
        base, size = d.handle[2], d.handle[1]
        if base <= t.data_ptr() < base + size:
            print(f"{name} is in a segment tagged {d.tag!r}")

Expected (what CUDA prints): kv is in the kv_cache pool and a kv_cache segment; rt is in the runtime pool and a runtime segment. The problem case is kv in the runtime pool, or rt in a segment tagged kv_cache.

Please post the output, so we can confirm this on hardware. Unless it shows otherwise, I'll skip the runtime pool around initialize_kv_cache on XPU (one line), so XPU keeps main's behaviour for that state. The runner-construction buffers stay under runtime either way.

🤖 Generated with Claude Code

verified on XPU and torch 2.14, claude is correct.

entries: [(<torch.xpu.memory.MemPool object at 0x75951a38dc10>, <torch.xpu.memory.XPUPluggableAllocator object at 0x75951a345310>)]
kv is in the 'runtime' pool
entries: [(<torch.xpu.memory.MemPool object at 0x75951a38dc70>, <torch.xpu.memory.XPUPluggableAllocator object at 0x75951a345730>)]
kv is in a segment tagged 'kv_cache'
entries: [(<torch.xpu.memory.MemPool object at 0x75951a38dc10>, <torch.xpu.memory.XPUPluggableAllocator object at 0x75951a345310>)]
rt is in the 'runtime' pool
entries: [(<torch.xpu.memory.MemPool object at 0x75951a38dc70>, <torch.xpu.memory.XPUPluggableAllocator object at 0x75951a345730>)]
rt is in a segment tagged 'kv_cache'

cc @yma11 please work with Intel torch team.

[mergify]

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

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

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #92580 for commit f3e87c387d78.