[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/buffer_utils.py

@@ -23,6 +23,10 @@ def set_default_max_concurrency(n: int) -> None:
     _DEFAULT_MAX_CONCURRENCY = max(2, n)
 
 
+def get_default_max_concurrency() -> int:
+    return _DEFAULT_MAX_CONCURRENCY
+
+
 def async_copy_to_gpu(
     x: torch.Tensor | np.ndarray,
     out: torch.Tensor | None = None,

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,143 @@
+# 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. It runs no language model: no KV cache, no sampler, no CUDA graphs.
+"""
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.config.compilation import CUDAGraphMode
+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.buffer_utils import get_default_max_concurrency
+from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
+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."
+        )
+        assert self.dp_size == 1, "An encoder-only instance does not support DP."
+
+        depth = get_default_max_concurrency()
+        # `UvaBufferPool` recycles a slot every `depth` steps and the device
+        # reads the pooled host buffers in place. A sampling step is ordered by
+        # the device wait in `AsyncOutput.get_output()`; this one never waits, so
+        # it keeps one event per slot generation. Blocking: no driver busy-poll.
+        self._input_reuse_events = [torch.Event(blocking=True) for _ in range(depth)]
+        self._input_reuse_idx = 0
+
+    def _wait_for_input_reuse(self) -> None:
+        """Wait for the step that last wrote the slots this step will reuse."""
+        self._input_reuse_events[self._input_reuse_idx].synchronize()
+
+    def _mark_input_reuse(self) -> None:
+        """Mark this step's last host write into the pooled input buffers."""
+        self._input_reuse_events[self._input_reuse_idx].record()
+        self._input_reuse_idx = (self._input_reuse_idx + 1) % len(
+            self._input_reuse_events
+        )
+
+    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, _ = self.gather_batch_req_state(scheduler_output, False)
+        assert batch_req_state is not None
+        # No CUDA graph, and no DP peer to agree a padded shape with.
+        self.prepare_inputs(
+            scheduler_output,
+            batch_req_state,
+            BatchExecutionDescriptor(
+                cg_mode=CUDAGraphMode.NONE,
+                num_tokens=batch_req_state.num_tokens,
+                num_reqs=None,
+            ),
+        )
+        # Before the encoder, not after: it reads no pooled metadata, and
+        # waiting on it costs ~22% throughput.
+        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,
+            )
+
+        # `prepare_inputs_embeds` would build an inputs_embeds nobody reads and
+        # raise "Encoder cache miss" for items 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)
+
+        return ModelRunnerOutput.with_ec_conn_output(
+            make_empty_encoder_model_runner_output(scheduler_output),
+            ec_connector_output,
+        )