[Core] Simplify KVBlockZeroer index tensor handling (#48399)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Nick Hill committed
Jul 23, 2026 at 11:12 UTC
521aa80f719bb11bf973d2d51873ca966afe373d
4 files changed
+6
-55
tests/v1/worker/test_kv_block_zeroer.py
+4
-7
@@ -14,14 +14,10 @@ def test_block_ids_are_not_overwritten_while_copy_is_in_flight():
14
page_size_el = 4
15
storage = torch.ones((num_blocks, page_size_el), dtype=torch.int32, device=device)
16
17
- # Build the minimal zeroer state directly so the test can focus on ID-buffer
18
- # lifetime without constructing model attention groups.
17
+ # Build the minimal zeroer state directly so the test can focus on the
18
+ # in-flight copy behavior without constructing model attention groups.
19
zeroer = KVBlockZeroer.__new__(KVBlockZeroer)
20
zeroer.device = device
21
- zeroer.pin_memory = True
22
- zeroer.max_concurrency = 2
23
- zeroer._id_cap = 8
24
- zeroer._allocate_id_buffers()
21
zeroer._meta = (
22
torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device),
23
page_size_el,
@@ -32,7 +28,8 @@ def test_block_ids_are_not_overwritten_while_copy_is_in_flight():
28
stream = torch.cuda.Stream()
29
with torch.cuda.stream(stream):
30
# Keep the first nonblocking H2D copy pending while the host submits the
35
- # second call. A single shared pinned source would be overwritten here.
31
+ # second call. Each call must stage from its own pinned source so the
32
+ # first copy is not corrupted before it runs.
33
torch.cuda._sleep(10_000_000)
34
zeroer.zero_block_ids([1])
35
zeroer.zero_block_ids([2])
vllm/v1/worker/gpu/model_runner.py
+1
-3
@@ -51,7 +51,7 @@ from vllm.sequence import IntermediateTensors
51
from vllm.tasks import SupportedTask
52
from vllm.utils.math_utils import cdiv
53
from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
54
-from vllm.utils.torch_utils import PIN_MEMORY, STR_DTYPE_TO_TORCH_DTYPE
54
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
55
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
56
from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
57
from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
@@ -521,12 +521,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
521
"""Build KV-block zeroing metadata; invoked from gpu_worker."""
522
self.kv_block_zeroer = KVBlockZeroer(
523
self.device,
524
- pin_memory=PIN_MEMORY,
524
attn_groups_iter=(g for groups in self.attn_groups for g in groups),
525
kernel_block_sizes=self.kernel_block_sizes,
526
cache_dtype=self.cache_config.cache_dtype,
527
static_forward_context=self.compilation_config.static_forward_context,
529
- max_concurrency=self.vllm_config.max_concurrent_batches,
528
)
529
530
@torch.inference_mode()
vllm/v1/worker/gpu_model_runner.py
-2
@@ -1134,13 +1134,11 @@ class GPUModelRunner(
1134
"""
1135
self._kv_block_zeroer = KVBlockZeroer(
1136
self.device,
1137
- pin_memory=PIN_MEMORY,
1137
attn_groups_iter=self._kv_cache_spec_attn_group_iterator(),
1138
kernel_block_sizes=self._kernel_block_sizes,
1139
cache_dtype=self.cache_config.cache_dtype,
1140
runner_only_attn_layers=self.runner_only_attn_layers,
1141
static_forward_context=self.compilation_config.static_forward_context,
1143
- max_concurrency=self.vllm_config.max_concurrent_batches,
1142
)
1143
1144
def _zero_block_ids(self, block_ids: list[int]) -> None:
vllm/v1/worker/utils.py
+1
-43
@@ -91,13 +91,11 @@ class KVBlockZeroer:
91
def __init__(
92
self,
93
device: torch.device,
94
- pin_memory: bool,
94
attn_groups_iter: Iterable["AttentionGroup"],
95
kernel_block_sizes: list[int],
96
cache_dtype: str,
97
static_forward_context: dict[str, Any],
98
runner_only_attn_layers: set[str] | None = None,
100
- max_concurrency: int = 1,
99
) -> None:
100
"""Precompute the absolute-address table for the Triton zeroing kernel.
101
@@ -112,15 +110,7 @@ class KVBlockZeroer:
110
Only AttentionSpec layers are processed; Mamba layers are skipped.
111
"""
112
self.device = device
115
- self.pin_memory = pin_memory
116
- if max_concurrency < 1:
117
- raise ValueError("max_concurrency must be at least 1")
118
- self.max_concurrency = max_concurrency
113
self._meta: tuple[torch.Tensor, int, int, int] | None = None
120
- self._id_cap: int = 0
121
- self._ids_pinned: list[torch.Tensor] = []
122
- self._ids_gpu: list[torch.Tensor] = []
123
- self._id_buffer_index = 0
114
115
if runner_only_attn_layers is None:
116
runner_only_attn_layers = set()
@@ -182,8 +172,6 @@ class KVBlockZeroer:
172
return
173
174
blk_size = min(largest_power_of_2_divisor(page_size_el), 1024)
185
- self._id_cap = 8192
186
- self._allocate_id_buffers()
175
self._meta = (
176
torch.tensor(seg_addrs, dtype=torch.uint64, device=self.device),
177
page_size_el,
@@ -191,43 +179,13 @@ class KVBlockZeroer:
179
len(seg_addrs),
180
)
181
194
- def _allocate_id_buffers(self) -> None:
195
- self._ids_pinned = [
196
- torch.empty(
197
- self._id_cap,
198
- dtype=torch.int64,
199
- pin_memory=self.pin_memory,
200
- )
201
- for _ in range(self.max_concurrency)
202
- ]
203
- self._ids_gpu = [
204
- torch.empty(self._id_cap, dtype=torch.int64, device=self.device)
205
- for _ in range(self.max_concurrency)
206
- ]
207
- self._id_buffer_index = 0
208
-
182
def zero_block_ids(self, block_ids: list[int]) -> None:
183
"""Zero the KV cache memory for the given block IDs."""
184
if not block_ids or self._meta is None:
185
return
186
seg_addrs, page_size_el, blk_size, n_segs = self._meta
187
n_blocks = len(block_ids)
215
- if n_blocks > self._id_cap:
216
- # The old pinned buffers may still be the source of an in-flight
217
- # nonblocking copy. Growing is rare, so we don't mind the sync overhead
218
- torch.accelerator.synchronize()
219
- self._id_cap = n_blocks * 2
220
- self._allocate_id_buffers()
221
-
222
- # The H2D copy is nonblocking, so its pinned source must not be mutated
223
- # while this batch is in flight. Rotate through as many buffers as concurrent
224
- # in-flight batches, to avoid collisions.
225
- buffer_index = self._id_buffer_index
226
- self._id_buffer_index = (buffer_index + 1) % self.max_concurrency
227
- ids_pinned = self._ids_pinned[buffer_index]
228
- ids_pinned[:n_blocks].numpy()[:] = block_ids
229
- idx = self._ids_gpu[buffer_index][:n_blocks]
230
- idx.copy_(ids_pinned[:n_blocks], non_blocking=True)
188
+ idx = async_tensor_h2d(block_ids, device=self.device, dtype=torch.int64)
189
grid = (n_blocks * n_segs * (page_size_el // blk_size),)
190
_zero_kv_blocks_kernel[grid](
191
seg_addrs,