[Bugfix][EPD][Model Runner V2] Skip gather mm embeddings for encoder only instance (#51222)
Signed-off-by: Tianyu Guo <guoty@inferact.ai>
Tianyu Guo committed
Aug 7, 2026 at 11:40 UTC
dd856e48bbf969e3f0e561e8c76f6e92c76e0795
5 files changed
+101
-11
tests/v1/core/test_scheduler.py
+31
@@ -5832,3 +5832,34 @@ def test_encoder_instance_finishes_request_once_prompt_is_consumed():
5832
assert request.status == RequestStatus.FINISHED_STOPPED
5833
# The encoder instance publishes an embedding, not tokens.
5834
assert request.num_output_tokens == 0
5835
+
5836
+
5837
+@pytest.mark.parametrize("ec_role", ["ec_producer", "ec_consumer"])
5838
+def test_encoder_input_skipped_when_connector_already_has_the_item(ec_role: str):
5839
+ """Neither role re-encodes what the connector already holds.
5840
+
5841
+ For a consumer the item is loaded; for a producer there is nothing left to
5842
+ do at all -- it published that embedding earlier (or a sibling encoder did),
5843
+ so a second ViT pass would be pure waste. Reached on any repeat: a second
5844
+ chat turn re-sending its image, a sibling encoder behind the proxy's
5845
+ round-robin, or a restart that kept the shared storage.
5846
+
5847
+ Pinned because the obvious "fix" for the encoder-instance crash this used to
5848
+ cause is to make the producer encode anyway; the crash belongs to the worker
5849
+ (an encoder instance must not gather embeddings it never needed), and paying
5850
+ for it here would cost every deployment a redundant encode.
5851
+ """
5852
+ scheduler = create_scheduler(
5853
+ max_num_seqs=8,
5854
+ max_num_batched_tokens=1024,
5855
+ use_ec_connector=True,
5856
+ ec_role=ec_role,
5857
+ )
5858
+ request = _make_encoder_instance_request(scheduler)
5859
+ req_id = request.request_id
5860
+ scheduler.ec_connector.has_cache_item = lambda *a, **k: True
5861
+
5862
+ output = scheduler.schedule()
5863
+
5864
+ assert output.num_scheduled_tokens[req_id] > 0
5865
+ assert not output.scheduled_encoder_inputs.get(req_id)
tests/v1/worker/test_encoder_runner.py
+43
@@ -9,6 +9,8 @@ and tolerated (token-embedding fallback) when it is not, while a miss within
9
the processed range still fails loudly.
10
"""
11
12
+from unittest.mock import MagicMock
13
+
14
import numpy as np
15
import pytest
16
import torch
@@ -16,6 +18,7 @@ import torch
18
from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
19
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
20
from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
21
+from vllm.v1.worker.gpu.model_states.interface import ModelState
22
23
pytestmark = pytest.mark.cpu_test
24
@@ -181,3 +184,43 @@ def test_gather_preserves_mixed_modalities():
184
assert len(mm_embeds) == 2
185
assert [e.modality for e in mm_embeds] == ["video", "audio"]
186
assert int(is_mm_embed.sum()) == 8
187
+
188
+
189
+def test_execute_mm_encoder_caches_outputs_without_gathering():
190
+ """An encoder instance encodes and publishes, and must stop there.
191
+
192
+ `ModelState.execute_mm_encoder` is the half of `get_mm_embeddings` that an
193
+ EPD encoder instance needs: it runs no language model, so gathering would
194
+ build an `inputs_embeds` nobody reads -- and the gather raises
195
+ `Encoder cache miss` for any scheduled item absent from the local cache,
196
+ which on a producer takes the whole engine down (the scheduler hands it
197
+ items the connector already holds, and a producer has no load path).
198
+ """
199
+ cache = EncoderCache()
200
+ state = MagicMock()
201
+ state.encoder_cache = cache
202
+ embedding = torch.ones(2, HIDDEN)
203
+ # (mm_hashes, [(modality, kwargs item), ...]), as prepare_mm_inputs returns.
204
+ state.encoder_runner.prepare_mm_inputs.return_value = (
205
+ ["hash0"],
206
+ [("image", MagicMock())],
207
+ )
208
+ state.encoder_runner.execute_mm_encoder.return_value = [embedding]
209
+
210
+ ModelState.execute_mm_encoder(state, {"req0": [0]})
211
+
212
+ assert cache.encoder_outputs == {"hash0": embedding}
213
+ state.encoder_runner.gather_mm_embeddings.assert_not_called()
214
+
215
+
216
+def test_execute_mm_encoder_is_a_noop_without_scheduled_items():
217
+ """A step that schedules no encoder input must not touch the encoder."""
218
+ cache = EncoderCache()
219
+ state = MagicMock()
220
+ state.encoder_cache = cache
221
+ state.encoder_runner.prepare_mm_inputs.return_value = ([], [])
222
+
223
+ ModelState.execute_mm_encoder(state, {})
224
+
225
+ assert not cache.encoder_outputs
226
+ state.encoder_runner.execute_mm_encoder.assert_not_called()
vllm/v1/worker/gpu/model_runner.py
+11
-3
@@ -1380,9 +1380,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1380
with self.ec_connector.maybe_get_output(
1381
scheduler_output
1382
) as ec_connector_output:
1383
- inputs_embeds = self.model_state.get_mm_embeddings(
1384
- scheduled_encoder_inputs, input_batch, self.req_states
1385
- )
1383
+ if self.is_encoder_only:
1384
+ # Encode and publish, nothing else: this instance runs no
1385
+ # language model, so the gather inside get_mm_embeddings
1386
+ # would build an inputs_embeds nobody reads -- and it
1387
+ # raises "Encoder cache miss" for any scheduled item this
1388
+ # instance did not encode, taking the engine down with it.
1389
+ self.model_state.execute_mm_encoder(scheduled_encoder_inputs)
1390
+ else:
1391
+ inputs_embeds = self.model_state.get_mm_embeddings(
1392
+ scheduled_encoder_inputs, input_batch, self.req_states
1393
+ )
1394
if inputs_embeds is not None and not self.model.requires_raw_input_tokens:
1395
input_ids = None
1396
vllm/v1/worker/gpu/model_states/default.py
+1
-8
@@ -70,14 +70,7 @@ class DefaultModelState(ModelState):
70
input_batch: InputBatch,
71
req_states: RequestState,
72
) -> torch.Tensor:
73
- mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs(
74
- scheduled_encoder_inputs
75
- )
76
- if mm_kwargs:
77
- # Execute the multimodal encoder.
78
- encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)
79
- # Cache the encoder outputs by mm_hash
80
- self.encoder_cache.encoder_outputs.update(zip(mm_hashes, encoder_outputs))
73
+ self.execute_mm_encoder(scheduled_encoder_inputs)
74
75
mm_embeds, is_mm_embed = super().gather_mm_embeddings(input_batch)
76
if self.mm_pruner is not None and mm_embeds:
vllm/v1/worker/gpu/model_states/interface.py
+15
@@ -138,6 +138,21 @@ class ModelState(ABC):
138
"""Pre-allocated inputs_embeds buffer for dummy runs (contents unused)."""
139
return None
140
141
+ def execute_mm_encoder(
142
+ self, scheduled_encoder_inputs: dict[str, list[int]]
143
+ ) -> None:
144
+ """Run the multi-modal encoder and cache its outputs by `mm_hash`.
145
+
146
+ The encode half of `get_mm_embeddings`, without the gather, for callers
147
+ that run no language model.
148
+ """
149
+ mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs(
150
+ scheduled_encoder_inputs
151
+ )
152
+ if mm_kwargs:
153
+ encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)
154
+ self.encoder_cache.encoder_outputs.update(zip(mm_hashes, encoder_outputs))
155
+
156
def gather_mm_embeddings(
157
self, input_batch: InputBatch, draft_lookahead: int = 0
158
) -> tuple[list[torch.Tensor], torch.Tensor]: