From 3cecee40f34cec995b4705b7ec928e8ce3327f97 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Tue, 30 Jun 2026 13:25:32 -0500 Subject: [PATCH] [Model Runner V2][Spec Decode] Fix stale values in idx_mapping from CG num reqs padding (#47066) --- vllm/v1/worker/gpu/sample/gumbel.py | 9 ++++++--- vllm/v1/worker/gpu/spec_decode/dflash/speculator.py | 8 +++++--- vllm/v1/worker/gpu/spec_decode/speculator.py | 4 ++++ 3 files changed, 15 insertions(+), 6 deletions(-) diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index 4b0a1694f7..74ee8abb2c 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -100,7 +100,10 @@ def gumbel_block_argmax( PER_TOKEN_COL: tl.constexpr = False, ): req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx).to(tl.int64) - temp = tl.load(temp_ptr + req_state_idx).to(tl.float32) + is_valid_req = req_state_idx >= 0 + temp = tl.load(temp_ptr + req_state_idx, mask=is_valid_req, other=0.0).to( + tl.float32 + ) if temp != 0.0 and APPLY_TEMPERATURE: # Apply temperature. # NOTE(woosuk): Match the behavior of _temperature_kernel. @@ -122,7 +125,7 @@ def gumbel_block_argmax( + col * vocab_size + block, logits, - mask=mask, + mask=mask & is_valid_req, ) # fp32 is the default reduction dtype; fp64 is ~1/32–1/64x the throughput @@ -131,7 +134,7 @@ def gumbel_block_argmax( logits = logits.to(tl.float64) if temp != 0.0: # Calculate the seed for gumbel noise. - seed = tl.load(seeds_ptr + req_state_idx) + seed = tl.load(seeds_ptr + req_state_idx, mask=is_valid_req, other=0) pos = tl.load(pos_ptr + token_idx) gumbel_seed = tl.randint(seed, pos) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index e458396749..7978fc5133 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -366,7 +366,7 @@ class DFlashSpeculator(DraftModelSpeculator): self.query_cudagraph_manager.run_fullgraph(batch_desc) else: self._generate_draft( - num_reqs_padded, + num_reqs, num_tokens_padded, draft_attn_metadata, draft_slot_mappings_by_layer, @@ -496,7 +496,9 @@ def _prepare_dflash_inputs_kernel( mask = block < max_num_reqs tl.store(out_seq_lens_ptr + block, 0, mask=mask) # Padded sample slots point at query index 0 (a valid row in - # last_hidden_states) so CG replay never reads OOB. + # last_hidden_states) so CG replay never reads OOB. Padded + # sample idx mappings point to -1, which is ignored during + # sampling to prevent writing stale values to draft logits. pad_start = num_reqs * num_speculative_steps pad_end = max_num_reqs * num_speculative_steps for i in range(pad_start, pad_end, BLOCK_SIZE): @@ -504,7 +506,7 @@ def _prepare_dflash_inputs_kernel( mask = block < pad_end tl.store(out_sample_indices_ptr + block, 0, mask=mask) tl.store(out_sample_pos_ptr + block, 0, mask=mask) - tl.store(out_sample_idx_mapping_ptr + block, 0, mask=mask) + tl.store(out_sample_idx_mapping_ptr + block, -1, mask=mask) # Pad query slot mappings past num_query_tokens with PAD so the # captured CG sees PAD slots (no K/V write) for replay sizes # larger than the current request count. diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 341ed715c7..9e36c7cc65 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -304,3 +304,7 @@ class DraftModelSpeculator(BaseSpeculator): self.temperature.copy_(temperature) self.seeds.copy_(seeds) self.idx_mapping[:num_reqs].copy_(idx_mapping) + if self.draft_logits is not None: + # 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)