A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 08:30 UTC
[Spec Decode] fix returning size mismatch on extract hidden states proposer (#38610)
Signed-off-by: Jaebok Lee <jaebok9541@naver.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
zzaebok committed
Apr 10, 2026 at 04:39 UTC
edee96519a1b3485bacea2d14b0ffb6c83cae871
1 file changed
+4
-1
vllm/v1/spec_decode/extract_hidden_states.py
+4
-1
@@ -145,7 +145,10 @@ class ExtractHiddenStatesProposer:
145
146
# Return the sampled tokens as "draft" tokens
147
# Shape: [batch_size, 1] to match num_speculative_tokens=1
148
- return sampled_token_ids
148
+ # On decode steps with spec tokens, sampled_token_ids may have
149
+ # shape [batch_size, 2] (target + spec verification); slice to
150
+ # return only the target-sampled column.
151
+ return sampled_token_ids[:, :1]
152
153
def _get_slot_mapping(
154
self,