[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>
Ronald committed
Aug 28, 2026 at 06:08 UTC
479eeb32d2b432dbb4e442fbd3f94ca2eca35d67
9 files changed
+228
-146
tests/basic_correctness/test_mem.py
+47
-4
@@ -6,6 +6,8 @@ import asyncio
6
import pytest
7
import torch
8
9
+import vllm.device_allocator.cumem as cumem
10
+import vllm.envs as envs
11
from vllm import LLM, AsyncEngineArgs, AsyncLLMEngine, SamplingParams
12
from vllm.device_allocator import get_mem_allocator_instance
13
from vllm.platforms import current_platform
@@ -98,10 +100,8 @@ def test_discard_tags():
100
# Weights are still usable
101
assert torch.allclose(weights, torch.ones_like(weights))
102
101
- # Wake up and verify kv_cache is remapped (zeroed content)
103
+ # Wake up and verify kv_cache is remapped; discarded contents are undefined.
104
allocator.wake_up()
103
- # After wake_up the VA is remapped; content is not preserved
104
- # but the allocation is valid
105
assert kv.shape == (512, 512)
106
107
# Full sleep/wake cycle still works after discard
@@ -377,7 +377,13 @@ def test_deep_sleep_async():
377
378
379
@requires_fp8
380
-def test_deep_sleep_fp8_kvcache():
380
+def test_deep_sleep_fp8_kvcache_mrv1(monkeypatch: pytest.MonkeyPatch):
381
+ # Regression test for https://github.com/vllm-project/vllm/pull/28783.
382
+ # In particular, verify that MRV1 does not rely on post_kv_cache_wake_up()
383
+ # to restore correct output after level-2 sleep.
384
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
385
+ envs.disable_envs_cache()
386
+
387
model = "Qwen/Qwen2-0.5B"
388
used_bytes_baseline = current_platform.get_current_memory_usage()
389
@@ -409,3 +415,40 @@ def test_deep_sleep_fp8_kvcache():
415
416
# cmp output
417
assert output[0].outputs[0].text == output2[0].outputs[0].text
418
+
419
+
420
+@requires_fp8
421
+def test_deep_sleep_fp8_kvcache_mrv1_with_undefined_remap(
422
+ monkeypatch: pytest.MonkeyPatch,
423
+):
424
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
425
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
426
+ envs.disable_envs_cache()
427
+
428
+ llm = LLM(
429
+ "Qwen/Qwen2-0.5B",
430
+ enable_sleep_mode=True,
431
+ kv_cache_dtype="fp8",
432
+ )
433
+ prompt = "How are you?"
434
+ sampling_params = SamplingParams(temperature=0, max_tokens=10)
435
+ expected = llm.generate(prompt, sampling_params)
436
+
437
+ llm.sleep(level=2)
438
+ llm.wake_up(tags=["weights"])
439
+ llm.collective_rpc("reload_weights")
440
+
441
+ original_create_and_map = cumem.create_and_map
442
+
443
+ def create_and_map_with_poison(handle) -> None:
444
+ original_create_and_map(handle)
445
+ _, size, ptr, _ = handle
446
+ cumem.libcudart.cudaMemset(ptr, 0xA5, size)
447
+
448
+ monkeypatch.setattr(cumem, "create_and_map", create_and_map_with_poison)
449
+
450
+ # New requests must overwrite undefined remapped KV bytes before reading them.
451
+ llm.wake_up(tags=["kv_cache"])
452
+ actual = llm.generate(prompt, sampling_params)
453
+
454
+ assert expected[0].outputs[0].text == actual[0].outputs[0].text
tests/v1/worker/test_gpu_model_runner.py
-56
@@ -1752,59 +1752,3 @@ def test_mamba_cache_raises_when_max_num_seqs_exceeds_blocks():
1752
1753
with pytest.raises(ValueError, match="max_num_seqs"):
1754
runner.initialize_kv_cache(kv_cache_config)
1755
-
1756
-
1757
-class TestInitFp8KvScalesHybridModels:
1758
- """Verify init_fp8_kv_scales handles heterogeneous kv_caches entries.
1759
-
1760
- Hybrid models (Mamba, DeltaNet) store per-layer state as a list of tensors
1761
- rather than a single tensor. init_fp8_kv_scales must iterate both forms.
1762
- """
1763
-
1764
- @staticmethod
1765
- def _make_runner_stub(kv_caches):
1766
- runner = Mock(spec=GPUModelRunner)
1767
- runner.cache_config = SimpleNamespace(cache_dtype="fp8_e4m3")
1768
- runner.kv_caches = kv_caches
1769
- runner.compilation_config = SimpleNamespace(static_forward_context={})
1770
- runner.init_fp8_kv_scales = GPUModelRunner.init_fp8_kv_scales.__get__(
1771
- runner, GPUModelRunner
1772
- )
1773
- return runner
1774
-
1775
- def test_zeroes_both_tensor_and_list_entries(self):
1776
- single_tensor = torch.ones(4, 8)
1777
- list_tensors = [torch.ones(2, 4), torch.ones(3, 6)]
1778
-
1779
- runner = self._make_runner_stub([single_tensor, list_tensors])
1780
- runner.init_fp8_kv_scales()
1781
-
1782
- assert (single_tensor == 0).all()
1783
- assert all((t == 0).all() for t in list_tensors)
1784
-
1785
- def test_skips_none_entries(self):
1786
- tensor = torch.ones(4, 8)
1787
- runner = self._make_runner_stub([None, tensor, None])
1788
- runner.init_fp8_kv_scales()
1789
-
1790
- assert (tensor == 0).all()
1791
-
1792
- def test_noop_when_kv_cache_not_quantized(self):
1793
- tensor = torch.ones(4, 8)
1794
- runner = self._make_runner_stub([tensor])
1795
- runner.cache_config.cache_dtype = "auto"
1796
- runner.init_fp8_kv_scales()
1797
-
1798
- assert (tensor == 1).all()
1799
-
1800
- def test_mixed_none_tensor_and_list(self):
1801
- t1 = torch.ones(2, 2)
1802
- t2 = torch.ones(3, 3)
1803
- list_entry = [torch.ones(1, 1), torch.ones(1, 1)]
1804
-
1805
- runner = self._make_runner_stub([None, t1, list_entry, None, t2])
1806
- runner.init_fp8_kv_scales()
1807
-
1808
- assert (t1 == 0).all()
1809
- assert (t2 == 0).all()
1810
- assert all((t == 0).all() for t in list_entry)
tests/v1/worker/test_kv_cache_allocation_scope.py
new
+135
@@ -0,0 +1,135 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from contextlib import AbstractContextManager
5
+from types import SimpleNamespace
6
+from typing import Any, cast
7
+
8
+import torch
9
+
10
+import vllm.v1.worker.gpu.attn_utils as attn_utils
11
+import vllm.v1.worker.gpu_model_runner as gpu_model_runner
12
+from vllm.v1.worker.gpu_worker import Worker
13
+
14
+
15
+class _AllocationScope(AbstractContextManager):
16
+ def __init__(self) -> None:
17
+ self.active = False
18
+
19
+ def __enter__(self):
20
+ assert not self.active
21
+ self.active = True
22
+ return self
23
+
24
+ def __exit__(self, *args: Any) -> None:
25
+ assert self.active
26
+ self.active = False
27
+
28
+
29
+def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None:
30
+ scope = _AllocationScope()
31
+ kv_caches = {"layer": torch.empty(0)}
32
+
33
+ def allocate(*args, **kwargs):
34
+ assert scope.active
35
+ return kv_caches
36
+
37
+ def bind(*args, **kwargs):
38
+ assert not scope.active
39
+
40
+ monkeypatch.setattr(attn_utils, "allocate_kv_cache", allocate)
41
+ monkeypatch.setattr(attn_utils, "bind_kv_cache", bind)
42
+ monkeypatch.setattr(attn_utils, "get_shared_kv_cache_layers", lambda config: {})
43
+
44
+ config = SimpleNamespace(
45
+ cache_config=SimpleNamespace(get_resolved_kv_cache_layout=lambda: None),
46
+ model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type="test")),
47
+ )
48
+ result = attn_utils.init_kv_cache(
49
+ [],
50
+ {},
51
+ object(),
52
+ torch.device("cpu"),
53
+ [],
54
+ config,
55
+ kv_cache_allocation_context=scope,
56
+ )
57
+
58
+ assert result is kv_caches
59
+ assert not scope.active
60
+
61
+
62
+def test_mrv1_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None:
63
+ scope = _AllocationScope()
64
+ kv_caches = {"layer": torch.empty(0)}
65
+
66
+ def allocate(*args, **kwargs):
67
+ assert scope.active
68
+ return kv_caches
69
+
70
+ def bind(*args, **kwargs):
71
+ assert not scope.active
72
+
73
+ monkeypatch.setattr(gpu_model_runner, "allocate_kv_cache", allocate)
74
+ monkeypatch.setattr(gpu_model_runner, "bind_kv_cache", bind)
75
+
76
+ runner = SimpleNamespace(
77
+ device=torch.device("cpu"),
78
+ cache_config=SimpleNamespace(get_resolved_kv_cache_layout=lambda: None),
79
+ shared_kv_cache_layers={},
80
+ model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type="test")),
81
+ compilation_config=SimpleNamespace(static_forward_context={}),
82
+ kv_caches=[],
83
+ )
84
+ result = gpu_model_runner.GPUModelRunner.initialize_kv_cache_tensors(
85
+ runner,
86
+ object(),
87
+ [],
88
+ kv_cache_allocation_context=scope,
89
+ )
90
+
91
+ assert result is kv_caches
92
+ assert not scope.active
93
+
94
+
95
+def test_kv_wake_does_not_run_model_runner_recovery() -> None:
96
+ model = torch.nn.Module()
97
+ model.register_buffer("_k_scale", torch.tensor(0.5))
98
+ model.register_buffer("_v_scale", torch.tensor(0.25))
99
+
100
+ class Runner:
101
+ def __init__(self) -> None:
102
+ self.model = model
103
+ self.layout_tensors = tuple(torch.tensor([i]) for i in range(5))
104
+ self.recovery_calls = 0
105
+
106
+ def post_kv_cache_wake_up(self) -> None:
107
+ self.recovery_calls += 1
108
+ self.model.get_buffer("_k_scale").fill_(1.0)
109
+ self.model.get_buffer("_v_scale").fill_(1.0)
110
+ self.layout_tensors = tuple(torch.tensor([i]) for i in range(5))
111
+
112
+ runner = Runner()
113
+ worker = cast(
114
+ Worker,
115
+ SimpleNamespace(
116
+ _get_sleep_mode_backend=lambda: SimpleNamespace(resume=lambda tags: None),
117
+ _sleep_saved_buffers={},
118
+ _sleep_saved_draft_buffers={},
119
+ model_runner=runner,
120
+ synchronize_device=lambda: None,
121
+ ),
122
+ )
123
+ layout_tensors = runner.layout_tensors
124
+ layout_ptrs = tuple(t.data_ptr() for t in layout_tensors)
125
+
126
+ Worker.wake_up(worker, tags=["kv_cache"])
127
+
128
+ assert runner.recovery_calls == 0
129
+ assert model.get_buffer("_k_scale").item() == 0.5
130
+ assert model.get_buffer("_v_scale").item() == 0.25
131
+ assert all(
132
+ actual is expected
133
+ for actual, expected in zip(runner.layout_tensors, layout_tensors, strict=True)
134
+ )
135
+ assert tuple(t.data_ptr() for t in runner.layout_tensors) == layout_ptrs
vllm/v1/worker/cpu_model_runner.py
+7
-2
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import sys
4
-from contextlib import contextmanager
4
+from contextlib import AbstractContextManager, contextmanager
5
from typing import Any
6
7
import torch
@@ -157,8 +157,13 @@ class CPUModelRunner(GPUModelRunner):
157
self,
158
kv_cache_config: KVCacheConfig,
159
is_profiling: bool = False,
160
+ kv_cache_allocation_context: AbstractContextManager | None = None,
161
) -> None:
161
- super().initialize_kv_cache(kv_cache_config, is_profiling)
162
+ super().initialize_kv_cache(
163
+ kv_cache_config,
164
+ is_profiling,
165
+ kv_cache_allocation_context=kv_cache_allocation_context,
166
+ )
167
168
if self.speculative_config:
169
if self.speculative_config.use_eagle():
vllm/v1/worker/gpu/attn_utils.py
+10
-6
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Mapping, Sequence
4
+from contextlib import AbstractContextManager, nullcontext
5
from dataclasses import dataclass
6
from typing import Any, cast
7
@@ -208,13 +209,16 @@ def init_kv_cache(
209
device: torch.device,
210
kernel_block_sizes: list[int],
211
vllm_config: VllmConfig,
212
+ kv_cache_allocation_context: AbstractContextManager | None = None,
213
) -> dict[str, Any]:
212
- kv_caches = allocate_kv_cache(
213
- kv_cache_config,
214
- device,
215
- vllm_config.cache_config.get_resolved_kv_cache_layout(),
216
- kernel_block_sizes,
217
- )
214
+ allocation_context = kv_cache_allocation_context or nullcontext()
215
+ with allocation_context:
216
+ kv_caches = allocate_kv_cache(
217
+ kv_cache_config,
218
+ device,
219
+ vllm_config.cache_config.get_resolved_kv_cache_layout(),
220
+ kernel_block_sizes,
221
+ )
222
for layer_name, target in get_shared_kv_cache_layers(vllm_config).items():
223
kv_caches[layer_name] = kv_caches[target]
224
# Dual-attention models (e.g. LongCat-Flash) put two Attention modules per
vllm/v1/worker/gpu/block_table.py
-5
@@ -86,11 +86,6 @@ class BlockTables:
86
)
87
88
def init_block_table_layout_tensors(self) -> None:
89
- # Called at init and after a CuMem kv_cache wake-up. The ptr tensors
90
- # cache raw data_ptr() values that go stale once the underlying tensors
91
- # are reallocated on wake; the size tensors need re-populating because
92
- # their storage lives under the kv_cache pool tag and comes back with
93
- # undefined contents.
89
self.block_table_ptrs = self._make_ptr_tensor(
90
[b.gpu for b in self.block_tables]
91
)
vllm/v1/worker/gpu/model_runner.py
+6
-4
@@ -20,6 +20,7 @@ instead of embedding feature-specific logic directly.
20
import functools
21
import gc
22
import time
23
+from contextlib import AbstractContextManager
24
from copy import deepcopy
25
from typing import Any, NamedTuple
26
@@ -524,7 +525,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
525
return get_kv_cache_spec(self.vllm_config)
526
527
def initialize_kv_cache(
527
- self, kv_cache_config: KVCacheConfig, is_profiling: bool = False
528
+ self,
529
+ kv_cache_config: KVCacheConfig,
530
+ is_profiling: bool = False,
531
+ kv_cache_allocation_context: AbstractContextManager | None = None,
532
) -> None:
533
kv_cache_config = deepcopy(kv_cache_config)
534
self.kv_cache_config = kv_cache_config
@@ -656,6 +660,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
660
self.device,
661
self.kernel_block_sizes,
662
self.vllm_config,
663
+ kv_cache_allocation_context=kv_cache_allocation_context,
664
)
665
if is_profiling:
666
self.kv_connector = NO_OP_KV_CONNECTOR
@@ -863,9 +868,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
868
self.reset_encoder_cache()
869
gc.collect()
870
866
- def post_kv_cache_wake_up(self) -> None:
867
- self.block_tables.init_block_table_layout_tensors()
868
-
871
def reset_mm_cache(self) -> None:
872
if self.encoder_cache is not None:
873
self.encoder_cache.reset_mm_cache()
vllm/v1/worker/gpu_model_runner.py
+17
-64
@@ -59,7 +59,7 @@ from vllm.forward_context import (
59
)
60
from vllm.logger import init_logger
61
from vllm.lora.layers import BaseLayerWithLoRA, LoRAMapping, LoRAMappingType
62
-from vllm.model_executor.layers.attention import Attention, MLAAttention
62
+from vllm.model_executor.layers.attention import Attention
63
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
64
from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
65
from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
@@ -132,7 +132,6 @@ from vllm.utils.torch_utils import (
132
PIN_MEMORY,
133
async_tensor_h2d,
134
current_stream,
135
- is_quantized_kv_cache,
135
kv_cache_dtype_str_to_dtype,
136
)
137
from vllm.v1.attention.backend import (
@@ -1027,60 +1026,6 @@ class GPUModelRunner(
1026
self.encoder_cache.clear()
1027
self.late_interaction_runner.clear()
1028
1030
- def post_kv_cache_wake_up(self) -> None:
1031
- self.init_fp8_kv_scales()
1032
-
1033
- @torch.inference_mode()
1034
- def init_fp8_kv_scales(self) -> None:
1035
- """
1036
- Re-initialize the KV cache and FP8 scales after waking from sleep.
1037
- 1. Zero out the KV cache tensors to remove garbage data from re-allocation.
1038
- 2. Reset Attention layer scaling factors (_k_scale, _v_scale) to 1.0.
1039
- If these are left at 0.0 (default after wake_up), all KV cache values
1040
- become effectively zero, causing gibberish output.
1041
- """
1042
- if not is_quantized_kv_cache(self.cache_config.cache_dtype):
1043
- return
1044
-
1045
- kv_caches = getattr(self, "kv_caches", [])
1046
- for cache_entry in kv_caches:
1047
- if cache_entry is None:
1048
- continue
1049
- # Hybrid models (Mamba, DeltaNet) store per-layer state as a
1050
- # list of tensors rather than a single tensor.
1051
- if isinstance(cache_entry, list):
1052
- for t in cache_entry:
1053
- t.zero_()
1054
- else:
1055
- cache_entry.zero_()
1056
-
1057
- k_attr_names = ("_k_scale", "k_scale")
1058
- v_attr_names = ("_v_scale", "v_scale")
1059
-
1060
- attn_layers = self.compilation_config.static_forward_context
1061
- for name, module in attn_layers.items():
1062
- if isinstance(module, (Attention, MLAAttention)):
1063
- # TODO: Generally, scale is 1.0 if user uses on-the-fly fp8
1064
- # kvcache quant. However, to get better accuracy, compression
1065
- # frameworks like llm-compressors allow users to tune the
1066
- # scale. We may need to restore the specific calibrated scales
1067
- # here in the future.
1068
- k_scale_val, v_scale_val = 1.0, 1.0
1069
-
1070
- # Processing K Scale
1071
- for attr in k_attr_names:
1072
- if hasattr(module, attr):
1073
- param = getattr(module, attr)
1074
- if isinstance(param, torch.Tensor):
1075
- param.fill_(k_scale_val)
1076
-
1077
- # Processing V Scale
1078
- for attr in v_attr_names:
1079
- if hasattr(module, attr):
1080
- param = getattr(module, attr)
1081
- if isinstance(param, torch.Tensor):
1082
- param.fill_(v_scale_val)
1083
-
1029
def _get_positions(self, num_tokens: Any):
1030
if isinstance(num_tokens, int):
1031
if self.uses_mrope:
@@ -7443,7 +7388,10 @@ class GPUModelRunner(
7388
yield from attn_groups
7389
7390
def initialize_kv_cache_tensors(
7446
- self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]
7391
+ self,
7392
+ kv_cache_config: KVCacheConfig,
7393
+ kernel_block_sizes: list[int],
7394
+ kv_cache_allocation_context: AbstractContextManager | None = None,
7395
) -> dict[str, torch.Tensor]:
7396
"""
7397
Initialize the memory buffer for KV cache.
@@ -7457,12 +7405,14 @@ class GPUModelRunner(
7405
corresponding memory buffer for KV cache.
7406
"""
7407
7460
- kv_caches = allocate_kv_cache(
7461
- kv_cache_config,
7462
- self.device,
7463
- self.cache_config.get_resolved_kv_cache_layout(),
7464
- kernel_block_sizes,
7465
- )
7408
+ allocation_context = kv_cache_allocation_context or nullcontext()
7409
+ with allocation_context:
7410
+ kv_caches = allocate_kv_cache(
7411
+ kv_cache_config,
7412
+ self.device,
7413
+ self.cache_config.get_resolved_kv_cache_layout(),
7414
+ kernel_block_sizes,
7415
+ )
7416
7417
# Set up cross-layer KV cache sharing
7418
for layer_name, target_layer_name in self.shared_kv_cache_layers.items():
@@ -7512,6 +7462,7 @@ class GPUModelRunner(
7462
self,
7463
kv_cache_config: KVCacheConfig,
7464
is_profiling: bool = False,
7465
+ kv_cache_allocation_context: AbstractContextManager | None = None,
7466
) -> None:
7467
"""
7468
Initialize KV cache based on `kv_cache_config`.
@@ -7544,7 +7495,9 @@ class GPUModelRunner(
7495
# Reinitialize need to after initialize_attn_backend
7496
self.may_reinitialize_input_batch(kv_cache_config, kernel_block_sizes)
7497
kv_caches = self.initialize_kv_cache_tensors(
7547
- kv_cache_config, kernel_block_sizes
7498
+ kv_cache_config,
7499
+ kernel_block_sizes,
7500
+ kv_cache_allocation_context=kv_cache_allocation_context,
7501
)
7502
7503
if (
vllm/v1/worker/gpu_worker.py
+6
-5
@@ -256,9 +256,6 @@ class Worker(WorkerBase):
256
buffer.data.copy_(self._sleep_saved_draft_buffers[name].data)
257
self._sleep_saved_draft_buffers = {}
258
259
- if tags is None or "kv_cache" in tags:
260
- self.model_runner.post_kv_cache_wake_up()
261
-
259
self.synchronize_device()
260
261
def checkpoint_prepare(self) -> None:
@@ -689,8 +686,12 @@ class Worker(WorkerBase):
686
# related to kv cache connector (e.g. kv cache sharing layers).
687
ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config)
688
692
- with self._maybe_get_memory_pool_context(tag="kv_cache"):
693
- self.model_runner.initialize_kv_cache(kv_cache_config)
689
+ self.model_runner.initialize_kv_cache(
690
+ kv_cache_config,
691
+ kv_cache_allocation_context=self._maybe_get_memory_pool_context(
692
+ tag="kv_cache"
693
+ ),
694
+ )
695
696
if self.model_config.enable_return_routed_experts:
697
self.model_runner.init_routed_experts_capturer()