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

Signed-off-by: Tianyu Guo <guoty@inferact.ai>

tests/v1/worker/test_gpu_warmup_blocks.py

@@ -77,7 +77,6 @@ def _make_runner(
         num_speculative_steps=num_spec_steps,
         decode_query_len=num_spec_steps + 1,
         is_pooling_model=False,
-        is_encoder_only=False,
         is_encoder_decoder=False,
         is_last_pp_rank=True,
         max_num_reqs=4,
@@ -88,7 +87,9 @@ def _make_runner(
         kv_cache_config=SimpleNamespace(
             kv_cache_groups=kv_cache_groups, num_blocks=1024
         ),
-        vllm_config=SimpleNamespace(num_lookahead_tokens=num_lookahead_tokens),
+        vllm_config=SimpleNamespace(
+            num_lookahead_tokens=num_lookahead_tokens, is_mm_encoder_only=False
+        ),
         kv_block_zeroer=None,
         kv_connector=SimpleNamespace(set_disabled=lambda disabled: None),
     )

vllm/config/vllm.py

@@ -572,7 +572,7 @@ class VllmConfig:
         )
 
     @property
-    def is_encoder_only(self) -> bool:
+    def is_mm_encoder_only(self) -> bool:
         mm_config = (
             self.model_config.multimodal_config
             if self.model_config is not None
@@ -1782,6 +1782,16 @@ class VllmConfig:
         # before the HMA check below, which inspects the connector class.
         self._post_init_kv_transfer_config()
 
+        if self.is_mm_encoder_only and self.cache_config.enable_prefix_caching:
+            # Such an 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 runs the "
+                "multi-modal 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/core/sched/scheduler.py

@@ -102,7 +102,7 @@ class Scheduler(SchedulerInterface):
             )
         self.structured_output_manager = structured_output_manager
         self.is_encoder_decoder = vllm_config.model_config.is_encoder_decoder
-        self.is_encoder_only = vllm_config.is_encoder_only
+        self.is_mm_encoder_only = vllm_config.is_mm_encoder_only
 
         # include_finished_set controls whether a separate set of finished
         # request ids should be included in the EngineCoreOutputs returned
@@ -1904,7 +1904,7 @@ class Scheduler(SchedulerInterface):
                 request.status = RequestStatus.FINISHED_STOPPED
                 stopped = True
             elif (
-                self.is_encoder_only
+                self.is_mm_encoder_only
                 and request.num_computed_tokens >= request.num_prompt_tokens
             ):
                 # An encoder instance runs the encoder and publishes the

vllm/v1/worker/gpu/model_runner.py

@@ -69,7 +69,6 @@ from vllm.v1.outputs import (
     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 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
         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 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         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:
@@ -646,9 +642,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         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."
@@ -809,12 +802,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                     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
         )
@@ -851,9 +838,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
     @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
@@ -1575,26 +1559,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 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,7 +209,7 @@ def warmup_kernels(
     We must call the provided worker's execute_model for pipeline parallel
     coordination.
     """
-    if model_runner.is_encoder_only:
+    if model_runner.vllm_config.is_mm_encoder_only:
         return
 
     num_spec_steps = model_runner.num_speculative_steps

vllm/v1/worker/gpu_worker.py

@@ -423,9 +423,14 @@ class Worker(WorkerBase):
 
         # 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_mm_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

vllm/v1/worker/mm_encoder_model_runner.py

@@ -0,0 +1,141 @@
+# 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 collections.abc import Iterator
+from contextlib import contextmanager
+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.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 = vllm_config.max_concurrent_batches
+        # Blocking (sleep) events: busy-polling the driver lock can make this
+        # rank a straggler under contention.
+        self._input_events = [torch.Event(blocking=True) for _ in range(depth)]
+        self._input_event_idx = 0
+
+    @contextmanager
+    def input_tensor_semaphore(self) -> Iterator[None]:
+        """Guard the host writes into the pooled input buffers.
+
+        `UvaBufferPool` recycles a slot every `max_concurrent_batches` 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 entering blocks until the step that last wrote these slots is
+        done with them and leaving records this step's last write.
+        """
+        idx = self._input_event_idx
+        input_event = self._input_events[idx]
+        input_event.synchronize()
+        try:
+            yield
+        finally:
+            input_event.record()
+            self._input_event_idx = (idx + 1) % len(self._input_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."
+
+        with self.input_tensor_semaphore():
+            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:
+                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.
+            batch_desc = BatchExecutionDescriptor(
+                cg_mode=CUDAGraphMode.NONE,
+                num_tokens=batch_req_state.num_tokens,
+                num_reqs=None,
+            )
+            self.prepare_inputs(scheduler_output, batch_req_state, batch_desc)
+
+        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,
+            )
+
+        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,
+        )