| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Model runner for instances that only run the multi-modal encoder. |
| 4 | |
| 5 | An encoder-only instance -- `--mm-encoder-only`, or the producer side of |
| 6 | encoder-cache disaggregation -- encodes the multi-modal items and publishes the |
| 7 | embeddings. It runs no language model: no KV cache, sampler, or decoder graphs. |
| 8 | """ |
| 9 | |
| 10 | import time |
| 11 | from collections.abc import Iterator |
| 12 | from contextlib import contextmanager |
| 13 | from typing import TYPE_CHECKING, Any |
| 14 | |
| 15 | import torch |
| 16 | |
| 17 | from vllm.config import VllmConfig |
| 18 | from vllm.config.compilation import CUDAGraphMode |
| 19 | from vllm.logger import init_logger |
| 20 | from vllm.sequence import IntermediateTensors |
| 21 | from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture |
| 22 | from vllm.v1.kv_cache_interface import KVCacheSpec |
| 23 | from vllm.v1.outputs import ( |
| 24 | ModelRunnerOutput, |
| 25 | make_empty_encoder_model_runner_output, |
| 26 | ) |
| 27 | from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor |
| 28 | from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras |
| 29 | from vllm.v1.worker.gpu.model_runner import GPUModelRunner |
| 30 | from vllm.v1.worker.workspace import lock_workspace |
| 31 | |
| 32 | if TYPE_CHECKING: |
| 33 | from vllm.v1.core.sched.output import SchedulerOutput |
| 34 | |
| 35 | logger = init_logger(__name__) |
| 36 | |
| 37 | |
| 38 | class MMEncoderModelRunner(GPUModelRunner): |
| 39 | """Encoder-only variant of the V2 GPU model runner.""" |
| 40 | |
| 41 | def __init__(self, vllm_config: VllmConfig, device: torch.device): |
| 42 | super().__init__(vllm_config, device) |
| 43 | assert self.supports_mm_inputs, ( |
| 44 | "An encoder-only instance must serve a multi-modal model." |
| 45 | ) |
| 46 | assert self.dp_size == 1, "An encoder-only instance does not support DP." |
| 47 | |
| 48 | depth = vllm_config.max_concurrent_batches |
| 49 | # Blocking (sleep) events: busy-polling the driver lock can make this |
| 50 | # rank a straggler under contention. |
| 51 | self._input_events = [torch.Event(blocking=True) for _ in range(depth)] |
| 52 | self._input_event_idx = 0 |
| 53 | |
| 54 | @contextmanager |
| 55 | def input_tensor_semaphore(self) -> Iterator[None]: |
| 56 | """Guard the host writes into the pooled input buffers. |
| 57 | |
| 58 | `UvaBufferPool` recycles a slot every `max_concurrent_batches` steps and |
| 59 | the device reads the pooled host buffers in place. A sampling step is |
| 60 | ordered by the device wait in `AsyncOutput.get_output()`; this one never |
| 61 | waits, so entering blocks until the step that last wrote these slots is |
| 62 | done with them and leaving records this step's last write. |
| 63 | """ |
| 64 | idx = self._input_event_idx |
| 65 | input_event = self._input_events[idx] |
| 66 | input_event.synchronize() |
| 67 | try: |
| 68 | yield |
| 69 | finally: |
| 70 | input_event.record() |
| 71 | self._input_event_idx = (idx + 1) % len(self._input_events) |
| 72 | |
| 73 | def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: |
| 74 | return {} |
| 75 | |
| 76 | @torch.inference_mode() |
| 77 | def capture_model(self, *, profile_only: bool = False) -> int: |
| 78 | encoder_runner = self.model_state.encoder_runner |
| 79 | if not encoder_runner.has_cudagraph(): |
| 80 | return 0 |
| 81 | |
| 82 | start_time = time.perf_counter() |
| 83 | with freeze_gc_for_cudagraph_capture(): |
| 84 | torch.accelerator.empty_cache() |
| 85 | start_free_memory = torch.accelerator.get_memory_info()[0] |
| 86 | encoder_runner.capture() |
| 87 | graph_memory = start_free_memory - torch.accelerator.get_memory_info()[0] |
| 88 | |
| 89 | if not profile_only: |
| 90 | lock_workspace() |
| 91 | |
| 92 | manager = encoder_runner.cudagraph_manager |
| 93 | assert manager is not None |
| 94 | logger.info( |
| 95 | "Encoder graph capturing finished in %.0f secs, took %.2f GiB (%d graphs)", |
| 96 | time.perf_counter() - start_time, |
| 97 | graph_memory / (1 << 30), |
| 98 | manager.get_num_graphs_to_capture(), |
| 99 | ) |
| 100 | return graph_memory |
| 101 | |
| 102 | def needs_cudagraph_capture(self) -> bool: |
| 103 | return False |
| 104 | |
| 105 | def _dummy_run( |
| 106 | self, *args: Any, **kwargs: Any |
| 107 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 108 | empty = torch.empty(0, device=self.device) |
| 109 | return empty, empty |
| 110 | |
| 111 | def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None: |
| 112 | return |
| 113 | |
| 114 | def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None: |
| 115 | return |
| 116 | |
| 117 | def _no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput: |
| 118 | return self._merge_ec_connector_no_forward( |
| 119 | scheduler_output, self.kv_connector.no_forward(scheduler_output) |
| 120 | ) |
| 121 | |
| 122 | @torch.inference_mode() |
| 123 | def execute_model( |
| 124 | self, |
| 125 | scheduler_output: "SchedulerOutput", |
| 126 | intermediate_tensors: IntermediateTensors | None = None, |
| 127 | dummy_run: bool = False, |
| 128 | skip_attn_for_dummy_run: bool = False, |
| 129 | is_profile: bool = False, |
| 130 | context_len: int = 0, |
| 131 | ) -> ModelRunnerOutput: |
| 132 | assert not dummy_run, "An encoder-only instance runs no dummy batch." |
| 133 | |
| 134 | with self.input_tensor_semaphore(): |
| 135 | self.update_pp_decode_requests() |
| 136 | self.finish_requests(scheduler_output) |
| 137 | self.free_states(scheduler_output) |
| 138 | self.add_requests(scheduler_output) |
| 139 | self.update_requests(scheduler_output) |
| 140 | self.block_tables.apply_staged_writes() |
| 141 | if scheduler_output.total_num_scheduled_tokens == 0: |
| 142 | return self._no_forward(scheduler_output) |
| 143 | |
| 144 | batch_req_state, _ = self.gather_batch_req_state(scheduler_output, False) |
| 145 | assert batch_req_state is not None |
| 146 | # No decoder graph or DP padding; encoder graphs run independently. |
| 147 | batch_desc = BatchExecutionDescriptor( |
| 148 | cg_mode=CUDAGraphMode.NONE, |
| 149 | num_tokens=batch_req_state.num_tokens, |
| 150 | num_reqs=None, |
| 151 | ) |
| 152 | self.prepare_inputs( |
| 153 | scheduler_output, batch_req_state, batch_desc, num_active_loras=0 |
| 154 | ) |
| 155 | |
| 156 | scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs |
| 157 | if self.lora_config is not None: |
| 158 | set_active_mm_loras( |
| 159 | model=self.model, |
| 160 | lora_manager=self.lora_manager, |
| 161 | encoder_cache=self.encoder_cache, |
| 162 | req_id_to_index=self.req_states.req_id_to_index, |
| 163 | lora_state=self.lora_state, |
| 164 | scheduled_encoder_inputs=scheduled_encoder_inputs, |
| 165 | ) |
| 166 | |
| 167 | with self.ec_connector.maybe_get_output( |
| 168 | scheduler_output |
| 169 | ) as ec_connector_output: |
| 170 | self.model_state.execute_mm_encoder(scheduled_encoder_inputs) |
| 171 | |
| 172 | return ModelRunnerOutput.with_ec_conn_output( |
| 173 | make_empty_encoder_model_runner_output(scheduler_output), |
| 174 | ec_connector_output, |
| 175 | ) |