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 14, 05:30 UTC
[Model] Define `merge_by_field_config` MM interface (#25676)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Sep 26, 2025 at 01:13 UTC
0ea80c87d9764a27f1cbdec446e38470a3f8c15b
5 files changed
+44
-12
tests/models/multimodal/processing/test_tensor_schema.py
+18
-5
@@ -19,6 +19,8 @@ from vllm.distributed import (cleanup_dist_env_and_memory,
19
init_distributed_environment,
20
initialize_model_parallel)
21
from vllm.model_executor.model_loader.utils import set_default_torch_dtype
22
+from vllm.model_executor.models.interfaces import (SupportsMultiModal,
23
+ supports_multimodal)
24
from vllm.multimodal import MULTIMODAL_REGISTRY, BatchedTensorInputs
25
from vllm.multimodal.processing import (BaseMultiModalProcessor,
26
InputProcessingContext)
@@ -88,6 +90,7 @@ def resize_mm_data(
90
91
92
def create_batched_mm_kwargs(
93
+ model_cls: type[SupportsMultiModal],
94
model_config: ModelConfig,
95
processor: BaseMultiModalProcessor,
96
size_factors: tuple[float, ...] = (1.0, 0.5, 0.25),
@@ -127,16 +130,22 @@ def create_batched_mm_kwargs(
130
mm_data=resized_mm_data,
131
hf_processor_mm_kwargs=processor_inputs.hf_processor_mm_kwargs,
132
tokenization_kwargs=processor_inputs.tokenization_kwargs,
130
- )["mm_kwargs"]
133
+ )["mm_kwargs"].require_data()
134
items = [
135
item for modality in supported_mm_limits
136
for item in mm_kwargs[modality]
137
]
135
- return group_mm_kwargs_by_modality(items)
138
+ return group_mm_kwargs_by_modality(
139
+ items,
140
+ merge_by_field_config=model_cls.merge_by_field_config,
141
+ )
142
143
144
@contextmanager
139
-def initialize_dummy_model(model_cls: nn.Module, model_config: ModelConfig):
145
+def initialize_dummy_model(
146
+ model_cls: type[nn.Module],
147
+ model_config: ModelConfig,
148
+):
149
temp_file = tempfile.mkstemp()[1]
150
init_distributed_environment(
151
world_size=1,
@@ -198,8 +207,12 @@ def test_model_tensor_schema(model_arch: str, model_id: str):
207
hf_overrides=hf_overrides_fn,
208
skip_tokenizer_init=model_info.skip_tokenizer_init,
209
enforce_eager=model_info.enforce_eager,
201
- dtype=model_info.dtype)
210
+ dtype=model_info.dtype,
211
+ )
212
+
213
model_cls = MULTIMODAL_REGISTRY._get_model_cls(model_config)
214
+ assert supports_multimodal(model_cls)
215
+
216
factories = MULTIMODAL_REGISTRY._processor_factories[model_cls]
217
218
inputs_parse_methods = []
@@ -228,7 +241,7 @@ def test_model_tensor_schema(model_arch: str, model_id: str):
241
242
with initialize_dummy_model(model_cls, model_config) as model:
243
for modality, _, mm_kwargs in create_batched_mm_kwargs(
231
- model_config, processor):
244
+ model_cls, model_config, processor):
245
for method_name in inputs_parse_methods:
246
print(f"Testing `{method_name}` with modality={modality} "
247
f"and mm_kwargs{list(mm_kwargs.keys())}")
vllm/config/model.py
+2
-3
@@ -63,13 +63,12 @@ ConvertType = Literal["none", "embed", "classify", "reward"]
63
ConvertOption = Literal["auto", ConvertType]
64
TaskOption = Literal["auto", "generate", "embedding", "embed", "classify",
65
"score", "reward", "transcription", "draft"]
66
-_ResolvedTask = Literal["generate", "transcription", "encode", "embed",
67
- "classify", "reward", "draft"]
66
TokenizerMode = Literal["auto", "slow", "mistral", "custom"]
67
ModelDType = Literal["auto", "half", "float16", "bfloat16", "float", "float32"]
68
LogprobsMode = Literal["raw_logits", "raw_logprobs", "processed_logits",
69
"processed_logprobs"]
72
-HfOverrides = Union[dict[str, Any], Callable[[type], type]]
70
+HfOverrides = Union[dict[str, Any], Callable[[PretrainedConfig],
71
+ PretrainedConfig]]
72
ModelImpl = Literal["auto", "vllm", "transformers", "terratorch"]
73
74
_RUNNER_TASKS: dict[RunnerType, list[TaskOption]] = {
vllm/model_executor/models/interfaces.py
+6
@@ -64,6 +64,12 @@ class SupportsMultiModal(Protocol):
64
`multimodal_config.mm_encoder_tp_mode="data"`.
65
"""
66
67
+ merge_by_field_config: ClassVar[bool] = False
68
+ """
69
+ A flag that indicates which implementation of
70
+ `vllm.multimodal.utils.group_mm_kwargs_by_modality` to use.
71
+ """
72
+
73
@classmethod
74
def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:
75
"""
vllm/v1/worker/gpu_model_runner.py
+11
-2
@@ -40,7 +40,8 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
40
from vllm.model_executor.layers.mamba.abstract import MambaBase
41
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
42
from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader
43
-from vllm.model_executor.models.interfaces import (is_mixture_of_experts,
43
+from vllm.model_executor.models.interfaces import (SupportsMultiModal,
44
+ is_mixture_of_experts,
45
supports_eagle3,
46
supports_mrope,
47
supports_transcription)
@@ -777,11 +778,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
778
mm_kwargs.append(feature.data)
779
780
# Input all modalities at once
781
+ model = cast(SupportsMultiModal, self.model)
782
mm_kwargs_combined: BatchedTensorInputs = {}
783
for _, _, mm_kwargs_group in group_mm_kwargs_by_modality(
784
mm_kwargs,
785
device=self.device,
786
pin_memory=self.pin_memory,
787
+ merge_by_field_config=model.merge_by_field_config,
788
):
789
mm_kwargs_combined.update(mm_kwargs_group)
790
@@ -1525,11 +1528,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1528
# in the same batch while still being able to benefit from batching
1529
# multimodal inputs. The proper solution should be reordering the
1530
# encoder outputs.
1531
+ model = cast(SupportsMultiModal, self.model)
1532
encoder_outputs = []
1533
for _, num_items, mm_kwargs_group in group_mm_kwargs_by_modality(
1534
mm_kwargs,
1535
device=self.device,
1536
pin_memory=self.pin_memory,
1537
+ merge_by_field_config=model.merge_by_field_config,
1538
):
1539
# Run the encoder.
1540
# `curr_group_outputs` is either of the following:
@@ -1538,7 +1543,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1543
# 2. A list or tuple (length: num_items) of tensors, each of shape
1544
# (feature_size, hidden_size) in case the feature size is dynamic
1545
# depending on the input multimodal items.
1541
- curr_group_outputs = self.model.get_multimodal_embeddings(
1546
+ curr_group_outputs = model.get_multimodal_embeddings(
1547
**mm_kwargs_group)
1548
1549
sanity_check_mm_encoder_outputs(
@@ -1623,11 +1628,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1628
return {}
1629
1630
# Group MM kwargs by modality and extract features
1631
+ model = cast(SupportsMultiModal, self.model)
1632
encoder_features = {}
1633
for _, _, mm_kwargs_group in group_mm_kwargs_by_modality(
1634
mm_kwargs,
1635
device=self.device,
1636
pin_memory=self.pin_memory,
1637
+ merge_by_field_config=model.merge_by_field_config,
1638
):
1639
# Add the grouped features to encoder_features dict
1640
# This allows the model to receive them as kwargs (e.g.,
@@ -2839,11 +2846,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2846
dummy_mm_item = dummy_mm_data[modality][0]
2847
dummy_mm_items = [dummy_mm_item] * max_items_per_batch
2848
2849
+ model = cast(SupportsMultiModal, self.model)
2850
return next(mm_kwargs_group
2851
for _, _, mm_kwargs_group in group_mm_kwargs_by_modality(
2852
dummy_mm_items,
2853
device=self.device,
2854
pin_memory=self.pin_memory,
2855
+ merge_by_field_config=model.merge_by_field_config,
2856
))
2857
2858
@torch.inference_mode()
vllm/v1/worker/tpu_model_runner.py
+7
-2
@@ -30,7 +30,8 @@ from vllm.logger import init_logger
30
from vllm.lora.layers import BaseLayerWithLoRA
31
from vllm.model_executor.model_loader import get_model_loader
32
from vllm.model_executor.model_loader.tpu import TPUModelLoader
33
-from vllm.model_executor.models.interfaces import supports_transcription
33
+from vllm.model_executor.models.interfaces import (SupportsMultiModal,
34
+ supports_transcription)
35
from vllm.model_executor.models.interfaces_base import (
36
is_pooling_model, is_text_generation_model)
37
from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -834,11 +835,13 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
835
# in the same batch while still being able to benefit from batching
836
# multimodal inputs. The proper solution should be reordering the
837
# encoder outputs.
838
+ model = cast(SupportsMultiModal, self.model)
839
encoder_outputs = []
840
for _, num_items, mm_kwargs_group in group_mm_kwargs_by_modality(
841
mm_kwargs,
842
device=self.device,
843
pin_memory=self.pin_memory,
844
+ merge_by_field_config=model.merge_by_field_config,
845
):
846
# Run the encoder.
847
# `curr_group_outputs` is either of the following:
@@ -848,7 +851,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
851
# (feature_size, hidden_size) in case the feature size is dynamic
852
# depending on the input multimodal items.
853
torch_xla.sync(wait=False)
851
- curr_group_outputs = self.model.get_multimodal_embeddings(
854
+ curr_group_outputs = model.get_multimodal_embeddings(
855
**mm_kwargs_group)
856
torch_xla.sync(wait=False)
857
@@ -1805,11 +1808,13 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1808
dummy_mm_item = dummy_mm_data[modality][0]
1809
dummy_mm_items = [dummy_mm_item] * max_items_per_batch
1810
1811
+ model = cast(SupportsMultiModal, self.model)
1812
return next(grouped_mm_kwargs
1813
for _, _, grouped_mm_kwargs in group_mm_kwargs_by_modality(
1814
dummy_mm_items,
1815
device=self.device,
1816
pin_memory=self.pin_memory,
1817
+ merge_by_field_config=model.merge_by_field_config,
1818
))
1819
1820