"docs/vscode:/vscode.git/clone" did not exist on "b7df53cd42f3eab007b4f287c151960858e949df"
Unverified Commit edee9651 authored by zzaebok's avatar zzaebok Committed by GitHub
Browse files

[Spec Decode] fix returning size mismatch on extract hidden states proposer (#38610)


Signed-off-by: default avatarJaebok Lee <jaebok9541@naver.com>
Co-authored-by: default avatarmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent adaabb8a
...@@ -145,7 +145,10 @@ class ExtractHiddenStatesProposer: ...@@ -145,7 +145,10 @@ class ExtractHiddenStatesProposer:
# Return the sampled tokens as "draft" tokens # Return the sampled tokens as "draft" tokens
# Shape: [batch_size, 1] to match num_speculative_tokens=1 # Shape: [batch_size, 1] to match num_speculative_tokens=1
return sampled_token_ids # On decode steps with spec tokens, sampled_token_ids may have
# shape [batch_size, 2] (target + spec verification); slice to
# return only the target-sampled column.
return sampled_token_ids[:, :1]
def _get_slot_mapping( def _get_slot_mapping(
self, self,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment