A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 00:15 UTC
[Model Runner V2] Support VLM (#32546)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jan 18, 2026 at 16:58 UTC
bb1848cd62e5e5f327f307bee26e2b947f00396e
6 files changed
+263
-15
vllm/v1/worker/gpu/cudagraph_utils.py
+7
@@ -76,6 +76,7 @@ class CudaGraphManager:
76
model: nn.Module,
77
input_buffers: InputBuffers,
78
mrope_positions: torch.Tensor | None,
79
+ inputs_embeds: torch.Tensor | None,
80
block_tables: BlockTables,
81
attn_metadata_builders: list[AttentionMetadataBuilder],
82
kv_cache_config: KVCacheConfig,
@@ -86,6 +87,8 @@ class CudaGraphManager:
87
if self.uses_mrope:
88
assert mrope_positions is not None
89
positions = mrope_positions[:, :num_tokens]
90
+ if inputs_embeds is not None:
91
+ inputs_embeds = inputs_embeds[:num_tokens]
92
attn_metadata = prepare_inputs_to_capture(
93
num_reqs,
94
num_tokens,
@@ -108,6 +111,7 @@ class CudaGraphManager:
111
hidden_states = model(
112
input_ids=input_ids,
113
positions=positions,
114
+ inputs_embeds=inputs_embeds,
115
)
116
if self.hidden_states is None:
117
self.hidden_states = torch.empty_like(hidden_states)
@@ -128,6 +132,7 @@ class CudaGraphManager:
132
hidden_states = model(
133
input_ids=input_ids,
134
positions=positions,
135
+ inputs_embeds=inputs_embeds,
136
)
137
self.hidden_states[:num_tokens] = hidden_states
138
self.graphs[num_tokens] = graph
@@ -138,6 +143,7 @@ class CudaGraphManager:
143
model: nn.Module,
144
input_buffers: InputBuffers,
145
mrope_positions: torch.Tensor | None,
146
+ inputs_embeds: torch.Tensor | None,
147
block_tables: BlockTables,
148
attn_metadata_builders: list[AttentionMetadataBuilder],
149
kv_cache_config: KVCacheConfig,
@@ -149,6 +155,7 @@ class CudaGraphManager:
155
model=model,
156
input_buffers=input_buffers,
157
mrope_positions=mrope_positions,
158
+ inputs_embeds=inputs_embeds,
159
block_tables=block_tables,
160
attn_metadata_builders=attn_metadata_builders,
161
kv_cache_config=kv_cache_config,
vllm/v1/worker/gpu/input_batch.py
+3
-3
@@ -15,9 +15,6 @@ class InputBuffers:
15
self,
16
max_num_reqs: int,
17
max_num_tokens: int,
18
- inputs_embeds_size: int,
19
- vocab_size: int,
20
- dtype: torch.dtype,
18
device: torch.device,
19
):
20
self.max_num_reqs = max_num_reqs
@@ -64,6 +61,8 @@ class InputBatch:
61
positions: torch.Tensor
62
# [3, num_tokens_after_padding]
63
mrope_positions: torch.Tensor | None
64
+ # [num_tokens_after_padding, hidden_size]
65
+ inputs_embeds: torch.Tensor | None
66
67
# layer_name -> Metadata
68
attn_metadata: dict[str, Any]
@@ -132,6 +131,7 @@ class InputBatch:
131
input_ids=input_ids,
132
positions=positions,
133
mrope_positions=None,
134
+ inputs_embeds=None,
135
attn_metadata=None, # type: ignore
136
logits_indices=logits_indices,
137
cu_num_logits=cu_num_logits,
vllm/v1/worker/gpu/mm/encoder_runner.py
new
+184
@@ -0,0 +1,184 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import numpy as np
4
+import torch
5
+
6
+from vllm.model_executor.models.interfaces import SupportsMultiModal
7
+from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalKwargsItem
8
+from vllm.multimodal.utils import group_mm_kwargs_by_modality
9
+from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool
10
+from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
11
+
12
+
13
+class EncoderRunner:
14
+ def __init__(
15
+ self,
16
+ max_num_tokens: int,
17
+ hidden_size: int,
18
+ dtype: torch.dtype,
19
+ device: torch.device,
20
+ ):
21
+ self.max_num_tokens = max_num_tokens
22
+ self.hidden_size = hidden_size
23
+ self.dtype = dtype
24
+ self.device = device
25
+
26
+ self.inputs_embeds = torch.zeros(
27
+ max_num_tokens,
28
+ hidden_size,
29
+ dtype=dtype,
30
+ device=device,
31
+ )
32
+ self.req_id_to_mm_features: dict[str, list[MultiModalFeatureSpec]] = {}
33
+ self.encoder_cache: dict[str, torch.Tensor] = {}
34
+
35
+ self.tmp_is_mm_embed = UvaBufferPool(max_num_tokens, torch.bool)
36
+
37
+ def add_request(self, req_id: str, mm_features: list[MultiModalFeatureSpec]):
38
+ self.req_id_to_mm_features[req_id] = mm_features
39
+
40
+ def free_encoder_cache(self, mm_hash: str) -> None:
41
+ self.encoder_cache.pop(mm_hash, None)
42
+
43
+ def remove_request(self, req_id: str) -> None:
44
+ self.req_id_to_mm_features.pop(req_id, None)
45
+
46
+ def prepare_mm_inputs(
47
+ self,
48
+ scheduled_encoder_inputs: dict[str, list[int]],
49
+ ) -> tuple[list[str], list[MultiModalKwargsItem]]:
50
+ mm_hashes: list[str] = []
51
+ mm_kwargs: list[MultiModalKwargsItem] = []
52
+ for req_id, encoder_input_ids in scheduled_encoder_inputs.items():
53
+ mm_features = self.req_id_to_mm_features[req_id]
54
+ for mm_input_id in encoder_input_ids:
55
+ mm_feature = mm_features[mm_input_id]
56
+ if mm_feature.data is None:
57
+ continue
58
+ mm_hashes.append(mm_feature.identifier)
59
+ mm_kwargs.append(mm_feature.data)
60
+ return mm_hashes, mm_kwargs
61
+
62
+ @torch.inference_mode()
63
+ def execute_mm_encoder(
64
+ self,
65
+ model: SupportsMultiModal,
66
+ mm_hashes: list[str],
67
+ mm_kwargs: list[MultiModalKwargsItem],
68
+ ) -> list[torch.Tensor]:
69
+ if not mm_hashes:
70
+ return []
71
+
72
+ encoder_outputs: list[torch.Tensor] = []
73
+ for modality, num_items, mm_kwargs_group in group_mm_kwargs_by_modality(
74
+ mm_kwargs,
75
+ device=self.device,
76
+ pin_memory=False,
77
+ ):
78
+ curr_group_outputs = model.embed_multimodal(**mm_kwargs_group)
79
+ sanity_check_mm_encoder_outputs(
80
+ curr_group_outputs,
81
+ expected_num_items=num_items,
82
+ )
83
+ encoder_outputs.extend(curr_group_outputs)
84
+
85
+ # Cache the encoder outputs by mm_hash
86
+ for mm_hash, output in zip(mm_hashes, encoder_outputs):
87
+ self.encoder_cache[mm_hash] = output
88
+ return encoder_outputs
89
+
90
+ def gather_mm_embeddings(
91
+ self,
92
+ req_ids: list[str],
93
+ total_num_scheduled_tokens: int,
94
+ num_scheduled_tokens: np.ndarray,
95
+ query_start_loc: np.ndarray,
96
+ prefill_lens: np.ndarray,
97
+ computed_prefill_lens: np.ndarray,
98
+ ) -> tuple[list[torch.Tensor], torch.Tensor]:
99
+ is_prefilling = (computed_prefill_lens < prefill_lens).tolist()
100
+ all_decode = not any(is_prefilling)
101
+ if all_decode:
102
+ # All decode requests, so no need to gather any embeddings.
103
+ return [], torch.zeros(
104
+ total_num_scheduled_tokens,
105
+ dtype=torch.bool,
106
+ device=self.device,
107
+ )
108
+
109
+ query_start = computed_prefill_lens.tolist()
110
+ query_end = (computed_prefill_lens + num_scheduled_tokens).tolist()
111
+
112
+ mm_embeds: list[torch.Tensor] = []
113
+ is_mm_embed = torch.zeros(
114
+ total_num_scheduled_tokens,
115
+ dtype=torch.bool,
116
+ device="cpu",
117
+ pin_memory=False,
118
+ )
119
+ for i, req_id in enumerate(req_ids):
120
+ if not is_prefilling[i]:
121
+ # OPTIMIZATION: Skip decode requests.
122
+ continue
123
+
124
+ mm_features = self.req_id_to_mm_features[req_id]
125
+ for mm_feature in mm_features:
126
+ pos_info = mm_feature.mm_position
127
+ start_pos = pos_info.offset
128
+ num_encoder_tokens = pos_info.length
129
+
130
+ if start_pos >= query_end[i]:
131
+ # The encoder output is not needed in this step.
132
+ break
133
+ if start_pos + num_encoder_tokens <= query_start[i]:
134
+ # The encoder output is already processed and stored
135
+ # in the decoder's KV cache.
136
+ continue
137
+
138
+ start_idx = max(query_start[i] - start_pos, 0)
139
+ end_idx = min(query_end[i] - start_pos, num_encoder_tokens)
140
+ assert start_idx < end_idx
141
+ curr_embeds_start, curr_embeds_end = (
142
+ pos_info.get_embeds_indices_in_range(start_idx, end_idx)
143
+ )
144
+ # If there are no embeddings in the current range, we skip
145
+ # gathering the embeddings.
146
+ if curr_embeds_start == curr_embeds_end:
147
+ continue
148
+
149
+ mm_hash = mm_feature.identifier
150
+ encoder_output = self.encoder_cache.get(mm_hash, None)
151
+ assert encoder_output is not None, f"Encoder cache miss for {mm_hash}."
152
+
153
+ if (is_embed := pos_info.is_embed) is not None:
154
+ is_embed = is_embed[start_idx:end_idx]
155
+ mm_embeds_item = encoder_output[curr_embeds_start:curr_embeds_end]
156
+ else:
157
+ mm_embeds_item = encoder_output[start_idx:end_idx]
158
+
159
+ req_start_pos = query_start_loc[i] + start_pos - query_start[i]
160
+ is_mm_embed[req_start_pos + start_idx : req_start_pos + end_idx] = (
161
+ True if is_embed is None else is_embed
162
+ )
163
+ mm_embeds.append(mm_embeds_item)
164
+
165
+ # Copy the is_mm_embed tensor to the GPU.
166
+ is_mm_embed = self.tmp_is_mm_embed.copy_to_gpu(is_mm_embed)
167
+ return mm_embeds, is_mm_embed
168
+
169
+ @torch.inference_mode()
170
+ def get_inputs_embeds(
171
+ self,
172
+ model: SupportsMultiModal,
173
+ input_ids: torch.Tensor,
174
+ mm_embeds: list[torch.Tensor],
175
+ is_mm_embed: torch.Tensor,
176
+ ) -> torch.Tensor:
177
+ x = model.embed_input_ids(
178
+ input_ids,
179
+ multimodal_embeddings=mm_embeds,
180
+ is_multimodal=is_mm_embed,
181
+ )
182
+ # Copy to the pre-allocated buffer for CUDA graphs.
183
+ self.inputs_embeds[: x.shape[0]] = x
184
+ return self.inputs_embeds
vllm/v1/worker/gpu/mm/mrope_utils.py
+3
-5
@@ -23,7 +23,7 @@ class MRopeState:
23
# NOTE(woosuk): This tensor can be extremely large (e.g., several GBs)
24
# wasting a lot of CPU memory.
25
self.prefill_mrope_positions = StagedWriteTensor(
26
- (max_num_reqs, 3 * max_model_len),
26
+ (max_num_reqs * 3, max_model_len),
27
dtype=torch.int32,
28
device=device,
29
uva_instead_of_gpu=True,
@@ -58,9 +58,7 @@ class MRopeState:
58
)
59
for i in range(3):
60
pos = prefill_mrope_positions[i].tolist()
61
- self.prefill_mrope_positions.stage_write(
62
- req_idx, i * self.max_model_len, pos
63
- )
61
+ self.prefill_mrope_positions.stage_write(3 * req_idx + i, 0, pos)
62
self.prefill_mrope_delta.np[req_idx] = prefill_mrope_delta
63
64
def apply_staged_writes(self) -> None:
@@ -79,7 +77,7 @@ class MRopeState:
77
self.mrope_positions,
78
self.mrope_positions.stride(0),
79
self.prefill_mrope_positions.gpu,
82
- self.prefill_mrope_positions.gpu.stride(0),
80
+ 3 * self.max_model_len,
81
self.max_model_len,
82
self.prefill_mrope_delta.gpu,
83
idx_mapping,
vllm/v1/worker/gpu/model_runner.py
+66
-4
@@ -14,6 +14,7 @@ from vllm.config.compilation import CUDAGraphMode
14
from vllm.forward_context import set_forward_context
15
from vllm.logger import init_logger
16
from vllm.model_executor.model_loader import get_model_loader
17
+from vllm.multimodal import MULTIMODAL_REGISTRY
18
from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
19
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
20
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
@@ -47,6 +48,7 @@ from vllm.v1.worker.gpu.input_batch import (
48
prepare_pos_seq_lens,
49
prepare_prefill_inputs,
50
)
51
+from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
52
from vllm.v1.worker.gpu.mm.mrope_utils import MRopeState
53
from vllm.v1.worker.gpu.sample.logprob import compute_prompt_logprobs
54
from vllm.v1.worker.gpu.sample.output import SamplerOutput
@@ -95,6 +97,17 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
97
self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()
98
99
# Multimodal
100
+ self.mm_registry = MULTIMODAL_REGISTRY
101
+ self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(
102
+ self.model_config
103
+ )
104
+ if self.supports_mm_inputs:
105
+ self.encoder_runner = EncoderRunner(
106
+ max_num_tokens=self.max_num_tokens,
107
+ hidden_size=self.inputs_embeds_size,
108
+ dtype=self.dtype,
109
+ device=self.device,
110
+ )
111
self.uses_mrope = self.model_config.uses_mrope
112
if self.uses_mrope:
113
self.mrope_states = MRopeState(
@@ -134,9 +147,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
147
self.input_buffers = InputBuffers(
148
max_num_reqs=self.max_num_reqs,
149
max_num_tokens=self.max_num_tokens,
137
- inputs_embeds_size=self.inputs_embeds_size,
138
- vocab_size=self.vocab_size,
139
- dtype=self.dtype,
150
device=self.device,
151
)
152
self.sampler = Sampler(
@@ -289,6 +299,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
299
input_batch.mrope_positions = self.mrope_states.mrope_positions[
300
:, :num_tokens
301
]
302
+ if self.supports_mm_inputs:
303
+ input_batch.inputs_embeds = self.encoder_runner.inputs_embeds[:num_tokens]
304
if not skip_attn:
305
self.prepare_dummy_attn_metadata(input_batch)
306
@@ -314,6 +326,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
326
hidden_states = self.model(
327
input_ids=input_batch.input_ids,
328
positions=positions,
329
+ inputs_embeds=input_batch.inputs_embeds,
330
)
331
sample_hidden_states = hidden_states[input_batch.logits_indices]
332
return hidden_states, sample_hidden_states
@@ -378,10 +391,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
391
mrope_positions = None
392
if self.uses_mrope:
393
mrope_positions = self.mrope_states.mrope_positions
394
+ inputs_embeds = None
395
+ if self.supports_mm_inputs:
396
+ inputs_embeds = self.encoder_runner.inputs_embeds
397
self.cudagraph_manager.capture(
398
model=self.model,
399
input_buffers=self.input_buffers,
400
mrope_positions=mrope_positions,
401
+ inputs_embeds=inputs_embeds,
402
block_tables=self.block_tables,
403
attn_metadata_builders=self.attn_metadata_builders,
404
kv_cache_config=self.kv_cache_config,
@@ -412,8 +429,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
429
if scheduler_output.preempted_req_ids is not None:
430
for req_id in scheduler_output.preempted_req_ids:
431
self.req_states.remove_request(req_id)
432
+ if self.supports_mm_inputs:
433
+ self.encoder_runner.remove_request(req_id)
434
for req_id in scheduler_output.finished_req_ids:
435
self.req_states.remove_request(req_id)
436
+ if self.supports_mm_inputs:
437
+ self.encoder_runner.remove_request(req_id)
438
+
439
+ if self.supports_mm_inputs:
440
+ for mm_hash in scheduler_output.free_encoder_mm_hashes:
441
+ self.encoder_runner.free_encoder_cache(mm_hash)
442
443
# Add new requests.
444
for new_req_data in scheduler_output.scheduled_new_reqs:
@@ -432,13 +457,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
457
)
458
req_index = self.req_states.req_id_to_index[req_id]
459
460
+ if self.supports_mm_inputs:
461
+ self.encoder_runner.add_request(req_id, new_req_data.mm_features)
462
+
463
# Pre-compute M-RoPE positions for prefill.
464
if self.uses_mrope:
465
self.mrope_states.init_prefill_mrope_positions(
466
req_index,
467
self.model, # type: ignore
468
new_req_data.prefill_token_ids,
441
- mm_features=[], # TODO
469
+ mm_features=new_req_data.mm_features,
470
)
471
472
self.block_tables.append_block_ids(
@@ -632,12 +660,33 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
660
input_ids=input_ids,
661
positions=positions,
662
mrope_positions=mrope_positions,
663
+ inputs_embeds=None,
664
attn_metadata=attn_metadata,
665
logits_indices=logits_indices,
666
cu_num_logits=cu_num_logits,
667
cu_num_logits_np=cu_num_logits_np,
668
)
669
670
+ @torch.inference_mode()
671
+ def get_mm_embeddings(
672
+ self,
673
+ scheduled_encoder_inputs: dict[str, list[int]],
674
+ input_batch: InputBatch,
675
+ ) -> tuple[list[torch.Tensor], torch.Tensor]:
676
+ mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs(
677
+ scheduled_encoder_inputs
678
+ )
679
+ self.encoder_runner.execute_mm_encoder(self.model, mm_hashes, mm_kwargs)
680
+ mm_embeds, is_mm_embed = self.encoder_runner.gather_mm_embeddings(
681
+ input_batch.req_ids,
682
+ input_batch.num_tokens,
683
+ input_batch.num_scheduled_tokens,
684
+ input_batch.query_start_loc_np,
685
+ self.req_states.prefill_len.np[input_batch.idx_mapping_np],
686
+ self.req_states.num_computed_prefill_tokens[input_batch.idx_mapping_np],
687
+ )
688
+ return mm_embeds, is_mm_embed
689
+
690
def sample(
691
self,
692
hidden_states: torch.Tensor,
@@ -930,6 +979,18 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
979
input_batch.num_scheduled_tokens,
980
)
981
self._set_active_loras(*lora_inputs)
982
+
983
+ if self.supports_mm_inputs:
984
+ # Execute the multimodal encoder.
985
+ mm_embeds, is_mm_embed = self.get_mm_embeddings(
986
+ scheduler_output.scheduled_encoder_inputs, input_batch
987
+ )
988
+ inputs_embeds = self.encoder_runner.get_inputs_embeds(
989
+ self.model, input_batch.input_ids, mm_embeds, is_mm_embed
990
+ )
991
+ input_batch.inputs_embeds = inputs_embeds[
992
+ : input_batch.num_tokens_after_padding
993
+ ]
994
else:
995
# No actual tokens to run. A dummy run for DP.
996
num_reqs = min(num_tokens_after_padding, self.max_num_reqs)
@@ -970,6 +1031,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1031
hidden_states = self.model(
1032
input_ids=input_batch.input_ids,
1033
positions=positions,
1034
+ inputs_embeds=input_batch.inputs_embeds,
1035
)
1036
1037
self.execute_model_state = hidden_states, input_batch
vllm/v1/worker/gpu/spec_decode/eagle.py
-3
@@ -48,9 +48,6 @@ class EagleSpeculator:
48
self.input_buffers = InputBuffers(
49
max_num_reqs=self.max_num_reqs,
50
max_num_tokens=self.max_num_tokens,
51
- inputs_embeds_size=self.inputs_embeds_size,
52
- vocab_size=self.vocab_size,
53
- dtype=self.dtype,
51
device=device,
52
)
53
self.hidden_states = torch.zeros(