@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 4a3447d200

[Bugfix][Model Runner V2] Restore multimodal draft capability detection (#50417)

Signed-off-by: Raphael Rialland <raphael.rialland@mistral.ai>

Raphaël Rialland committed Aug 3, 2026 at 23:18 UTC 4a3447d200e5aa428d68d1a00aa00f1a19a1a729
12 files changed +317 -58
tests/v1/worker/test_gpu_autoregressive_speculator.py
+199 -1
@@ -4,18 +4,32 @@
4 from contextlib import nullcontext
5 from types import SimpleNamespace
6
7 +import pytest
8 import torch
9
10 from vllm.config.compilation import CUDAGraphMode
11 +from vllm.model_executor.models import supports_multimodal_embeddings
12 +from vllm.model_executor.models.exaone4_5_mtp import Exaone4_5_MTP
13 +from vllm.model_executor.models.llama4_eagle import EagleLlama4ForCausalLM
14 +from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
15 +from vllm.model_executor.models.mistral_eagle import EagleMistralForCausalLM
16 +from vllm.model_executor.models.mistral_large_3_eagle import (
17 + EagleMistralLarge3ForCausalLM,
18 +)
19 +from vllm.v1.worker.gpu.spec_decode import speculator as base_spec_module
20 from vllm.v1.worker.gpu.spec_decode.autoregressive import speculator as spec_module
21 from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import (
22 AutoRegressiveSpeculator,
23 )
24 +from vllm.v1.worker.gpu.spec_decode.multi_module_mtp.speculator import (
25 + MultiModuleMTPSpeculator,
26 +)
27 +from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
28
29
30 class _TestSpeculator(AutoRegressiveSpeculator):
31 def load_draft_model(self, target_model, target_attn_layer_names):
18 - raise NotImplementedError
32 + return self.test_draft_model
33
34
35 class _DraftModel(torch.nn.Module):
@@ -27,6 +41,43 @@ class _DraftModel(torch.nn.Module):
41 return self.output
42
43
44 +class _MultimodalDraftModel(torch.nn.Module):
45 + supports_multimodal_embeddings = True
46 +
47 + def embed_input_ids(
48 + self,
49 + input_ids,
50 + multimodal_embeddings=None,
51 + *,
52 + is_multimodal=None,
53 + ):
54 + raise AssertionError("embed_input_ids should not be called during loading")
55 +
56 +
57 +class _TextOnlyDraftModel(torch.nn.Module):
58 + def embed_input_ids(
59 + self,
60 + input_ids,
61 + multimodal_embeddings=None,
62 + *,
63 + is_multimodal=None,
64 + ):
65 + raise AssertionError("embed_input_ids should not be called during loading")
66 +
67 +
68 +def _mock_base_model_load(monkeypatch):
69 + monkeypatch.setattr(
70 + base_spec_module,
71 + "get_layers_from_vllm_config",
72 + lambda *args, **kwargs: {},
73 + )
74 + monkeypatch.setattr(
75 + DraftModelSpeculator,
76 + "_validate_local_argmax_reduction",
77 + lambda self: None,
78 + )
79 +
80 +
81 def _make_speculator(
82 monkeypatch,
83 output: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
@@ -49,6 +100,153 @@ def _make_speculator(
100 return speculator
101
102
103 +def test_mm_support_configured_after_model_load(monkeypatch):
104 + target_model_config = object()
105 + draft_model_config = object()
106 + vllm_config = SimpleNamespace(model_config=target_model_config)
107 + draft_model = _MultimodalDraftModel()
108 +
109 + def init_base(speculator, vllm_config, device):
110 + speculator.vllm_config = vllm_config
111 + speculator.device = device
112 + speculator.max_num_tokens = 4
113 + speculator.max_num_reqs = 2
114 + speculator.hidden_size = 3
115 + speculator.dtype = torch.float32
116 + speculator.draft_model_config = draft_model_config
117 + speculator.supports_mm_inputs = False
118 +
119 + checked_configs = []
120 +
121 + def supports_multimodal_inputs(model_config):
122 + checked_configs.append(model_config)
123 + return True
124 +
125 + monkeypatch.setattr(DraftModelSpeculator, "__init__", init_base)
126 + _mock_base_model_load(monkeypatch)
127 + monkeypatch.setattr(
128 + base_spec_module.MULTIMODAL_REGISTRY,
129 + "supports_multimodal_inputs",
130 + supports_multimodal_inputs,
131 + )
132 +
133 + speculator = _TestSpeculator(vllm_config, torch.device("cpu"))
134 +
135 + assert checked_configs == []
136 + assert not speculator.supports_mm_inputs
137 + assert speculator.inputs_embeds is None
138 +
139 + speculator.test_draft_model = draft_model
140 + speculator.load_model(torch.nn.Module())
141 +
142 + assert checked_configs == [target_model_config]
143 + assert speculator.supports_mm_inputs
144 + assert speculator.inputs_embeds is not None
145 + assert speculator.inputs_embeds.shape == (4, 3)
146 +
147 +
148 +def test_load_model_keeps_mm_support_for_capable_drafter(monkeypatch):
149 + speculator = object.__new__(_TestSpeculator)
150 + speculator.supports_mm_inputs = False
151 + speculator.inputs_embeds = None
152 + speculator.vllm_config = SimpleNamespace(model_config=object())
153 + speculator.max_num_tokens = 4
154 + speculator.hidden_size = 3
155 + speculator.dtype = torch.float32
156 + speculator.device = torch.device("cpu")
157 + draft_model = _MultimodalDraftModel()
158 + speculator.test_draft_model = draft_model
159 + _mock_base_model_load(monkeypatch)
160 + monkeypatch.setattr(
161 + base_spec_module.MULTIMODAL_REGISTRY,
162 + "supports_multimodal_inputs",
163 + lambda model_config: True,
164 + )
165 +
166 + speculator.load_model(torch.nn.Module())
167 +
168 + assert speculator.supports_mm_inputs
169 + assert speculator.inputs_embeds is not None
170 +
171 +
172 +def test_load_model_disables_mm_support_for_text_only_drafter(monkeypatch):
173 + speculator = object.__new__(_TestSpeculator)
174 + speculator.supports_mm_inputs = False
175 + speculator.inputs_embeds = None
176 + speculator.vllm_config = SimpleNamespace(model_config=object())
177 + draft_model = _TextOnlyDraftModel()
178 + speculator.test_draft_model = draft_model
179 + warning_messages = []
180 + _mock_base_model_load(monkeypatch)
181 + monkeypatch.setattr(
182 + base_spec_module.MULTIMODAL_REGISTRY,
183 + "supports_multimodal_inputs",
184 + lambda model_config: True,
185 + )
186 + monkeypatch.setattr(
187 + base_spec_module.logger,
188 + "warning_once",
189 + lambda message, *args: warning_messages.append(message % args),
190 + )
191 +
192 + speculator.load_model(torch.nn.Module())
193 +
194 + assert not speculator.supports_mm_inputs
195 + assert warning_messages == [
196 + "Draft model _TextOnlyDraftModel does not support external multimodal "
197 + "embeddings. Embeddings from the target model will not be passed to the "
198 + "drafter; using text-only draft inputs instead."
199 + ]
200 +
201 +
202 +def test_multi_module_mm_support_configured_after_model_load(monkeypatch):
203 + speculator = object.__new__(MultiModuleMTPSpeculator)
204 + speculator.supports_mm_inputs = False
205 + speculator.inputs_embeds = None
206 + speculator.cached_draft_input_embeds = None
207 + speculator.vllm_config = SimpleNamespace(model_config=object())
208 + speculator.max_num_tokens = 4
209 + speculator.max_num_reqs = 2
210 + speculator.num_speculative_steps = 3
211 + speculator.hidden_size = 3
212 + speculator.dtype = torch.float32
213 + speculator.device = torch.device("cpu")
214 + draft_model = _MultimodalDraftModel()
215 + _mock_base_model_load(monkeypatch)
216 + monkeypatch.setattr(
217 + MultiModuleMTPSpeculator,
218 + "load_draft_model",
219 + lambda self, target_model, target_attn_layer_names: draft_model,
220 + )
221 + monkeypatch.setattr(
222 + base_spec_module.MULTIMODAL_REGISTRY,
223 + "supports_multimodal_inputs",
224 + lambda model_config: True,
225 + )
226 +
227 + speculator.load_model(torch.nn.Module())
228 +
229 + assert speculator.supports_mm_inputs
230 + assert speculator.inputs_embeds is not None
231 + assert speculator.inputs_embeds.shape == (4, 3)
232 + assert speculator.cached_draft_input_embeds is not None
233 + assert speculator.cached_draft_input_embeds.shape == (2, 2, 3)
234 +
235 +
236 +@pytest.mark.parametrize(
237 + ("model_cls", "expected"),
238 + [
239 + (EagleLlama4ForCausalLM, True),
240 + (EagleMistralForCausalLM, True),
241 + (EagleMistralLarge3ForCausalLM, True),
242 + (Exaone4_5_MTP, True),
243 + (Eagle3LlamaForCausalLM, False),
244 + ],
245 +)
246 +def test_draft_model_multimodal_embedding_capability(model_cls, expected):
247 + assert supports_multimodal_embeddings(model_cls) is expected
248 +
249 +
250 def test_run_model_unpacks_tuple_return_for_mtp(monkeypatch):
251 logits_hidden = torch.full((4, 3), 1.0)
252 feedback_hidden = torch.full((4, 3), 2.0)
vllm/model_executor/models/__init__.py
+4
@@ -6,12 +6,14 @@ from .interfaces import (
6 SupportsLoRA,
7 SupportsMRoPE,
8 SupportsMultiModal,
9 + SupportsMultiModalEmbeddings,
10 SupportsPP,
11 SupportsTranscription,
12 has_inner_state,
13 supports_lora,
14 supports_mrope,
15 supports_multimodal,
16 + supports_multimodal_embeddings,
17 supports_pp,
18 supports_transcription,
19 )
@@ -35,6 +37,8 @@ __all__ = [
37 "supports_lora",
38 "SupportsMultiModal",
39 "supports_multimodal",
40 + "SupportsMultiModalEmbeddings",
41 + "supports_multimodal_embeddings",
42 "SupportsMRoPE",
43 "supports_mrope",
44 "SupportsPP",
vllm/model_executor/models/interfaces.py
+34 -1
@@ -110,7 +110,22 @@ _language_model_by_module = dict[nn.Module, "VllmModel"]()
110
111
112 @runtime_checkable
113 -class SupportsMultiModal(Protocol):
113 +class SupportsMultiModalEmbeddings(Protocol):
114 + """The interface for models that can merge external multimodal embeddings."""
115 +
116 + supports_multimodal_embeddings: ClassVar[Literal[True]] = True
117 +
118 + def embed_input_ids(
119 + self,
120 + input_ids: Tensor,
121 + multimodal_embeddings: MultiModalEmbeddings | None = None,
122 + *,
123 + is_multimodal: Tensor | None = None,
124 + ) -> Tensor: ...
125 +
126 +
127 +@runtime_checkable
128 +class SupportsMultiModal(SupportsMultiModalEmbeddings, Protocol):
129 """The interface required for all multi-modal models."""
130
131 supports_multimodal: ClassVar[Literal[True]] = True
@@ -488,6 +503,24 @@ def supports_multimodal(
503 return getattr(model, "supports_multimodal", False)
504
505
506 +@overload
507 +def supports_multimodal_embeddings(
508 + model: type[object],
509 +) -> TypeIs[type[SupportsMultiModalEmbeddings]]: ...
510 +
511 +
512 +@overload
513 +def supports_multimodal_embeddings(
514 + model: object,
515 +) -> TypeIs[SupportsMultiModalEmbeddings]: ...
516 +
517 +
518 +def supports_multimodal_embeddings(
519 + model: type[object] | object,
520 +) -> TypeIs[type[SupportsMultiModalEmbeddings]] | TypeIs[SupportsMultiModalEmbeddings]:
521 + return getattr(model, "supports_multimodal_embeddings", False)
522 +
523 +
524 def supports_multimodal_raw_input_only(model: type[object] | object) -> bool:
525 return getattr(model, "supports_multimodal_raw_input_only", False)
526
vllm/model_executor/models/llama4_eagle.py
+2 -2
@@ -35,7 +35,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
35 from vllm.model_executor.models.llama4 import Llama4DecoderLayer, Llama4ForCausalLM
36 from vllm.model_executor.models.utils import extract_layer_index
37
38 -from .interfaces import SupportsMultiModal
38 +from .interfaces import SupportsMultiModal, SupportsMultiModalEmbeddings
39 from .utils import (
40 AutoWeightsLoader,
41 WeightsMapper,
@@ -152,7 +152,7 @@ class LlamaModel(nn.Module):
152 }
153
154
155 -class EagleLlama4ForCausalLM(Llama4ForCausalLM):
155 +class EagleLlama4ForCausalLM(Llama4ForCausalLM, SupportsMultiModalEmbeddings):
156 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
157 nn.Module.__init__(self)
158 self.config = vllm_config.speculative_config.draft_model_config.hf_config
vllm/model_executor/models/mistral_eagle.py
+5 -2
@@ -14,7 +14,10 @@ from vllm.model_executor.layers.linear import RowParallelLinear
14 from vllm.model_executor.layers.logits_processor import LogitsProcessor
15 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
16 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
17 -from vllm.model_executor.models.interfaces import MultiModalEmbeddings
17 +from vllm.model_executor.models.interfaces import (
18 + MultiModalEmbeddings,
19 + SupportsMultiModalEmbeddings,
20 +)
21 from vllm.model_executor.models.llama import LlamaConfig
22 from vllm.model_executor.models.mistral import (
23 MistralDecoderLayer,
@@ -109,7 +112,7 @@ class EagleMistralModel(MistralModel):
112 return hidden_states, hidden_states
113
114
112 -class EagleMistralForCausalLM(MistralForCausalLM):
115 +class EagleMistralForCausalLM(MistralForCausalLM, SupportsMultiModalEmbeddings):
116 mistral_mapping = MistralForCausalLM.mistral_mapping | {
117 "eagle_linear": "model.fc",
118 }
vllm/model_executor/models/mistral_large_3_eagle.py
+4 -2
@@ -22,7 +22,7 @@ from vllm.model_executor.models.deepseek_v2 import (
22 )
23 from vllm.model_executor.models.mistral_large_3 import MistralLarge3ForCausalLM
24
25 -from .interfaces import SupportsMultiModal
25 +from .interfaces import SupportsMultiModal, SupportsMultiModalEmbeddings
26 from .utils import WeightsMapper, make_empty_intermediate_tensors_factory, maybe_prefix
27
28 logger = init_logger(__name__)
@@ -107,7 +107,9 @@ class EagleMistralLarge3Model(DeepseekV2Model):
107 return output
108
109
110 -class EagleMistralLarge3ForCausalLM(MistralLarge3ForCausalLM):
110 +class EagleMistralLarge3ForCausalLM(
111 + MistralLarge3ForCausalLM, SupportsMultiModalEmbeddings
112 +):
113 hf_to_vllm_mapper = MistralLarge3ForCausalLM.hf_to_vllm_mapper | WeightsMapper(
114 orig_to_new_regex={
115 regex.compile(r"\Aeagle_linear\.weight\Z"): r"model.fc.weight",
vllm/models/inkling/amd/mtp.py
+2 -1
@@ -29,6 +29,7 @@ from vllm.model_executor.model_loader.mtp_validation import (
29 is_mtp_completeness_check_enabled,
30 )
31 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
32 +from vllm.model_executor.models.interfaces import SupportsMultiModalEmbeddings
33 from vllm.model_executor.models.utils import maybe_prefix
34 from vllm.sequence import IntermediateTensors
35
@@ -213,7 +214,7 @@ class InklingMultiTokenPredictor(nn.Module):
214 return hidden
215
216
216 -class InklingMTP(nn.Module):
217 +class InklingMTP(nn.Module, SupportsMultiModalEmbeddings):
218 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
219 super().__init__()
220 assert vllm_config.speculative_config is not None
vllm/models/inkling/nvidia/mtp.py
+2 -1
@@ -33,6 +33,7 @@ from vllm.model_executor.model_loader.mtp_validation import (
33 is_mtp_completeness_check_enabled,
34 )
35 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
36 +from vllm.model_executor.models.interfaces import SupportsMultiModalEmbeddings
37 from vllm.model_executor.models.utils import maybe_prefix
38 from vllm.sequence import IntermediateTensors
39
@@ -242,7 +243,7 @@ class InklingMultiTokenPredictor(nn.Module):
243 return hidden
244
245
245 -class InklingMTP(nn.Module):
246 +class InklingMTP(nn.Module, SupportsMultiModalEmbeddings):
247 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
248 super().__init__()
249 assert vllm_config.speculative_config is not None
vllm/v1/spec_decode/llm_base_proposer.py
+14 -13
@@ -25,7 +25,10 @@ from vllm.forward_context import set_forward_context
25 from vllm.logger import init_logger
26 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
27 from vllm.model_executor.model_loader import get_model
28 -from vllm.model_executor.models import supports_multimodal
28 +from vllm.model_executor.models import (
29 + supports_multimodal,
30 + supports_multimodal_embeddings,
31 +)
32 from vllm.model_executor.models.deepseek_eagle3 import Eagle3DeepseekV2ForCausalLM
33 from vllm.model_executor.models.interfaces import SupportsMultiModal
34 from vllm.model_executor.models.laguna_dflash import DFlashLagunaForCausalLM
@@ -1349,18 +1352,16 @@ class SpecDecodeBaseProposer:
1352 if all_attn_layers[name].get_kv_cache_spec(self.vllm_config) is not None
1353 }
1354
1352 - if self.supports_mm_inputs:
1353 - # Even if the target model is multimodal, we can also use
1354 - # text-only draft models
1355 - try:
1356 - dummy_input_ids = torch.tensor([[1]], device=self.input_ids.device)
1357 - self.model.embed_input_ids(dummy_input_ids, multimodal_embeddings=None)
1358 - except (NotImplementedError, AttributeError, TypeError):
1359 - logger.warning(
1360 - "Draft model does not support multimodal inputs, "
1361 - "falling back to text-only mode"
1362 - )
1363 - self.supports_mm_inputs = False
1355 + # Even if the target model is multimodal, we can also use
1356 + # text-only draft models
1357 + if self.supports_mm_inputs and not supports_multimodal_embeddings(self.model):
1358 + logger.warning_once(
1359 + "Draft model %s does not support external multimodal embeddings. "
1360 + "Embeddings from the target model will not be passed to the "
1361 + "drafter; using text-only draft inputs instead.",
1362 + type(self.model).__name__,
1363 + )
1364 + self.supports_mm_inputs = False
1365
1366 if supports_multimodal(target_model):
1367 # handle multimodality
vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
+15 -8
@@ -3,12 +3,12 @@
3 from typing import Any
4
5 import torch
6 +import torch.nn as nn
7
8 from vllm.config import VllmConfig
9 from vllm.config.compilation import CUDAGraphMode
10 from vllm.forward_context import BatchDescriptor, set_forward_context
11 from vllm.logger import init_logger
11 -from vllm.multimodal import MULTIMODAL_REGISTRY
12 from vllm.triton_utils import tl, triton
13 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
14 from vllm.v1.worker.gpu.cudagraph_utils import (
@@ -37,17 +37,23 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
37 self.max_num_reqs, dtype=torch.int64, device=device
38 )
39
40 - self.supports_mm_inputs = MULTIMODAL_REGISTRY.supports_multimodal_inputs(
41 - self.draft_model_config
42 - )
43 - if self.supports_mm_inputs:
44 - self.inputs_embeds = torch.zeros(
45 - self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=device
46 - )
40 + self.inputs_embeds: torch.Tensor | None = None
41
42 self.prefill_cudagraph_manager: SpeculatorCudaGraphManager | None = None
43 self.decode_cudagraph_manager: SpeculatorCudaGraphManager | None = None
44
45 + def load_model(self, target_model: nn.Module) -> None:
46 + super().load_model(target_model)
47 + if not self.supports_mm_inputs:
48 + return
49 +
50 + self.inputs_embeds = torch.zeros(
51 + self.max_num_tokens,
52 + self.hidden_size,
53 + dtype=self.dtype,
54 + device=self.device,
55 + )
56 +
57 @property
58 def advance_draft_positions(self) -> bool:
59 """
@@ -295,6 +301,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
301 ):
302 inputs_embeds = None
303 if self.supports_mm_inputs:
304 + assert self.inputs_embeds is not None
305 # Merge multimodal embeddings with input ids.
306 mm_embeds, is_mm_embed = mm_inputs or (None, None)
307 num_input_tokens = (
vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py
+19 -27
@@ -9,7 +9,6 @@ from vllm.config import VllmConfig
9 from vllm.config.compilation import CUDAGraphMode
10 from vllm.forward_context import BatchDescriptor, set_forward_context
11 from vllm.logger import init_logger
12 -from vllm.multimodal import MULTIMODAL_REGISTRY
12 from vllm.triton_utils import tl, triton
13 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
14 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
@@ -39,25 +38,7 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):
38 self.max_num_reqs, dtype=torch.int64, device=device
39 )
40
42 - self.supports_mm_inputs = MULTIMODAL_REGISTRY.supports_multimodal_inputs(
43 - self.draft_model_config
44 - )
45 - # HACK: the Inkling MTP draft has no MM processor of its own (its draft
46 - # config is flattened text-only), but it consumes the target's merged
47 - # embeddings at draft prefill — treat it as MM-capable whenever the
48 - # target is.
49 - if (
50 - not self.supports_mm_inputs
51 - and self.draft_model_config.hf_config.model_type == "inkling_mtp"
52 - ):
53 - self.supports_mm_inputs = MULTIMODAL_REGISTRY.supports_multimodal_inputs(
54 - vllm_config.model_config
55 - )
41 self.inputs_embeds: torch.Tensor | None = None
57 - if self.supports_mm_inputs:
58 - self.inputs_embeds = torch.zeros(
59 - self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=device
60 - )
42
43 # Input id overrides for the last num_speculative_steps - 1 draft steps.
44 # Used by chunked-prefilling requests to swap in the future prefill token
@@ -80,14 +61,6 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):
61 device=self.device,
62 )
63 self.cached_draft_input_embeds: torch.Tensor | None = None
83 - if self.supports_mm_inputs:
84 - self.cached_draft_input_embeds = torch.zeros(
85 - self.max_num_reqs,
86 - self.num_speculative_steps - 1,
87 - self.hidden_size,
88 - dtype=self.dtype,
89 - device=self.device,
90 - )
64 self.cached_target_hidden_states = torch.zeros(
65 self.max_num_reqs,
66 self.num_speculative_steps - 1,
@@ -105,6 +78,25 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):
78 ) -> nn.Module:
79 return load_eagle_model(target_model, self.vllm_config)
80
81 + def load_model(self, target_model: nn.Module) -> None:
82 + super().load_model(target_model)
83 + if not self.supports_mm_inputs:
84 + return
85 +
86 + self.inputs_embeds = torch.zeros(
87 + self.max_num_tokens,
88 + self.hidden_size,
89 + dtype=self.dtype,
90 + device=self.device,
91 + )
92 + self.cached_draft_input_embeds = torch.zeros(
93 + self.max_num_reqs,
94 + self.num_speculative_steps - 1,
95 + self.hidden_size,
96 + dtype=self.dtype,
97 + device=self.device,
98 + )
99 +
100 def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
101 # TODO(TheEpicDolphin): Support piecewise cudagraph for multi-module MTP.
102 if cudagraph_mode.has_piecewise_cudagraphs():
vllm/v1/worker/gpu/spec_decode/speculator.py
+17
@@ -13,6 +13,8 @@ from vllm.config.compilation import CUDAGraphMode
13 from vllm.distributed.eplb.eplb_state import EplbState
14 from vllm.logger import init_logger
15 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
16 +from vllm.model_executor.models import supports_multimodal_embeddings
17 +from vllm.multimodal import MULTIMODAL_REGISTRY
18 from vllm.v1.kv_cache_interface import KVCacheConfig
19 from vllm.v1.worker.gpu.attn_utils import (
20 build_attn_metadata,
@@ -137,6 +139,8 @@ class DraftModelSpeculator(BaseSpeculator):
139 device=device,
140 )
141
142 + self.supports_mm_inputs = False
143 +
144 @abstractmethod
145 def load_draft_model(
146 self,
@@ -164,6 +168,19 @@ class DraftModelSpeculator(BaseSpeculator):
168 )
169 self.draft_attn_layer_names = all_attn_layers - target_attn_layer_names
170
171 + target_supports_mm = MULTIMODAL_REGISTRY.supports_multimodal_inputs(
172 + self.vllm_config.model_config
173 + )
174 + draft_supports_mm = supports_multimodal_embeddings(self.model)
175 + self.supports_mm_inputs = target_supports_mm and draft_supports_mm
176 + if target_supports_mm and not draft_supports_mm:
177 + logger.warning_once(
178 + "Draft model %s does not support external multimodal embeddings. "
179 + "Embeddings from the target model will not be passed to the "
180 + "drafter; using text-only draft inputs instead.",
181 + type(self.model).__name__,
182 + )
183 +
184 def set_eplb_state(self, eplb_state: EplbState) -> None:
185 """Inject EPLB state after construction."""
186 self.eplb_state = eplb_state