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 Glm4vImageInputs, Glm4vVideoInputs to TensorSchema (#21678)
Signed-off-by: Benji Beck <benjibeck@meta.com> Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Benji Beck committed
Jul 27, 2025 at 22:36 UTC
88e46c7c8dfa7c761259795773b758114721169a
2 files changed
+69
-66
tests/standalone_tests/test_tensor_schema.py
+23
-1
@@ -5,6 +5,7 @@ import pytest
5
import torch
6
7
from vllm.model_executor.models.fuyu import FuyuImagePatchInputs
8
+from vllm.model_executor.models.glm4_1v import Glm4vImageEmbeddingInputs
9
from vllm.model_executor.models.phi3v import Phi3VImagePixelInputs
10
11
@@ -145,4 +146,25 @@ def test_tensor_schema_with_list_of_symbolic_dim_mismatch_in_length():
146
FuyuImagePatchInputs(
147
flat_data=flat_data,
148
patches_per_image=patches_per_image,
148
- )
\ No newline at end of file
149
+ )
150
+
151
+
152
+def test_valid_tensor_schema_with_static_last_dim():
153
+ image_embeds = torch.randn(256, 1024)
154
+ image_grid_thw = torch.randint(0, 4, (2, 3))
155
+
156
+ Glm4vImageEmbeddingInputs(
157
+ image_embeds=image_embeds,
158
+ image_grid_thw=image_grid_thw,
159
+ )
160
+
161
+
162
+def test_invalid_tensor_schema_with_static_last_dim():
163
+ image_embeds = torch.randn(256, 1024)
164
+ image_grid_thw = torch.randint(0, 4, (2, 4)) # Wrong last dim
165
+
166
+ with pytest.raises(ValueError, match="dim\\[1\\] expected 3, got 4"):
167
+ Glm4vImageEmbeddingInputs(
168
+ image_embeds=image_embeds,
169
+ image_grid_thw=image_grid_thw,
170
+ )
vllm/model_executor/models/glm4_1v.py
+46
-65
@@ -29,7 +29,7 @@
29
import math
30
from collections.abc import Iterable, Mapping, Sequence
31
from functools import partial
32
-from typing import Any, Callable, Literal, Optional, TypedDict, Union
32
+from typing import Annotated, Any, Callable, Literal, Optional, Union
33
34
import numpy as np
35
import torch
@@ -70,6 +70,7 @@ from vllm.multimodal.profiling import BaseDummyInputsBuilder
70
from vllm.platforms import _Backend
71
from vllm.sequence import IntermediateTensors
72
from vllm.transformers_utils.config import uses_mrope
73
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
74
75
from ..layers.activation import SiluAndMul
76
from .interfaces import (MultiModalEmbeddings, SupportsLoRA,
@@ -88,80 +89,68 @@ _MAX_FRAMES_PER_VIDEO = 600
89
# === Vision Inputs === #
90
91
91
-class Glm4vImagePixelInputs(TypedDict):
92
- type: Literal["pixel_values"]
93
- pixel_values: torch.Tensor
94
- """Shape:
95
- `(num_patches, num_channels * patch_size * patch_size)`
92
+class Glm4vImagePixelInputs(TensorSchema):
93
"""
97
-
98
- image_grid_thw: torch.Tensor
99
- """Shape: `(num_images, 3)`
100
- This should be in `(grid_t, grid_h, grid_w)` format.
94
+ Dimensions:
95
+ - np: Number of patches
96
+ - cpp: Number of channels * patch_size * patch_size
97
+ - ni: Number of images
98
+ - g: Grid dimensions (3 for grid_t, grid_h, grid_w)
99
"""
100
+ type: Literal["pixel_values"] = "pixel_values"
101
102
+ pixel_values: Annotated[torch.Tensor, TensorShape("np", "cpp")]
103
+ image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]
104
104
-class Glm4vImageEmbeddingInputs(TypedDict):
105
- type: Literal["image_embeds"]
106
- image_embeds: torch.Tensor
107
- """Supported types:
108
- - List[`torch.Tensor`]: A list of tensors holding all images' features.
109
- Each tensor holds an image's features.
110
- - `torch.Tensor`: A tensor holding all images' features
111
- (concatenation of all images' feature tensors).
105
113
- Tensor shape: `(num_image_features, hidden_size)`
114
- - `num_image_features` varies based on
115
- the number and resolution of the images.
116
- - `hidden_size` must match the hidden size of language model backbone.
106
+class Glm4vImageEmbeddingInputs(TensorSchema):
107
"""
118
-
119
- image_grid_thw: torch.Tensor
120
- """Shape: `(num_images, 3)`
121
- This should be in `(grid_t, grid_h, grid_w)` format.
108
+ Dimensions:
109
+ - f: Number of image features (varies based on image resolution)
110
+ - h: Hidden size (must match language model backbone)
111
+ - n: Number of images
112
+ - g: Grid dimensions (3 for grid_t, grid_h, grid_w)
113
"""
114
+ type: Literal["image_embeds"] = "image_embeds"
115
+
116
+ image_embeds: Annotated[torch.Tensor, TensorShape("f", "h")]
117
+ image_grid_thw: Annotated[torch.Tensor, TensorShape("n", 3)]
118
119
120
Glm4vImageInputs = Union[Glm4vImagePixelInputs, Glm4vImageEmbeddingInputs]
121
122
128
-class Glm4vVideoPixelInputs(TypedDict):
129
- type: Literal["pixel_values_videos"]
130
- pixel_values_videos: torch.Tensor
131
- """Shape:
132
- `(num_patches,
133
- num_channels * temporal_patch_size * patch_size * patch_size)`
123
+class Glm4vVideoPixelInputs(TensorSchema):
124
"""
135
- # video_metadata: Union[list[VideoMetadata], list[dict]]
136
- video_grid_thw: Union[list[torch.Tensor], torch.Tensor]
137
- """Shape: `(num_videos, num_frames, 3)` or `(1, num_frames, 3)`
138
- for single video.
139
- Each entry represents [grid_t, grid_h, grid_w] format where:
140
- - grid_t: Temporal grid size (usually 1 for processed video)
141
- - grid_h: Height grid size
142
- - grid_w: Width grid size
143
- This describes the grid structure of the video patches.
125
+ Dimensions:
126
+ - np: Number of patches
127
+ - ctpp: Number of channels * temporal_patch_size *
128
+ patch_size * patch_size
129
+ - nv: Number of videos
130
+ - f: Number of frames
131
+ - g: Grid dimensions (3 for grid_t which is usually 1 for processed
132
+ video, grid_h, grid_w)
133
"""
134
+ type: Literal["pixel_values_videos"] = "pixel_values_videos"
135
136
+ pixel_values_videos: Annotated[torch.Tensor, TensorShape("np", "ctpp")]
137
+ # video_metadata: Union[list[VideoMetadata], list[dict]]
138
+ video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", "f", 3)]
139
147
-class Glm4vVideoEmbeddingInputs(TypedDict):
148
- type: Literal["video_embeds"]
140
150
- video_embeds: torch.Tensor
141
+class Glm4vVideoEmbeddingInputs(TensorSchema):
142
"""
152
- Tensor shape: `(num_video_patches, hidden_size)`
153
- - `num_video_patches`: Total number of video patches across all frames
154
- - `hidden_size`: Must match the hidden size of language model backbone
143
+ Dimensions:
144
+ - p: Number of video patches across all frames
145
+ - h: Hidden size (must match language model backbone)
146
+ - n: Number of videos
147
+ - g: Grid dimensions (3 for grid_t which is usually 1 for processed
148
+ video, grid_h, grid_w)
149
"""
150
+ type: Literal["video_embeds"] = "video_embeds"
151
157
- video_grid_thw: torch.Tensor
158
- """Shape: `(num_videos, 1, 3)` or `(1, 1, 3)` for single video
159
- Each entry represents [grid_t, grid_h, grid_w] format where:
160
- - grid_t: Temporal grid size (usually 1 for processed video)
161
- - grid_h: Height grid size
162
- - grid_w: Width grid size
163
- This describes the grid structure of the video patches.
164
- """
152
+ video_embeds: Annotated[torch.Tensor, TensorShape("p", "h")]
153
+ video_grid_thw: Annotated[torch.Tensor, TensorShape("n", 1, 3)]
154
155
156
Glm4vVideoInputs = Union[Glm4vVideoPixelInputs, Glm4vVideoEmbeddingInputs]
@@ -1324,10 +1313,6 @@ class Glm4vForConditionalGeneration(nn.Module, SupportsMultiModal,
1313
image_grid_thw = self._validate_and_reshape_mm_tensor(
1314
image_grid_thw, "image grid_thw")
1315
1327
- if not isinstance(pixel_values, (torch.Tensor, list)):
1328
- raise ValueError("Incorrect type of image pixel values. "
1329
- f"Got type: {type(pixel_values)}")
1330
-
1316
return Glm4vImagePixelInputs(
1317
type="pixel_values",
1318
pixel_values=pixel_values,
@@ -1340,9 +1325,6 @@ class Glm4vForConditionalGeneration(nn.Module, SupportsMultiModal,
1325
image_grid_thw = self._validate_and_reshape_mm_tensor(
1326
image_grid_thw, "image grid_thw")
1327
1343
- if not isinstance(image_embeds, torch.Tensor):
1344
- raise ValueError("Incorrect type of image embeddings. "
1345
- f"Got type: {type(image_embeds)}")
1328
return Glm4vImageEmbeddingInputs(
1329
type="image_embeds",
1330
image_embeds=image_embeds,
@@ -1354,8 +1336,10 @@ class Glm4vForConditionalGeneration(nn.Module, SupportsMultiModal,
1336
pixel_values_videos = kwargs.pop("pixel_values_videos", None)
1337
video_embeds = kwargs.pop("video_embeds", None)
1338
video_grid_thw = kwargs.pop("video_grid_thw", None)
1339
+
1340
if pixel_values_videos is None and video_embeds is None:
1341
return None
1342
+
1343
if pixel_values_videos is not None:
1344
pixel_values_videos = self._validate_and_reshape_mm_tensor(
1345
pixel_values_videos, "video pixel values")
@@ -1375,9 +1359,6 @@ class Glm4vForConditionalGeneration(nn.Module, SupportsMultiModal,
1359
video_grid_thw = self._validate_and_reshape_mm_tensor(
1360
video_grid_thw, "video grid_thw")
1361
1378
- if not isinstance(video_embeds, torch.Tensor):
1379
- raise ValueError("Incorrect type of video embeddings. "
1380
- f"Got type: {type(video_embeds)}")
1362
return Glm4vVideoEmbeddingInputs(
1363
type="video_embeds",
1364
video_embeds=video_embeds,