[ModelRunner V2] Support realtime embeddings (#46762)
Nick Hill committed
Jun 26, 2026 at 19:42 UTC
c6dd32a810aa8c4eda5696722c807e53d9f595a5
6 files changed
+52
-36
tests/v1/worker/test_encoder_runner.py
+3
-3
@@ -53,14 +53,14 @@ def _make_runner(
53
54
55
def _gather(runner: EncoderRunner, *, num_scheduled: int, draft_lookahead: int):
56
- # Single prefilling request, computed_prefill=0, prefill_len large.
56
+ # Single prefilling request, num_computed_tokens=0, prefill_len large.
57
return runner.gather_mm_embeddings(
58
req_ids=["req0"],
59
total_num_scheduled_tokens=num_scheduled,
60
num_scheduled_tokens=np.array([num_scheduled]),
61
query_start_loc=np.array([0]),
62
prefill_lens=np.array([1000]),
63
- computed_prefill_lens=np.array([0]),
63
+ num_computed_tokens=np.array([0]),
64
draft_lookahead=draft_lookahead,
65
)
66
@@ -147,7 +147,7 @@ def test_multi_request_batch_gathers_per_request(draft_lookahead):
147
num_scheduled_tokens=np.array([8, 8]),
148
query_start_loc=np.array([0, 8]),
149
prefill_lens=np.array([1000, 1000]),
150
- computed_prefill_lens=np.array([0, 0]),
150
+ num_computed_tokens=np.array([0, 0]),
151
draft_lookahead=draft_lookahead,
152
)
153
vllm/model_executor/models/diffusion_gemma.py
+1
-1
@@ -878,7 +878,7 @@ class DiffusionGemmaModelState(ModelState):
878
scheduled_encoder_inputs: dict[str, list[int]],
879
input_batch: InputBatch,
880
req_states: RequestState,
881
- ) -> torch.Tensor:
881
+ ) -> torch.Tensor | None:
882
if not self.supports_mm_inputs:
883
return None
884
vllm/v1/worker/gpu/mm/encoder_runner.py
+20
-17
@@ -3,7 +3,7 @@
3
import numpy as np
4
import torch
5
6
-from vllm.model_executor.models.interfaces import SupportsMultiModal
6
+from vllm.model_executor.models.interfaces import SupportsMultiModal, supports_realtime
7
from vllm.multimodal.inputs import MultiModalKwargsItem
8
from vllm.multimodal.utils import get_mm_features_in_window, group_and_batch_mm_kwargs
9
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
@@ -26,6 +26,7 @@ class EncoderRunner:
26
self.encoder_cache = encoder_cache
27
self.dtype = dtype
28
self.device = device
29
+ self.is_realtime = supports_realtime(model)
30
31
self.inputs_embeds = torch.zeros(
32
max_num_tokens, hidden_size, dtype=dtype, device=device
@@ -67,30 +68,32 @@ class EncoderRunner:
68
num_scheduled_tokens: np.ndarray,
69
query_start_loc: np.ndarray,
70
prefill_lens: np.ndarray,
70
- computed_prefill_lens: np.ndarray,
71
+ num_computed_tokens: np.ndarray,
72
draft_lookahead: int = 0,
73
) -> tuple[list[torch.Tensor], torch.Tensor]:
74
if draft_lookahead:
74
- computed_prefill_lens = computed_prefill_lens + draft_lookahead
75
+ num_computed_tokens = num_computed_tokens + draft_lookahead
76
76
- is_prefilling_np = computed_prefill_lens < prefill_lens
77
- if not is_prefilling_np.any():
78
- # All decode requests, so no need to gather any embeddings.
79
- return [], torch.zeros(
80
- total_num_scheduled_tokens, dtype=torch.bool, device=self.device
81
- )
82
-
83
- is_prefilling = is_prefilling_np.tolist()
84
- query_start = computed_prefill_lens.tolist()
85
- query_end = (computed_prefill_lens + num_scheduled_tokens).tolist()
86
-
87
- mm_embeds: list[torch.Tensor] = []
77
is_mm_embed = torch.zeros(
78
total_num_scheduled_tokens, dtype=torch.bool, device="cpu"
79
)
80
+
81
+ # Whether to gather media embeddings this step.
82
+ exclude_embeddings: list[bool] | None = None
83
+ if not self.is_realtime:
84
+ # Non-realtime models only have media embeddings within the prompt.
85
+ is_decode = num_computed_tokens >= prefill_lens
86
+ if is_decode.all():
87
+ # All decode requests, so no need to gather any embeddings.
88
+ return [], is_mm_embed
89
+ exclude_embeddings = is_decode.tolist()
90
+
91
+ query_start = num_computed_tokens.tolist()
92
+ query_end = (num_computed_tokens + num_scheduled_tokens).tolist()
93
+
94
+ mm_embeds: list[torch.Tensor] = []
95
for i, req_id in enumerate(req_ids):
92
- if not is_prefilling[i]:
93
- # OPTIMIZATION: Skip decode requests.
96
+ if exclude_embeddings is not None and exclude_embeddings[i]:
97
continue
98
99
cur_query_start = query_start[i]
vllm/v1/worker/gpu/model_runner.py
+18
-13
@@ -1218,20 +1218,25 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1218
if self.supports_mm_inputs and self.is_first_pp_rank:
1219
# Run MM encoder (if needed) and get multimodal embeddings.
1220
# Only first PP rank prepares multimodal embeddings.
1221
- # NOTE(woosuk): We must call get_mm_embeddings even during dummy runs
1222
- # to obtain inputs_embeds, because the compiled model expects this input.
1223
- if self.lora_config is not None:
1224
- set_active_mm_loras(
1225
- model=self.model,
1226
- lora_manager=self.lora_manager,
1227
- encoder_cache=self.encoder_cache,
1228
- req_id_to_index=self.req_states.req_id_to_index,
1229
- lora_state=self.lora_state,
1230
- scheduled_encoder_inputs=scheduler_output.scheduled_encoder_inputs,
1221
+ if dummy_run:
1222
+ # Obtain mm embeddings of correct shape for compiled model.
1223
+ inputs_embeds = self.model_state.dummy_inputs_embeds(
1224
+ input_batch.num_tokens_after_padding
1225
+ )
1226
+ else:
1227
+ scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
1228
+ if self.lora_config is not None:
1229
+ set_active_mm_loras(
1230
+ model=self.model,
1231
+ lora_manager=self.lora_manager,
1232
+ encoder_cache=self.encoder_cache,
1233
+ req_id_to_index=self.req_states.req_id_to_index,
1234
+ lora_state=self.lora_state,
1235
+ scheduled_encoder_inputs=scheduled_encoder_inputs,
1236
+ )
1237
+ inputs_embeds = self.model_state.get_mm_embeddings(
1238
+ scheduled_encoder_inputs, input_batch, self.req_states
1239
)
1232
- inputs_embeds = self.model_state.get_mm_embeddings(
1233
- scheduler_output.scheduled_encoder_inputs, input_batch, self.req_states
1234
- )
1240
if inputs_embeds is not None and not self.model.requires_raw_input_tokens:
1241
input_ids = None
1242
vllm/v1/worker/gpu/model_states/default.py
+4
@@ -57,6 +57,10 @@ class DefaultModelState(ModelState):
57
if self.rope_state is not None:
58
self.rope_state.apply_staged_writes()
59
60
+ def dummy_inputs_embeds(self, num_tokens: int) -> torch.Tensor:
61
+ """Pre-allocated inputs_embeds buffer for dummy runs (contents unused)."""
62
+ return self.encoder_runner.inputs_embeds[:num_tokens]
63
+
64
def get_mm_embeddings(
65
self,
66
scheduled_encoder_inputs: dict[str, list[int]],
vllm/v1/worker/gpu/model_states/interface.py
+6
-2
@@ -109,17 +109,21 @@ class ModelState(ABC):
109
) -> torch.Tensor | None:
110
raise NotImplementedError
111
112
+ def dummy_inputs_embeds(self, num_tokens: int) -> torch.Tensor | None:
113
+ """Pre-allocated inputs_embeds buffer for dummy runs (contents unused)."""
114
+ return None
115
+
116
def gather_mm_embeddings(
117
self, input_batch: InputBatch, draft_lookahead: int = 0
118
) -> tuple[list[torch.Tensor], torch.Tensor]:
115
- """Gather cached multimodal embeddings for a speculator's draft forward."""
119
+ """Gather cached multimodal embeddings."""
120
return self.encoder_runner.gather_mm_embeddings(
121
input_batch.req_ids,
122
input_batch.num_tokens,
123
input_batch.num_scheduled_tokens,
124
input_batch.query_start_loc_np,
125
input_batch.prefill_len_np,
122
- input_batch.num_computed_prefill_tokens_np,
126
+ input_batch.num_computed_tokens_np,
127
draft_lookahead=draft_lookahead,
128
)
129