diff --git a/tests/basic_correctness/test_mem.py b/tests/basic_correctness/test_mem.py index 696d2f66bc5b..6fed9b6c8c84 100644 --- a/tests/basic_correctness/test_mem.py +++ b/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 diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 94eb8054848c..bd8fc7dc0109 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/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) diff --git a/tests/v1/worker/test_kv_cache_allocation_scope.py b/tests/v1/worker/test_kv_cache_allocation_scope.py new file mode 100644 index 000000000000..81afe44275fe --- /dev/null +++ b/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 diff --git a/vllm/v1/worker/cpu_model_runner.py b/vllm/v1/worker/cpu_model_runner.py index 3e7f0d4c1e88..67ddf83baec5 100644 --- a/vllm/v1/worker/cpu_model_runner.py +++ b/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(): diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index d4ab1e506623..05febec328d4 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/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 diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index dae1630a2c19..5e09383ee6ab 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/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] ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index a58c86166df2..b2d56a57d500 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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() diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 0daabc978dab..a1cab24fbef5 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -59,7 +59,7 @@ ) 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 @@ 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 @@ def reset_encoder_cache(self) -> None: 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 @@ def _kv_cache_spec_attn_group_iterator(self) -> Iterator[AttentionGroup]: 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 @@ def initialize_kv_cache_tensors( 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 @@ def initialize_kv_cache( 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 @@ def initialize_kv_cache( # 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 ( diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 990fae6db2c1..65797bc2a1b0 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -256,9 +256,6 @@ def wake_up(self, tags: list[str] | None = None) -> None: 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 @@ def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None: # 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()