[Bugfix][MRV2] Isolate sleep-mode KV allocations (#53508)

Signed-off-by: Ronald1995 <ronaldautomobile@163.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: aoshen02 <aoshen@inferact.ai>

tests/basic_correctness/test_mem.py

@@ -6,6 +6,8 @@ import asyncio
 import pytest
 import torch
 
+import vllm.device_allocator.cumem as cumem
+import vllm.envs as envs
 from vllm import LLM, AsyncEngineArgs, AsyncLLMEngine, SamplingParams
 from vllm.device_allocator import get_mem_allocator_instance
 from vllm.platforms import current_platform
@@ -98,10 +100,8 @@ def test_discard_tags():
     # Weights are still usable
     assert torch.allclose(weights, torch.ones_like(weights))
 
-    # Wake up and verify kv_cache is remapped (zeroed content)
+    # Wake up and verify kv_cache is remapped; discarded contents are undefined.
     allocator.wake_up()
-    # After wake_up the VA is remapped; content is not preserved
-    # but the allocation is valid
     assert kv.shape == (512, 512)
 
     # Full sleep/wake cycle still works after discard
@@ -377,7 +377,13 @@ def test_deep_sleep_async():
 
 
 @requires_fp8
-def test_deep_sleep_fp8_kvcache():
+def test_deep_sleep_fp8_kvcache_mrv1(monkeypatch: pytest.MonkeyPatch):
+    # Regression test for https://github.com/vllm-project/vllm/pull/28783.
+    # In particular, verify that MRV1 does not rely on post_kv_cache_wake_up()
+    # to restore correct output after level-2 sleep.
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
+    envs.disable_envs_cache()
+
     model = "Qwen/Qwen2-0.5B"
     used_bytes_baseline = current_platform.get_current_memory_usage()
 
@@ -409,3 +415,40 @@ def test_deep_sleep_fp8_kvcache():
 
     # cmp output
     assert output[0].outputs[0].text == output2[0].outputs[0].text
+
+
+@requires_fp8
+def test_deep_sleep_fp8_kvcache_mrv1_with_undefined_remap(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    envs.disable_envs_cache()
+
+    llm = LLM(
+        "Qwen/Qwen2-0.5B",
+        enable_sleep_mode=True,
+        kv_cache_dtype="fp8",
+    )
+    prompt = "How are you?"
+    sampling_params = SamplingParams(temperature=0, max_tokens=10)
+    expected = llm.generate(prompt, sampling_params)
+
+    llm.sleep(level=2)
+    llm.wake_up(tags=["weights"])
+    llm.collective_rpc("reload_weights")
+
+    original_create_and_map = cumem.create_and_map
+
+    def create_and_map_with_poison(handle) -> None:
+        original_create_and_map(handle)
+        _, size, ptr, _ = handle
+        cumem.libcudart.cudaMemset(ptr, 0xA5, size)
+
+    monkeypatch.setattr(cumem, "create_and_map", create_and_map_with_poison)
+
+    # New requests must overwrite undefined remapped KV bytes before reading them.
+    llm.wake_up(tags=["kv_cache"])
+    actual = llm.generate(prompt, sampling_params)
+
+    assert expected[0].outputs[0].text == actual[0].outputs[0].text

tests/v1/worker/test_gpu_model_runner.py

@@ -1752,59 +1752,3 @@ def test_mamba_cache_raises_when_max_num_seqs_exceeds_blocks():
 
         with pytest.raises(ValueError, match="max_num_seqs"):
             runner.initialize_kv_cache(kv_cache_config)
-
-
-class TestInitFp8KvScalesHybridModels:
-    """Verify init_fp8_kv_scales handles heterogeneous kv_caches entries.
-
-    Hybrid models (Mamba, DeltaNet) store per-layer state as a list of tensors
-    rather than a single tensor. init_fp8_kv_scales must iterate both forms.
-    """
-
-    @staticmethod
-    def _make_runner_stub(kv_caches):
-        runner = Mock(spec=GPUModelRunner)
-        runner.cache_config = SimpleNamespace(cache_dtype="fp8_e4m3")
-        runner.kv_caches = kv_caches
-        runner.compilation_config = SimpleNamespace(static_forward_context={})
-        runner.init_fp8_kv_scales = GPUModelRunner.init_fp8_kv_scales.__get__(
-            runner, GPUModelRunner
-        )
-        return runner
-
-    def test_zeroes_both_tensor_and_list_entries(self):
-        single_tensor = torch.ones(4, 8)
-        list_tensors = [torch.ones(2, 4), torch.ones(3, 6)]
-
-        runner = self._make_runner_stub([single_tensor, list_tensors])
-        runner.init_fp8_kv_scales()
-
-        assert (single_tensor == 0).all()
-        assert all((t == 0).all() for t in list_tensors)
-
-    def test_skips_none_entries(self):
-        tensor = torch.ones(4, 8)
-        runner = self._make_runner_stub([None, tensor, None])
-        runner.init_fp8_kv_scales()
-
-        assert (tensor == 0).all()
-
-    def test_noop_when_kv_cache_not_quantized(self):
-        tensor = torch.ones(4, 8)
-        runner = self._make_runner_stub([tensor])
-        runner.cache_config.cache_dtype = "auto"
-        runner.init_fp8_kv_scales()
-
-        assert (tensor == 1).all()
-
-    def test_mixed_none_tensor_and_list(self):
-        t1 = torch.ones(2, 2)
-        t2 = torch.ones(3, 3)
-        list_entry = [torch.ones(1, 1), torch.ones(1, 1)]
-
-        runner = self._make_runner_stub([None, t1, list_entry, None, t2])
-        runner.init_fp8_kv_scales()
-
-        assert (t1 == 0).all()
-        assert (t2 == 0).all()
-        assert all((t == 0).all() for t in list_entry)

tests/v1/worker/test_kv_cache_allocation_scope.py

@@ -0,0 +1,135 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from contextlib import AbstractContextManager
+from types import SimpleNamespace
+from typing import Any, cast
+
+import torch
+
+import vllm.v1.worker.gpu.attn_utils as attn_utils
+import vllm.v1.worker.gpu_model_runner as gpu_model_runner
+from vllm.v1.worker.gpu_worker import Worker
+
+
+class _AllocationScope(AbstractContextManager):
+    def __init__(self) -> None:
+        self.active = False
+
+    def __enter__(self):
+        assert not self.active
+        self.active = True
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        assert self.active
+        self.active = False
+
+
+def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None:
+    scope = _AllocationScope()
+    kv_caches = {"layer": torch.empty(0)}
+
+    def allocate(*args, **kwargs):
+        assert scope.active
+        return kv_caches
+
+    def bind(*args, **kwargs):
+        assert not scope.active
+
+    monkeypatch.setattr(attn_utils, "allocate_kv_cache", allocate)
+    monkeypatch.setattr(attn_utils, "bind_kv_cache", bind)
+    monkeypatch.setattr(attn_utils, "get_shared_kv_cache_layers", lambda config: {})
+
+    config = SimpleNamespace(
+        cache_config=SimpleNamespace(get_resolved_kv_cache_layout=lambda: None),
+        model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type="test")),
+    )
+    result = attn_utils.init_kv_cache(
+        [],
+        {},
+        object(),
+        torch.device("cpu"),
+        [],
+        config,
+        kv_cache_allocation_context=scope,
+    )
+
+    assert result is kv_caches
+    assert not scope.active
+
+
+def test_mrv1_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None:
+    scope = _AllocationScope()
+    kv_caches = {"layer": torch.empty(0)}
+
+    def allocate(*args, **kwargs):
+        assert scope.active
+        return kv_caches
+
+    def bind(*args, **kwargs):
+        assert not scope.active
+
+    monkeypatch.setattr(gpu_model_runner, "allocate_kv_cache", allocate)
+    monkeypatch.setattr(gpu_model_runner, "bind_kv_cache", bind)
+
+    runner = SimpleNamespace(
+        device=torch.device("cpu"),
+        cache_config=SimpleNamespace(get_resolved_kv_cache_layout=lambda: None),
+        shared_kv_cache_layers={},
+        model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type="test")),
+        compilation_config=SimpleNamespace(static_forward_context={}),
+        kv_caches=[],
+    )
+    result = gpu_model_runner.GPUModelRunner.initialize_kv_cache_tensors(
+        runner,
+        object(),
+        [],
+        kv_cache_allocation_context=scope,
+    )
+
+    assert result is kv_caches
+    assert not scope.active
+
+
+def test_kv_wake_does_not_run_model_runner_recovery() -> None:
+    model = torch.nn.Module()
+    model.register_buffer("_k_scale", torch.tensor(0.5))
+    model.register_buffer("_v_scale", torch.tensor(0.25))
+
+    class Runner:
+        def __init__(self) -> None:
+            self.model = model
+            self.layout_tensors = tuple(torch.tensor([i]) for i in range(5))
+            self.recovery_calls = 0
+
+        def post_kv_cache_wake_up(self) -> None:
+            self.recovery_calls += 1
+            self.model.get_buffer("_k_scale").fill_(1.0)
+            self.model.get_buffer("_v_scale").fill_(1.0)
+            self.layout_tensors = tuple(torch.tensor([i]) for i in range(5))
+
+    runner = Runner()
+    worker = cast(
+        Worker,
+        SimpleNamespace(
+            _get_sleep_mode_backend=lambda: SimpleNamespace(resume=lambda tags: None),
+            _sleep_saved_buffers={},
+            _sleep_saved_draft_buffers={},
+            model_runner=runner,
+            synchronize_device=lambda: None,
+        ),
+    )
+    layout_tensors = runner.layout_tensors
+    layout_ptrs = tuple(t.data_ptr() for t in layout_tensors)
+
+    Worker.wake_up(worker, tags=["kv_cache"])
+
+    assert runner.recovery_calls == 0
+    assert model.get_buffer("_k_scale").item() == 0.5
+    assert model.get_buffer("_v_scale").item() == 0.25
+    assert all(
+        actual is expected
+        for actual, expected in zip(runner.layout_tensors, layout_tensors, strict=True)
+    )
+    assert tuple(t.data_ptr() for t in runner.layout_tensors) == layout_ptrs

