[Refactor][Model Runner V2][Multimodal] Move the encoder-only path out of the shared runner

vllm/config/vllm.py

@@ -1752,6 +1752,13 @@ def has_blocked_weights():
         # before the HMA check below, which inspects the connector class.
         self._post_init_kv_transfer_config()
 
+        if self.is_encoder_only and self.cache_config.enable_prefix_caching:
+            # An encoder-only instance publishes encoder embeddings and runs no
+            # language model, so it holds no KV cache for prefix caching to
+            # reuse and its coordinator would have no group to manage.
+            logger.info("Disabling prefix caching: this instance is encoder-only.")
+            self.cache_config.enable_prefix_caching = False
+
         # Hybrid KV cache manager (HMA) runtime rules:
         # - Explicit enable (--no-disable-kv-cache-manager): error if runtime
         #   disables it

vllm/v1/worker/gpu/model_runner.py

@@ -69,7 +69,6 @@
     ECConnectorOutput,
     ModelRunnerOutput,
     RoutedExpertsTensors,
-    make_empty_encoder_model_runner_output,
 )
 from vllm.v1.worker.block_table import get_block_table_width
 from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
@@ -175,7 +174,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
 
         self.device = device
         self.dtype = self.model_config.dtype
-        self.is_encoder_only = vllm_config.is_encoder_only
         self.kv_cache_dtype = self.dtype
         if self.cache_config.cache_dtype != "auto":
             # Quantized KV cache.
@@ -499,8 +497,6 @@ def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]:
         return encoder_runner.get_encoder_timing_stats()
 
     def get_kv_cache_spec(self):
-        if self.is_encoder_only:
-            return {}
         return get_kv_cache_spec(self.vllm_config)
 
     def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
