A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 00:15 UTC
[Model Runner V2] Fix `_compute_slot_mappings_kernel` for chunked prefill (#36580)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Mar 10, 2026 at 00:23 UTC
9efc3bdcd6749f6d0ba26b12aee27cc8829c6f93
1 file changed
+5
-4
vllm/v1/worker/gpu/block_table.py
+5
-4
@@ -138,10 +138,8 @@ class BlockTables:
138
num_tokens_padded: int,
139
) -> torch.Tensor:
140
num_reqs = idx_mapping.shape[0]
141
- num_tokens = positions.shape[0]
141
num_groups = self.num_kv_cache_groups
142
_compute_slot_mappings_kernel[(num_groups, num_reqs + 1)](
144
- num_tokens,
143
self.max_num_batched_tokens,
144
idx_mapping,
145
query_start_loc,
@@ -213,7 +211,6 @@ def _gather_block_tables_kernel(
211
212
@triton.jit
213
def _compute_slot_mappings_kernel(
216
- num_tokens,
214
max_num_tokens,
215
idx_mapping, # [num_reqs]
216
query_start_loc, # [num_reqs + 1]
@@ -236,7 +233,11 @@ def _compute_slot_mappings_kernel(
233
234
if batch_idx == tl.num_programs(1) - 1:
235
# Pad remaining slots to -1. This is needed for CUDA graphs.
239
- for i in range(num_tokens, max_num_tokens, TRITON_BLOCK_SIZE):
236
+ # Start from actual token count (not padded) to cover the gap
237
+ # between actual tokens and padded tokens that can contain stale
238
+ # valid slot IDs from previous chunks during chunked prefill.
239
+ actual_num_tokens = tl.load(query_start_loc + batch_idx)
240
+ for i in range(actual_num_tokens, max_num_tokens, TRITON_BLOCK_SIZE):
241
offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
242
tl.store(slot_mapping_ptr + offset, PAD_ID, mask=offset < max_num_tokens)
243
return