[Bugfix][MRV2] Isolate sleep-mode KV allocations

[Ronald1995]

Purpose

Part of #53343. Split from #53344.

The Model Runner V2 KV-cache memory-pool context currently covers initialization work beyond the discardable cache backing allocation. Persistent metadata created in that scope can inherit the kv_cache tag and lose its contents during level-1 sleep while retaining the same virtual address.

This PR is intentionally MRV2-only. It passes the allocation context through the MRV2 model runner and enters it only around allocate_kv_cache. MRV1 keeps its existing behavior.

Changes

• restrict the MRV2 kv_cache pool to backing allocations;
• keep cache sharing, binding, and runtime setup outside the pool;
• add an MRV2 allocation-boundary test;
• add deterministic poisoned-remap allocator and CUDA Graph reproducers.

Dependency

Independent. Can merge directly into main.

Test Plan

pytest -q tests/v1/worker/test_kv_cache_allocation_scope.py
pytest -q tests/basic_correctness/test_mem.py -k "tagged_ordinary_tensor or cudagraph_replays"

test result:

===================================================================================== warnings summary =====================================================================================
.venv/lib/python3.12/site-packages/torch/jit/_script.py:365: 14 warnings
  /home/liurong/vllm/.venv/lib/python3.12/site-packages/torch/jit/_script.py:365: DeprecationWarning: `torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`.
    warnings.warn(

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
1 passed, 14 warnings in 1.00s
..                                                                                                                                                                                   [100%]
===================================================================================== warnings summary =====================================================================================
.venv/lib/python3.12/site-packages/torch/jit/_script.py:365: 14 warnings
  /home/liurong/vllm/.venv/lib/python3.12/site-packages/torch/jit/_script.py:365: DeprecationWarning: `torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`.
    warnings.warn(

tests/basic_correctness/test_mem.py::test_tagged_ordinary_tensor_is_discarded_with_kv_cache
tests/basic_correctness/test_mem.py::test_cudagraph_replays_with_corrupted_tagged_constant
  /home/liurong/vllm/tests/utils.py:1759: DeprecationWarning: This process (pid=4125262) is multi-threaded, use of fork() may lead to deadlocks in the child.
    pid = os.fork()

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
2 passed, 11 deselected, 16 warnings in 34.01s

Changed files pass Ruff, git diff --check, and Python compilation. CUDA tests require a supported Linux CUDA runtime and were not executed on the Windows development host.

Follow-up: reevaluate post_kv_cache_wake_up

A follow-up PR will evaluate removing the MRV2 post_kv_cache_wake_up() implementation that rebuilds the BlockTable layout tensors.

That hook currently compensates for the previous broad allocation scope: BlockTables and its pointer, stride, and block-size metadata were created while the kv_cache pool was active, so level-1 sleep discarded their contents and wake-up had to reconstruct them.

After this PR, those metadata tensors are created outside the discardable kv_cache pool. They should therefore remain resident and no longer require wake-time reconstruction. Removing the workaround would also avoid replacing tensor objects after wake-up, which is preferable for consumers that rely on stable storage addresses.

The cleanup is intentionally deferred until a follow-up verifies pointer/value stability across sleep and wake-up and confirms behavior across supported platforms and sleep-mode backends.

[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.

[mergify]

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

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

[aoshen02]

I reproduced the sleep/wake failure and tested the follow-up in Ronald1995#9.

Environment

• Node: 1x H200 node, 4 GPUs used
• Image: vllm/vllm-openai:nightly@sha256:3578c1fa6a9676e1de068b9d75c777cc865d251fadfbe6175ae82278739c6674
• Image vLLM commit: a9a17e7095a66ef6c6685a1c7ddd657781a78d3c
• PR patch rebased onto vllm-project/vllm@5e379a361e3ea8bb82b7efd768c36f39a0cf32fd
• Model: Qwen/Qwen3.5-35B-A3B
• Model snapshot: 59d61f3ce65a6d9863b86d2e96597125219dc754
• No source bind mount: the modified Python files were copied into the container.

Engine configuration

os.environ["VLLM_USE_V2_MODEL_RUNNER"] = "1"

llm = LLM(
    model="Qwen/Qwen3.5-35B-A3B",
    tensor_parallel_size=4,
    speculative_config={"method": "mtp", "num_speculative_tokens": 3},
    compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.FULL),
    enable_sleep_mode=True,
    enable_prefix_caching=False,
    max_model_len=256,
    max_num_seqs=1,
    limit_mm_per_prompt={"image": 0, "video": 0},
    trust_remote_code=True,
)

Additional environment:

HF_HUB_OFFLINE=1
TRANSFORMERS_OFFLINE=1
VLLM_ALLOW_INSECURE_SERIALIZATION=1

Reproduction sequence

1. Verify all four workers use fused MTP decode.
2. Record the five block-table layout tensor data_ptr() values on every worker.
3. Generate four tokens for The capital of France is.
4. llm.sleep(level=2).
5. llm.wake_up(tags=["weights"]).
6. llm.collective_rpc("reload_weights").
7. llm.wake_up(tags=["kv_cache"]).
8. Record layout pointers again and repeat the same deterministic decode.

Results

PR #53508 as-is

• sleep: 0.503149 s
• wake weights: 0.407355 s
• wake KV: 1.497417 s
• first post-wake decode: failed on all TP workers with CUDA error: an illegal memory access was encountered

The remaining problem is GPUModelRunner.post_kv_cache_wake_up() calling init_block_table_layout_tensors(): it replaces tensors referenced by the fused multi-step CUDA graph even though #53508 already moved this metadata outside the sleepable KV allocation scope.

#53508 + remove MRV2 post-wake reconstruction

• layout_ptrs_stable=true
• tokens_equal=true
• sleep: 0.507078 s
• wake weights: 0.454879 s
• wake KV: 1.132419 s
• first post-wake decode: 0.120755 s
• no CUDA error

The legacy Model Runner callback is retained because it still restores FP8 KV state; only MRV2 skips the callback. Pre-commit, mypy, and DCO checks pass.

[aoshen02]

Strict review: the allocation-scope direction is correct, but this revision is not merge-ready.

1. Blocking runtime failure: MRV2 still calls post_kv_cache_wake_up(), which recreates the block-table layout tensors. Those tensors are now intentionally resident outside the sleepable KV pool, so rebuilding them is unnecessary and changes addresses captured by fused multi-step CUDA graphs. In the 4x H200 / Qwen3.5-35B-A3B / TP4 / fused-MTP / full-CUDA-Graph reproducer, PR #53508 as-is fails on the first post-wake decode on every TP worker with CUDA error: an illegal memory access was encountered. Removing the MRV2 reconstruction makes layout pointers stable and preserves identical output. The minimal follow-up is Ronald1995#9.

2. Blocking type-check failure: Worker.model_runner is typed as the legacy GPUModelRunner, so the new kv_cache_allocation_context= argument at vllm/v1/worker/gpu_worker.py:691 fails mypy with Unexpected keyword argument ... [call-arg]. The changed files must pass pre-commit after rebasing.

3. Regression coverage is insufficient: test_kv_cache_allocation_scope.py calls attn_utils.init_kv_cache() directly and the allocator tests only demonstrate generic corruption. Neither test exercises Worker.initialize_from_config() followed by the real MRV2 sleep/wake callback. Consequently all proposed tests pass while the production fused-MTP path still crashes. Add a regression that asserts MRV2 does not rebuild graph-visible block-table metadata and that pointer identity survives KV wake-up.

4. The PR currently conflicts with main; rebase and rerun the complete changed-file checks.

Once the MRV2 wake reconstruction is removed, the type boundary is fixed, and the regression covers the lifecycle path, the core design is preferable to an in-place refresh workaround.

[aoshen02]

Correction to item 4 in my strict review: GitHub currently reports this commit as mergeable, so the earlier Mergify conflict notice is stale. The accurate validation blocker is that is failed and is skipped; the revision still needs a green required-check run after the three substantive fixes. The runtime failure, mypy failure, and regression-coverage findings are unchanged.

[github-actions]

✅ @Ronald1995, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[aoshen02]

Re-reviewed after the follow-up fix was merged and the branch was rebased onto latest main. The sleep/wakeup allocation scope and stable block-table refresh are now integrated; changed-file pre-commit passes locally and the remote pre-run-check/pre-commit workflow is green. The H200 MRV2 fused-MTP sleep/wakeup regression also passed with stable metadata pointers and identical tokens.

[aoshen02]

MRV1 support is now included in commit 9d2d78745. Both model runners pass the memory-pool context down to the backing KV allocation only; MRV1 keeps its FP8 KV-cache wake-up recovery. Validation: tests/v1/worker/test_kv_cache_allocation_scope.py passes for MRV1 and MRV2 (2 passed), full mypy 3.13 passes, changed-file pre-commit passes, and the updated remote pre-commit is green. The PR title should now be [Bugfix][MRV1][MRV2] Isolate sleep-mode KV allocations.

[aoshen02]

Re-reviewed after extending the allocation boundary to Model Runner V1. The shared worker path now scopes only the actual KV backing allocation for both runners, while MRV1 retains its FP8 wake recovery. MRV1/MRV2 unit tests, full mypy 3.13, local pre-commit, and remote pre-commit all pass.

[Ronald1995]

Why MRV1 no longer needs post_kv_cache_wake_up

Commit c84879845c removes the MRV1 post-KV wake-up hook, along with the now-unreachable post_kv_cache_wake_up() / init_fp8_kv_scales() implementation.

init_fp8_kv_scales() bundled two operations:

1. zeroing the remapped FP8 KV-cache backing tensors;
2. resetting _k_scale and _v_scale to 1.0.

Neither operation is required by the current sleep/wake lifecycle:

• _k_scale and _v_scale are now registered model buffers. Before level-2 sleep, Worker.sleep() saves model.named_buffers() to CPU, and Worker.wake_up(tags=["weights"]) restores those buffers. The host-side scale values are not part of the discarded KV allocation either. Therefore FP8 scales no longer depend on a KV-cache wake-up callback.
• KV-cache contents discarded by level-2 sleep are no longer logically valid. After wake-up, cache slots used by new requests are written before they are read. Their initial byte contents therefore do not need to be zero. This invariant is dtype-independent, which is also why BF16 KV cache has never required an equivalent zeroing step.
• The zeroing added in [Bugfix][sleepmode][fp8 kv cache]: Fix FP8 KV cache + sleep(level=2) gibberish output #28783 was coupled with the scale-reset workaround. The correctness issue was the lost FP8 scale state, rather than FP8 requiring an all-zero cache allocation.

A regression test now explicitly sets VLLM_USE_V2_MODEL_RUNNER=0, performs FP8 generation, level-2 sleep, weights/KV wake-up, and deterministic generation again, then verifies that the pre/post-wake outputs are identical. This directly covers the #28783 failure mode without relying on post_kv_cache_wake_up().

Ruff and changed-file static checks pass locally. The FP8 CUDA regression requires a supported GPU and is left to GPU CI.

[Ronald1995]

test result for test_deep_sleep_fp8_kvcache_mrv1 merged by #28783

===================================================================================== warnings summary =====================================================================================
.venv/lib/python3.12/site-packages/torch/jit/_script.py:365: 14 warnings
  /home/liurong/vllm/.venv/lib/python3.12/site-packages/torch/jit/_script.py:365: DeprecationWarning: `torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`.
    warnings.warn(

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
============================================================================= 1 passed, 14 warnings in 31.23s ==============================================================================

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #85729 for commit 093ae6c23453.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #85735 for commit 093ae6c23453.

[njhill]

Thanks @Ronald1995