[Bugfix][Spec Decode] Restrict embedding-width share guard to EAGLE drafts (#47953)

Signed-off-by: Evan Takahashi <evan10takahashi@gmail.com>
Co-authored-by: Michael Goin <mgoin64@gmail.com>
This commit is contained in:
Evan Takahashi
2026-07-21 10:53:37 -07:00
committed by GitHub
co-authored by Michael Goin
parent de6ec294ef
commit b2b8f679d0
+6 -2
View File
@@ -1477,9 +1477,13 @@ class SpecDecodeBaseProposer:
"Sharing target model embedding weights with the draft model."
)
if share_embeddings:
if share_embeddings and hasattr(self.model, "has_own_embed_tokens"):
# EAGLE drafts consume input embeddings at their own hidden
# size, so only share when the widths match. MTP drafts
# project target-width embeddings (e.g. Gemma4 MTP's
# pre_projection takes 2 * backbone_hidden_size), so the
# width check does not apply to them.
draft_embed = self.model.model.embed_tokens
# Only share when both models use the same embedding width.
# Guard with isinstance so non-Tensor weights (e.g. in tests)
# are not affected — mirrors the weight-equality check above.
if isinstance(target_embed_tokens.weight, torch.Tensor) and isinstance(