[Model Runner V2][Spec Decode] Fix stale values in idx_mapping from CG num reqs padding (#47066)
Giancarlo Delfin committed
Jun 30, 2026 at 13:25 UTC
3cecee40f34cec995b4705b7ec928e8ce3327f97
3 files changed
+15
-6
vllm/v1/worker/gpu/sample/gumbel.py
+6
-3
@@ -100,7 +100,10 @@ def gumbel_block_argmax(
100
PER_TOKEN_COL: tl.constexpr = False,
101
):
102
req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx).to(tl.int64)
103
- temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
103
+ is_valid_req = req_state_idx >= 0
104
+ temp = tl.load(temp_ptr + req_state_idx, mask=is_valid_req, other=0.0).to(
105
+ tl.float32
106
+ )
107
if temp != 0.0 and APPLY_TEMPERATURE:
108
# Apply temperature.
109
# NOTE(woosuk): Match the behavior of _temperature_kernel.
@@ -122,7 +125,7 @@ def gumbel_block_argmax(
125
+ col * vocab_size
126
+ block,
127
logits,
125
- mask=mask,
128
+ mask=mask & is_valid_req,
129
)
130
131
# fp32 is the default reduction dtype; fp64 is ~1/32–1/64x the throughput
@@ -131,7 +134,7 @@ def gumbel_block_argmax(
134
logits = logits.to(tl.float64)
135
if temp != 0.0:
136
# Calculate the seed for gumbel noise.
134
- seed = tl.load(seeds_ptr + req_state_idx)
137
+ seed = tl.load(seeds_ptr + req_state_idx, mask=is_valid_req, other=0)
138
pos = tl.load(pos_ptr + token_idx)
139
gumbel_seed = tl.randint(seed, pos)
140
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+5
-3
@@ -366,7 +366,7 @@ class DFlashSpeculator(DraftModelSpeculator):
366
self.query_cudagraph_manager.run_fullgraph(batch_desc)
367
else:
368
self._generate_draft(
369
- num_reqs_padded,
369
+ num_reqs,
370
num_tokens_padded,
371
draft_attn_metadata,
372
draft_slot_mappings_by_layer,
@@ -496,7 +496,9 @@ def _prepare_dflash_inputs_kernel(
496
mask = block < max_num_reqs
497
tl.store(out_seq_lens_ptr + block, 0, mask=mask)
498
# Padded sample slots point at query index 0 (a valid row in
499
- # last_hidden_states) so CG replay never reads OOB.
499
+ # last_hidden_states) so CG replay never reads OOB. Padded
500
+ # sample idx mappings point to -1, which is ignored during
501
+ # sampling to prevent writing stale values to draft logits.
502
pad_start = num_reqs * num_speculative_steps
503
pad_end = max_num_reqs * num_speculative_steps
504
for i in range(pad_start, pad_end, BLOCK_SIZE):
@@ -504,7 +506,7 @@ def _prepare_dflash_inputs_kernel(
506
mask = block < pad_end
507
tl.store(out_sample_indices_ptr + block, 0, mask=mask)
508
tl.store(out_sample_pos_ptr + block, 0, mask=mask)
507
- tl.store(out_sample_idx_mapping_ptr + block, 0, mask=mask)
509
+ tl.store(out_sample_idx_mapping_ptr + block, -1, mask=mask)
510
# Pad query slot mappings past num_query_tokens with PAD so the
511
# captured CG sees PAD slots (no K/V write) for replay sizes
512
# larger than the current request count.
vllm/v1/worker/gpu/spec_decode/speculator.py
+4
@@ -304,3 +304,7 @@ class DraftModelSpeculator(BaseSpeculator):
304
self.temperature.copy_(temperature)
305
self.seeds.copy_(seeds)
306
self.idx_mapping[:num_reqs].copy_(idx_mapping)
307
+ if self.draft_logits is not None:
308
+ # idx_mapping for CG padded requests points to -1, which is ignored
309
+ # during sampling to prevent writing stale values to draft logits.
310
+ self.idx_mapping[num_reqs:].fill_(-1)