[SpecDecode] Skip embedding sharing when target/draft hidden dims differ - #44631
Closed
chaojun-zhang wants to merge 1 commit into
Closed
chaojun-zhang wants to merge 1 commit into
chaojun-zhang wants to merge 1 commit into
Conversation
chaojun-zhang
requested review from
MatthewBonanni,
benchislett and
luccafong
as code owners
June 5, 2026 08:04
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Contributor
|
Hi @chaojun-zhang, the pre-commit checks have failed. Please run: uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, Tip Is
|
_maybe_share_embeddings blindly replaces the draft model's embed_tokens with the target's when the draft checkpoint has no own embedding weights. This breaks EAGLE3 variants paired with a target of different hidden size (e.g. EAGLE3-LLaMA3.1-8B draft + MiniMax-M2 target): the first decoder layer does cat([embed(tokens), hidden_states]) expecting (N,4096)+(N,4096), but receives (N,3072)+(N,4096) after the bad share. Add a dimension check before sharing: if target embedding output dim != draft hidden size, skip sharing and log a warning. The draft model falls back to random-initialized embeddings, which is the correct behavior when dimensions are incompatible. Fixes Eagle3MiniMaxM2ForCausalLM initialization test. Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
chaojun-zhang
force-pushed
the
fix/eagle3-embedding-sharing
branch
from
June 5, 2026 08:40
7008908 to
fcb68e9
Compare
Contributor
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fix
_maybe_share_embeddingsinSpecDecodeBaseProposer: when a draft checkpoint has no ownembed_tokens(e.g.yuhuili/EAGLE3-LLaMA3.1-Instruct-8B), the target embedding was blindly shared to the draft model even if their hidden dimensions differ. With an EAGLE3-LLaMA3.1-8B draft (hidden_size=4096) paired with a MiniMax-M2 target (hidden_size=3072), the first decoder layer doescat([embed(tokens), hidden_states])expecting dim 8192 but receives 7168, triggering atorch._assertfailure duringtorch.compiletracing and crashing engine initialization.Add a dimension guard before sharing; skip and log a warning when dims differ.
Test
Without this fix the test fails with
RuntimeError: Engine core initialization failed(root cause: shape mismatch incatduring compile). With this fix the test passes.