[Refactor][Model Runner V2][Multimodal] Move the encoder-only path out of the shared runner (#53176)
Signed-off-by: Tianyu Guo <guoty@inferact.ai>
Tianyu Guo committed
Aug 22, 2026 at 13:11 UTC
7ca49fbe4bab019e55d57cdc4b7fd3d55c67c1a6
7 files changed
+169
-42
tests/v1/worker/test_gpu_warmup_blocks.py
+3
-2
@@ -77,7 +77,6 @@ def _make_runner(
77
num_speculative_steps=num_spec_steps,
78
decode_query_len=num_spec_steps + 1,
79
is_pooling_model=False,
80
- is_encoder_only=False,
80
is_encoder_decoder=False,
81
is_last_pp_rank=True,
82
max_num_reqs=4,
@@ -88,7 +87,9 @@ def _make_runner(
87
kv_cache_config=SimpleNamespace(
88
kv_cache_groups=kv_cache_groups, num_blocks=1024
89
),
91
- vllm_config=SimpleNamespace(num_lookahead_tokens=num_lookahead_tokens),
90
+ vllm_config=SimpleNamespace(
91
+ num_lookahead_tokens=num_lookahead_tokens, is_mm_encoder_only=False
92
+ ),
93
kv_block_zeroer=None,
94
kv_connector=SimpleNamespace(set_disabled=lambda disabled: None),
95
)
vllm/config/vllm.py
+11
-1
@@ -572,7 +572,7 @@ class VllmConfig:
572
)
573
574
@property
575
- def is_encoder_only(self) -> bool:
575
+ def is_mm_encoder_only(self) -> bool:
576
mm_config = (
577
self.model_config.multimodal_config
578
if self.model_config is not None
@@ -1782,6 +1782,16 @@ class VllmConfig:
1782
# before the HMA check below, which inspects the connector class.
1783
self._post_init_kv_transfer_config()
1784
1785
+ if self.is_mm_encoder_only and self.cache_config.enable_prefix_caching:
1786
+ # Such an instance publishes encoder embeddings and runs no language
1787
+ # model, so it holds no KV cache for prefix caching to reuse and its
1788
+ # coordinator would have no group to manage.
1789
+ logger.info(
1790
+ "Disabling prefix caching: this instance runs the "
1791
+ "multi-modal encoder only."
1792
+ )
1793
+ self.cache_config.enable_prefix_caching = False
1794
+
1795
# Hybrid KV cache manager (HMA) runtime rules:
1796
# - Explicit enable (--no-disable-kv-cache-manager): error if runtime
1797
# disables it
vllm/v1/core/sched/scheduler.py
+2
-2
@@ -102,7 +102,7 @@ class Scheduler(SchedulerInterface):
102
)
103
self.structured_output_manager = structured_output_manager
104
self.is_encoder_decoder = vllm_config.model_config.is_encoder_decoder
105
- self.is_encoder_only = vllm_config.is_encoder_only
105
+ self.is_mm_encoder_only = vllm_config.is_mm_encoder_only
106
107
# include_finished_set controls whether a separate set of finished
108
# request ids should be included in the EngineCoreOutputs returned
@@ -1904,7 +1904,7 @@ class Scheduler(SchedulerInterface):
1904
request.status = RequestStatus.FINISHED_STOPPED
1905
stopped = True
1906
elif (
1907
- self.is_encoder_only
1907
+ self.is_mm_encoder_only
1908
and request.num_computed_tokens >= request.num_prompt_tokens
1909
):
1910
# An encoder instance runs the encoder and publishes the
vllm/v1/worker/gpu/model_runner.py
+3
-33
@@ -69,7 +69,6 @@ from vllm.v1.outputs import (
69
ECConnectorOutput,
70
ModelRunnerOutput,
71
RoutedExpertsTensors,
72
- make_empty_encoder_model_runner_output,
72
)
73
from vllm.v1.worker.block_table import get_block_table_width
74
from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
@@ -175,7 +174,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
174
175
self.device = device
176
self.dtype = self.model_config.dtype
178
- self.is_encoder_only = vllm_config.is_encoder_only
177
self.kv_cache_dtype = self.dtype
178
if self.cache_config.cache_dtype != "auto":
179
# Quantized KV cache.
@@ -499,8 +497,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
497
return encoder_runner.get_encoder_timing_stats()
498
499
def get_kv_cache_spec(self):
502
- if self.is_encoder_only:
503
- return {}
500
return get_kv_cache_spec(self.vllm_config)
501
502
def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
@@ -646,9 +642,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
642
is_profile: bool = False,
643
**kwargs,
644
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
649
- if self.is_encoder_only:
650
- empty = torch.empty(0, device=self.device)
651
- return empty, empty
645
if skip_attn and not is_profile:
646
raise ValueError(
647
"skip_attn must only be True for initial memory profiling."
@@ -809,12 +802,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
802
dummy_mm_inputs, mm_budget
803
)
804
812
- if self.is_encoder_only:
813
- torch.accelerator.synchronize()
814
- self.reset_encoder_cache()
815
- gc.collect()
816
- return
817
-
805
hidden_states, sample_hidden_states = self._dummy_run(
806
self.max_num_tokens, skip_attn=True, is_profile=True
807
)
@@ -851,9 +838,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
838
839
@torch.inference_mode()
840
def capture_model(self) -> int:
854
- if self.is_encoder_only:
855
- return 0
856
-
841
assert self.cudagraph_manager is not None
842
capture_encoder = (
843
self.model_state.supports_mm_inputs
@@ -1575,26 +1559,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1559
with self.ec_connector.maybe_get_output(
1560
scheduler_output
1561
) as ec_connector_output:
1578
- if self.is_encoder_only:
1579
- # Encode and publish, nothing else: this instance runs no
1580
- # language model, so the gather inside prepare_inputs_embeds
1581
- # would build an inputs_embeds nobody reads -- and it
1582
- # raises "Encoder cache miss" for any scheduled item this
1583
- # instance did not encode, taking the engine down with it.
1584
- self.model_state.execute_mm_encoder(scheduled_encoder_inputs)
1585
- else:
1586
- inputs_embeds = self.model_state.prepare_inputs_embeds(
1587
- scheduled_encoder_inputs, input_batch, self.req_states
1588
- )
1562
+ inputs_embeds = self.model_state.prepare_inputs_embeds(
1563
+ scheduled_encoder_inputs, input_batch, self.req_states
1564
+ )
1565
if inputs_embeds is not None and not requires_raw_input_tokens(self.model):
1566
input_ids = None
1567
1592
- if self.is_encoder_only:
1593
- return ModelRunnerOutput.with_ec_conn_output(
1594
- make_empty_encoder_model_runner_output(scheduler_output),
1595
- ec_connector_output,
1596
- )
1597
-
1568
model_inputs = {
1569
"input_ids": input_ids,
1570
"positions": input_batch.positions,
vllm/v1/worker/gpu/warmup.py
+1
-1
@@ -209,7 +209,7 @@ def warmup_kernels(
209
We must call the provided worker's execute_model for pipeline parallel
210
coordination.
211
"""
212
- if model_runner.is_encoder_only:
212
+ if model_runner.vllm_config.is_mm_encoder_only:
213
return
214
215
num_spec_steps = model_runner.num_speculative_steps
vllm/v1/worker/gpu_worker.py
+8
-3
@@ -423,9 +423,14 @@ class Worker(WorkerBase):
423
424
# Construct the model runner
425
if self.use_v2_model_runner:
426
- from vllm.v1.worker.gpu.model_runner import (
427
- GPUModelRunner as GPUModelRunnerV2,
428
- )
426
+ if self.vllm_config.is_mm_encoder_only:
427
+ from vllm.v1.worker.mm_encoder_model_runner import (
428
+ MMEncoderModelRunner as GPUModelRunnerV2,
429
+ )
430
+ else:
431
+ from vllm.v1.worker.gpu.model_runner import ( # type: ignore[assignment]
432
+ GPUModelRunner as GPUModelRunnerV2,
433
+ )
434
435
# HACK(woosuk): This is a temporary fix to avoid type errors.
436
self.model_runner: GPUModelRunner = GPUModelRunnerV2( # type: ignore
vllm/v1/worker/mm_encoder_model_runner.py
new
+141
@@ -0,0 +1,141 @@
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, no sampler, no CUDA graphs.
8
+"""
9
+
10
+from collections.abc import Iterator
11
+from contextlib import contextmanager
12
+from typing import TYPE_CHECKING, Any
13
+
14
+import torch
15
+
16
+from vllm.config import VllmConfig
17
+from vllm.config.compilation import CUDAGraphMode
18
+from vllm.sequence import IntermediateTensors
19
+from vllm.v1.kv_cache_interface import KVCacheSpec
20
+from vllm.v1.outputs import (
21
+ ModelRunnerOutput,
22
+ make_empty_encoder_model_runner_output,
23
+)
24
+from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
25
+from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras
26
+from vllm.v1.worker.gpu.model_runner import GPUModelRunner
27
+
28
+if TYPE_CHECKING:
29
+ from vllm.v1.core.sched.output import SchedulerOutput
30
+
31
+
32
+class MMEncoderModelRunner(GPUModelRunner):
33
+ """Encoder-only variant of the V2 GPU model runner."""
34
+
35
+ def __init__(self, vllm_config: VllmConfig, device: torch.device):
36
+ super().__init__(vllm_config, device)
37
+ assert self.supports_mm_inputs, (
38
+ "An encoder-only instance must serve a multi-modal model."
39
+ )
40
+ assert self.dp_size == 1, "An encoder-only instance does not support DP."
41
+
42
+ depth = vllm_config.max_concurrent_batches
43
+ # Blocking (sleep) events: busy-polling the driver lock can make this
44
+ # rank a straggler under contention.
45
+ self._input_events = [torch.Event(blocking=True) for _ in range(depth)]
46
+ self._input_event_idx = 0
47
+
48
+ @contextmanager
49
+ def input_tensor_semaphore(self) -> Iterator[None]:
50
+ """Guard the host writes into the pooled input buffers.
51
+
52
+ `UvaBufferPool` recycles a slot every `max_concurrent_batches` steps and
53
+ the device reads the pooled host buffers in place. A sampling step is
54
+ ordered by the device wait in `AsyncOutput.get_output()`; this one never
55
+ waits, so entering blocks until the step that last wrote these slots is
56
+ done with them and leaving records this step's last write.
57
+ """
58
+ idx = self._input_event_idx
59
+ input_event = self._input_events[idx]
60
+ input_event.synchronize()
61
+ try:
62
+ yield
63
+ finally:
64
+ input_event.record()
65
+ self._input_event_idx = (idx + 1) % len(self._input_events)
66
+
67
+ def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
68
+ return {}
69
+
70
+ def capture_model(self) -> int:
71
+ return 0
72
+
73
+ def _dummy_run(
74
+ self, *args: Any, **kwargs: Any
75
+ ) -> tuple[torch.Tensor, torch.Tensor]:
76
+ empty = torch.empty(0, device=self.device)
77
+ return empty, empty
78
+
79
+ def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
80
+ return
81
+
82
+ def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None:
83
+ return
84
+
85
+ def _no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
86
+ return self._merge_ec_connector_no_forward(
87
+ scheduler_output, self.kv_connector.no_forward(scheduler_output)
88
+ )
89
+
90
+ @torch.inference_mode()
91
+ def execute_model(
92
+ self,
93
+ scheduler_output: "SchedulerOutput",
94
+ intermediate_tensors: IntermediateTensors | None = None,
95
+ dummy_run: bool = False,
96
+ skip_attn_for_dummy_run: bool = False,
97
+ is_profile: bool = False,
98
+ context_len: int = 0,
99
+ ) -> ModelRunnerOutput:
100
+ assert not dummy_run, "An encoder-only instance runs no dummy batch."
101
+
102
+ with self.input_tensor_semaphore():
103
+ self.update_pp_decode_requests()
104
+ self.finish_requests(scheduler_output)
105
+ self.free_states(scheduler_output)
106
+ self.add_requests(scheduler_output)
107
+ self.update_requests(scheduler_output)
108
+ self.block_tables.apply_staged_writes()
109
+ if scheduler_output.total_num_scheduled_tokens == 0:
110
+ return self._no_forward(scheduler_output)
111
+
112
+ batch_req_state, _ = self.gather_batch_req_state(scheduler_output, False)
113
+ assert batch_req_state is not None
114
+ # No CUDA graph, and no DP peer to agree a padded shape with.
115
+ batch_desc = BatchExecutionDescriptor(
116
+ cg_mode=CUDAGraphMode.NONE,
117
+ num_tokens=batch_req_state.num_tokens,
118
+ num_reqs=None,
119
+ )
120
+ self.prepare_inputs(scheduler_output, batch_req_state, batch_desc)
121
+
122
+ scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
123
+ if self.lora_config is not None:
124
+ set_active_mm_loras(
125
+ model=self.model,
126
+ lora_manager=self.lora_manager,
127
+ encoder_cache=self.encoder_cache,
128
+ req_id_to_index=self.req_states.req_id_to_index,
129
+ lora_state=self.lora_state,
130
+ scheduled_encoder_inputs=scheduled_encoder_inputs,
131
+ )
132
+
133
+ with self.ec_connector.maybe_get_output(
134
+ scheduler_output
135
+ ) as ec_connector_output:
136
+ self.model_state.execute_mm_encoder(scheduled_encoder_inputs)
137
+
138
+ return ModelRunnerOutput.with_ec_conn_output(
139
+ make_empty_encoder_model_runner_output(scheduler_output),
140
+ ec_connector_output,
141
+ )