[Bugfix][Kernel] Handle kernel block sizes in V2 DCP slot mapping (#51031)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Nicolò Lucchesi committed
Aug 24, 2026 at 09:28 UTC
0ecc284790e5403f74b899524ef82ecb69f83cb3
2 files changed
+73
-15
tests/v1/worker/test_gpu_block_table.py
+44
@@ -132,6 +132,50 @@ def test_block_tables_apply_staged_writes_single_group():
132
)
133
134
135
+@pytest.mark.parametrize("cp_rank", range(4))
136
+def test_dcp_slot_mapping_with_smaller_kernel_blocks(cp_rank: int):
137
+ """DCP interleave is expressed in logical-block token coordinates."""
138
+ device = torch.device("cuda")
139
+ block_tables = BlockTables(
140
+ block_sizes=[128],
141
+ max_num_reqs=1,
142
+ max_num_batched_tokens=1024,
143
+ max_num_blocks_per_group=[2],
144
+ device=device,
145
+ kernel_block_sizes=[64],
146
+ cp_size=4,
147
+ cp_rank=cp_rank,
148
+ cp_interleave=128,
149
+ )
150
+ block_tables.append_block_ids(
151
+ req_index=0,
152
+ new_block_ids=([5, 9],),
153
+ overwrite=True,
154
+ )
155
+ block_tables.apply_staged_writes()
156
+
157
+ idx_mapping = torch.zeros(1, dtype=torch.int32, device=device)
158
+ query_start_loc = torch.tensor([0, 1024], dtype=torch.int32, device=device)
159
+ positions = torch.arange(1024, dtype=torch.int64, device=device)
160
+ actual = block_tables.compute_slot_mappings(
161
+ idx_mapping,
162
+ query_start_loc,
163
+ positions,
164
+ num_tokens_padded=1024,
165
+ )[0]
166
+
167
+ expected = torch.full((1024,), -1, dtype=torch.int64, device=device)
168
+ first_start = cp_rank * 128
169
+ second_start = 512 + first_start
170
+ expected[first_start : first_start + 128] = torch.arange(
171
+ 5 * 128, 6 * 128, dtype=torch.int64, device=device
172
+ )
173
+ expected[second_start : second_start + 128] = torch.arange(
174
+ 9 * 128, 10 * 128, dtype=torch.int64, device=device
175
+ )
176
+ assert torch.equal(actual, expected)
177
+
178
+
179
def test_v1_block_table_move_row_clears_vacated_row():
180
"""condense() moves the last row into a freed slot; the vacated row must
181
not keep stale block ids. Padded dummy-run batches dereference stale rows
vllm/v1/worker/gpu/block_table.py
+29
-15
@@ -88,9 +88,9 @@ class BlockTables:
88
def init_block_table_layout_tensors(self) -> None:
89
# Called at init and after a CuMem kv_cache wake-up. The ptr tensors
90
# cache raw data_ptr() values that go stale once the underlying tensors
91
- # are reallocated on wake; block_sizes_tensor needs re-populating
92
- # because its storage lives under the kv_cache pool tag and comes back
93
- # with undefined contents.
91
+ # are reallocated on wake; the size tensors need re-populating because
92
+ # their storage lives under the kv_cache pool tag and comes back with
93
+ # undefined contents.
94
self.block_table_ptrs = self._make_ptr_tensor(
95
[b.gpu for b in self.block_tables]
96
)
@@ -100,6 +100,9 @@ class BlockTables:
100
device=self.device,
101
)
102
self.block_sizes_tensor = torch.tensor(
103
+ self.block_sizes, dtype=torch.int32, device=self.device
104
+ )
105
+ self.kernel_block_sizes_tensor = torch.tensor(
106
self.kernel_block_sizes, dtype=torch.int32, device=self.device
107
)
108
self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables)
@@ -203,6 +206,7 @@ class BlockTables:
206
self.block_table_ptrs,
207
self.block_table_strides,
208
self.block_sizes_tensor,
209
+ self.kernel_block_sizes_tensor,
210
slot_mappings,
211
slot_mappings.stride(0),
212
self.cp_rank,
@@ -274,6 +278,7 @@ def _compute_slot_mappings_kernel(
278
block_table_ptrs, # [num_kv_cache_groups]
279
block_table_strides, # [num_kv_cache_groups]
280
block_sizes, # [num_kv_cache_groups]
281
+ kernel_block_sizes, # [num_kv_cache_groups]
282
slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens]
283
slot_mappings_stride,
284
cp_rank,
@@ -300,7 +305,8 @@ def _compute_slot_mappings_kernel(
305
306
block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32)
307
block_table_stride = tl.load(block_table_strides + group_id)
303
- block_size = tl.load(block_sizes + group_id)
308
+ kv_block_size = tl.load(block_sizes + group_id)
309
+ kernel_block_size = tl.load(kernel_block_sizes + group_id)
310
311
req_state_idx = tl.load(idx_mapping + batch_idx)
312
start_idx = tl.load(query_start_loc + batch_idx)
@@ -309,22 +315,30 @@ def _compute_slot_mappings_kernel(
315
offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
316
positions = tl.load(pos + offset, mask=offset < end_idx, other=0)
317
312
- block_indices = positions // (block_size * CP_SIZE)
313
- block_offsets = positions % (block_size * CP_SIZE)
314
- block_numbers = tl.load(
315
- block_table_ptr + req_state_idx * block_table_stride + block_indices
316
- )
317
-
318
if CP_SIZE == 1:
319
# Common case: Context parallelism is not used.
320
- slot_ids = block_numbers * block_size + block_offsets
320
+ local_positions = positions
321
+ is_local = True
322
else:
323
# Context parallelism is used.
323
- is_local = block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank
324
- rounds = block_offsets // (CP_INTERLEAVE * CP_SIZE)
325
- remainder = block_offsets % CP_INTERLEAVE
324
+ virtual_block_size = kv_block_size * CP_SIZE
325
+ virtual_block_indices = positions // virtual_block_size
326
+ virtual_block_offsets = positions % virtual_block_size
327
+ is_local = virtual_block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank
328
+ rounds = virtual_block_offsets // (CP_INTERLEAVE * CP_SIZE)
329
+ remainder = virtual_block_offsets % CP_INTERLEAVE
330
local_offsets = rounds * CP_INTERLEAVE + remainder
327
- slot_ids = block_numbers * block_size + local_offsets
331
+ local_positions = virtual_block_indices * kv_block_size + local_offsets
332
+
333
+ block_indices = local_positions // kernel_block_size
334
+ block_offsets = local_positions % kernel_block_size
335
+ block_numbers = tl.load(
336
+ block_table_ptr + req_state_idx * block_table_stride + block_indices,
337
+ mask=is_local,
338
+ other=0,
339
+ )
340
+ slot_ids = block_numbers * kernel_block_size + block_offsets
341
+ if CP_SIZE != 1:
342
slot_ids = tl.where(is_local, slot_ids, PAD_ID)
343
344
tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx)