Support llama3 eagle3 head with llama4 verifier #117
Conversation
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
rahul-tuli
force-pushed
the
support-llama3-eagle3-head-with-llama4-verifier
branch
10 times, most recently
from
cf02c8d to
1695608
Compare
rahul-tuli
force-pushed
the
support-llama3-eagle3-head-with-llama4-verifier
branch
3 times, most recently
from
1f6fd40 to
5e93541
Compare
Support configuring eagle_aux_hidden_state_layer_ids and inference_type in the Eagle3 speculator configuration. This allows users to specify which verifier layers should output auxiliary hidden states for the drafter to consume during speculative decoding. Signed-off-by: rahul-tuli <[email protected]> Signed-off-by: Rahul Tuli <[email protected]>
Add documentation explaining that get_eagle3_aux_hidden_state_layers() provides default layer selection and that the GPU model runner can override this with values from speculative config for dynamic configuration. Signed-off-by: rahul-tuli <[email protected]> Signed-off-by: Rahul Tuli <[email protected]>
Add Eagle3 support to Llama4ForConditionalGeneration by implementing set_aux_hidden_state_layers() and get_eagle3_aux_hidden_state_layers() methods. Both methods delegate to the underlying Llama4ForCausalLM language model, enabling Eagle3 speculative decoding with Llama4 multimodal verifier models. This allows text-only Eagle3 drafters to work with Llama4 multimodal verifiers by consuming auxiliary hidden states from specified layers. Signed-off-by: rahul-tuli <[email protected]> Signed-off-by: Rahul Tuli <[email protected]>
Implement custom get_input_embeddings() in Eagle3LlamaForCausalLM that accepts multimodal parameters but only processes text embeddings. This ensures the Llama3-based Eagle3 drafter correctly handles text inputs while remaining compatible with multimodal verifier interfaces. The drafter receives multimodal context through auxiliary hidden states from the verifier rather than processing multimodal inputs directly. Signed-off-by: rahul-tuli <[email protected]> Signed-off-by: Rahul Tuli <[email protected]>
Implement _get_eagle3_aux_layers_from_config() helper method to extract auxiliary layer IDs from the draft model's speculative config. The GPU model runner now prefers config-specified layers over model defaults, with fallback to model's get_eagle3_aux_hidden_state_layers() when not configured. Changes: - Refactor auxiliary layer setup with early return pattern for errors - Add config extraction with proper error handling - Log only when using non-default layer configuration - Enable dynamic layer configuration per deployment Signed-off-by: rahul-tuli <[email protected]> Signed-off-by: Rahul Tuli <[email protected]>
Signed-off-by: Rahul Tuli <[email protected]>
Signed-off-by: Rahul Tuli <[email protected]>
rahul-tuli
force-pushed
the
support-llama3-eagle3-head-with-llama4-verifier
branch
from
cac1941 to
1037b36
Compare
mgoin
force-pushed
the
main
branch
2 times, most recently
from
0340f45 to
2f7dbc9
Compare
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.
Summary
This PR enables Eagle3 speculative decoding with Llama3 drafter and Llama4 multimodal verifier support, with configurable auxiliary hidden state layers.
Key Features
eagle_aux_hidden_state_layer_idsin the speculator config, allowing non-default layer selection for optimal performance across different model architecturesConfiguration
Auxiliary layer indices can be set in the Eagle3 draft model config:
{ "eagle_aux_hidden_state_layer_ids": [1, 23, 44] }This enables using hidden states from non-default layers (e.g., layers 1, 23, 44 instead of default 2, 23, 44) for cross-architecture scenarios where different layer combinations may work better.
Testing
Command:
Results: