diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index e79a7afbd81..d2a2a6aedbf 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -47,18 +47,7 @@ class BlockTables: device=device, ) self.block_tables.append(block_table) - self.block_table_ptrs = self._make_ptr_tensor( - [b.gpu for b in self.block_tables] - ) - self.block_table_strides = torch.tensor( - [b.gpu.stride(0) for b in self.block_tables], - dtype=torch.int64, - device=self.device, - ) - self.block_sizes_tensor = torch.tensor( - self.block_sizes, dtype=torch.int32, device=self.device - ) self.num_blocks = UvaBackedTensor( (self.num_kv_cache_groups, self.max_num_reqs), dtype=torch.int32, @@ -69,7 +58,6 @@ class BlockTables: self.input_block_tables: list[torch.Tensor] = [ torch.zeros_like(b.gpu) for b in self.block_tables ] - self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables) self.slot_mappings = torch.zeros( self.num_kv_cache_groups, @@ -78,12 +66,33 @@ class BlockTables: device=self.device, ) + self.init_block_table_layout_tensors() + def _make_ptr_tensor(self, x: Iterable[torch.Tensor]) -> torch.Tensor: # NOTE(woosuk): Use uint64 instead of int64 to cover all possible addresses. return torch.tensor( [t.data_ptr() for t in x], dtype=torch.uint64, device=self.device ) + 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; block_sizes_tensor needs re-populating + # because its 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] + ) + self.block_table_strides = torch.tensor( + [b.gpu.stride(0) for b in self.block_tables], + dtype=torch.int64, + device=self.device, + ) + self.block_sizes_tensor = torch.tensor( + self.block_sizes, dtype=torch.int32, device=self.device + ) + self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables) + def append_block_ids( self, req_index: int, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index a3eb1589a3f..dfe50cb135d 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -558,6 +558,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): del hidden_states, sample_hidden_states 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 e2af34eecb9..53197a5c81a 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -884,6 +884,9 @@ 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: """ diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 19d0a68142e..d955f401d87 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -46,7 +46,7 @@ from vllm.tasks import SupportedTask from vllm.tracing import instrument from vllm.utils.mem_constants import GiB_bytes from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling -from vllm.utils.torch_utils import is_quantized_kv_cache, set_random_seed +from vllm.utils.torch_utils import set_random_seed from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec from vllm.v1.outputs import ( @@ -192,15 +192,8 @@ class Worker(WorkerBase): buffer.data.copy_(self._sleep_saved_buffers[name].data) self._sleep_saved_buffers = {} - # If the KV cache has just been woken up, - # the internal state of cache_engine must be reset, - # especially the FP8 scaling factor. - if ( - (tags is None or "kv_cache" in tags) - and is_quantized_kv_cache(self.cache_config.cache_dtype) - and hasattr(self.model_runner, "init_fp8_kv_scales") - ): - self.model_runner.init_fp8_kv_scales() + if tags is None or "kv_cache" in tags: + self.model_runner.post_kv_cache_wake_up() def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager: if not self.vllm_config.model_config.enable_sleep_mode: