[Bugfix] Bound KV block zeroing launch geometry (#52058)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Lucas Wilkinson committed
Aug 12, 2026 at 23:42 UTC
79f3183f86b89c3bda05d467041bf3ef9ef60426
2 files changed
+86
-37
tests/v1/worker/test_kv_block_zeroer.py
+69
-8
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import math
5
from types import SimpleNamespace
6
7
import pytest
@@ -10,6 +11,7 @@ from vllm.v1.kv_cache_interface import (
11
ChunkedLocalAttentionSpec,
12
SlidingWindowSpec,
13
)
14
+from vllm.v1.worker import utils as worker_utils
15
from vllm.v1.worker.utils import (
16
AttentionGroup,
17
KVBlockZeroer,
@@ -121,10 +123,7 @@ def test_non_uniform_page_sizes():
123
seg_page_sizes = [page_size_a, page_size_b]
124
max_ps = max(seg_page_sizes)
125
124
- def largest_power_of_2_divisor(n):
125
- return n & -n
126
-
127
- blk_size = min(min(largest_power_of_2_divisor(ps) for ps in seg_page_sizes), 1024)
126
+ blk_size = min(1 << (max_ps - 1).bit_length(), 1024)
127
128
zeroer._meta = (
129
torch.tensor(
@@ -134,7 +133,7 @@ def test_non_uniform_page_sizes():
133
),
134
torch.tensor(seg_page_sizes, dtype=torch.int64, device=device),
135
torch.tensor(seg_page_sizes, dtype=torch.int64, device=device),
137
- max_ps // blk_size,
136
+ (max_ps + blk_size - 1) // blk_size,
137
blk_size,
138
2,
139
)
@@ -186,12 +185,74 @@ def test_packed_segment_zeros_only_its_last_block_page():
185
assert torch.equal(backing, expected)
186
187
188
+def test_large_dsv4_launch_geometry(monkeypatch):
189
+ """Keep the failing DSV4 shape efficient and within launch limits."""
190
+ device = torch.device("cpu")
191
+ n_blocks, n_segs = 6870, 181
192
+ layer_names = [f"layer.{i}" for i in range(n_segs)]
193
+ page_sizes = [9344 if i % 2 == 0 else 292 for i in range(n_segs)]
194
+ spec = SlidingWindowSpec(
195
+ block_size=1,
196
+ num_kv_heads=1,
197
+ head_size=1,
198
+ dtype=torch.int32,
199
+ sliding_window=1,
200
+ )
201
+ storages = {
202
+ name: torch.ones((1, page_size), dtype=torch.int32)
203
+ for name, page_size in zip(layer_names, page_sizes)
204
+ }
205
+ zeroer = KVBlockZeroer(
206
+ device,
207
+ attn_groups_iter=[
208
+ AttentionGroup(
209
+ _BlockFirstBackend, # type: ignore[arg-type]
210
+ [name],
211
+ spec,
212
+ group_id,
213
+ )
214
+ for group_id, name in enumerate(layer_names)
215
+ ],
216
+ kernel_block_sizes=[1] * n_segs,
217
+ cache_dtype="auto",
218
+ static_forward_context={
219
+ name: SimpleNamespace(kv_cache=storage)
220
+ for name, storage in storages.items()
221
+ },
222
+ )
223
+
224
+ assert zeroer._meta is not None
225
+ _, _, seg_page_sizes, max_chunks, blk_size, n_segs = zeroer._meta
226
+ assert seg_page_sizes.tolist() == page_sizes
227
+ assert (max_chunks, blk_size, n_segs) == (10, 1024, 181)
228
+
229
+ captured_grids = []
230
+
231
+ class FakeKernel:
232
+ def __getitem__(self, grid):
233
+ captured_grids.append(grid)
234
+ return lambda *args, **kwargs: None
235
+
236
+ monkeypatch.setattr(worker_utils, "_zero_kv_blocks_kernel", FakeKernel())
237
+ monkeypatch.setattr(
238
+ worker_utils,
239
+ "async_tensor_h2d",
240
+ lambda values, **kwargs: torch.tensor(values, dtype=torch.int64),
241
+ )
242
+
243
+ zeroer.zero_block_ids(list(range(n_blocks)))
244
+
245
+ old_max_chunks = max(page_sizes) // 4
246
+ assert math.prod((n_blocks, n_segs, old_max_chunks)) > 2**31 - 1
247
+ assert captured_grids == [(n_blocks, n_segs, max_chunks)]
248
+
249
+
250
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
190
-def test_warmup_compiles_every_n_blocks_specialization():
251
+def test_warmup_compiles_for_all_block_counts():
252
"""After warmup, no launch should trigger a first-request JIT compile.
253
193
- ``n_blocks`` is ``do_not_specialize``, so a single warmup launch must
194
- cover every block count.
254
+ The block count is carried by the launch grid, so changing it must reuse
255
+ the warmup's compiled kernel.
256
"""
257
device = torch.device("cuda")
258
num_blocks = 64
vllm/v1/worker/utils.py
+17
-29
@@ -17,7 +17,6 @@ from vllm.model_executor.models.interfaces import MultiModalEmbeddings
17
from vllm.model_executor.models.utils import extract_layer_index
18
from vllm.platforms import current_platform
19
from vllm.triton_utils import tl, triton
20
-from vllm.utils.math_utils import largest_power_of_2_divisor
20
from vllm.utils.mem_utils import MemorySnapshot, format_gib
21
from vllm.utils.torch_utils import async_tensor_h2d
22
from vllm.v1.attention.backend import (
@@ -52,15 +51,12 @@ def raise_if_nan_logits(num_nans_in_logits: Mapping[str, int]) -> None:
51
raise RuntimeError(f"NaNs detected in logits: {corrupted_requests}")
52
53
55
-@triton.jit(do_not_specialize=["n_blocks"])
54
+@triton.jit
55
def _zero_kv_blocks_kernel(
56
seg_addrs_ptr,
57
seg_block_strides_ptr,
58
seg_page_sizes_ptr,
59
block_ids_ptr,
61
- n_blocks,
62
- N_SEGS: tl.constexpr,
63
- MAX_CHUNKS: tl.constexpr,
60
BLOCK_SIZE: tl.constexpr,
61
):
62
"""Zero KV cache blocks across all segments in a single launch.
@@ -78,29 +74,27 @@ def _zero_kv_blocks_kernel(
74
seg_addrs_ptr holds absolute byte addresses (int64) for each segment,
75
allowing segments to live in different CUDA allocations.
76
81
- Programs are mapped as (block_index, seg_index, chunk_index).
77
+ Programs are mapped directly onto a 3-D grid as
78
+ (block_index, seg_index, chunk_index).
79
"""
83
- pid = tl.program_id(0)
84
- work_per_block = N_SEGS * MAX_CHUNKS
85
- block_index = pid // work_per_block
86
- if block_index >= n_blocks:
87
- return
88
- remainder = pid % work_per_block
89
- seg_index = remainder // MAX_CHUNKS
90
- chunk_index = remainder % MAX_CHUNKS
80
+ block_index = tl.program_id(0)
81
+ seg_index = tl.program_id(1)
82
+ chunk_index = tl.program_id(2)
83
block_stride_el = tl.load(seg_block_strides_ptr + seg_index)
84
page_size_el = tl.load(seg_page_sizes_ptr + seg_index)
93
- if chunk_index >= page_size_el // BLOCK_SIZE:
85
+ chunk_offset = chunk_index.to(tl.int64) * BLOCK_SIZE
86
+ if chunk_offset >= page_size_el:
87
return
88
block_id = tl.load(block_ids_ptr + block_index)
89
seg_addr = tl.load(seg_addrs_ptr + seg_index)
90
ptr = tl.cast(seg_addr, tl.pointer_type(tl.int32))
98
- offset = (
99
- block_id.to(tl.int64) * block_stride_el.to(tl.int64)
100
- + chunk_index.to(tl.int64) * BLOCK_SIZE
91
+ block_offset = block_id.to(tl.int64) * block_stride_el.to(tl.int64)
92
+ cols = chunk_offset + tl.arange(0, BLOCK_SIZE).to(tl.int64)
93
+ tl.store(
94
+ ptr + block_offset + cols,
95
+ tl.zeros([BLOCK_SIZE], dtype=tl.int32),
96
+ mask=cols < page_size_el,
97
)
102
- cols = tl.arange(0, BLOCK_SIZE).to(tl.int64)
103
- tl.store(ptr + offset + cols, tl.zeros([BLOCK_SIZE], dtype=tl.int32))
98
99
100
class KVBlockZeroer:
@@ -204,15 +198,12 @@ class KVBlockZeroer:
198
return
199
200
max_page_size_el = max(seg_page_sizes)
207
- blk_size = min(
208
- min(largest_power_of_2_divisor(ps) for ps in seg_page_sizes),
209
- 1024,
210
- )
201
+ blk_size = min(1 << (max_page_size_el - 1).bit_length(), 1024)
202
self._meta = (
203
torch.tensor(seg_addrs, dtype=torch.uint64, device=self.device),
204
torch.tensor(seg_block_strides, dtype=torch.int64, device=self.device),
205
torch.tensor(seg_page_sizes, dtype=torch.int64, device=self.device),
215
- max_page_size_el // blk_size,
206
+ (max_page_size_el + blk_size - 1) // blk_size,
207
blk_size,
208
len(seg_addrs),
209
)
@@ -231,15 +222,12 @@ class KVBlockZeroer:
222
) = self._meta
223
n_blocks = len(block_ids)
224
idx = async_tensor_h2d(block_ids, device=self.device, dtype=torch.int64)
234
- grid = (n_blocks * n_segs * max_chunks,)
225
+ grid = (n_blocks, n_segs, max_chunks)
226
_zero_kv_blocks_kernel[grid](
227
seg_addrs,
228
seg_block_strides,
229
seg_page_sizes,
230
idx,
240
- n_blocks,
241
- N_SEGS=n_segs,
242
- MAX_CHUNKS=max_chunks,
231
BLOCK_SIZE=blk_size,
232
)
233