Use get_input_embeddings

Signed-off-by: Rahul Tuli <[email protected]>

Author: rahul-tuli

vllm/model_executor/models/llama_eagle3.py

@@ -248,8 +248,7 @@ def get_input_embeddings(
         multimodal_embeddings: Optional[NestedTensors] = None,
         is_multimodal: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        # The llama3 drafter only processes text embeddings
-        return self.model.embed_tokens(input_ids)
+        return self.model.get_input_embeddings(input_ids)
 
     def forward(
         self,

vllm/model_executor/models/mllama4.py

@@ -64,8 +64,12 @@
 from vllm.sequence import IntermediateTensors
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
 
-from .interfaces import (MultiModalEmbeddings, SupportsEagle3,
-                         SupportsMultiModal, SupportsPP)
+from .interfaces import (
+    MultiModalEmbeddings,
+    SupportsEagle3,
+    SupportsMultiModal,
+    SupportsPP,
+)
 from .llama4 import Llama4ForCausalLM
 from .utils import AutoWeightsLoader, flatten_bn, maybe_prefix
 from .vision import run_dp_sharded_vision_model
@@ -773,7 +777,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
     def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
         """Set which layers should output auxiliary hidden states for EAGLE3."""
         # Delegate to underlying language model (Llama4ForCausalLM)
-        assert hasattr(self.language_model, 'set_aux_hidden_state_layers')
+        assert hasattr(self.language_model, "set_aux_hidden_state_layers")
         self.language_model.set_aux_hidden_state_layers(layers)
 
     def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
@@ -783,9 +787,7 @@ def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
         the speculative config if available, providing dynamic configuration.
         """
         # Delegate to underlying language model (Llama4ForCausalLM)
-        assert hasattr(
-            self.language_model, "get_eagle3_aux_hidden_state_layers"
-        )
+        assert hasattr(self.language_model, "get_eagle3_aux_hidden_state_layers")
         return self.language_model.get_eagle3_aux_hidden_state_layers()
 
     def _parse_and_validate_image_input(

vllm/transformers_utils/configs/speculators/algos.py

@@ -34,4 +34,5 @@ def update_eagle3(config_dict: dict, vllm_config: dict) -> None:
     vllm_config["architectures"] = ["Eagle3LlamaForCausalLM"]
     if config_dict.get("eagle_aux_hidden_state_layer_ids"):
         vllm_config["eagle_aux_hidden_state_layer_ids"] = config_dict[
-            "eagle_aux_hidden_state_layer_ids"]
+            "eagle_aux_hidden_state_layer_ids"
+        ]

vllm/v1/worker/gpu_model_runner.py

@@ -3026,12 +3026,11 @@ def _get_eagle3_aux_layers_from_config(self) -> Optional[tuple[int, ...]]:
             Tuple of layer indices if found in draft model config,
             None otherwise.
         """
-        if not (self.speculative_config
-                and self.speculative_config.draft_model_config):
+        if not (self.speculative_config and self.speculative_config.draft_model_config):
             return None
 
         hf_config = self.speculative_config.draft_model_config.hf_config
-        if not hasattr(hf_config, 'eagle_aux_hidden_state_layer_ids'):
+        if not hasattr(hf_config, "eagle_aux_hidden_state_layer_ids"):
             return None
 
         layer_ids = hf_config.eagle_aux_hidden_state_layer_ids