[Bugfix] Support non-uniform page sizes in KVBlockZeroer (#49704)
Signed-off-by: Elvir Crncevic <elvircrn@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Elvir Crnčević committed
Jul 24, 2026 at 22:38 UTC
972848f2764e2f79deb7853320ccf0094be452aa
3 files changed
+92
-31
tests/v1/worker/test_kv_block_zeroer.py
+50
-3
@@ -20,9 +20,10 @@ def test_block_ids_are_not_overwritten_while_copy_is_in_flight():
20
zeroer.device = device
21
zeroer._meta = (
22
torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device),
23
- page_size_el,
24
- page_size_el,
25
- 1,
23
+ torch.tensor([page_size_el], dtype=torch.int64, device=device),
24
+ page_size_el // page_size_el, # max_chunks = 1
25
+ page_size_el, # blk_size
26
+ 1, # n_segs
27
)
28
29
stream = torch.cuda.Stream()
@@ -39,3 +40,49 @@ def test_block_ids_are_not_overwritten_while_copy_is_in_flight():
40
assert torch.all(storage[1] == 0)
41
assert torch.all(storage[2] == 0)
42
assert torch.all(storage[3] == 1)
43
+
44
+
45
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
46
+def test_non_uniform_page_sizes():
47
+ """Two segments with different page sizes (e.g. MLA + DSA indexer)."""
48
+ device = torch.device("cuda")
49
+ num_blocks = 4
50
+ page_size_a = 10496 # int32 elements
51
+ page_size_b = 2112
52
+
53
+ storage_a = torch.ones((num_blocks, page_size_a), dtype=torch.int32, device=device)
54
+ storage_b = torch.ones((num_blocks, page_size_b), dtype=torch.int32, device=device)
55
+
56
+ zeroer = KVBlockZeroer.__new__(KVBlockZeroer)
57
+ zeroer.device = device
58
+
59
+ seg_page_sizes = [page_size_a, page_size_b]
60
+ max_ps = max(seg_page_sizes)
61
+
62
+ def largest_power_of_2_divisor(n):
63
+ return n & -n
64
+
65
+ blk_size = min(min(largest_power_of_2_divisor(ps) for ps in seg_page_sizes), 1024)
66
+
67
+ zeroer._meta = (
68
+ torch.tensor(
69
+ [storage_a.data_ptr(), storage_b.data_ptr()],
70
+ dtype=torch.uint64,
71
+ device=device,
72
+ ),
73
+ torch.tensor(seg_page_sizes, dtype=torch.int64, device=device),
74
+ max_ps // blk_size,
75
+ blk_size,
76
+ 2,
77
+ )
78
+
79
+ stream = torch.cuda.Stream()
80
+ with torch.cuda.stream(stream):
81
+ zeroer.zero_block_ids([1, 2])
82
+ stream.synchronize()
83
+
84
+ for storage in (storage_a, storage_b):
85
+ assert torch.all(storage[0] == 1)
86
+ assert torch.all(storage[1] == 0)
87
+ assert torch.all(storage[2] == 0)
88
+ assert torch.all(storage[3] == 1)
vllm/model_executor/warmup/qwen_triton_warmup.py
+10
-6
@@ -36,7 +36,8 @@ _FLA_POST_CONV_WARMUP_LENGTHS = (1, 2, 16)
36
37
@dataclass(frozen=True)
38
class _ZeroKvWarmupConfig:
39
- page_size_el: int
39
+ seg_page_sizes: torch.Tensor
40
+ max_chunks: int
41
block_size: int
42
n_segs: int
43
@@ -160,9 +161,10 @@ def _zero_kv_warmup_config(runner: object) -> _ZeroKvWarmupConfig | None:
161
if meta is None:
162
return None
163
163
- _, page_size_el, block_size, n_segs = meta
164
+ _, seg_page_sizes, max_chunks, block_size, n_segs = meta
165
return _ZeroKvWarmupConfig(
165
- page_size_el=int(page_size_el),
166
+ seg_page_sizes=seg_page_sizes,
167
+ max_chunks=int(max_chunks),
168
block_size=int(block_size),
169
n_segs=int(n_segs),
170
)
@@ -185,8 +187,9 @@ def _warm_zero_kv_blocks_kernel(
187
from vllm.v1.worker.utils import _zero_kv_blocks_kernel
188
189
max_n_blocks = max(_ZERO_KV_N_BLOCKS)
190
+ max_page_size = int(config.seg_page_sizes.max().item())
191
scratch = torch.empty(
189
- max_n_blocks * config.page_size_el,
192
+ max_n_blocks * max_page_size,
193
dtype=torch.int32,
194
device=device,
195
)
@@ -198,13 +201,14 @@ def _warm_zero_kv_blocks_kernel(
201
202
for n_blocks in _ZERO_KV_N_BLOCKS:
203
block_ids = torch.arange(n_blocks, dtype=torch.int64, device=device)
201
- grid = (n_blocks * config.n_segs * (config.page_size_el // config.block_size),)
204
+ grid = (n_blocks * config.n_segs * config.max_chunks,)
205
_zero_kv_blocks_kernel[grid](
206
seg_addrs,
207
+ config.seg_page_sizes,
208
block_ids,
209
n_blocks,
210
N_SEGS=config.n_segs,
207
- PAGE_SIZE_EL=config.page_size_el,
211
+ MAX_CHUNKS=config.max_chunks,
212
BLOCK_SIZE=config.block_size,
213
)
214
vllm/v1/worker/utils.py
+32
-22
@@ -43,10 +43,11 @@ logger = init_logger(__name__)
43
@triton.jit
44
def _zero_kv_blocks_kernel(
45
seg_addrs_ptr,
46
+ seg_page_sizes_ptr,
47
block_ids_ptr,
48
n_blocks,
49
N_SEGS: tl.constexpr,
49
- PAGE_SIZE_EL: tl.constexpr,
50
+ MAX_CHUNKS: tl.constexpr,
51
BLOCK_SIZE: tl.constexpr,
52
):
53
"""Zero KV cache blocks across all segments in a single launch.
@@ -56,25 +57,33 @@ def _zero_kv_blocks_kernel(
57
buffer. For backends where K/V is outermost (block_dim=1) there are
58
two segments per buffer (one for K, one for V).
59
60
+ Segments may have different page sizes (e.g. models with multiple KV
61
+ cache groups like MLA + DSA indexer). Each segment's page size is
62
+ read from seg_page_sizes_ptr; programs whose chunk_index falls beyond
63
+ their segment's page size early-exit.
64
+
65
seg_addrs_ptr holds absolute byte addresses (int64) for each segment,
66
allowing segments to live in different CUDA allocations.
67
68
Programs are mapped as (block_index, seg_index, chunk_index).
69
"""
70
pid = tl.program_id(0)
65
- chunks = PAGE_SIZE_EL // BLOCK_SIZE
66
- work_per_block = N_SEGS * chunks
71
+ work_per_block = N_SEGS * MAX_CHUNKS
72
block_index = pid // work_per_block
73
if block_index >= n_blocks:
74
return
75
remainder = pid % work_per_block
71
- seg_index = remainder // chunks
72
- chunk_index = remainder % chunks
76
+ seg_index = remainder // MAX_CHUNKS
77
+ chunk_index = remainder % MAX_CHUNKS
78
+ page_size_el = tl.load(seg_page_sizes_ptr + seg_index)
79
+ if chunk_index >= page_size_el // BLOCK_SIZE:
80
+ return
81
block_id = tl.load(block_ids_ptr + block_index)
82
seg_addr = tl.load(seg_addrs_ptr + seg_index)
83
ptr = tl.cast(seg_addr, tl.pointer_type(tl.int32))
84
offset = (
77
- block_id.to(tl.int64) * PAGE_SIZE_EL + chunk_index.to(tl.int64) * BLOCK_SIZE
85
+ block_id.to(tl.int64) * page_size_el.to(tl.int64)
86
+ + chunk_index.to(tl.int64) * BLOCK_SIZE
87
)
88
cols = tl.arange(0, BLOCK_SIZE).to(tl.int64)
89
tl.store(ptr + offset + cols, tl.zeros([BLOCK_SIZE], dtype=tl.int32))
@@ -104,19 +113,19 @@ class KVBlockZeroer:
113
114
Block IDs from the scheduler reference logical blocks whose size
115
may differ from the kernel block size (virtual block splitting).
107
- PAGE_SIZE_EL accounts for this ratio so that
108
- ``block_id * PAGE_SIZE_EL`` lands at the correct offset.
116
+ Each segment's page_size_el accounts for this ratio so that
117
+ ``block_id * page_size_el`` lands at the correct offset.
118
119
Only AttentionSpec layers are processed; Mamba layers are skipped.
120
"""
121
self.device = device
113
- self._meta: tuple[torch.Tensor, int, int, int] | None = None
122
+ self._meta: tuple[torch.Tensor, torch.Tensor, int, int, int] | None = None
123
124
if runner_only_attn_layers is None:
125
runner_only_attn_layers = set()
126
seen_ptrs: set[int] = set()
127
seg_addrs: list[int] = []
119
- page_size_el: int | None = None
128
+ seg_page_sizes: list[int] = []
129
130
for group in attn_groups_iter:
131
spec = group.kv_cache_spec
@@ -149,12 +158,6 @@ class KVBlockZeroer:
158
assert cur_bytes % 4 == 0
159
kernel_block_el = cur_bytes // 4
160
cur_page_el = kernel_block_el * ratio
152
- if page_size_el is None:
153
- page_size_el = cur_page_el
154
- else:
155
- assert page_size_el == cur_page_el, (
156
- f"Non-uniform page sizes: {page_size_el} vs {cur_page_el}"
157
- )
161
162
block_stride_bytes = cur_bytes
163
outer_dims = [
@@ -166,15 +169,21 @@ class KVBlockZeroer:
169
for outer in iprod(*(range(kv.shape[d]) for d in outer_dims)):
170
off_bytes = sum(i * s for i, s in zip(outer, outer_strides))
171
seg_addrs.append(dp + off_bytes)
172
+ seg_page_sizes.append(cur_page_el)
173
170
- if not seg_addrs or page_size_el is None:
174
+ if not seg_addrs:
175
self._meta = None
176
return
177
174
- blk_size = min(largest_power_of_2_divisor(page_size_el), 1024)
178
+ max_page_size_el = max(seg_page_sizes)
179
+ blk_size = min(
180
+ min(largest_power_of_2_divisor(ps) for ps in seg_page_sizes),
181
+ 1024,
182
+ )
183
self._meta = (
184
torch.tensor(seg_addrs, dtype=torch.uint64, device=self.device),
177
- page_size_el,
185
+ torch.tensor(seg_page_sizes, dtype=torch.int64, device=self.device),
186
+ max_page_size_el // blk_size,
187
blk_size,
188
len(seg_addrs),
189
)
@@ -183,16 +192,17 @@ class KVBlockZeroer:
192
"""Zero the KV cache memory for the given block IDs."""
193
if not block_ids or self._meta is None:
194
return
186
- seg_addrs, page_size_el, blk_size, n_segs = self._meta
195
+ seg_addrs, seg_page_sizes, max_chunks, blk_size, n_segs = self._meta
196
n_blocks = len(block_ids)
197
idx = async_tensor_h2d(block_ids, device=self.device, dtype=torch.int64)
189
- grid = (n_blocks * n_segs * (page_size_el // blk_size),)
198
+ grid = (n_blocks * n_segs * max_chunks,)
199
_zero_kv_blocks_kernel[grid](
200
seg_addrs,
201
+ seg_page_sizes,
202
idx,
203
n_blocks,
204
N_SEGS=n_segs,
195
- PAGE_SIZE_EL=page_size_el,
205
+ MAX_CHUNKS=max_chunks,
206
BLOCK_SIZE=blk_size,
207
)
208