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 (R-T) (#26260)
Signed-off-by: Ayush Satyam <ayushsatyam146@gmail.com> Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: DarkLight1337 <tlleungac@connect.ust.hk>
Ayush Satyam committed
Oct 7, 2025 at 13:40 UTC
de342585ffce0c8d55e9658f9bb6ca52a3fbd592
3 files changed
+46
-45
vllm/model_executor/models/step3_vl.py
+36
-33
@@ -4,7 +4,7 @@ import math
4
from collections.abc import Iterable, Mapping, Sequence
5
from itertools import product
6
from math import ceil, sqrt
7
-from typing import Any, Literal, Optional, TypedDict, Union
7
+from typing import Annotated, Any, Literal, Optional, Union
8
9
import numpy as np
10
import torch
@@ -44,28 +44,48 @@ from vllm.multimodal.profiling import BaseDummyInputsBuilder
44
from vllm.sequence import IntermediateTensors
45
from vllm.transformers_utils.configs import Step3VisionEncoderConfig
46
from vllm.transformers_utils.tokenizer import AnyTokenizer
47
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
48
49
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
50
from .utils import (
51
AutoWeightsLoader,
52
WeightsMapper,
52
- flatten_bn,
53
init_vllm_registered_model,
54
maybe_prefix,
55
)
56
from .vision import run_dp_sharded_vision_model
57
58
59
-class Step3VLImagePixelInputs(TypedDict):
59
+class Step3VLImagePixelInputs(TensorSchema):
60
+ """
61
+ Dimensions:
62
+ - bn: Batch size * number of images
63
+ - c: Number of channels (3)
64
+ - h: Height
65
+ - w: Width
66
+ - bnp: Batch size * number of images * number of patches
67
+ - hp: Height of patch
68
+ - wp: Width of patch
69
+ """
70
+
71
type: Literal["pixel_values"]
61
- pixel_values: torch.Tensor
62
- patch_pixel_values: Optional[torch.Tensor]
63
- num_patches: list[int]
72
+ pixel_values: Annotated[torch.Tensor, TensorShape("bn", 3, "h", "w")]
73
+ patch_pixel_values: Annotated[
74
+ Optional[torch.Tensor], TensorShape("bnp", 3, "hp", "wp")
75
+ ]
76
+ num_patches: Annotated[torch.Tensor, TensorShape("bn")]
77
+
78
79
+class Step3VLImageEmbeddingInputs(TensorSchema):
80
+ """
81
+ Dimensions:
82
+ - bn: Batch size * number of images
83
+ - f: Image feature size
84
+ - h: Hidden size (must match the hidden size of language model backbone)
85
+ """
86
66
-class Step3VLImageEmbeddingInputs(TypedDict):
67
- type: Literal["image_embeds"]
68
- image_embeds: torch.Tensor
87
+ type: Literal["image_embeds"] = "image_embeds"
88
+ data: Annotated[torch.Tensor, TensorShape("bn", "f", "h")]
89
90
91
Step3VLImageInputs = Union[Step3VLImagePixelInputs, Step3VLImageEmbeddingInputs]
@@ -895,6 +915,8 @@ class Step3VisionTransformer(nn.Module):
915
dummy_inputs=Step3VLDummyInputsBuilder,
916
)
917
class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
918
+ merge_by_field_config = True
919
+
920
hf_to_vllm_mapper = WeightsMapper(
921
orig_to_new_prefix={
922
"model.": "language_model.model.",
@@ -982,41 +1004,22 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)
1004
return None
1005
1006
if pixel_values is not None:
985
- pixel_values = flatten_bn(pixel_values, concat=True)
986
- if pixel_values.dim() >= 3:
987
- pixel_values = pixel_values.view(-1, *pixel_values.shape[-3:])
988
- if patch_pixel_values is not None:
989
- patch_pixel_values = flatten_bn(patch_pixel_values, concat=True)
990
- patch_pixel_values = patch_pixel_values.view(
991
- -1, *patch_pixel_values.shape[-3:]
992
- )
993
- # Handle empty patch_pixel_values by setting to None
994
- if patch_pixel_values.shape[0] == 0:
995
- patch_pixel_values = None
996
- num_patches = flatten_bn(num_patches, concat=True).tolist()
997
-
1007
return Step3VLImagePixelInputs(
1008
type="pixel_values",
1000
- pixel_values=pixel_values.to(self.dtype).to(self.device),
1001
- patch_pixel_values=patch_pixel_values.to(self.dtype).to(self.device)
1009
+ pixel_values=pixel_values.to(self.dtype),
1010
+ patch_pixel_values=patch_pixel_values.to(self.dtype)
1011
if patch_pixel_values is not None
1012
else None,
1013
num_patches=num_patches,
1014
)
1015
1016
if image_embeds is not None:
1008
- if image_embeds.dim() == 2 or image_embeds.dim() >= 3:
1009
- image_embeds = image_embeds.view(-1, image_embeds.shape[-1])
1010
- else:
1011
- raise ValueError(
1012
- f"Unexpected shape for image_embeds: {image_embeds.shape}"
1013
- )
1014
-
1017
return Step3VLImageEmbeddingInputs(
1018
type="image_embeds",
1017
- image_embeds=image_embeds.to(self.dtype).to(self.device),
1019
+ image_embeds=image_embeds.to(self.dtype),
1020
)
1019
- return None
1021
+
1022
+ raise AssertionError("This line should be unreachable.")
1023
1024
def _process_image_features(self, image_features: torch.Tensor) -> torch.Tensor:
1025
B, P = image_features.shape[:2]
vllm/model_executor/models/tarsier.py
+1
-5
@@ -47,11 +47,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape
47
from .clip import CLIPVisionModel
48
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
49
from .siglip import SiglipVisionModel
50
-from .utils import (
51
- AutoWeightsLoader,
52
- init_vllm_registered_model,
53
- maybe_prefix,
54
-)
50
+from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
51
from .vision import (
52
VisionEncoderInfo,
53
get_num_selected_vision_tokens,
vllm/model_executor/models/terratorch.py
+9
-7
@@ -87,12 +87,10 @@ def _terratorch_field_factory(
87
if input.type == InputTypeEnum.tensor:
88
fields[input_name] = "image"
89
90
- mm_fields_config = {}
91
- for field_name, field_modality in fields.items():
92
- mm_fields_config[field_name] = MultiModalFieldConfig.shared(
93
- batch_size=1, modality=field_modality
94
- )
95
- return mm_fields_config
90
+ return {
91
+ field_name: MultiModalFieldConfig.batched(modality=field_modality)
92
+ for field_name, field_modality in fields.items()
93
+ }
94
95
return _terratorch_field_config
96
@@ -192,9 +190,12 @@ class TerratorchMultiModalProcessor(BaseMultiModalProcessor):
190
) -> MultiModalInputs:
191
if "image" in mm_data:
192
image_data = mm_data["image"]
193
+ image_data = {k: v.unsqueeze(0) for k, v in image_data.items()}
194
else:
195
image_data = mm_data
197
- mm_data = {"image": mm_data}
196
+ image_data = {k: v.unsqueeze(0) for k, v in image_data.items()}
197
+
198
+ mm_data = {"image": image_data}
199
200
mm_items = self._to_mm_items(mm_data)
201
tokenization_kwargs = tokenization_kwargs or {}
@@ -226,6 +227,7 @@ class TerratorchMultiModalProcessor(BaseMultiModalProcessor):
227
dummy_inputs=TerratorchInputBuilder,
228
)
229
class Terratorch(nn.Module, IsAttentionFree, SupportsMultiModal):
230
+ merge_by_field_config = True
231
supports_multimodal_raw_input_only = True
232
is_pooling_model = True
233