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 FuyuImagePatchInputs to TensorSchema (#21662)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Jul 26, 2025 at 19:34 UTC
3339cba3ff4de0d4a516b31076dc73673510c227
3 files changed
+57
-45
tests/standalone_tests/test_tensor_schema.py
+22
@@ -4,6 +4,7 @@
4
import pytest
5
import torch
6
7
+from vllm.model_executor.models.fuyu import FuyuImagePatchInputs
8
from vllm.model_executor.models.phi3v import Phi3VImagePixelInputs
9
10
@@ -124,3 +125,24 @@ def test_tensor_schema_with_invalid_resolve_binding_dims():
125
"w": 336
126
},
127
)
128
+
129
+
130
+def test_tensor_schema_with_list_of_symbolic_dim():
131
+ flat_data = torch.stack([torch.randn(768) for _ in range(3)]) # (bn=3, fn)
132
+ patches_per_image = [64, 64, 64] # len = bn = 3
133
+
134
+ FuyuImagePatchInputs(
135
+ flat_data=flat_data,
136
+ patches_per_image=patches_per_image,
137
+ )
138
+
139
+
140
+def test_tensor_schema_with_list_of_symbolic_dim_mismatch_in_length():
141
+ flat_data = torch.stack([torch.randn(768) for _ in range(4)]) # (bn=4, fn)
142
+ patches_per_image = [64, 64, 64] # len = 3 ≠ bn
143
+
144
+ with pytest.raises(ValueError, match="expected 'bn'=4, got 3"):
145
+ FuyuImagePatchInputs(
146
+ flat_data=flat_data,
147
+ patches_per_image=patches_per_image,
148
+ )
\ No newline at end of file
vllm/model_executor/models/fuyu.py
+19
-36
@@ -19,7 +19,7 @@
19
""" PyTorch Fuyu model."""
20
import math
21
from collections.abc import Iterable, Mapping, Sequence
22
-from typing import Literal, Optional, TypedDict
22
+from typing import Annotated, Literal, Optional
23
24
import torch
25
import torch.nn as nn
@@ -40,6 +40,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
40
PromptUpdate, PromptUpdateDetails)
41
from vllm.multimodal.profiling import BaseDummyInputsBuilder
42
from vllm.sequence import IntermediateTensors
43
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
44
45
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
46
from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn, maybe_prefix,
@@ -50,18 +51,24 @@ _IMAGE_TOKEN_ID = 71011
51
_NEWLINE_TOKEN_ID = 71019
52
53
53
-class FuyuImagePatchInputs(TypedDict):
54
- type: Literal["image_patches"]
55
- flat_data: torch.Tensor
54
+class FuyuImagePatchInputs(TensorSchema):
55
"""
57
- Shape:
58
- `(batch_size * num_patches, patch_size_x * patch_size_y * num_channels)`
56
+ Dimensions:
57
+ - bn: Batch size * number of images
58
+ - fn: Num channels * patch_size_x * patch_size_y
59
"""
60
61
- patches_per_image: list[int]
61
+ type: Literal["image_patches"] = "image_patches"
62
+
63
+ flat_data: Annotated[
64
+ torch.Tensor,
65
+ TensorShape("bn", "fn"),
66
+ ]
67
+
68
+ patches_per_image: Annotated[list[int], TensorShape("bn")]
69
"""
70
The number of total patches for each image in the batch.
64
-
71
+
72
This is used to split the embeddings which has the first two dimensions
73
flattened just like `flat_data`.
74
"""
@@ -297,42 +304,18 @@ class FuyuForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
304
self.make_empty_intermediate_tensors = (
305
self.language_model.make_empty_intermediate_tensors)
306
300
- def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:
301
-
302
- h = w = self.config.patch_size
303
- num_channels = self.config.num_channels
304
- expected_dims = num_channels * h * w
305
-
306
- def _validate_shape(d: torch.Tensor):
307
- actual_dims = d.size(-1)
308
-
309
- if actual_dims != expected_dims:
310
- expected_expr = str(expected_dims)
311
- raise ValueError(
312
- "The expected shape of pixel values per image per batch "
313
- f"per patch is {expected_expr}. "
314
- f"You supplied {tuple(d.shape)}.")
315
-
316
- for d in data:
317
- _validate_shape(d)
318
-
319
- return data.to(self.vision_embed_tokens.weight.dtype)
320
-
307
def _parse_and_validate_image_input(
308
self, **kwargs: object) -> Optional[FuyuImagePatchInputs]:
309
image_patches = kwargs.pop("image_patches", None)
310
if image_patches is not None:
325
- if not isinstance(image_patches, (torch.Tensor, list)):
326
- raise ValueError("Incorrect type of image patches. "
327
- f"Got type: {type(image_patches)}")
328
-
311
image_patches_flat = flatten_bn(image_patches)
330
-
312
+ flat_data = flatten_bn(image_patches, concat=True).data.to(
313
+ self.vision_embed_tokens.weight.dtype)
314
return FuyuImagePatchInputs(
315
type="image_patches",
333
- flat_data=self._validate_pixel_values(
334
- flatten_bn(image_patches_flat, concat=True)),
316
+ flat_data=flat_data,
317
patches_per_image=[x.size(0) for x in image_patches_flat],
318
+ resolve_bindings={"fn": self.image_feature_size},
319
)
320
321
return None
vllm/utils/tensor_schema.py
+16
-9
@@ -86,9 +86,6 @@ class TensorSchema:
86
expected_shape: tuple[Union[int, str], ...],
87
dynamic_dims: set[str, ...]) -> tuple[int, ...]:
88
"""Validate a list/tuple of tensors and return the actual shape."""
89
- if not value:
90
- raise ValueError(f"{field_name} is an empty list")
91
-
89
# Ensure all tensors in the list have the same
90
# shape, besides dynamic dimensions
91
first = value[0]
@@ -117,6 +114,7 @@ class TensorSchema:
114
int],
115
dynamic_dims: set[str, ...]) -> None:
116
"""Validate that the actual tensor shape matches the expected shape."""
117
+
118
if len(actual_shape) != len(expected_shape):
119
raise ValueError(f"{field_name} has rank {len(actual_shape)} "
120
f"but expected {len(expected_shape)}")
@@ -160,12 +158,11 @@ class TensorSchema:
158
# Skip validation when Union contains None
159
if type(None) in args:
160
continue
163
- # If not optional, raise error
161
+ # Otherwise field is required, raise error
162
raise ValueError(f"Required field '{field_name}' is missing")
163
164
# Field exists, proceed with validation
165
value = getattr(self, field_name)
168
-
166
if get_origin(field_type) is not None:
167
args = get_args(field_type)
168
@@ -173,13 +170,23 @@ class TensorSchema:
170
if isinstance(arg, TensorShape):
171
expected_shape = arg.resolve(**self._resolve_bindings)
172
if isinstance(value, (list, tuple)):
176
- actual_shape = self._validate_nested_tensors(
177
- value, field_name, expected_shape,
178
- arg.dynamic_dims)
179
-
173
+ # list/tuple of Tensors → shape = (len(value), ...)
174
+ if value and isinstance(value[0], torch.Tensor):
175
+ actual_shape = self._validate_nested_tensors(
176
+ value, field_name, expected_shape,
177
+ arg.dynamic_dims)
178
+ elif value:
179
+ # list/tuple of scalars → shape = (len(value),)
180
+ actual_shape = (len(value), )
181
+ else:
182
+ raise ValueError(
183
+ f"{field_name} is an empty list")
184
+
185
+ # Tensor → shape = tensor.shape
186
elif isinstance(value, torch.Tensor):
187
actual_shape = value.shape
188
189
+ # Otherwise, it's an unsupported type
190
else:
191
type_names = []
192
for arg in args: