A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 12:30 UTC
[Model Runner v2] Fix block table IMA issue (#40648)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Wentao Ye committed
Apr 29, 2026 at 11:30 UTC
51fda1ba44ff3fd08e9202ce4f404cf3a1feaec1
4 files changed
+30
-22
vllm/v1/worker/gpu/block_table.py
+21
-12
@@ -47,18 +47,7 @@ class BlockTables:
47
device=device,
48
)
49
self.block_tables.append(block_table)
50
- self.block_table_ptrs = self._make_ptr_tensor(
51
- [b.gpu for b in self.block_tables]
52
- )
53
- self.block_table_strides = torch.tensor(
54
- [b.gpu.stride(0) for b in self.block_tables],
55
- dtype=torch.int64,
56
- device=self.device,
57
- )
50
59
- self.block_sizes_tensor = torch.tensor(
60
- self.block_sizes, dtype=torch.int32, device=self.device
61
- )
51
self.num_blocks = UvaBackedTensor(
52
(self.num_kv_cache_groups, self.max_num_reqs),
53
dtype=torch.int32,
@@ -69,7 +58,6 @@ class BlockTables:
58
self.input_block_tables: list[torch.Tensor] = [
59
torch.zeros_like(b.gpu) for b in self.block_tables
60
]
72
- self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables)
61
62
self.slot_mappings = torch.zeros(
63
self.num_kv_cache_groups,
@@ -78,12 +66,33 @@ class BlockTables:
66
device=self.device,
67
)
68
69
+ self.init_block_table_layout_tensors()
70
+
71
def _make_ptr_tensor(self, x: Iterable[torch.Tensor]) -> torch.Tensor:
72
# NOTE(woosuk): Use uint64 instead of int64 to cover all possible addresses.
73
return torch.tensor(
74
[t.data_ptr() for t in x], dtype=torch.uint64, device=self.device
75
)
76
77
+ def init_block_table_layout_tensors(self) -> None:
78
+ # Called at init and after a CuMem kv_cache wake-up. The ptr tensors
79
+ # cache raw data_ptr() values that go stale once the underlying tensors
80
+ # are reallocated on wake; block_sizes_tensor needs re-populating
81
+ # because its storage lives under the kv_cache pool tag and comes back
82
+ # with undefined contents.
83
+ self.block_table_ptrs = self._make_ptr_tensor(
84
+ [b.gpu for b in self.block_tables]
85
+ )
86
+ self.block_table_strides = torch.tensor(
87
+ [b.gpu.stride(0) for b in self.block_tables],
88
+ dtype=torch.int64,
89
+ device=self.device,
90
+ )
91
+ self.block_sizes_tensor = torch.tensor(
92
+ self.block_sizes, dtype=torch.int32, device=self.device
93
+ )
94
+ self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables)
95
+
96
def append_block_ids(
97
self,
98
req_index: int,
vllm/v1/worker/gpu/model_runner.py
+3
@@ -558,6 +558,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
558
del hidden_states, sample_hidden_states
559
gc.collect()
560
561
+ def post_kv_cache_wake_up(self) -> None:
562
+ self.block_tables.init_block_table_layout_tensors()
563
+
564
def reset_mm_cache(self) -> None:
565
if self.encoder_cache is not None:
566
self.encoder_cache.reset_mm_cache()
vllm/v1/worker/gpu_model_runner.py
+3
@@ -884,6 +884,9 @@ class GPUModelRunner(
884
self.encoder_cache.clear()
885
self.late_interaction_runner.clear()
886
887
+ def post_kv_cache_wake_up(self) -> None:
888
+ self.init_fp8_kv_scales()
889
+
890
@torch.inference_mode()
891
def init_fp8_kv_scales(self) -> None:
892
"""
vllm/v1/worker/gpu_worker.py
+3
-10
@@ -46,7 +46,7 @@ from vllm.tasks import SupportedTask
46
from vllm.tracing import instrument
47
from vllm.utils.mem_constants import GiB_bytes
48
from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling
49
-from vllm.utils.torch_utils import is_quantized_kv_cache, set_random_seed
49
+from vllm.utils.torch_utils import set_random_seed
50
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
51
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
52
from vllm.v1.outputs import (
@@ -192,15 +192,8 @@ class Worker(WorkerBase):
192
buffer.data.copy_(self._sleep_saved_buffers[name].data)
193
self._sleep_saved_buffers = {}
194
195
- # If the KV cache has just been woken up,
196
- # the internal state of cache_engine must be reset,
197
- # especially the FP8 scaling factor.
198
- if (
199
- (tags is None or "kv_cache" in tags)
200
- and is_quantized_kv_cache(self.cache_config.cache_dtype)
201
- and hasattr(self.model_runner, "init_fp8_kv_scales")
202
- ):
203
- self.model_runner.init_fp8_kv_scales()
195
+ if tags is None or "kv_cache" in tags:
196
+ self.model_runner.post_kv_cache_wake_up()
197
198
def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:
199
if not self.vllm_config.model_config.enable_sleep_mode: