diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index e77156eef9..5c245d1c4d 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -5,6 +5,7 @@ import pytest import torch from vllm.platforms import current_platform +from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.worker.gpu.block_table import BlockTables pytestmark = pytest.mark.skipif( @@ -273,3 +274,41 @@ def test_get_dummy_block_tables_returns_zeroed_rows(): assert (dummy[0] == 0).all() # CUDA graph invariant: same persistent tensor, not a fresh allocation. assert dummy[0].data_ptr() == block_tables.input_block_tables[0].data_ptr() + + +def test_dummy_request_slot_mapping_is_pad(): + """idx_mapping == -1 marks a dummy (or CUDA-graph padding) request. + + A dummy draft decode step must not resolve slots through the persistent + block-table row of a real request slot, which may be stale and point at + blocks that are now in the prefix cache. + """ + device = torch.device("cuda") + block_tables = BlockTables( + block_sizes=[4], + max_num_reqs=2, + max_num_batched_tokens=16, + max_num_blocks_per_group=[4], + device=device, + kernel_block_sizes=[4], + ) + block_tables.append_block_ids(req_index=0, new_block_ids=([7, 8],), overwrite=True) + block_tables.apply_staged_writes() + query_start_loc = torch.tensor([0, 3], dtype=torch.int32, device=device) + positions = torch.tensor([1, 2, 3], dtype=torch.int64, device=device) + + real = block_tables.compute_slot_mappings( + torch.tensor([0], dtype=torch.int32, device=device), + query_start_loc, + positions, + num_tokens_padded=3, + ) + assert real[0].tolist() == [7 * 4 + 1, 7 * 4 + 2, 7 * 4 + 3] + + dummy = block_tables.compute_slot_mappings( + torch.tensor([-1], dtype=torch.int32, device=device), + query_start_loc, + positions, + num_tokens_padded=3, + ) + assert dummy[0].tolist() == [PAD_SLOT_ID] * 3 diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index c01895e635..6563e30c34 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -315,6 +315,9 @@ def _compute_slot_mappings_kernel( mapping_enabled = tl.load(slot_mapping_enabled + group_id) req_state_idx = tl.load(idx_mapping + batch_idx) + # idx_mapping == -1 marks a dummy (or CUDA-graph padding) request that owns + # no blocks: never read its block-table row and emit PAD for its tokens. + is_real_req = req_state_idx >= 0 start_idx = tl.load(query_start_loc + batch_idx) end_idx = tl.load(query_start_loc + batch_idx + 1) for i in range(start_idx, end_idx, TRITON_BLOCK_SIZE): @@ -342,12 +345,12 @@ def _compute_slot_mappings_kernel( block_offsets = local_positions % kernel_block_size block_numbers = tl.load( block_table_ptr + req_state_idx * block_table_stride + block_indices, - mask=is_local, + mask=is_local & is_real_req, other=0, ) slot_ids = block_numbers * kernel_block_size + block_offsets if CP_SIZE != 1: slot_ids = tl.where(is_local, slot_ids, PAD_ID) - slot_ids = tl.where(mapping_enabled, slot_ids, PAD_ID) + slot_ids = tl.where(mapping_enabled & is_real_req, slot_ids, PAD_ID) tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx) diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index f166d63e8e..c787efa07f 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -254,6 +254,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): input_batch.idx_mapping, temperature, seeds, + dummy_run=dummy_run, ) # Get the input ids and last token indices for the speculator. diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py index 1c6ae315cf..fc0924e3a3 100644 --- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py @@ -166,6 +166,7 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator): input_batch.idx_mapping, temperature, seeds, + dummy_run=dummy_run, ) num_tokens = input_batch.num_tokens diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 299cac6416..534b6ef8eb 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -481,6 +481,7 @@ class DraftModelSpeculator(BaseSpeculator): temperature: torch.Tensor, # [max_num_reqs] seeds: torch.Tensor, + dummy_run: bool = False, ) -> None: # Copy temperature, seeds, and idx mapping to the pre-allocated buffers. # NOTE(woosuk): For draft sampling, we only consider the temperature @@ -490,10 +491,15 @@ class DraftModelSpeculator(BaseSpeculator): # affect the output distribution after rejection sampling. self.temperature.copy_(temperature) self.seeds.copy_(seeds) - self.idx_mapping[:num_reqs].copy_(idx_mapping) - # idx_mapping for CG padded requests points to -1, which is ignored - # during sampling to prevent writing stale values to draft logits. - self.idx_mapping[num_reqs:].fill_(-1) + # idx_mapping == -1 marks a row the drafter must not act on: sampling + # skips it and compute_slot_mappings emits PAD. CUDA-graph padded rows + # always get it; a dummy batch gets it for every row, since its arange + # idx_mapping names request-state slots with stale block tables. + if dummy_run: + self.idx_mapping.fill_(-1) + else: + self.idx_mapping[:num_reqs].copy_(idx_mapping) + self.idx_mapping[num_reqs:].fill_(-1) def _build_uniform_batch_dp_sync( self,