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, 23:45 UTC
Migrate InternVLImageInputs and InternVLVideoInputs to TensorSchema (#21684)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Jul 28, 2025 at 22:09 UTC
f1e2c095ecee01db02d0b63aae26d039b940d894
1 file changed
+49
-62
vllm/model_executor/models/internvl.py
+49
-62
@@ -9,7 +9,7 @@
9
# --------------------------------------------------------
10
from abc import ABC, abstractmethod
11
from collections.abc import Iterable, Mapping, Sequence
12
-from typing import Any, Literal, Optional, TypedDict, TypeVar, Union
12
+from typing import Annotated, Any, Literal, Optional, TypeVar, Union
13
14
import numpy.typing as npt
15
import torch
@@ -37,6 +37,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
37
from vllm.multimodal.profiling import BaseDummyInputsBuilder
38
from vllm.sequence import IntermediateTensors
39
from vllm.transformers_utils.tokenizer import AnyTokenizer
40
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
41
42
from .interfaces import (MultiModalEmbeddings, SupportsLoRA,
43
SupportsMultiModal, SupportsPP)
@@ -51,54 +52,60 @@ IMAGENET_MEAN = (0.485, 0.456, 0.406)
52
IMAGENET_STD = (0.229, 0.224, 0.225)
53
54
54
-class InternVLImagePixelInputs(TypedDict):
55
- type: Literal["pixel_values"]
56
- pixel_values_flat: torch.Tensor
55
+class InternVLImagePixelInputs(TensorSchema):
56
"""
58
- Shape:
59
- `(batch_size * num_images * (1 + num_patches), num_channels, height, width)`
57
+ Dimensions:
58
+ - bn: Batch size * number of images
59
+ - bnp: Batch size * number of images * (1 + num_patches)
60
+ - c: Number of channels (3)
61
+ - h: Height of each image patch
62
+ - w: Width of each image patch
63
"""
64
+ type: Literal["pixel_values"]
65
+ pixel_values_flat: Annotated[torch.Tensor, TensorShape("bnp", 3, "h", "w")]
66
+ num_patches: Annotated[torch.Tensor, TensorShape("bn")]
67
62
- num_patches: torch.Tensor
63
- """Shape: `(batch_size * num_images)`"""
64
-
65
-
66
-class InternVLImageEmbeddingInputs(TypedDict):
67
- type: Literal["image_embeds"]
68
- data: Union[torch.Tensor, list[torch.Tensor]]
69
- """
70
- A tensor of shape `(num_images, total_image_feature_size, hidden_size)`
71
- or a list of tensors of shape `(total_image_feature_size, hidden_size)`
68
73
- `hidden_size` must match the hidden size of language model backbone.
69
+class InternVLImageEmbeddingInputs(TensorSchema):
70
+ """
71
+ Dimensions:
72
+ - n: Number of images
73
+ - f: Total image feature size
74
+ - h: Hidden size (must match the hidden size of language model backbone)
75
"""
76
+ type: Literal["image_embeds"]
77
+ data: Annotated[Union[torch.Tensor, list[torch.Tensor]],
78
+ TensorShape("n", "f", "h")]
79
80
81
InternVLImageInputs = Union[InternVLImagePixelInputs,
82
InternVLImageEmbeddingInputs]
83
84
81
-class InternVLVideoPixelInputs(TypedDict):
82
- type: Literal["pixel_values_videos"]
83
- pixel_values_flat: torch.Tensor
85
+class InternVLVideoPixelInputs(TensorSchema):
86
"""
85
- Shape:
86
- `(batch_size * num_video * num_frames, num_channels, height, width)`
87
+ Dimensions:
88
+ - bvf: Batch size * number of videos * num_frames
89
+ - bn: Batch size * number of images
90
+ - c: Number of channels (3)
91
+ - h: Height of each video frame
92
+ - w: Width of each video frame
93
"""
94
+ type: Literal["pixel_values_videos"]
95
+ pixel_values_flat: Annotated[torch.Tensor, TensorShape("bvf", 3, "h", "w")]
96
+ num_patches: Annotated[torch.Tensor, TensorShape("bn")]
97
89
- num_patches: torch.Tensor
90
- """Shape: `(batch_size * num_images)`"""
91
-
92
-
93
-class InternVLVideoEmbeddingInputs(TypedDict):
94
- type: Literal["video_embeds"]
95
- data: Union[torch.Tensor, list[torch.Tensor]]
96
- """
97
- A tensor of shape `(num_videos, total_video_feature_size, hidden_size)`
98
- or a list of tensors of shape `(total_video_feature_size, hidden_size)`
98
100
- `hidden_size` must match the hidden size of language model backbone.
99
+class InternVLVideoEmbeddingInputs(TensorSchema):
100
+ """
101
+ Dimensions:
102
+ - n: Number of videos
103
+ - f: Total video feature size
104
+ - h: Hidden size (must match the hidden size of language model backbone)
105
"""
106
+ type: Literal["video_embeds"]
107
+ data: Annotated[Union[torch.Tensor, list[torch.Tensor]],
108
+ TensorShape("n", "f", "h")]
109
110
111
InternVLVideoInputs = Union[InternVLVideoPixelInputs,
@@ -1151,26 +1158,6 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP,
1158
vit_embeds = self.mlp1(vit_embeds)
1159
return vit_embeds
1160
1154
- def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:
1155
-
1156
- h = w = self.config.vision_config.image_size
1157
- expected_dims = (3, h, w)
1158
-
1159
- def _validate_shape(d: torch.Tensor):
1160
- actual_dims = tuple(d.shape)
1161
-
1162
- if actual_dims != expected_dims:
1163
- expected_expr = str(expected_dims)
1164
- raise ValueError(
1165
- "The expected shape of pixel values per image per batch "
1166
- f" per patch is {expected_expr}. "
1167
- f"You supplied {tuple(d.shape)}.")
1168
-
1169
- for d in data:
1170
- _validate_shape(d)
1171
-
1172
- return data
1173
-
1161
def _parse_and_validate_image_input(
1162
self, **kwargs: object) -> Optional[InternVLImageInputs]:
1163
pixel_values_flat = kwargs.pop("pixel_values_flat", None)
@@ -1205,12 +1192,14 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP,
1192
1193
pixel_values_flat = flatten_bn(pixel_values_flat, concat=True)
1194
image_num_patches = flatten_bn(image_num_patches, concat=True)
1195
+ expected_h = expected_w = self.config.vision_config.image_size
1196
+ resolve_bindings = {"h": expected_h, "w": expected_w}
1197
1198
return InternVLImagePixelInputs(
1199
type="pixel_values",
1211
- pixel_values_flat=self._validate_pixel_values(
1212
- pixel_values_flat),
1200
+ pixel_values_flat=pixel_values_flat,
1201
num_patches=image_num_patches,
1202
+ resolve_bindings=resolve_bindings,
1203
)
1204
1205
raise AssertionError("This line should be unreachable.")
@@ -1225,11 +1214,7 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP,
1214
return None
1215
1216
if video_embeds is not None:
1228
- if not isinstance(video_embeds, (torch.Tensor, list)):
1229
- raise ValueError("Incorrect type of video embeddings. "
1230
- f"Got type: {type(video_embeds)}")
1231
-
1232
- return InternVLImageEmbeddingInputs(
1217
+ return InternVLVideoEmbeddingInputs(
1218
type="video_embeds",
1219
data=flatten_bn(video_embeds),
1220
)
@@ -1250,12 +1235,14 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP,
1235
pixel_values_flat_video = flatten_bn(pixel_values_flat_video,
1236
concat=True)
1237
video_num_patches = flatten_bn(video_num_patches, concat=True)
1238
+ expected_h = expected_w = self.config.vision_config.image_size
1239
+ resolve_bindings = {"h": expected_h, "w": expected_w}
1240
1241
return InternVLVideoPixelInputs(
1242
type="pixel_values_videos",
1256
- pixel_values_flat=self._validate_pixel_values(
1257
- pixel_values_flat_video),
1243
+ pixel_values_flat=pixel_values_flat_video,
1244
num_patches=video_num_patches,
1245
+ resolve_bindings=resolve_bindings,
1246
)
1247
1248
raise AssertionError("This line should be unreachable.")