[Model Runner V2][Spec Decode] Fix stale values in idx_mapping from CG num reqs padding (#47066)

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)
 

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.

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)