vllm/v1/worker/cpu_model_runner.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import sys
-from contextlib import contextmanager
+from contextlib import AbstractContextManager, contextmanager
 from typing import Any
 
 import torch
@@ -157,8 +157,13 @@ class CPUModelRunner(GPUModelRunner):
         self,
         kv_cache_config: KVCacheConfig,
         is_profiling: bool = False,
+        kv_cache_allocation_context: AbstractContextManager | None = None,
     ) -> None:
-        super().initialize_kv_cache(kv_cache_config, is_profiling)
+        super().initialize_kv_cache(
+            kv_cache_config,
+            is_profiling,
+            kv_cache_allocation_context=kv_cache_allocation_context,
+        )
 
         if self.speculative_config:
             if self.speculative_config.use_eagle():

vllm/v1/worker/gpu/attn_utils.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Mapping, Sequence
+from contextlib import AbstractContextManager, nullcontext
 from dataclasses import dataclass
 from typing import Any, cast
 
@@ -208,13 +209,16 @@ def init_kv_cache(
     device: torch.device,
     kernel_block_sizes: list[int],
     vllm_config: VllmConfig,
+    kv_cache_allocation_context: AbstractContextManager | None = None,
 ) -> dict[str, Any]:
-    kv_caches = allocate_kv_cache(
-        kv_cache_config,
-        device,
-        vllm_config.cache_config.get_resolved_kv_cache_layout(),
-        kernel_block_sizes,
-    )
+    allocation_context = kv_cache_allocation_context or nullcontext()
+    with allocation_context:
+        kv_caches = allocate_kv_cache(
+            kv_cache_config,
+            device,
+            vllm_config.cache_config.get_resolved_kv_cache_layout(),
+            kernel_block_sizes,
+        )
     for layer_name, target in get_shared_kv_cache_layers(vllm_config).items():
         kv_caches[layer_name] = kv_caches[target]
     # Dual-attention models (e.g. LongCat-Flash) put two Attention modules per

vllm/v1/worker/gpu/block_table.py

@@ -86,11 +86,6 @@ class BlockTables:
         )
 
     def init_block_table_layout_tensors(self) -> None:
-        # Called at init and after a CuMem kv_cache wake-up. The ptr tensors
-        # cache raw data_ptr() values that go stale once the underlying tensors
-        # are reallocated on wake; the size tensors need re-populating because
-        # their storage lives under the kv_cache pool tag and comes back with
-        # undefined contents.
         self.block_table_ptrs = self._make_ptr_tensor(
             [b.gpu for b in self.block_tables]
         )

vllm/v1/worker/gpu/model_runner.py

@@ -20,6 +20,7 @@ instead of embedding feature-specific logic directly.
 import functools
 import gc
 import time
+from contextlib import AbstractContextManager
 from copy import deepcopy
 from typing import Any, NamedTuple
 
@@ -524,7 +525,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         return get_kv_cache_spec(self.vllm_config)
 
     def initialize_kv_cache(
-        self, kv_cache_config: KVCacheConfig, is_profiling: bool = False
+        self,
+        kv_cache_config: KVCacheConfig,
+        is_profiling: bool = False,
+        kv_cache_allocation_context: AbstractContextManager | None = None,
     ) -> None:
         kv_cache_config = deepcopy(kv_cache_config)
         self.kv_cache_config = kv_cache_config
