[Bugfix][Spec Decode] Fix multimodal embedding merge in Llama 4 and Mistral Large 3 EAGLE drafters (#57633)

Signed-off-by: Nishant Revur <kumrfni@amazon.com>
Co-authored-by: Claude <noreply@anthropic.com>

tests/v1/worker/test_gpu_target_dependent_ar_speculator.py

@@ -337,6 +337,30 @@ def test_draft_model_multimodal_embedding_capability(model_cls, expected):
     assert supports_multimodal_embeddings(model_cls) is expected
 
 
+@pytest.mark.parametrize(
+    "model_cls", [EagleLlama4ForCausalLM, EagleMistralLarge3ForCausalLM]
+)
+def test_mm_capable_drafter_merges_multimodal_embeddings(model_cls):
+    """Drafters advertising multimodal embeddings must embed text ids and
+    scatter the multimodal embeddings themselves; they do not inherit
+    SupportsMultiModal's helpers."""
+    table = torch.arange(8, dtype=torch.float32).repeat(3, 1).T  # [8, 3]
+    drafter = object.__new__(model_cls)
+    drafter.model = SimpleNamespace(embed_input_ids=lambda ids: table[ids])
+    input_ids = torch.tensor([1, 5, 6, 2])
+
+    text_only = model_cls.embed_input_ids(drafter, input_ids)
+    assert torch.equal(text_only, table[input_ids])
+
+    mm_embeds = [torch.full((2, 3), -1.0)]
+    is_multimodal = torch.tensor([False, True, True, False])
+    merged = model_cls.embed_input_ids(
+        drafter, input_ids, mm_embeds, is_multimodal=is_multimodal
+    )
+    assert torch.equal(merged[[0, 3]], table[[1, 2]])
+    assert torch.equal(merged[[1, 2]], mm_embeds[0])
+
+
 def test_run_model_unpacks_tuple_return_for_mtp(monkeypatch):
     logits_hidden = torch.full((4, 3), 1.0)
     feedback_hidden = torch.full((4, 3), 2.0)

vllm/model_executor/models/llama4_eagle.py

@@ -36,10 +36,15 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
 from vllm.model_executor.models.llama4 import Llama4DecoderLayer, Llama4ForCausalLM
 from vllm.model_executor.models.utils import extract_layer_index
 
-from .interfaces import SupportsMultiModal, SupportsMultiModalEmbeddings
+from .interfaces import (
+    MultiModalEmbeddings,
+    SupportsMultiModalEmbeddings,
+    _require_is_multimodal,
+)
 from .utils import (
     AutoWeightsLoader,
     WeightsMapper,
+    _merge_multimodal_embeddings,
     maybe_prefix,
     process_eagle_weight,
 )
@@ -204,7 +209,21 @@ class EagleLlama4ForCausalLM(_EagleLlama4ForCausalLMBase, SupportsMultiModalEmbe
     def get_language_model(self) -> torch.nn.Module:
         return self.model
 
-    embed_input_ids = SupportsMultiModal.embed_input_ids  # type: ignore
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: MultiModalEmbeddings | None = None,
+        *,
+        is_multimodal: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        inputs_embeds = self.model.embed_input_ids(input_ids)
+        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
+            return inputs_embeds
+        return _merge_multimodal_embeddings(
+            inputs_embeds=inputs_embeds,
+            multimodal_embeddings=multimodal_embeddings,
+            is_multimodal=_require_is_multimodal(is_multimodal),
+        )
 
     def forward(
         self,

vllm/model_executor/models/mistral_large_3_eagle.py

@@ -23,8 +23,17 @@ from vllm.model_executor.models.deepseek_v2 import (
 )
 from vllm.model_executor.models.mistral_large_3 import MistralLarge3ForCausalLM
 
-from .interfaces import SupportsMultiModal, SupportsMultiModalEmbeddings
-from .utils import WeightsMapper, make_empty_intermediate_tensors_factory, maybe_prefix
+from .interfaces import (
+    MultiModalEmbeddings,
+    SupportsMultiModalEmbeddings,
+    _require_is_multimodal,
+)
+from .utils import (
+    WeightsMapper,
+    _merge_multimodal_embeddings,
+    make_empty_intermediate_tensors_factory,
+    maybe_prefix,
+)
 
 logger = init_logger(__name__)
 
@@ -148,7 +157,21 @@ class EagleMistralLarge3ForCausalLM(
     def get_language_model(self) -> torch.nn.Module:
         return self.model
 
-    embed_input_ids = SupportsMultiModal.embed_input_ids  # type: ignore
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: MultiModalEmbeddings | None = None,
+        *,
+        is_multimodal: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        inputs_embeds = self.model.embed_input_ids(input_ids)
+        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
+            return inputs_embeds
+        return _merge_multimodal_embeddings(
+            inputs_embeds=inputs_embeds,
+            multimodal_embeddings=multimodal_embeddings,
+            is_multimodal=_require_is_multimodal(is_multimodal),
+        )
 
     def forward(
         self,