[Bugfix][Spec Decode] Stop dummy draft decode steps from writing KV through stale block-table rows (#56734)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Yifan Qiao committed
Sep 20, 2026 at 18:06 UTC
d2983f2f163249acbe3053c073cb9830290984a0
5 files changed
+56
-6
tests/v1/worker/test_gpu_block_table.py
+39
@@ -5,6 +5,7 @@ import pytest
5
import torch
6
7
from vllm.platforms import current_platform
8
+from vllm.v1.attention.backends.utils import PAD_SLOT_ID
9
from vllm.v1.worker.gpu.block_table import BlockTables
10
11
pytestmark = pytest.mark.skipif(
@@ -273,3 +274,41 @@ def test_get_dummy_block_tables_returns_zeroed_rows():
274
assert (dummy[0] == 0).all()
275
# CUDA graph invariant: same persistent tensor, not a fresh allocation.
276
assert dummy[0].data_ptr() == block_tables.input_block_tables[0].data_ptr()
277
+
278
+
279
+def test_dummy_request_slot_mapping_is_pad():
280
+ """idx_mapping == -1 marks a dummy (or CUDA-graph padding) request.
281
+
282
+ A dummy draft decode step must not resolve slots through the persistent
283
+ block-table row of a real request slot, which may be stale and point at
284
+ blocks that are now in the prefix cache.
285
+ """
286
+ device = torch.device("cuda")
287
+ block_tables = BlockTables(
288
+ block_sizes=[4],
289
+ max_num_reqs=2,
290
+ max_num_batched_tokens=16,
291
+ max_num_blocks_per_group=[4],
292
+ device=device,
293
+ kernel_block_sizes=[4],
294
+ )
295
+ block_tables.append_block_ids(req_index=0, new_block_ids=([7, 8],), overwrite=True)
296
+ block_tables.apply_staged_writes()
297
+ query_start_loc = torch.tensor([0, 3], dtype=torch.int32, device=device)
298
+ positions = torch.tensor([1, 2, 3], dtype=torch.int64, device=device)
299
+
300
+ real = block_tables.compute_slot_mappings(
301
+ torch.tensor([0], dtype=torch.int32, device=device),
302
+ query_start_loc,
303
+ positions,
304
+ num_tokens_padded=3,
305
+ )
306
+ assert real[0].tolist() == [7 * 4 + 1, 7 * 4 + 2, 7 * 4 + 3]
307
+
308
+ dummy = block_tables.compute_slot_mappings(
309
+ torch.tensor([-1], dtype=torch.int32, device=device),
310
+ query_start_loc,
311
+ positions,
312
+ num_tokens_padded=3,
313
+ )
314
+ assert dummy[0].tolist() == [PAD_SLOT_ID] * 3
vllm/v1/worker/gpu/block_table.py
+5
-2
@@ -315,6 +315,9 @@ def _compute_slot_mappings_kernel(
315
mapping_enabled = tl.load(slot_mapping_enabled + group_id)
316
317
req_state_idx = tl.load(idx_mapping + batch_idx)
318
+ # idx_mapping == -1 marks a dummy (or CUDA-graph padding) request that owns
319
+ # no blocks: never read its block-table row and emit PAD for its tokens.
320
+ is_real_req = req_state_idx >= 0
321
start_idx = tl.load(query_start_loc + batch_idx)
322
end_idx = tl.load(query_start_loc + batch_idx + 1)
323
for i in range(start_idx, end_idx, TRITON_BLOCK_SIZE):
@@ -342,12 +345,12 @@ def _compute_slot_mappings_kernel(
345
block_offsets = local_positions % kernel_block_size
346
block_numbers = tl.load(
347
block_table_ptr + req_state_idx * block_table_stride + block_indices,
345
- mask=is_local,
348
+ mask=is_local & is_real_req,
349
other=0,
350
)
351
slot_ids = block_numbers * kernel_block_size + block_offsets
352
if CP_SIZE != 1:
353
slot_ids = tl.where(is_local, slot_ids, PAD_ID)
354
352
- slot_ids = tl.where(mapping_enabled, slot_ids, PAD_ID)
355
+ slot_ids = tl.where(mapping_enabled & is_real_req, slot_ids, PAD_ID)
356
tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx)
vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
+1
@@ -254,6 +254,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
254
input_batch.idx_mapping,
255
temperature,
256
seeds,
257
+ dummy_run=dummy_run,
258
)
259
260
# Get the input ids and last token indices for the speculator.
vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py
+1
@@ -166,6 +166,7 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):
166
input_batch.idx_mapping,
167
temperature,
168
seeds,
169
+ dummy_run=dummy_run,
170
)
171
172
num_tokens = input_batch.num_tokens
vllm/v1/worker/gpu/spec_decode/speculator.py
+10
-4
@@ -481,6 +481,7 @@ class DraftModelSpeculator(BaseSpeculator):
481
temperature: torch.Tensor,
482
# [max_num_reqs]
483
seeds: torch.Tensor,
484
+ dummy_run: bool = False,
485
) -> None:
486
# Copy temperature, seeds, and idx mapping to the pre-allocated buffers.
487
# NOTE(woosuk): For draft sampling, we only consider the temperature
@@ -490,10 +491,15 @@ class DraftModelSpeculator(BaseSpeculator):
491
# affect the output distribution after rejection sampling.
492
self.temperature.copy_(temperature)
493
self.seeds.copy_(seeds)
493
- self.idx_mapping[:num_reqs].copy_(idx_mapping)
494
- # idx_mapping for CG padded requests points to -1, which is ignored
495
- # during sampling to prevent writing stale values to draft logits.
496
- self.idx_mapping[num_reqs:].fill_(-1)
494
+ # idx_mapping == -1 marks a row the drafter must not act on: sampling
495
+ # skips it and compute_slot_mappings emits PAD. CUDA-graph padded rows
496
+ # always get it; a dummy batch gets it for every row, since its arange
497
+ # idx_mapping names request-state slots with stale block tables.
498
+ if dummy_run:
499
+ self.idx_mapping.fill_(-1)
500
+ else:
501
+ self.idx_mapping[:num_reqs].copy_(idx_mapping)
502
+ self.idx_mapping[num_reqs:].fill_(-1)
503
504
def _build_uniform_batch_dp_sync(
505
self,