[Spec Decode] Avoid redundant hidden-states gather in draft prefill (#46968)
Woosuk Kwon committed
Jun 28, 2026 at 17:04 UTC
0472436541c842ecda6d249411f1d35649291a79
1 file changed
+4
-1
vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
+4
-1
@@ -360,7 +360,10 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
360
self.current_draft_step,
361
self.draft_logits,
362
)
363
- self.hidden_states[:num_reqs] = hidden_states[last_token_indices]
363
+ if last_hidden_states is hidden_states:
364
+ self.hidden_states[:num_reqs] = sample_hidden_states
365
+ else:
366
+ self.hidden_states[:num_reqs] = hidden_states[last_token_indices]
367
self.input_buffers.positions[:num_reqs] = positions
368
369
def _multi_step_decode(