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, 03:15 UTC
[Model Runner V2] Gather multimodal embeddings before draft model postprocess (#37932)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Giancarlo Delfin committed
Mar 23, 2026 at 18:14 UTC
8f4824b664c7e67db2f4e06e991748d932474955
1 file changed
+18
-18
vllm/v1/worker/gpu/model_runner.py
+18
-18
@@ -1162,6 +1162,24 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1162
copy_event=self.output_copy_event,
1163
)
1164
1165
+ mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
1166
+ if self.speculator is not None and self.speculator.supports_mm_inputs:
1167
+ # Get cached multimodal embeddings for draft forward.
1168
+ # NOTE: This is done here because postprocess updates
1169
+ # num_computed_prefill_tokens.
1170
+ prefill_lens = self.req_states.prefill_len.np[input_batch.idx_mapping_np]
1171
+ computed_prefill_lens = self.req_states.num_computed_prefill_tokens[
1172
+ input_batch.idx_mapping_np
1173
+ ]
1174
+ mm_inputs = self.model_state.encoder_runner.gather_mm_embeddings(
1175
+ input_batch.req_ids,
1176
+ input_batch.num_tokens,
1177
+ input_batch.num_scheduled_tokens,
1178
+ input_batch.query_start_loc_np,
1179
+ prefill_lens,
1180
+ computed_prefill_lens + 1, # +1 to consider the skew in eagle
1181
+ )
1182
+
1183
# Postprocess results and update request states.
1184
# NOTE: This is intentionally done after creating the AsyncOutput,
1185
# ensuring that `copy_event` is recorded before calling postprocess.
@@ -1173,24 +1191,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1191
1192
if self.speculator is not None:
1193
assert self.sampler is not None
1176
- mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
1177
- if self.speculator.supports_mm_inputs:
1178
- # Get cached multimodal embeddings for draft forward.
1179
- prefill_lens = self.req_states.prefill_len.np[
1180
- input_batch.idx_mapping_np
1181
- ]
1182
- computed_prefill_lens = self.req_states.num_computed_prefill_tokens[
1183
- input_batch.idx_mapping_np
1184
- ]
1185
- mm_inputs = self.model_state.encoder_runner.gather_mm_embeddings(
1186
- input_batch.req_ids,
1187
- input_batch.num_tokens,
1188
- input_batch.num_scheduled_tokens,
1189
- input_batch.query_start_loc_np,
1190
- prefill_lens,
1191
- computed_prefill_lens + 1, # + 1 to consider the skew in eagle
1192
- )
1193
-
1194
draft_tokens = self.speculator.propose(
1195
input_batch,
1196
attn_metadata,