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 11, 00:00 UTC
[Chore] Simplify logic of `_execute_mm_encoder` (#31222)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Dec 24, 2025 at 10:15 UTC
ca6a95ba259aaa4c89eccdd254fa7922e31eddc2
1 file changed
+21
-25
vllm/v1/worker/gpu_model_runner.py
+21
-25
@@ -61,6 +61,7 @@ from vllm.model_executor.layers.rotary_embedding import (
61
)
62
from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader
63
from vllm.model_executor.models.interfaces import (
64
+ MultiModalEmbeddings,
65
SupportsMRoPE,
66
SupportsMultiModal,
67
SupportsXDRoPE,
@@ -78,11 +79,7 @@ from vllm.model_executor.models.interfaces_base import (
79
is_text_generation_model,
80
)
81
from vllm.multimodal import MULTIMODAL_REGISTRY
81
-from vllm.multimodal.inputs import (
82
- BatchedTensorInputs,
83
- MultiModalKwargsItem,
84
- PlaceholderRange,
85
-)
82
+from vllm.multimodal.inputs import BatchedTensorInputs, MultiModalKwargsItem
83
from vllm.multimodal.utils import group_mm_kwargs_by_modality
84
from vllm.pooling_params import PoolingParams
85
from vllm.sampling_params import SamplingType
@@ -2097,28 +2094,27 @@ class GPUModelRunner(
2094
]
2095
return logits_indices_padded
2096
2100
- def _batch_mm_kwargs_from_scheduler(
2097
+ def _batch_mm_inputs_from_scheduler(
2098
self,
2099
scheduler_output: "SchedulerOutput",
2103
- ) -> tuple[list[MultiModalKwargsItem], list[tuple[str, PlaceholderRange]]]:
2104
- """Batch multimodal kwargs from scheduled encoder inputs.
2100
+ ) -> tuple[list[str], list[MultiModalKwargsItem]]:
2101
+ """Batch multimodal inputs from scheduled encoder inputs.
2102
2103
Args:
2104
scheduler_output: The scheduler output containing scheduled encoder
2105
inputs.
2106
2107
Returns:
2111
- A tuple of (mm_kwargs, req_ids_pos) where:
2112
- - mm_kwargs: List of multimodal kwargs items to be batched
2113
- - mm_hashes_pos: List of (mm_hash, position_info) tuples
2108
+ A tuple of (mm_hashes, mm_kwargs) where:
2109
+ - mm_hashes: List of multimodal hashes for each item
2110
+ - mm_kwargs: List of multimodal kwargs for each item
2111
"""
2112
scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
2113
if not scheduled_encoder_inputs:
2114
return [], []
2118
- # Batch the multi-modal inputs.
2115
+
2116
+ mm_hashes = list[str]()
2117
mm_kwargs = list[MultiModalKwargsItem]()
2120
- # list of tuple (mm_hash, position_info)
2121
- mm_hashes_pos = list[tuple[str, PlaceholderRange]]()
2118
for req_id, encoder_input_ids in scheduled_encoder_inputs.items():
2119
req_state = self.requests[req_id]
2120
@@ -2126,19 +2122,16 @@ class GPUModelRunner(
2122
mm_feature = req_state.mm_features[mm_input_id]
2123
if mm_feature.data is None:
2124
continue
2129
- mm_hash = mm_feature.identifier
2125
+
2126
+ mm_hashes.append(mm_feature.identifier)
2127
mm_kwargs.append(mm_feature.data)
2131
- mm_hashes_pos.append((mm_hash, mm_feature.mm_position))
2128
2133
- return mm_kwargs, mm_hashes_pos
2129
+ return mm_hashes, mm_kwargs
2130
2131
def _execute_mm_encoder(
2132
self, scheduler_output: "SchedulerOutput"
2133
) -> list[torch.Tensor]:
2138
- # Batch the multi-modal inputs using the helper method.
2139
- mm_kwargs, mm_hashes_pos = self._batch_mm_kwargs_from_scheduler(
2140
- scheduler_output
2141
- )
2134
+ mm_hashes, mm_kwargs = self._batch_mm_inputs_from_scheduler(scheduler_output)
2135
2136
if not mm_kwargs:
2137
return []
@@ -2157,7 +2150,7 @@ class GPUModelRunner(
2150
device=self.device,
2151
pin_memory=self.pin_memory,
2152
):
2160
- curr_group_outputs: list[torch.Tensor] = []
2153
+ curr_group_outputs: MultiModalEmbeddings
2154
2155
# EVS-related change.
2156
# (ekhvedchenia): Temporary hack to limit peak memory usage when
@@ -2173,6 +2166,7 @@ class GPUModelRunner(
2166
and modality == "video"
2167
and num_items > 1
2168
):
2169
+ curr_group_outputs_lst = list[torch.Tensor]()
2170
for video_mm_kwargs_item in filter(
2171
lambda item: item.modality == "video", mm_kwargs
2172
):
@@ -2188,7 +2182,9 @@ class GPUModelRunner(
2182
**micro_batch_mm_inputs
2183
)
2184
2191
- curr_group_outputs.extend(micro_batch_outputs)
2185
+ curr_group_outputs_lst.extend(micro_batch_outputs)
2186
+
2187
+ curr_group_outputs = curr_group_outputs_lst
2188
else:
2189
# Run the encoder.
2190
# `curr_group_outputs` is either of the following:
@@ -2197,7 +2193,7 @@ class GPUModelRunner(
2193
# 2. A list or tuple (length: num_items) of tensors,
2194
# each of shape (feature_size, hidden_size) in case the feature
2195
# size is dynamic depending on the input multimodal items.
2200
- curr_group_outputs = model.embed_multimodal(**mm_kwargs_group) # type: ignore[assignment]
2196
+ curr_group_outputs = model.embed_multimodal(**mm_kwargs_group)
2197
2198
sanity_check_mm_encoder_outputs(
2199
curr_group_outputs,
@@ -2206,7 +2202,7 @@ class GPUModelRunner(
2202
encoder_outputs.extend(curr_group_outputs)
2203
2204
# Cache the encoder outputs by mm_hash
2209
- for (mm_hash, pos_info), output in zip(mm_hashes_pos, encoder_outputs):
2205
+ for mm_hash, output in zip(mm_hashes, encoder_outputs):
2206
self.encoder_cache[mm_hash] = output
2207
logger.debug("Finish execute for mm hash %s", mm_hash)
2208
self.maybe_save_ec_to_connector(self.encoder_cache, mm_hash)