From 581692cc0fd564546ca1f58475e795e2a154cb3f Mon Sep 17 00:00:00 2001 From: Nishant Revur Date: Thu, 8 Oct 2026 18:44:50 -0700 Subject: [PATCH] [Bugfix][Spec Decode] Fix multimodal embedding merge in Llama 4 and Mistral Large 3 EAGLE drafters (#57633) Signed-off-by: Nishant Revur Co-authored-by: Claude --- ...test_gpu_target_dependent_ar_speculator.py | 24 +++++++++++++++ vllm/model_executor/models/llama4_eagle.py | 23 +++++++++++++-- .../models/mistral_large_3_eagle.py | 29 +++++++++++++++++-- 3 files changed, 71 insertions(+), 5 deletions(-) diff --git a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py index 20e86f4a78..b63bd7a2ff 100644 --- a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py +++ b/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) diff --git a/vllm/model_executor/models/llama4_eagle.py b/vllm/model_executor/models/llama4_eagle.py index 975ee28dc2..8634f113d0 100644 --- a/vllm/model_executor/models/llama4_eagle.py +++ b/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, diff --git a/vllm/model_executor/models/mistral_large_3_eagle.py b/vllm/model_executor/models/mistral_large_3_eagle.py index d2e263f3e7..1398fd19fc 100644 --- a/vllm/model_executor/models/mistral_large_3_eagle.py +++ b/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,