@@ -649,9 +645,6 @@ def _dummy_run(
         is_profile: bool = False,
         **kwargs,
     ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
-        if self.is_encoder_only:
-            empty = torch.empty(0, device=self.device)
-            return empty, empty
         if skip_attn and not is_profile:
             raise ValueError(
                 "skip_attn must only be True for initial memory profiling."
@@ -812,12 +805,6 @@ def profile_run(self) -> None:
                     dummy_mm_inputs, mm_budget
                 )
 
-        if self.is_encoder_only:
-            torch.accelerator.synchronize()
-            self.reset_encoder_cache()
-            gc.collect()
-            return
-
         hidden_states, sample_hidden_states = self._dummy_run(
             self.max_num_tokens, skip_attn=True, is_profile=True
         )
@@ -854,9 +841,6 @@ def profile_cudagraph_memory(self) -> int:
 
     @torch.inference_mode()
     def capture_model(self) -> int:
-        if self.is_encoder_only:
-            return 0
-
         assert self.cudagraph_manager is not None
         capture_encoder = (
             self.model_state.supports_mm_inputs
@@ -1578,26 +1562,12 @@ def execute_model(
                 with self.ec_connector.maybe_get_output(
                     scheduler_output
                 ) as ec_connector_output:
-                    if self.is_encoder_only:
-                        # Encode and publish, nothing else: this instance runs no
-                        # language model, so the gather inside prepare_inputs_embeds
-                        # 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.prepare_inputs_embeds(
-                            scheduled_encoder_inputs, input_batch, self.req_states
-                        )
+                    inputs_embeds = self.model_state.prepare_inputs_embeds(
+                        scheduled_encoder_inputs, input_batch, self.req_states
+                    )
             if inputs_embeds is not None and not requires_raw_input_tokens(self.model):
                 input_ids = None
 
-        if self.is_encoder_only:
-            return ModelRunnerOutput.with_ec_conn_output(
-                make_empty_encoder_model_runner_output(scheduler_output),
-                ec_connector_output,
-            )
-
         model_inputs = {
             "input_ids": input_ids,
             "positions": input_batch.positions,

vllm/v1/worker/gpu/warmup.py

@@ -209,9 +209,6 @@ def warmup_kernels(
     We must call the provided worker's execute_model for pipeline parallel
     coordination.
     """
-    if model_runner.is_encoder_only:
-        return
-
     num_spec_steps = model_runner.num_speculative_steps
     decode_query_len = model_runner.decode_query_len
     # Use decode_query_len + 1 tokens so the prefill batch's per-request query

vllm/v1/worker/gpu_worker.py

@@ -422,9 +422,14 @@ def init_device(self):
 
         # Construct the model runner
         if self.use_v2_model_runner:
-            from vllm.v1.worker.gpu.model_runner import (
-                GPUModelRunner as GPUModelRunnerV2,
-            )
+            if self.vllm_config.is_encoder_only:
+                from vllm.v1.worker.mm_encoder_model_runner import (
+                    MMEncoderModelRunner as GPUModelRunnerV2,
+                )
+            else:
+                from vllm.v1.worker.gpu.model_runner import (  # type: ignore[assignment]
+                    GPUModelRunner as GPUModelRunnerV2,
+                )
 
             # HACK(woosuk): This is a temporary fix to avoid type errors.
             self.model_runner: GPUModelRunner = GPUModelRunnerV2(  # type: ignore
@@ -805,8 +810,9 @@ def compile_or_warm_up_model(self) -> CompilationTimes:
 
             maybe_save_startup_plan(self, kv_cache_memory_bytes_to_requested_limit)
 
-        if self.use_v2_model_runner:
+        if self.use_v2_model_runner and not self.vllm_config.is_encoder_only:
             # V2: Run full execute_model + sample_tokens to JIT compile triton kernels.
+            # An encoder-only instance runs no language model to warm up.
             warmup_kernels(self.model_runner, self.execute_model, self.sample_tokens)
         elif get_pp_group().is_last_rank:
             # V1: Warm up sampler and preallocate memory buffer for logits and other

vllm/v1/worker/mm_encoder_model_runner.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Model runner for instances that only run the multi-modal encoder.
+
+An encoder-only instance -- `--mm-encoder-only`, or the producer side of
+encoder-cache disaggregation -- encodes the multi-modal items and publishes the
+embeddings for a peer to consume. It runs no language model, holds no KV cache
+and samples no token, so most of a step does not apply to it: no attention
+metadata, no forward pass, no sampler, no CUDA graphs.
+
+Keeping that here rather than as `is_encoder_only` branches inside the shared
+runner keeps the exceptions in one place, and keeps invariants that only hold
+for a full step (see `execute_model` below) from being read as universal.
+"""
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.sequence import IntermediateTensors
+from vllm.v1.kv_cache_interface import KVCacheSpec
+from vllm.v1.outputs import (
+    ModelRunnerOutput,
+    make_empty_encoder_model_runner_output,
+)
+from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
+from vllm.v1.worker.gpu.lora_utils import get_num_active_loras_for_dispatch
+from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras
+from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+
+if TYPE_CHECKING:
+    from vllm.v1.core.sched.output import SchedulerOutput
+
+
+class MMEncoderModelRunner(GPUModelRunner):
+    """Encoder-only variant of the V2 GPU model runner."""
+
+    def __init__(self, vllm_config: VllmConfig, device: torch.device):
+        super().__init__(vllm_config, device)
+        assert self.supports_mm_inputs, (
+            "An encoder-only instance must serve a multi-modal model."
+        )
+
+        # The device reads the pooled host input buffers in place, and they are
+        # recycled every max_concurrent_batches steps. A sampling step waits on
+        # its output copy before its own slot comes round again; this runner
+        # returns a host-built output and never waits on the device (see
+        # `execute_model`), so it is the one that needs an explicit barrier.
+        self._input_reuse_event: torch.Event | None = None
+        if vllm_config.max_concurrent_batches > 1:
+            # Blocking (sleep) event: busy-polling the driver lock can make this
+            # rank a straggler under contention.
+            self._input_reuse_event = torch.Event(blocking=True)
+
+    def _wait_for_input_reuse(self) -> None:
+        """Wait for the step still reading the pooled input buffers."""
+        if self._input_reuse_event is not None:
+            self._input_reuse_event.synchronize()
+
+    def _mark_input_reuse(self) -> None:
+        """Mark this step's last host write into the pooled input buffers."""
+        if self._input_reuse_event is not None:
+            self._input_reuse_event.record()
+
+    def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
+        return {}
+
+    def capture_model(self) -> int:
+        return 0
+
+    def _dummy_run(
+        self, *args: Any, **kwargs: Any
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        empty = torch.empty(0, device=self.device)
+        return empty, empty
+
+    def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
+        return
+
+    def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None:
+        return
+
+    def _no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
+        return self._merge_ec_connector_no_forward(
+            scheduler_output, self.kv_connector.no_forward(scheduler_output)
+        )
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        scheduler_output: "SchedulerOutput",
+        intermediate_tensors: IntermediateTensors | None = None,
+        dummy_run: bool = False,
+        skip_attn_for_dummy_run: bool = False,
+        is_profile: bool = False,
+        context_len: int = 0,
+    ) -> ModelRunnerOutput:
+        assert not dummy_run, "An encoder-only instance runs no dummy batch."
+
+        self._wait_for_input_reuse()
+        self.update_pp_decode_requests()
+        self.finish_requests(scheduler_output)
+        self.free_states(scheduler_output)
+        self.add_requests(scheduler_output)
+        self.update_requests(scheduler_output)
+        self.block_tables.apply_staged_writes()
+        if scheduler_output.total_num_scheduled_tokens == 0:
+            self._mark_input_reuse()
+            return self._no_forward(scheduler_output)
+
+        batch_req_state, uniform_tok_count = self.gather_batch_req_state(
+            scheduler_output, False
+        )
+        assert batch_req_state is not None
+        num_active_loras = 0
+        if self.lora_config:
+            num_active_loras = get_num_active_loras_for_dispatch(
+                self.lora_config,
+                self.lora_state,
+                list(scheduler_output.num_scheduled_tokens.keys()),
+                False,
+            )
+        # This rank runs no compiled graph, but the DP peers size their padding
+        # from the shape agreed here, so it still has to take part.
+        batch_desc, _ = dispatch_cg_and_sync_dp(
+            self.cudagraph_manager,
+            len(scheduler_output.num_scheduled_tokens),
+            batch_req_state.num_tokens,
+            uniform_tok_count,
+            self.dp_size,
+            self.dp_rank,
+            max_query_len=max(scheduler_output.num_scheduled_tokens.values()),
+            num_active_loras=num_active_loras,
+        )
+        if batch_desc.num_tokens == 0:
+            self._mark_input_reuse()
+            return self._no_forward(scheduler_output)
+
+        self.prepare_inputs(scheduler_output, batch_req_state, batch_desc)
+        # Last host write into the pooled buffers. Recorded here rather than at
+        # function exit: after `execute_mm_encoder` the next step's wait would
+        # also have to wait for this step's ViT, which costs ~22% throughput
+        # and protects nothing -- the encoder reads no pooled metadata.
+        self._mark_input_reuse()
+
+        scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
+        if self.lora_config is not None:
+            set_active_mm_loras(
+                model=self.model,
+                lora_manager=self.lora_manager,
+                encoder_cache=self.encoder_cache,
+                req_id_to_index=self.req_states.req_id_to_index,
+                lora_state=self.lora_state,
+                scheduled_encoder_inputs=scheduled_encoder_inputs,
+            )
+
+        # Encode and publish, nothing else. `prepare_inputs_embeds` would build
+        # an inputs_embeds nobody reads, and would raise "Encoder cache miss"
+        # for any scheduled item this instance did not encode.
+        with self.ec_connector.maybe_get_output(
+            scheduler_output
+        ) as ec_connector_output:
+            self.model_state.execute_mm_encoder(scheduled_encoder_inputs)
+
+        # NOTE: This output is built on the host and carries no sampled token,
+        # so unlike a full step it never waits on the device. Anything the
+        # device still reads from a recycled host buffer must be ordered here
+        # explicitly -- a sampling step gets that ordering from its output copy.
+        return ModelRunnerOutput.with_ec_conn_output(
+            make_empty_encoder_model_runner_output(scheduler_output),
+            ec_connector_output,
+        )