[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>
Evan Takahashi committed
Jul 21, 2026 at 10:53 UTC
b2b8f679d0589f0c956f3e734cc70dab07b27b8a
1 file changed
+6
-2
vllm/v1/spec_decode/llm_base_proposer.py
+6
-2
@@ -1477,9 +1477,13 @@ class SpecDecodeBaseProposer:
1477
"Sharing target model embedding weights with the draft model."
1478
)
1479
1480
- if share_embeddings:
1480
+ if share_embeddings and hasattr(self.model, "has_own_embed_tokens"):
1481
+ # EAGLE drafts consume input embeddings at their own hidden
1482
+ # size, so only share when the widths match. MTP drafts
1483
+ # project target-width embeddings (e.g. Gemma4 MTP's
1484
+ # pre_projection takes 2 * backbone_hidden_size), so the
1485
+ # width check does not apply to them.
1486
draft_embed = self.model.model.embed_tokens
1482
- # Only share when both models use the same embedding width.
1487
# Guard with isinstance so non-Tensor weights (e.g. in tests)
1488
# are not affected — mirrors the weight-equality check above.
1489
if isinstance(target_embed_tokens.weight, torch.Tensor) and isinstance(