[Bugfix] Fix packed KV block zeroing stride (#50276)
Signed-off-by: wangxian001 <120719093+wangxian001@users.noreply.github.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
wangxian001 committed
Aug 7, 2026 at 04:49 UTC
d6af803f434397222674e8ea5cc6f25b3a208e62
2 files changed
+81
-17
tests/v1/worker/test_kv_block_zeroer.py
+39
@@ -21,6 +21,7 @@ def test_block_ids_are_not_overwritten_while_copy_is_in_flight():
21
zeroer._meta = (
22
torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device),
23
torch.tensor([page_size_el], dtype=torch.int64, device=device),
24
+ torch.tensor([page_size_el], dtype=torch.int64, device=device),
25
page_size_el // page_size_el, # max_chunks = 1
26
page_size_el, # blk_size
27
1, # n_segs
@@ -71,6 +72,7 @@ def test_non_uniform_page_sizes():
72
device=device,
73
),
74
torch.tensor(seg_page_sizes, dtype=torch.int64, device=device),
75
+ torch.tensor(seg_page_sizes, dtype=torch.int64, device=device),
76
max_ps // blk_size,
77
blk_size,
78
2,
@@ -88,6 +90,41 @@ def test_non_uniform_page_sizes():
90
assert torch.all(storage[3] == 1)
91
92
93
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
94
+def test_packed_segment_zeros_only_its_last_block_page():
95
+ """A packed KV segment steps by block stride but clears only its page."""
96
+ device = torch.device("cuda")
97
+ num_blocks = 4
98
+ block_stride_el = 12
99
+ page_size_el = 4
100
+ page_offset_el = 3
101
+ backing = torch.ones(
102
+ (num_blocks, block_stride_el), dtype=torch.int32, device=device
103
+ )
104
+
105
+ zeroer = KVBlockZeroer.__new__(KVBlockZeroer)
106
+ zeroer.device = device
107
+ zeroer._meta = (
108
+ torch.tensor(
109
+ [backing.data_ptr() + page_offset_el * backing.element_size()],
110
+ dtype=torch.uint64,
111
+ device=device,
112
+ ),
113
+ torch.tensor([block_stride_el], dtype=torch.int64, device=device),
114
+ torch.tensor([page_size_el], dtype=torch.int64, device=device),
115
+ 1,
116
+ page_size_el,
117
+ 1,
118
+ )
119
+
120
+ zeroer.zero_block_ids([num_blocks - 1])
121
+ torch.accelerator.synchronize()
122
+
123
+ expected = torch.ones_like(backing)
124
+ expected[-1, page_offset_el : page_offset_el + page_size_el] = 0
125
+ assert torch.equal(backing, expected)
126
+
127
+
128
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
129
def test_warmup_compiles_every_n_blocks_specialization():
130
"""After warmup, no launch should trigger a first-request JIT compile.
@@ -105,6 +142,7 @@ def test_warmup_compiles_every_n_blocks_specialization():
142
zeroer._meta = (
143
torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device),
144
torch.tensor([page_size_el], dtype=torch.int64, device=device),
145
+ torch.tensor([page_size_el], dtype=torch.int64, device=device),
146
1, # max_chunks
147
page_size_el, # blk_size
148
1, # n_segs
@@ -141,6 +179,7 @@ def test_warmup_respects_available_block_count():
179
zeroer._meta = (
180
torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device),
181
torch.tensor([page_size_el], dtype=torch.int64, device=device),
182
+ torch.tensor([page_size_el], dtype=torch.int64, device=device),
183
1,
184
page_size_el,
185
1,
vllm/v1/worker/utils.py
+42
-17
@@ -56,6 +56,7 @@ def raise_if_nan_logits(num_nans_in_logits: Mapping[str, int]) -> None:
56
@triton.jit(do_not_specialize=["n_blocks"])
57
def _zero_kv_blocks_kernel(
58
seg_addrs_ptr,
59
+ seg_block_strides_ptr,
60
seg_page_sizes_ptr,
61
block_ids_ptr,
62
n_blocks,
@@ -70,10 +71,10 @@ def _zero_kv_blocks_kernel(
71
buffer. For backends where K/V is outermost (block_dim=1) there are
72
two segments per buffer (one for K, one for V).
73
73
- Segments may have different page sizes (e.g. models with multiple KV
74
- cache groups like MLA + DSA indexer). Each segment's page size is
75
- read from seg_page_sizes_ptr; programs whose chunk_index falls beyond
76
- their segment's page size early-exit.
74
+ Segments may have different block strides and page sizes (e.g. packed
75
+ KV views or models with multiple KV cache groups like MLA + DSA
76
+ indexer). Each segment's block stride determines where a logical block
77
+ begins, while its page size determines how many elements are cleared.
78
79
seg_addrs_ptr holds absolute byte addresses (int64) for each segment,
80
allowing segments to live in different CUDA allocations.
@@ -88,6 +89,7 @@ def _zero_kv_blocks_kernel(
89
remainder = pid % work_per_block
90
seg_index = remainder // MAX_CHUNKS
91
chunk_index = remainder % MAX_CHUNKS
92
+ block_stride_el = tl.load(seg_block_strides_ptr + seg_index)
93
page_size_el = tl.load(seg_page_sizes_ptr + seg_index)
94
if chunk_index >= page_size_el // BLOCK_SIZE:
95
return
@@ -95,7 +97,7 @@ def _zero_kv_blocks_kernel(
97
seg_addr = tl.load(seg_addrs_ptr + seg_index)
98
ptr = tl.cast(seg_addr, tl.pointer_type(tl.int32))
99
offset = (
98
- block_id.to(tl.int64) * page_size_el.to(tl.int64)
100
+ block_id.to(tl.int64) * block_stride_el.to(tl.int64)
101
+ chunk_index.to(tl.int64) * BLOCK_SIZE
102
)
103
cols = tl.arange(0, BLOCK_SIZE).to(tl.int64)
@@ -126,18 +128,21 @@ class KVBlockZeroer:
128
129
Block IDs from the scheduler reference logical blocks whose size
130
may differ from the kernel block size (virtual block splitting).
129
- Each segment's page_size_el accounts for this ratio so that
130
- ``block_id * page_size_el`` lands at the correct offset.
131
+ Each virtual block is represented as an independent segment so its
132
+ physical block stride and zeroed page span remain independent.
133
134
Only AttentionSpec layers are processed; Mamba layers are skipped.
135
"""
136
self.device = device
135
- self._meta: tuple[torch.Tensor, torch.Tensor, int, int, int] | None = None
137
+ self._meta: (
138
+ tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int] | None
139
+ ) = None
140
141
if runner_only_attn_layers is None:
142
runner_only_attn_layers = set()
143
seen_ptrs: set[int] = set()
144
seg_addrs: list[int] = []
145
+ seg_block_strides: list[int] = []
146
seg_page_sizes: list[int] = []
147
148
for group in attn_groups_iter:
@@ -147,6 +152,7 @@ class KVBlockZeroer:
152
if group.kv_cache_group_id >= len(kernel_block_sizes):
153
continue
154
kernel_bs = kernel_block_sizes[group.kv_cache_group_id]
155
+ assert spec.block_size % kernel_bs == 0
156
ratio = spec.block_size // kernel_bs
157
block_dim = group.backend.get_kv_cache_block_dim(
158
kernel_bs,
@@ -167,22 +173,32 @@ class KVBlockZeroer:
173
seen_ptrs.add(dp)
174
175
el = kv.element_size()
170
- cur_bytes = kv.stride(block_dim) * el
171
- assert cur_bytes % 4 == 0
172
- kernel_block_el = cur_bytes // 4
173
- cur_page_el = kernel_block_el * ratio
174
-
175
- block_stride_bytes = cur_bytes
176
+ block_stride_bytes = kv.stride(block_dim) * el
177
+ assert block_stride_bytes % 4 == 0
178
+ assert kv.shape[block_dim] % ratio == 0
179
outer_dims = [
180
d
181
for d in range(block_dim)
182
if kv.stride(d) * el > block_stride_bytes
183
]
184
outer_strides = [kv.stride(d) * el for d in outer_dims]
185
+ inner_dims = [
186
+ d for d in range(kv.ndim) if d != block_dim and d not in outer_dims
187
+ ]
188
+ kernel_page_bytes = el + sum(
189
+ (kv.shape[d] - 1) * kv.stride(d) * el for d in inner_dims
190
+ )
191
+ assert kernel_page_bytes % 4 == 0
192
+ logical_block_stride_bytes = block_stride_bytes * ratio
193
for outer in iprod(*(range(kv.shape[d]) for d in outer_dims)):
194
off_bytes = sum(i * s for i, s in zip(outer, outer_strides))
184
- seg_addrs.append(dp + off_bytes)
185
- seg_page_sizes.append(cur_page_el)
195
+ assert (dp + off_bytes) % 4 == 0
196
+ for virtual_index in range(ratio):
197
+ seg_addrs.append(
198
+ dp + off_bytes + virtual_index * block_stride_bytes
199
+ )
200
+ seg_block_strides.append(logical_block_stride_bytes // 4)
201
+ seg_page_sizes.append(kernel_page_bytes // 4)
202
203
if not seg_addrs:
204
self._meta = None
@@ -195,6 +211,7 @@ class KVBlockZeroer:
211
)
212
self._meta = (
213
torch.tensor(seg_addrs, dtype=torch.uint64, device=self.device),
214
+ torch.tensor(seg_block_strides, dtype=torch.int64, device=self.device),
215
torch.tensor(seg_page_sizes, dtype=torch.int64, device=self.device),
216
max_page_size_el // blk_size,
217
blk_size,
@@ -205,12 +222,20 @@ class KVBlockZeroer:
222
"""Zero the KV cache memory for the given block IDs."""
223
if not block_ids or self._meta is None:
224
return
208
- seg_addrs, seg_page_sizes, max_chunks, blk_size, n_segs = self._meta
225
+ (
226
+ seg_addrs,
227
+ seg_block_strides,
228
+ seg_page_sizes,
229
+ max_chunks,
230
+ blk_size,
231
+ n_segs,
232
+ ) = self._meta
233
n_blocks = len(block_ids)
234
idx = async_tensor_h2d(block_ids, device=self.device, dtype=torch.int64)
235
grid = (n_blocks * n_segs * max_chunks,)
236
_zero_kv_blocks_kernel[grid](
237
seg_addrs,
238
+ seg_block_strides,
239
seg_page_sizes,
240
idx,
241
n_blocks,