@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 965bc71b04
Fork 0

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

Integrate TensorSchema with shape validation for Phi3VImagePixelInputs (#21232)

Signed-off-by: Benji Beck <benjibeck@meta.com>

Benji Beck committed Jul 24, 2025 at 21:43 UTC 965bc71b0445d7010bd40b3808d423d49ee68e58
3 files changed +372 -72
tests/standalone_tests/test_tensor_schema.py new
+126
@@ -0,0 +1,126 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +from vllm.model_executor.models.phi3v import Phi3VImagePixelInputs
8 +
9 +
10 +def test_tensor_schema_valid_tensor():
11 + Phi3VImagePixelInputs(
12 + data=torch.randn(16, 64, 3, 32, 32),
13 + image_sizes=torch.randint(0, 256, (16, 2)),
14 + )
15 +
16 +
17 +def test_tensor_schema_optional_fields():
18 + Phi3VImagePixelInputs(
19 + data=torch.randn(16, 64, 3, 32, 32),
20 + image_sizes=None,
21 + )
22 +
23 + Phi3VImagePixelInputs(data=torch.randn(16, 64, 3, 32, 32), )
24 +
25 +
26 +def test_tensor_schema_constant_dim_failure():
27 + with pytest.raises(ValueError, match="dim\\[2\\] expected 3, got 4"):
28 + Phi3VImagePixelInputs(
29 + data=torch.randn(16, 64, 4, 32, 32), # dim[2] = 4
30 + image_sizes=torch.randint(0, 256, (16, 2)),
31 + )
32 +
33 +
34 +def test_tensor_schema_symbolic_dim_mismatch():
35 + with pytest.raises(ValueError, match="expected 'bn'=12, got 16"):
36 + Phi3VImagePixelInputs(
37 + data=torch.randn(12, 64, 3, 32, 32),
38 + image_sizes=torch.randint(0, 256, (16, 2)),
39 + )
40 +
41 +
42 +def test_tensor_schema_list_tensor_valid():
43 + Phi3VImagePixelInputs(
44 + data=[torch.randn(64, 3, 32, 32) for _ in range(16)],
45 + image_sizes=torch.randint(0, 256, (16, 2)),
46 + )
47 +
48 +
49 +def test_tensor_schema_variable_patch_counts_valid():
50 + # Each image has a different number of patches (p)
51 + # Each tensor has shape (p, 3, 32, 32)
52 + data = [
53 + torch.randn(16, 3, 32, 32), # p = 16
54 + torch.randn(32, 3, 32, 32), # p = 32
55 + torch.randn(64, 3, 32, 32), # p = 64
56 + ]
57 + image_sizes = torch.randint(0, 256, (3, 2)) # bn = 3
58 + Phi3VImagePixelInputs(
59 + data=data,
60 + image_sizes=image_sizes,
61 + )
62 +
63 +
64 +def test_tensor_schema_tuple_tensor_valid():
65 + Phi3VImagePixelInputs(
66 + data=tuple(torch.randn(64, 3, 32, 32) for _ in range(16)),
67 + image_sizes=torch.randint(0, 256, (16, 2)),
68 + )
69 +
70 +
71 +def test_tensor_schema_inconsistent_shapes_in_list():
72 + with pytest.raises(ValueError, match="contains inconsistent shapes"):
73 + Phi3VImagePixelInputs(
74 + data=[torch.randn(64, 3, 32, 32),
75 + torch.randn(64, 3, 16, 16)] +
76 + [torch.randn(64, 3, 32, 32) for _ in range(14)],
77 + image_sizes=torch.randint(0, 256, (16, 2)),
78 + )
79 +
80 +
81 +def test_tensor_schema_empty_list():
82 + with pytest.raises(ValueError, match="is an empty list"):
83 + Phi3VImagePixelInputs(
84 + data=[],
85 + image_sizes=torch.randint(0, 256, (0, 2)),
86 + )
87 +
88 +
89 +def test_tensor_schema_validation_disabled_skips_shape_check():
90 + # This should NOT raise, because validation is turned off
91 + # This would normally fail (dim[2] should be 3, not 4)
92 + Phi3VImagePixelInputs(
93 + data=torch.randn(16, 64, 4, 32, 32),
94 + image_sizes=torch.randint(0, 256, (16, 2)),
95 + validate=False,
96 + )
97 +
98 +
99 +def test_tensor_schema_with_valid_resolve_binding_dims():
100 + data = torch.randn(16, 64, 3, 336, 336) # h=336, w=336
101 + image_sizes = torch.randint(0, 256, (16, 2))
102 +
103 + Phi3VImagePixelInputs(
104 + data=data,
105 + image_sizes=image_sizes,
106 + resolve_bindings={
107 + "h": 336,
108 + "w": 336
109 + },
110 + )
111 +
112 +
113 +def test_tensor_schema_with_invalid_resolve_binding_dims():
114 + data = torch.randn(16, 64, 3, 36, 36) # h=36, w=36
115 + image_sizes = torch.randint(0, 256, (16, 2))
116 +
117 + # Should raise because 'h' and 'w' don't match resolve bindings
118 + with pytest.raises(ValueError, match="dim\\[3\\] expected 336, got 36"):
119 + Phi3VImagePixelInputs(
120 + data=data,
121 + image_sizes=image_sizes,
122 + resolve_bindings={
123 + "h": 336,
124 + "w": 336
125 + },
126 + )
vllm/model_executor/models/phi3v.py
+36 -72
@@ -16,7 +16,7 @@
16 # See the License for the specific language governing permissions and
17 # limitations under the License.
18 from collections.abc import Iterable, Mapping, Sequence
19 -from typing import Any, Literal, Optional, TypedDict, Union
19 +from typing import Annotated, Any, Literal, Optional, Union
20
21 import regex as re
22 import torch
@@ -45,6 +45,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
45 from vllm.multimodal.profiling import BaseDummyInputsBuilder
46 from vllm.sequence import IntermediateTensors
47 from vllm.utils import is_list_of
48 +from vllm.utils.tensor_schema import TensorSchema, TensorShape
49
50 from .clip import CLIPVisionModel
51 from .interfaces import (MultiModalEmbeddings, SupportsMultiModal, SupportsPP,
@@ -93,32 +94,42 @@ def _init_img_processor(hf_config: PretrainedConfig,
94 return img_processor
95
96
96 -class Phi3VImagePixelInputs(TypedDict):
97 - type: Literal["pixel_values"]
98 - data: Union[torch.Tensor, list[torch.Tensor]]
97 +class Phi3VImagePixelInputs(TensorSchema):
98 """
100 - Shape:
101 - `(batch_size * num_images, 1 + num_patches, num_channels, height, width)`
102 -
103 - Note that `num_patches` may be different per batch and image,
104 - in which case the data is passed as a list instead of a batched tensor.
99 + Dimensions:
100 + - b: Batch size
101 + - n: Number of images
102 + - p: Number of patches
103 + - h: Height of each patch
104 + - w: Width of each patch
105 """
106
107 - image_sizes: torch.Tensor
108 - """
109 - Shape: `(batch_size * num_images, 2)`
107 + type: Literal["pixel_values", "image_embeds"] = "pixel_values"
108
111 - This should be in `(height, width)` format.
112 - """
109 + # Supports either a stacked tensor or a list of (p, 3, h, w) tensors
110 + data: Annotated[
111 + Union[torch.Tensor, list[torch.Tensor]],
112 + TensorShape("bn", "p", 3, "h", "w", dynamic_dims={"p"}
113 + ), # 'p' may vary across items
114 + ]
115
116 + # Stacked tensor with height and width for each image
117 + image_sizes: Annotated[Optional[torch.Tensor], TensorShape("bn", 2)]
118
115 -class Phi3VImageEmbeddingInputs(TypedDict):
116 - type: Literal["image_embeds"]
117 - data: Union[torch.Tensor, list[torch.Tensor]]
118 - """Shape: `(batch_size * num_images, image_feature_size, hidden_size)`
119
120 - `hidden_size` must match the hidden size of language model backbone.
120 +class Phi3VImageEmbeddingInputs(TensorSchema):
121 """
122 + Dimensions:
123 + - b: Batch size
124 + - n: Number of images
125 + - f: Image feature size (e.g., number of tokens per image)
126 + - h: Hidden size (must match language model backbone)
127 + """
128 + type: Literal["image_embeds"] = "image_embeds"
129 + data: Annotated[
130 + Union[torch.Tensor, list[torch.Tensor]],
131 + TensorShape("bn", "f", "h"),
132 + ]
133
134
135 Phi3VImageInputs = Union[Phi3VImagePixelInputs, Phi3VImageEmbeddingInputs]
@@ -563,44 +574,6 @@ class Phi3VForCausalLM(nn.Module, SupportsMultiModal, SupportsPP,
574 self.make_empty_intermediate_tensors = (
575 self.language_model.make_empty_intermediate_tensors)
576
566 - def _validate_image_sizes(self, data: torch.Tensor) -> torch.Tensor:
567 - expected_dims = (2, )
568 -
569 - def _validate_shape(d: torch.Tensor):
570 - actual_dims = tuple(d.shape)
571 -
572 - if actual_dims != expected_dims:
573 - expected_expr = str(expected_dims)
574 - raise ValueError(
575 - f"The expected shape of image sizes per image per batch "
576 - f"is {expected_expr}. You supplied {tuple(d.shape)}.")
577 -
578 - for d in data:
579 - _validate_shape(d)
580 -
581 - return data
582 -
583 - def _validate_pixel_values(
584 - self, data: Union[torch.Tensor, list[torch.Tensor]]
585 - ) -> Union[torch.Tensor, list[torch.Tensor]]:
586 -
587 - h = w = CLIP_VIT_LARGE_PATCH14_336_CONFIG.image_size
588 - expected_dims = (3, h, w)
589 -
590 - def _validate_shape(d: torch.Tensor):
591 - actual_dims = tuple(d.shape[1:])
592 -
593 - if actual_dims != expected_dims:
594 - expected_expr = ("num_patches", *map(str, expected_dims))
595 - raise ValueError(
596 - "The expected shape of pixel values per image per batch "
597 - f"is {expected_expr}. You supplied {tuple(d.shape)}.")
598 -
599 - for d in data:
600 - _validate_shape(d)
601 -
602 - return data
603 -
577 def _parse_and_validate_image_input(
578 self, **kwargs: object) -> Optional[Phi3VImageInputs]:
579 pixel_values = kwargs.pop("pixel_values", None)
@@ -611,25 +584,16 @@ class Phi3VForCausalLM(nn.Module, SupportsMultiModal, SupportsPP,
584 return None
585
586 if pixel_values is not None:
614 - if not isinstance(pixel_values, (torch.Tensor, list)):
615 - raise ValueError("Incorrect type of pixel values. "
616 - f"Got type: {type(pixel_values)}")
617 -
618 - if not isinstance(image_sizes, (torch.Tensor, list)):
619 - raise ValueError("Incorrect type of image sizes. "
620 - f"Got type: {type(image_sizes)}")
621 -
587 return Phi3VImagePixelInputs(
588 type="pixel_values",
624 - data=self._validate_pixel_values(flatten_bn(pixel_values)),
625 - image_sizes=self._validate_image_sizes(
626 - flatten_bn(image_sizes, concat=True)))
589 + data=flatten_bn(pixel_values),
590 + image_sizes=flatten_bn(image_sizes, concat=True),
591 + resolve_bindings={
592 + "h": CLIP_VIT_LARGE_PATCH14_336_CONFIG.image_size,
593 + "w": CLIP_VIT_LARGE_PATCH14_336_CONFIG.image_size
594 + })
595
596 if image_embeds is not None:
629 - if not isinstance(image_embeds, torch.Tensor):
630 - raise ValueError("Incorrect type of image embeddings. "
631 - f"Got type: {type(image_embeds)}")
632 -
597 return Phi3VImageEmbeddingInputs(
598 type="image_embeds",
599 data=flatten_bn(image_embeds),
vllm/utils/tensor_schema.py new
+210
@@ -0,0 +1,210 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Annotated, Any, Union, get_args, get_origin, get_type_hints
4 +
5 +import torch
6 +
7 +from vllm.logger import init_logger
8 +
9 +logger = init_logger(__name__)
10 +
11 +
12 +class TensorShape:
13 +
14 + def __init__(self,
15 + *dims: Union[int, str],
16 + dynamic_dims: set[str, ...] = None) -> None:
17 + self.dims = dims
18 + self.dynamic_dims = dynamic_dims if dynamic_dims else set()
19 +
20 + def resolve(self, **bindings: dict[str,
21 + int]) -> tuple[Union[int, str], ...]:
22 + resolved = []
23 + for dim in self.dims:
24 + if isinstance(dim, str) and dim in bindings:
25 + resolved.append(bindings[dim])
26 + else:
27 + resolved.append(dim)
28 + return tuple(resolved)
29 +
30 + def __str__(self) -> str:
31 + """Return a string representation of the tensor shape."""
32 + dim_strs = []
33 + for dim in self.dims:
34 + if isinstance(dim, str):
35 + if dim in self.dynamic_dims:
36 + dim_strs.append(
37 + f"{dim}*") # Mark dynamic dimensions with *
38 + else:
39 + dim_strs.append(dim)
40 + else:
41 + dim_strs.append(str(dim))
42 + return f"({', '.join(dim_strs)})"
43 +
44 +
45 +class TensorSchema:
46 +
47 + def __init__(self,
48 + *,
49 + validate: bool = True,
50 + resolve_bindings: dict[str, int] = None,
51 + **kwargs: Any) -> None:
52 + self._resolve_bindings = resolve_bindings if resolve_bindings else {}
53 +
54 + for key, value in kwargs.items():
55 + setattr(self, key, value)
56 +
57 + if validate:
58 + self.validate()
59 +
60 + def __getitem__(self, item) -> Any:
61 + return getattr(self, item)
62 +
63 + def _match_shape_with_dynamic(self, actual: tuple[int, ...],
64 + reference: tuple[int, ...],
65 + expected_shape: tuple[Union[int, str], ...],
66 + dynamic_dims: set[str, ...]) -> bool:
67 + if len(actual) != len(reference) or len(actual) > len(expected_shape):
68 + return False
69 +
70 + for i, (a, r) in enumerate(zip(actual, reference)):
71 + # When validating list inputs, we match shape suffixes only
72 + # (e.g. "p", 3, "h", "w"), assuming the list length corresponds
73 + # to the leading symbolic dim (e.g. "bn"). This allows comparing
74 + # only the trailing dimensions of each element in the list.
75 + dim = expected_shape[-len(actual) + i]
76 + # Skip this dimension if it's marked dynamic
77 + if dim in dynamic_dims:
78 + continue
79 + if a != r:
80 + return False
81 + return True
82 +
83 + def _validate_nested_tensors(
84 + self, value: Union[list[torch.Tensor, ...],
85 + tuple[torch.Tensor, ...]], field_name: str,
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 +
92 + # Ensure all tensors in the list have the same
93 + # shape, besides dynamic dimensions
94 + first = value[0]
95 + for i, v in enumerate(value):
96 + if not isinstance(v, torch.Tensor):
97 + raise ValueError(f"{field_name}[{i}] is not a "
98 + f"torch.Tensor")
99 + if not self._match_shape_with_dynamic(
100 + v.shape,
101 + first.shape,
102 + expected_shape,
103 + dynamic_dims,
104 + ):
105 + raise ValueError(f"{field_name} contains inconsistent "
106 + f"shapes: {first.shape} vs {v.shape} "
107 + f"at index {i}")
108 +
109 + # Treat the list as a stacked tensor:
110 + # shape = (len(list), *tensor.shape)
111 + return (len(value), ) + first.shape
112 +
113 + def _validate_tensor_shape_expected(self, actual_shape: tuple[int, ...],
114 + expected_shape: tuple[Union[int, str],
115 + ...],
116 + field_name: str, shape_env: dict[str,
117 + int],
118 + dynamic_dims: set[str, ...]) -> None:
119 + """Validate that the actual tensor shape matches the expected shape."""
120 + if len(actual_shape) != len(expected_shape):
121 + raise ValueError(f"{field_name} has rank {len(actual_shape)} "
122 + f"but expected {len(expected_shape)}")
123 +
124 + for i, dim in enumerate(expected_shape):
125 + if dim in dynamic_dims:
126 + continue
127 + elif isinstance(dim, int):
128 + if actual_shape[i] != dim:
129 + raise ValueError(f"{field_name} dim[{i}] expected "
130 + f"{dim}, got {actual_shape[i]}")
131 + elif isinstance(dim, str):
132 + if dim in shape_env:
133 + if actual_shape[i] != shape_env[dim]:
134 + raise ValueError(f"{field_name} dim[{i}] expected "
135 + f"'{dim}'={shape_env[dim]}, got "
136 + f"{actual_shape[i]}")
137 + else:
138 + shape_env[dim] = actual_shape[i]
139 + else:
140 + raise TypeError(f"{field_name} dim[{i}] has unsupported "
141 + f"type: {type(dim)}")
142 +
143 + def validate(self) -> None:
144 + type_hints = get_type_hints(self.__class__, include_extras=True)
145 + shape_env = {}
146 +
147 + for field_name, field_type in type_hints.items():
148 + # Check if field is missing
149 + if (not hasattr(self, field_name)
150 + or getattr(self, field_name) is None):
151 + # Check if field is marked as optional
152 + actual_type = field_type
153 + if get_origin(field_type) is Annotated:
154 + args = get_args(field_type)
155 + actual_type = args[0]
156 +
157 + # Check arg was provided as Union
158 + if get_origin(actual_type) is Union:
159 + args = get_args(actual_type)
160 + # Skip validation when Union contains None
161 + if type(None) in args:
162 + continue
163 + # If not optional, raise error
164 + raise ValueError(f"Required field '{field_name}' is missing")
165 +
166 + # Field exists, proceed with validation
167 + value = getattr(self, field_name)
168 +
169 + if get_origin(field_type) is not None:
170 + args = get_args(field_type)
171 +
172 + for arg in args:
173 + if isinstance(arg, TensorShape):
174 + expected_shape = arg.resolve(**self._resolve_bindings)
175 + if isinstance(value, (list, tuple)):
176 + actual_shape = self._validate_nested_tensors(
177 + value, field_name, expected_shape,
178 + arg.dynamic_dims)
179 +
180 + elif isinstance(value, torch.Tensor):
181 + actual_shape = value.shape
182 +
183 + else:
184 + type_names = []
185 + for arg in args:
186 + if hasattr(arg, "__name__"):
187 + type_names.append(str(arg.__name__))
188 + else:
189 + type_names.append(str(arg))
190 +
191 + expected_types = ", ".join(type_names)
192 + raise ValueError(
193 + f"{field_name} is not one of the expected "
194 + f"types: {expected_types}")
195 +
196 + self._validate_tensor_shape_expected(
197 + actual_shape, expected_shape, field_name,
198 + shape_env, arg.dynamic_dims)
199 +
200 + def print_shapes(self) -> None:
201 + """Print TensorShape annotations for debugging."""
202 + logger.debug("Shapes in %s:", self.__class__.__name__)
203 + type_hints = get_type_hints(self.__class__, include_extras=True)
204 +
205 + for field_name, field_type in type_hints.items():
206 + if get_origin(field_type) is not None:
207 + args = get_args(field_type)
208 + for arg in args:
209 + if isinstance(arg, TensorShape):
210 + logger.debug(" %s: %s", field_name, str(arg))