[MRV2] Keep block table metadata across KV wake-up

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]
         )

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()

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"