diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index fbc4f383cb35..387c9789360f 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; 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] ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index addc3b965f65..7bbed459e1f8 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -833,9 +833,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_worker.py b/vllm/v1/worker/gpu_worker.py index 7a86c3b68073..ef66808de7f8 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -256,7 +256,7 @@ 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: + if not self.use_v2_model_runner and (tags is None or "kv_cache" in tags): self.model_runner.post_kv_cache_wake_up() def checkpoint_prepare(self) -> None: @@ -688,7 +688,7 @@ def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None: ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config) if self.use_v2_model_runner: - self.model_runner.initialize_kv_cache( + self.model_runner.initialize_kv_cache( # type: ignore[call-arg] kv_cache_config, kv_cache_allocation_context=self._maybe_get_memory_pool_context( tag="kv_cache"