@@ -656,6 +660,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             self.device,
             self.kernel_block_sizes,
             self.vllm_config,
+            kv_cache_allocation_context=kv_cache_allocation_context,
         )
         if is_profiling:
             self.kv_connector = NO_OP_KV_CONNECTOR
@@ -863,9 +868,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         self.reset_encoder_cache()
         gc.collect()
 
-    def post_kv_cache_wake_up(self) -> None:
-        self.block_tables.init_block_table_layout_tensors()
-
     def reset_mm_cache(self) -> None:
         if self.encoder_cache is not None:
             self.encoder_cache.reset_mm_cache()

vllm/v1/worker/gpu_model_runner.py

@@ -59,7 +59,7 @@ from vllm.forward_context import (
 )
 from vllm.logger import init_logger
 from vllm.lora.layers import BaseLayerWithLoRA, LoRAMapping, LoRAMappingType
-from vllm.model_executor.layers.attention import Attention, MLAAttention
+from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
 from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
@@ -132,7 +132,6 @@ from vllm.utils.torch_utils import (
     PIN_MEMORY,
     async_tensor_h2d,
     current_stream,
-    is_quantized_kv_cache,
     kv_cache_dtype_str_to_dtype,
 )
 from vllm.v1.attention.backend import (
@@ -1027,60 +1026,6 @@ class GPUModelRunner(
         self.encoder_cache.clear()
         self.late_interaction_runner.clear()
 
-    def post_kv_cache_wake_up(self) -> None:
-        self.init_fp8_kv_scales()
-
-    @torch.inference_mode()
-    def init_fp8_kv_scales(self) -> None:
-        """
-        Re-initialize the KV cache and FP8 scales after waking from sleep.
-        1. Zero out the KV cache tensors to remove garbage data from re-allocation.
-        2. Reset Attention layer scaling factors (_k_scale, _v_scale) to 1.0.
-          If these are left at 0.0 (default after wake_up), all KV cache values
-          become effectively zero, causing gibberish output.
-        """
-        if not is_quantized_kv_cache(self.cache_config.cache_dtype):
-            return
-
-        kv_caches = getattr(self, "kv_caches", [])
-        for cache_entry in kv_caches:
-            if cache_entry is None:
-                continue
-            # Hybrid models (Mamba, DeltaNet) store per-layer state as a
-            # list of tensors rather than a single tensor.
-            if isinstance(cache_entry, list):
-                for t in cache_entry:
-                    t.zero_()
-            else:
-                cache_entry.zero_()
-
-        k_attr_names = ("_k_scale", "k_scale")
-        v_attr_names = ("_v_scale", "v_scale")
-
-        attn_layers = self.compilation_config.static_forward_context
-        for name, module in attn_layers.items():
-            if isinstance(module, (Attention, MLAAttention)):
-                # TODO: Generally, scale is 1.0 if user uses on-the-fly fp8
-                # kvcache quant. However, to get better accuracy, compression
-                # frameworks like llm-compressors allow users to tune the
-                # scale. We may need to restore the specific calibrated scales
-                # here in the future.
-                k_scale_val, v_scale_val = 1.0, 1.0
-
-                # Processing K Scale
-                for attr in k_attr_names:
-                    if hasattr(module, attr):
-                        param = getattr(module, attr)
-                        if isinstance(param, torch.Tensor):
-                            param.fill_(k_scale_val)
-
-                # Processing V Scale
-                for attr in v_attr_names:
-                    if hasattr(module, attr):
-                        param = getattr(module, attr)
-                        if isinstance(param, torch.Tensor):
-                            param.fill_(v_scale_val)
-
     def _get_positions(self, num_tokens: Any):
         if isinstance(num_tokens, int):
             if self.uses_mrope:
@@ -7443,7 +7388,10 @@ class GPUModelRunner(
             yield from attn_groups
 
     def initialize_kv_cache_tensors(
-        self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]
+        self,
+        kv_cache_config: KVCacheConfig,
+        kernel_block_sizes: list[int],
+        kv_cache_allocation_context: AbstractContextManager | None = None,
     ) -> dict[str, torch.Tensor]:
         """
         Initialize the memory buffer for KV cache.
@@ -7457,12 +7405,14 @@ class GPUModelRunner(
             corresponding memory buffer for KV cache.
         """
 
-        kv_caches = allocate_kv_cache(
-            kv_cache_config,
-            self.device,
-            self.cache_config.get_resolved_kv_cache_layout(),
-            kernel_block_sizes,
-        )
+        allocation_context = kv_cache_allocation_context or nullcontext()
+        with allocation_context:
+            kv_caches = allocate_kv_cache(
+                kv_cache_config,
+                self.device,
+                self.cache_config.get_resolved_kv_cache_layout(),
+                kernel_block_sizes,
+            )
 
         # Set up cross-layer KV cache sharing
         for layer_name, target_layer_name in self.shared_kv_cache_layers.items():
@@ -7512,6 +7462,7 @@ class GPUModelRunner(
         self,
         kv_cache_config: KVCacheConfig,
         is_profiling: bool = False,
+        kv_cache_allocation_context: AbstractContextManager | None = None,
     ) -> None:
         """
         Initialize KV cache based on `kv_cache_config`.
@@ -7544,7 +7495,9 @@ class GPUModelRunner(
         # Reinitialize need to after initialize_attn_backend
         self.may_reinitialize_input_batch(kv_cache_config, kernel_block_sizes)
         kv_caches = self.initialize_kv_cache_tensors(
-            kv_cache_config, kernel_block_sizes
+            kv_cache_config,
+            kernel_block_sizes,
+            kv_cache_allocation_context=kv_cache_allocation_context,
         )
 
         if (

vllm/v1/worker/gpu_worker.py

@@ -256,9 +256,6 @@ class Worker(WorkerBase):
                         buffer.data.copy_(self._sleep_saved_draft_buffers[name].data)
             self._sleep_saved_draft_buffers = {}
 
-        if tags is None or "kv_cache" in tags:
-            self.model_runner.post_kv_cache_wake_up()
-
         self.synchronize_device()
 
     def checkpoint_prepare(self) -> None:
@@ -689,8 +686,12 @@ class Worker(WorkerBase):
         # related to kv cache connector (e.g. kv cache sharing layers).
         ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config)
 
-        with self._maybe_get_memory_pool_context(tag="kv_cache"):
-            self.model_runner.initialize_kv_cache(kv_cache_config)
+        self.model_runner.initialize_kv_cache(
+            kv_cache_config,
+            kv_cache_allocation_context=self._maybe_get_memory_pool_context(
+                tag="kv_cache"
+            ),
+        )
 
         if self.model_config.enable_return_routed_experts:
             self.model_runner.init_routed_experts_capturer()