[Bugfix][EPD][Model Runner V2] Skip gather mm embeddings for encoder only instance (#51222)

Signed-off-by: Tianyu Guo <[email protected]>

tests/v1/core/test_scheduler.py

@@ -5832,3 +5832,34 @@ def test_encoder_instance_finishes_request_once_prompt_is_consumed():
     assert request.status == RequestStatus.FINISHED_STOPPED
     # The encoder instance publishes an embedding, not tokens.
     assert request.num_output_tokens == 0
+
+
+@pytest.mark.parametrize("ec_role", ["ec_producer", "ec_consumer"])
+def test_encoder_input_skipped_when_connector_already_has_the_item(ec_role: str):
+    """Neither role re-encodes what the connector already holds.
+
+    For a consumer the item is loaded; for a producer there is nothing left to
+    do at all -- it published that embedding earlier (or a sibling encoder did),
+    so a second ViT pass would be pure waste. Reached on any repeat: a second
+    chat turn re-sending its image, a sibling encoder behind the proxy's
+    round-robin, or a restart that kept the shared storage.
+
+    Pinned because the obvious "fix" for the encoder-instance crash this used to
+    cause is to make the producer encode anyway; the crash belongs to the worker
+    (an encoder instance must not gather embeddings it never needed), and paying
+    for it here would cost every deployment a redundant encode.
+    """
+    scheduler = create_scheduler(
+        max_num_seqs=8,
+        max_num_batched_tokens=1024,
+        use_ec_connector=True,
+        ec_role=ec_role,
+    )
+    request = _make_encoder_instance_request(scheduler)
+    req_id = request.request_id
+    scheduler.ec_connector.has_cache_item = lambda *a, **k: True
+
+    output = scheduler.schedule()
+
+    assert output.num_scheduled_tokens[req_id] > 0
+    assert not output.scheduled_encoder_inputs.get(req_id)

tests/v1/worker/test_encoder_runner.py

@@ -9,6 +9,8 @@ and tolerated (token-embedding fallback) when it is not, while a miss within
 the processed range still fails loudly.
 """
 
+from unittest.mock import MagicMock
+
 import numpy as np
 import pytest
 import torch
@@ -16,6 +18,7 @@ import torch
 from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
+from vllm.v1.worker.gpu.model_states.interface import ModelState
 
 pytestmark = pytest.mark.cpu_test
 
@@ -181,3 +184,43 @@ def test_gather_preserves_mixed_modalities():
     assert len(mm_embeds) == 2
     assert [e.modality for e in mm_embeds] == ["video", "audio"]
     assert int(is_mm_embed.sum()) == 8
+
+
+def test_execute_mm_encoder_caches_outputs_without_gathering():
+    """An encoder instance encodes and publishes, and must stop there.
+
+    `ModelState.execute_mm_encoder` is the half of `get_mm_embeddings` that an
+    EPD encoder instance needs: it runs no language model, so gathering would
+    build an `inputs_embeds` nobody reads -- and the gather raises
+    `Encoder cache miss` for any scheduled item absent from the local cache,
+    which on a producer takes the whole engine down (the scheduler hands it
+    items the connector already holds, and a producer has no load path).
+    """
+    cache = EncoderCache()
+    state = MagicMock()
+    state.encoder_cache = cache
+    embedding = torch.ones(2, HIDDEN)
+    # (mm_hashes, [(modality, kwargs item), ...]), as prepare_mm_inputs returns.
+    state.encoder_runner.prepare_mm_inputs.return_value = (
+        ["hash0"],
+        [("image", MagicMock())],
+    )
+    state.encoder_runner.execute_mm_encoder.return_value = [embedding]
+
+    ModelState.execute_mm_encoder(state, {"req0": [0]})
+
+    assert cache.encoder_outputs == {"hash0": embedding}
+    state.encoder_runner.gather_mm_embeddings.assert_not_called()
+
+
+def test_execute_mm_encoder_is_a_noop_without_scheduled_items():
+    """A step that schedules no encoder input must not touch the encoder."""
+    cache = EncoderCache()
+    state = MagicMock()
+    state.encoder_cache = cache
+    state.encoder_runner.prepare_mm_inputs.return_value = ([], [])
+
+    ModelState.execute_mm_encoder(state, {})
+
+    assert not cache.encoder_outputs
+    state.encoder_runner.execute_mm_encoder.assert_not_called()

vllm/v1/worker/gpu/model_runner.py

@@ -1380,9 +1380,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 with self.ec_connector.maybe_get_output(
                     scheduler_output
                 ) as ec_connector_output:
-                    inputs_embeds = self.model_state.get_mm_embeddings(
-                        scheduled_encoder_inputs, input_batch, self.req_states
-                    )
+                    if self.is_encoder_only:
+                        # Encode and publish, nothing else: this instance runs no
+                        # language model, so the gather inside get_mm_embeddings
+                        # would build an inputs_embeds nobody reads -- and it
+                        # raises "Encoder cache miss" for any scheduled item this
+                        # instance did not encode, taking the engine down with it.
+                        self.model_state.execute_mm_encoder(scheduled_encoder_inputs)
+                    else:
+                        inputs_embeds = self.model_state.get_mm_embeddings(
+                            scheduled_encoder_inputs, input_batch, self.req_states
+                        )
             if inputs_embeds is not None and not self.model.requires_raw_input_tokens:
                 input_ids = None
 

vllm/v1/worker/gpu/model_states/default.py

@@ -70,14 +70,7 @@ class DefaultModelState(ModelState):
         input_batch: InputBatch,
         req_states: RequestState,
     ) -> torch.Tensor:
-        mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs(
-            scheduled_encoder_inputs
-        )
-        if mm_kwargs:
-            # Execute the multimodal encoder.
-            encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)
-            # Cache the encoder outputs by mm_hash
-            self.encoder_cache.encoder_outputs.update(zip(mm_hashes, encoder_outputs))
+        self.execute_mm_encoder(scheduled_encoder_inputs)
 
         mm_embeds, is_mm_embed = super().gather_mm_embeddings(input_batch)
         if self.mm_pruner is not None and mm_embeds:

vllm/v1/worker/gpu/model_states/interface.py

@@ -138,6 +138,21 @@ class ModelState(ABC):
         """Pre-allocated inputs_embeds buffer for dummy runs (contents unused)."""
         return None
 
+    def execute_mm_encoder(
+        self, scheduled_encoder_inputs: dict[str, list[int]]
+    ) -> None:
+        """Run the multi-modal encoder and cache its outputs by `mm_hash`.
+
+        The encode half of `get_mm_embeddings`, without the gather, for callers
+        that run no language model.
+        """
+        mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs(
+            scheduled_encoder_inputs
+        )
+        if mm_kwargs:
+            encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)
+            self.encoder_cache.encoder_outputs.update(zip(mm_hashes, encoder_outputs))
+
     def gather_mm_embeddings(
         self, input_batch: InputBatch, draft_lookahead: int = 0
     ) -> tuple[list[torch.Tensor], torch.Tensor]: