diff --git a/tests/v1/worker/test_gpu_warmup_blocks.py b/tests/v1/worker/test_gpu_warmup_blocks.py index 7b3c5b011bbd..11be99f2089f 100644 --- a/tests/v1/worker/test_gpu_warmup_blocks.py +++ b/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), ) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index a3e027617685..acde48a931a2 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -572,7 +572,7 @@ def is_ec_producer_only(self) -> bool: ) @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 @@ def has_blocked_weights(): # 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 diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 26674154c648..b41727f273cc 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -102,7 +102,7 @@ def __init__( ) 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 @@ def update_from_output( 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 diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 6d8a28933f53..07e080a5bb9f 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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: @@ -646,9 +642,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." @@ -809,12 +802,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 ) @@ -851,9 +838,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 @@ -1575,26 +1559,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, diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 9d0085233743..e3c5e11e034d 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/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 diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index c35e22643f52..05c03af75eac 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -423,9 +423,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_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 diff --git a/vllm/v1/worker/mm_encoder_model_runner.py b/vllm/v1/worker/mm_encoder_model_runner.py new file mode 100644 index 000000000000..2a68ab37a937 --- /dev/null +++ b/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, + )