@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 224 lines 7.06 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from typing import Annotated, Literal
5
6 import pytest
7 import torch
8
9 from vllm.model_executor.models.glm4_1v import Glm4vImageEmbeddingInputs
10 from vllm.model_executor.models.granite_speech import GraniteSpeechAudioInputs
11 from vllm.model_executor.models.phi3v import Phi3VImagePixelInputs
12 from vllm.utils.tensor_schema import TensorSchema, TensorShape
13
14
15 class _DoubleNestedVideoPixelInputs(TensorSchema):
16 """Minimal schema exercising doubly-nested list[list[torch.Tensor]] inputs.
17
18 Dimensions:
19 - n: Number of videos
20 - f: Number of frames
21 - g: Number of grids
22 - c: Number of channels (3)
23 - h: Height
24 - w: Width
25 """
26
27 type: Literal["pixel_values_videos"] = "pixel_values_videos"
28 pixel_values_videos: Annotated[
29 list[list[torch.Tensor]],
30 TensorShape("n", "f", "g", 3, "h", "w", dynamic_dims={"f", "g"}),
31 ]
32
33
34 def test_tensor_schema_valid_tensor():
35 Phi3VImagePixelInputs(
36 pixel_values=torch.randn(16, 64, 3, 32, 32),
37 image_sizes=torch.randint(0, 256, (16, 2)),
38 )
39
40
41 def test_tensor_schema_optional_fields():
42 Phi3VImagePixelInputs(
43 pixel_values=torch.randn(16, 64, 3, 32, 32),
44 image_sizes=None,
45 )
46
47 Phi3VImagePixelInputs(pixel_values=torch.randn(16, 64, 3, 32, 32))
48
49
50 def test_tensor_schema_constant_dim_failure():
51 with pytest.raises(ValueError, match="dim\\[2\\] expected 3, got 4"):
52 Phi3VImagePixelInputs(
53 pixel_values=torch.randn(16, 64, 4, 32, 32), # dim[2] = 4
54 image_sizes=torch.randint(0, 256, (16, 2)),
55 )
56
57
58 def test_tensor_schema_invalid_types_in_list():
59 with pytest.raises(TypeError, match="is not one of the expected types"):
60 Phi3VImagePixelInputs(
61 pixel_values=[
62 torch.randn(64, 3, 32, 32),
63 "not_a_tensor",
64 torch.randn(64, 3, 32, 32),
65 ],
66 image_sizes=torch.randint(0, 256, (3, 2)),
67 )
68
69
70 def test_tensor_schema_rank_mismatch():
71 with pytest.raises(ValueError, match="has rank 3 but expected 5"):
72 Phi3VImagePixelInputs(
73 pixel_values=torch.randn(16, 64, 3),
74 image_sizes=torch.randint(0, 256, (16, 2)),
75 )
76
77
78 def test_tensor_schema_missing_required_field():
79 with pytest.raises(ValueError, match="Required field 'pixel_values' is missing"):
80 Phi3VImagePixelInputs(
81 image_sizes=torch.randint(0, 256, (16, 2)),
82 )
83
84
85 def test_tensor_schema_symbolic_dim_mismatch():
86 with pytest.raises(ValueError, match="expected 'bn'=12, got 16"):
87 Phi3VImagePixelInputs(
88 pixel_values=torch.randn(12, 64, 3, 32, 32),
89 image_sizes=torch.randint(0, 256, (16, 2)),
90 )
91
92
93 def test_tensor_schema_list_tensor_valid():
94 Phi3VImagePixelInputs(
95 pixel_values=[torch.randn(64, 3, 32, 32) for _ in range(16)],
96 image_sizes=torch.randint(0, 256, (16, 2)),
97 )
98
99
100 def test_tensor_schema_variable_patch_counts_valid():
101 # Each image has a different number of patches (p)
102 # Each tensor has shape (p, 3, 32, 32)
103 Phi3VImagePixelInputs(
104 pixel_values=[
105 torch.randn(16, 3, 32, 32), # p = 16
106 torch.randn(32, 3, 32, 32), # p = 32
107 torch.randn(64, 3, 32, 32), # p = 64
108 ],
109 image_sizes=torch.randint(0, 256, (3, 2)), # bn = 3
110 )
111
112
113 def test_tensor_schema_tuple_tensor_valid():
114 Phi3VImagePixelInputs(
115 pixel_values=tuple(torch.randn(64, 3, 32, 32) for _ in range(16)),
116 image_sizes=torch.randint(0, 256, (16, 2)),
117 )
118
119
120 def test_tensor_schema_double_nested_tensors():
121 x = torch.rand(4, 3, 32, 32)
122 y = torch.rand(2, 3, 32, 32)
123
124 _DoubleNestedVideoPixelInputs(pixel_values_videos=([x, y, x], [y], [x, y]))
125
126
127 def test_tensor_schema_inconsistent_shapes_in_list():
128 with pytest.raises(ValueError, match="contains inconsistent shapes"):
129 Phi3VImagePixelInputs(
130 pixel_values=[
131 torch.randn(64, 3, 32, 32),
132 torch.randn(64, 3, 16, 16),
133 *(torch.randn(64, 3, 32, 32) for _ in range(14)),
134 ],
135 image_sizes=torch.randint(0, 256, (16, 2)),
136 )
137
138
139 def test_tensor_schema_empty_list():
140 with pytest.raises(ValueError, match="is an empty sequence"):
141 Phi3VImagePixelInputs(
142 pixel_values=[],
143 image_sizes=torch.randint(0, 256, (0, 2)),
144 )
145
146
147 def test_tensor_schema_validation_disabled_skips_shape_check():
148 # This should NOT raise, because validation is turned off
149 # This would normally fail (dim[2] should be 3, not 4)
150 Phi3VImagePixelInputs(
151 pixel_values=torch.randn(16, 64, 4, 32, 32),
152 image_sizes=torch.randint(0, 256, (16, 2)),
153 validate=False,
154 )
155
156
157 def test_tensor_schema_with_valid_resolve_binding_dims():
158 pixel_values = torch.randn(16, 64, 3, 336, 336) # h=336, w=336
159 image_sizes = torch.randint(0, 256, (16, 2))
160
161 Phi3VImagePixelInputs(
162 pixel_values=pixel_values,
163 image_sizes=image_sizes,
164 resolve_bindings={"h": 336, "w": 336},
165 )
166
167
168 def test_tensor_schema_with_invalid_resolve_binding_dims():
169 pixel_values = torch.randn(16, 64, 3, 36, 36) # h=36, w=36
170 image_sizes = torch.randint(0, 256, (16, 2))
171
172 # Should raise because 'h' and 'w' don't match resolve bindings
173 with pytest.raises(ValueError, match="dim\\[3\\] expected 336, got 36"):
174 Phi3VImagePixelInputs(
175 pixel_values=pixel_values,
176 image_sizes=image_sizes,
177 resolve_bindings={"h": 336, "w": 336},
178 )
179
180
181 def test_tensor_schema_with_list_of_symbolic_dim():
182 input_features = torch.randn(3, 10, 160) # (b=3, fi=10, 160)
183 input_features_mask = torch.randn(3, 8) # (b=3, fo=8)
184 audio_embed_sizes = [8, 8, 8] # len = b = 3
185
186 GraniteSpeechAudioInputs(
187 input_features=input_features,
188 input_features_mask=input_features_mask,
189 audio_embed_sizes=audio_embed_sizes,
190 )
191
192
193 def test_tensor_schema_with_list_of_symbolic_dim_mismatch_in_length():
194 input_features = torch.randn(4, 10, 160) # (b=4, fi=10, 160)
195 input_features_mask = torch.randn(4, 8) # (b=4, fo=8)
196 audio_embed_sizes = [8, 8, 8] # len = 3 ≠ b
197
198 with pytest.raises(ValueError, match="expected 'b'=4, got 3"):
199 GraniteSpeechAudioInputs(
200 input_features=input_features,
201 input_features_mask=input_features_mask,
202 audio_embed_sizes=audio_embed_sizes,
203 )
204
205
206 def test_valid_tensor_schema_with_static_last_dim():
207 image_embeds = torch.randn(256, 1024)
208 image_grid_thw = torch.randint(0, 4, (2, 3))
209
210 Glm4vImageEmbeddingInputs(
211 image_embeds=image_embeds,
212 image_grid_thw=image_grid_thw,
213 )
214
215
216 def test_invalid_tensor_schema_with_static_last_dim():
217 image_embeds = torch.randn(256, 1024)
218 image_grid_thw = torch.randint(0, 4, (2, 4)) # Wrong last dim
219
220 with pytest.raises(ValueError, match="dim\\[1\\] expected 3, got 4"):
221 Glm4vImageEmbeddingInputs(
222 image_embeds=image_embeds,
223 image_grid_thw=image_grid_thw,
224 )