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

tests/basic_correctness/test_mem.py

@@ -6,6 +6,8 @@
 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 @@ async def test():
 
 
 @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 @@ def initialize_kv_cache(
         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 @@ def _make_ptr_tensor(self, x: Iterable[torch.Tensor]) -> torch.Tensor:
         )
 
     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 @@
 import functools
 import gc
 import time
+from contextlib import AbstractContextManager
 from copy import deepcopy
 from typing import Any, NamedTuple
 
@@ -524,7 +525,10 @@ def get_kv_cache_spec(self):
         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 @@ def initialize_kv_cache(
             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
@@ -861,9 +866,6 @@ def profile_run(self) -> None:
         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()