Add EAGLE-3 draft head#20149
Open
digantdesai wants to merge 1 commit into
Open
Conversation
Contributor
Author
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20149
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ⏳ No Failures, 3 PendingAs of commit f6c749f with merge base dc55469 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
This was referenced Jun 9, 2026
Contributor
There was a problem hiding this comment.
Pull request overview
This PR introduces an EAGLE-3 “draft head” module under examples/models/eagle3/ that can execute a single decoder-layer forward pass and load vLLM speculator-format safetensors checkpoints, along with unit tests covering forward execution, checkpoint loading (mono + sharded), vocab remapping tensors, RoPE precision behavior, and rejection of unsupported checkpoint variants.
Changes:
- Add
Eagle3Draft+ supporting modules/config, including a safetensors checkpoint loader for vLLM speculator checkpoints. - Add unit tests validating forward shapes, mapping behavior, checkpoint roundtrip (including sharded), and unsupported config variants.
- Register the new tests in
pytest.ini.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
pytest.ini |
Adds the new EAGLE-3 draft-head test file to pytest discovery. |
examples/models/eagle3/draft.py |
Implements the EAGLE-3 draft-head model and speculator safetensors checkpoint loading. |
examples/models/eagle3/test_draft.py |
Adds unit tests for forward pass, checkpoint loading, vocab mapping, and RoPE precision. |
examples/models/eagle3/__init__.py |
Declares the eagle3 examples package. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+174
to
+176
| def fuse(self, aux: torch.Tensor) -> torch.Tensor: | ||
| """Fuse concatenated target aux hidden states (B,T,3*D) -> feature (B,T,D).""" | ||
| return self.fc(aux) |
Comment on lines
+15
to
+17
| Draft ids map back to target ids with ``target_id = draft_id + d2t[draft_id]``. | ||
| Speculator checkpoints store the decoder layer under ``layers.0.*`` and may | ||
| include ``embed_tokens``, ``d2t``, and ``t2d``. |
Comment on lines
+254
to
+258
| # d2t/t2d are index/mask tensors (their checkpoint shape differs from the | ||
| # placeholder buffers); register them directly, load the rest strict. | ||
| model.register_buffer("d2t", state_dict.pop("d2t"), persistent=False) | ||
| model.register_buffer("t2d", state_dict.pop("t2d"), persistent=False) | ||
| model.load_state_dict(state_dict, strict=True, assign=True) |
Comment on lines
+156
to
+178
| @pytest.mark.parametrize("sharded", [False, True]) | ||
| def test_from_checkpoint_roundtrip(tmp_path, sharded): | ||
| cfg = tiny_config() | ||
| src, d2t, t2d = _write_checkpoint(str(tmp_path), cfg, sharded=sharded) | ||
|
|
||
| model, loaded_cfg = Eagle3Draft.from_checkpoint( | ||
| str(tmp_path), device="cpu", dtype=torch.float32 | ||
| ) | ||
| assert loaded_cfg.has_own_embed | ||
| assert loaded_cfg.aux_hidden_state_layers == cfg.aux_hidden_state_layers | ||
| assert loaded_cfg.target_vocab_size == cfg.target_vocab_size | ||
| torch.testing.assert_close( | ||
| model.midlayer.self_attn.q_proj.weight, src.midlayer.self_attn.q_proj.weight | ||
| ) | ||
| torch.testing.assert_close(model.fc.weight, src.fc.weight) | ||
| assert torch.equal(model.d2t, d2t) | ||
| assert torch.equal(model.t2d, t2d) | ||
| assert model.midlayer.self_attn.inv_freq.dtype == torch.float32 | ||
| T = 4 | ||
| feat = model.fuse(torch.randn(1, T, 3 * cfg.target_hidden_size)) | ||
| emb = model.embed(torch.randint(0, cfg.target_vocab_size, (T,))).unsqueeze(0) | ||
| logits, g = model(emb, feat, torch.arange(T)) | ||
| assert logits.shape == (1, T, cfg.draft_vocab_size) |
mergennachin
left a comment
mergennachin
left a comment
Contributor
There was a problem hiding this comment.
Codex is giving me this feedback
examples/models/eagle3/draft.py:227
This assumes every vLLM/speculators EAGLE-3 config stores RoPE as transformer_layer_config["rope_parameters"]["rope_theta"]. Real checkpoints often use top-level rope_theta instead, for example RedHatAI/Qwen3-8B-speculator.eagle3 and Llama-4 speculators. As written, from_checkpoint() raises KeyError before
loading weights. Please support both config formats, and either implement or explicitly reject unsupported rope_scaling.
examples/models/eagle3/draft.py:231
This requires eagle_aux_hidden_state_layer_ids, but some real EAGLE-3 speculator configs omit it or set it to null. Since the FC input size depends on the number of aux hidden states, the loader should either default to 3 like vLLM/speculators does, derive the count from fc.weight.shape, or fail with a
clear unsupported-checkpoint error instead of KeyError/TypeError.
examples/models/eagle3/draft.py:256
d2t/t2d are treated as mandatory, but full-vocabulary EAGLE-3 checkpoints can omit them by design when draft_vocab_size == target_vocab_size (for example the Llama-4 Maverick speculator). This currently raises KeyError: 'd2t'. Please load these tensors only when present and use identity mapping/no mapping
for equal-vocab checkpoints.
examples/models/eagle3/draft.py:72
This hardcodes default theta-only RoPE. Llama-3/Llama-4 style EAGLE-3 configs can include rope_scaling, so even after config parsing is fixed this will produce wrong positions for scaled-RoPE checkpoints. Please either implement the scaled RoPE variant or reject those configs explicitly during
from_checkpoint().
mergennachin
approved these changes
Jun 12, 2026
This was referenced Jun 17, 2026
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
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.
Adds the EAGLE-3 draft-head module, vLLM speculator checkpoint loading,
reduced-vocabulary id mapping, and tests for forward execution,
checkpoint loading, vocab remapping, RoPE precision, and unsupported
checkpoint variants.
Authored with assistance from Claude Code.