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 2, 06:00 UTC
[Model Runner V2] Support zeroing freshly allocated KV blocks for hybrid + fp8 KVCache (#43990)
Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com>
zhrrr committed
Jun 2, 2026 at 13:56 UTC
8a9eb40808dd1a78c322c93e4466e73842be5957
3 files changed
+45
-20
vllm/v1/worker/gpu/model_runner.py
+26
-3
@@ -47,6 +47,7 @@ from vllm.sequence import IntermediateTensors
47
from vllm.tasks import SupportedTask
48
from vllm.utils.math_utils import cdiv
49
from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
50
+from vllm.utils.platform_utils import is_pin_memory_available
51
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
52
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
53
from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
@@ -103,6 +104,7 @@ from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler
104
from vllm.v1.worker.gpu.states import RequestState
105
from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
106
from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
107
+from vllm.v1.worker.utils import KVBlockZeroer
108
109
logger = init_logger(__name__)
110
@@ -129,6 +131,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
131
self.cache_config.cache_dtype
132
]
133
134
+ # Lazily initialized in _init_kv_zero_meta() when the KV cache needs
135
+ # zeroing (e.g. hybrid models with fp8 KV cache).
136
+ self.kv_block_zeroer: KVBlockZeroer | None = None
137
+
138
self.vocab_size = self.model_config.get_vocab_size()
139
self.max_model_len = self.model_config.max_model_len
140
self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
@@ -393,7 +399,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
399
) + spec.num_speculative_blocks
400
max_num_blocks_per_group.append(max_num_blocks)
401
396
- self.attn_groups, attn_cg_support, kernel_block_sizes = init_attn_backend(
402
+ self.attn_groups, attn_cg_support, self.kernel_block_sizes = init_attn_backend(
403
self.kv_cache_config, self.vllm_config, self.device
404
)
405
self.block_tables = BlockTables(
@@ -402,7 +408,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
408
max_num_batched_tokens=self.max_num_tokens,
409
max_num_blocks_per_group=max_num_blocks_per_group,
410
device=self.device,
405
- kernel_block_sizes=kernel_block_sizes,
411
+ kernel_block_sizes=self.kernel_block_sizes,
412
cp_size=self.dcp_size,
413
cp_rank=self.dcp_rank,
414
cp_interleave=self.cp_interleave,
@@ -442,11 +448,22 @@ class GPUModelRunner(LoRAModelRunnerMixin):
448
self.attn_groups,
449
self.device,
450
self.cache_config.cache_dtype,
445
- kernel_block_sizes,
451
+ self.kernel_block_sizes,
452
self.vllm_config,
453
)
454
self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)
455
456
+ def _init_kv_zero_meta(self) -> None:
457
+ """Build KV-block zeroing metadata; invoked from gpu_worker."""
458
+ self.kv_block_zeroer = KVBlockZeroer(
459
+ self.device,
460
+ is_pin_memory_available(),
461
+ attn_groups_iter=(g for groups in self.attn_groups for g in groups),
462
+ kernel_block_sizes=self.kernel_block_sizes,
463
+ cache_dtype=self.cache_config.cache_dtype,
464
+ static_forward_context=self.compilation_config.static_forward_context,
465
+ )
466
+
467
@torch.inference_mode()
468
@step_eplb_after(is_dummy=True)
469
def _dummy_run(
@@ -753,6 +770,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
770
out=self.req_states.num_computed_prefill_tokens,
771
)
772
773
+ # Zero GPU memory for freshly allocated cache blocks to prevent
774
+ # stale NaN/data from corrupting attention or SSM computation.
775
+ if scheduler_output.new_block_ids_to_zero:
776
+ assert self.kv_block_zeroer is not None
777
+ self.kv_block_zeroer.zero_block_ids(scheduler_output.new_block_ids_to_zero)
778
+
779
def prepare_inputs(
780
self, scheduler_output: SchedulerOutput, batch_desc: BatchExecutionDescriptor
781
) -> InputBatch:
vllm/v1/worker/gpu_model_runner.py
+3
-3
@@ -1084,11 +1084,11 @@ class GPUModelRunner(
1084
def _init_kv_zero_meta(self) -> None:
1085
"""One-time precomputation for _zero_block_ids.
1086
1087
- Delegates to KVBlockZeroer.init_meta with the runner's state.
1087
Called from gpu_worker.py outside the CuMem pool context.
1088
"""
1090
- self._kv_block_zeroer = KVBlockZeroer(self.device, self.pin_memory)
1091
- self._kv_block_zeroer.init_meta(
1089
+ self._kv_block_zeroer = KVBlockZeroer(
1090
+ self.device,
1091
+ self.pin_memory,
1092
attn_groups_iter=self._kv_cache_spec_attn_group_iterator(),
1093
kernel_block_sizes=self._kernel_block_sizes,
1094
cache_dtype=self.cache_config.cache_dtype,
vllm/v1/worker/utils.py
+16
-14
@@ -80,30 +80,23 @@ def _zero_kv_blocks_kernel(
80
class KVBlockZeroer:
81
"""Manages efficient zeroing of KV cache blocks via a Triton kernel.
82
83
- Call :meth:`init_meta` once after KV caches are allocated to precompute
84
- segment addresses, then call :meth:`zero_block_ids` each step to zero
83
+ Construct once after KV caches are allocated to precompute segment
84
+ addresses, then call :meth:`zero_block_ids` each step to zero
85
newly-allocated blocks.
86
"""
87
88
- def __init__(self, device: torch.device, pin_memory: bool):
89
- self.device = device
90
- self.pin_memory = pin_memory
91
- self._meta: tuple[torch.Tensor, int, int, int] | None = None
92
- self._id_cap: int = 0
93
- self._ids_pinned: torch.Tensor | None = None
94
- self._ids_gpu: torch.Tensor | None = None
95
-
96
- def init_meta(
88
+ def __init__(
89
self,
90
+ device: torch.device,
91
+ pin_memory: bool,
92
attn_groups_iter: Iterable["AttentionGroup"],
93
kernel_block_sizes: list[int],
94
cache_dtype: str,
101
- runner_only_attn_layers: set[str],
95
static_forward_context: dict[str, Any],
96
+ runner_only_attn_layers: set[str] | None = None,
97
) -> None:
104
- """One-time precomputation for zero_block_ids.
98
+ """Precompute the absolute-address table for the Triton zeroing kernel.
99
106
- Builds absolute-address table for the Triton zeroing kernel.
100
Each entry is the absolute byte address of a segment start on the
101
GPU, so segments in different CUDA allocations work correctly.
102
@@ -114,6 +107,15 @@ class KVBlockZeroer:
107
108
Only AttentionSpec layers are processed; Mamba layers are skipped.
109
"""
110
+ self.device = device
111
+ self.pin_memory = pin_memory
112
+ self._meta: tuple[torch.Tensor, int, int, int] | None = None
113
+ self._id_cap: int = 0
114
+ self._ids_pinned: torch.Tensor | None = None
115
+ self._ids_gpu: torch.Tensor | None = None
116
+
117
+ if runner_only_attn_layers is None:
118
+ runner_only_attn_layers = set()
119
seen_ptrs: set[int] = set()
120
seg_addrs: list[int] = []
121
page_size_el: int | None = None