A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 05:45 UTC
[Model Runner V2] Support multi-modal embeddings for spec decode model (#36097)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>
Giancarlo Delfin committed
Mar 22, 2026 at 02:48 UTC
b3e846017d9f6f48d196185e1d125e421ec7ab6b
2 files changed
+58
vllm/v1/worker/gpu/model_runner.py
+31
@@ -430,6 +430,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
430
# dummy run the eagle speculator's propose to ensure DP/EP sync.
431
if self.speculator is not None:
432
assert self.sampler is not None
433
+ mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
434
+ if self.speculator.supports_mm_inputs:
435
+ mm_inputs = (
436
+ [],
437
+ torch.zeros(
438
+ input_batch.num_tokens,
439
+ dtype=torch.bool,
440
+ device=self.device,
441
+ ),
442
+ )
443
self.speculator.propose(
444
input_batch=input_batch,
445
attn_metadata=attn_metadata,
@@ -449,6 +459,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
459
num_tokens_across_dp=num_tokens_across_dp,
460
dummy_run=True,
461
skip_attn_for_dummy_run=skip_attn,
462
+ mm_inputs=mm_inputs,
463
)
464
465
assert hidden_states is not None # Last PP rank always has hidden_states
@@ -1142,8 +1153,27 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1153
self.postprocess(
1154
input_batch, sampler_output.sampled_token_ids, num_sampled, num_rejected
1155
)
1156
+
1157
if self.speculator is not None:
1158
assert self.sampler is not None
1159
+ mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
1160
+ if self.speculator.supports_mm_inputs:
1161
+ # Get cached multimodal embeddings for draft forward.
1162
+ prefill_lens = self.req_states.prefill_len.np[
1163
+ input_batch.idx_mapping_np
1164
+ ]
1165
+ computed_prefill_lens = self.req_states.num_computed_prefill_tokens[
1166
+ input_batch.idx_mapping_np
1167
+ ]
1168
+ mm_inputs = self.model_state.encoder_runner.gather_mm_embeddings(
1169
+ input_batch.req_ids,
1170
+ input_batch.num_tokens,
1171
+ input_batch.num_scheduled_tokens,
1172
+ input_batch.query_start_loc_np,
1173
+ prefill_lens,
1174
+ computed_prefill_lens + 1, # + 1 to consider the skew in eagle
1175
+ )
1176
+
1177
draft_tokens = self.speculator.propose(
1178
input_batch,
1179
attn_metadata,
@@ -1157,6 +1187,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1187
self.sampler.sampling_states.temperature.gpu,
1188
self.sampler.sampling_states.seeds.gpu,
1189
num_tokens_across_dp=num_tokens_across_dp,
1190
+ mm_inputs=mm_inputs,
1191
)
1192
self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
1193
self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens)
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+27
@@ -10,6 +10,7 @@ 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.model_executor.layers.attention_layer_base import AttentionLayerBase
13
+from vllm.multimodal import MULTIMODAL_REGISTRY
14
from vllm.triton_utils import tl, triton
15
from vllm.v1.kv_cache_interface import KVCacheConfig
16
from vllm.v1.worker.gpu.attn_utils import (
@@ -76,6 +77,14 @@ class EagleSpeculator:
77
device=device,
78
)
79
80
+ self.supports_mm_inputs = MULTIMODAL_REGISTRY.supports_multimodal_inputs(
81
+ self.draft_model_config
82
+ )
83
+ if self.supports_mm_inputs:
84
+ self.inputs_embeds = torch.zeros(
85
+ self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=device
86
+ )
87
+
88
cache_draft_logits = self.speculative_config.rejection_sample_method != "strict"
89
self.draft_logits: torch.Tensor | None = None
90
if cache_draft_logits:
@@ -138,6 +147,7 @@ class EagleSpeculator:
147
slot_mappings: dict[str, torch.Tensor] | None,
148
num_tokens_across_dp: torch.Tensor | None,
149
cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
150
+ mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
151
) -> tuple[torch.Tensor, torch.Tensor]:
152
batch_descriptor = BatchDescriptor(num_tokens=num_tokens)
153
with set_forward_context(
@@ -149,10 +159,25 @@ class EagleSpeculator:
159
slot_mapping=slot_mappings,
160
batch_descriptor=batch_descriptor,
161
):
162
+ inputs_embeds = None
163
+ if self.supports_mm_inputs:
164
+ # Merge multimodal embeddings with input ids.
165
+ mm_embeds, is_mm_embed = mm_inputs or (None, None)
166
+ num_input_tokens = (
167
+ is_mm_embed.shape[0] if is_mm_embed is not None else num_tokens
168
+ )
169
+ self.inputs_embeds[:num_input_tokens] = self.model.embed_input_ids(
170
+ self.input_buffers.input_ids[:num_input_tokens],
171
+ multimodal_embeddings=mm_embeds,
172
+ is_multimodal=is_mm_embed,
173
+ )
174
+ inputs_embeds = self.inputs_embeds[:num_tokens]
175
+
176
ret_hidden_states = self.model(
177
input_ids=self.input_buffers.input_ids[:num_tokens],
178
positions=self.input_buffers.positions[:num_tokens],
179
hidden_states=self.hidden_states[:num_tokens],
180
+ inputs_embeds=inputs_embeds,
181
)
182
if self.method == "mtp":
183
last_hidden_states = ret_hidden_states
@@ -254,6 +279,7 @@ class EagleSpeculator:
279
num_tokens_across_dp: torch.Tensor | None = None,
280
dummy_run: bool = False,
281
skip_attn_for_dummy_run: bool = False,
282
+ mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
283
) -> torch.Tensor:
284
# NOTE(woosuk): To avoid CPU-GPU synchronization without CPU knowing the
285
# number of rejected tokens, we maintain the size of eagle's input_ids and
@@ -288,6 +314,7 @@ class EagleSpeculator:
314
attn_metadata,
315
slot_mappings,
316
num_tokens_across_dp=num_tokens_across_dp,
317
+ mm_inputs=mm_inputs,
318
)
319
sample_hidden_states = last_hidden_states[last_token_indices]
320
logits = self.model.compute_logits(sample_hidden_states)