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 11, 05:00 UTC
Add hidden dimension validation for multimodal embedding inputs (#30968)
Signed-off-by: Wenqi Glantz <wglantz@nvidia.com>
Wenqi Glantz committed
Dec 19, 2025 at 02:59 UTC
4924ac582c5007fbf7b15d719708812e06961009
4 files changed
+589
-10
tests/entrypoints/openai/test_embedding_shape_validation.py
new
+223
@@ -0,0 +1,223 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Embedding shape validation in multimodal APIs.
5
+
6
+Tests verify that embeddings with correct ndim but incorrect hidden_size
7
+are rejected before they can cause crashes during model inference.
8
+
9
+Validation is performed by the parser (MultiModalDataParser) and EmbeddingItems
10
+classes, not by CompletionRenderer or MediaIO classes.
11
+"""
12
+
13
+import pytest
14
+import torch
15
+
16
+from vllm.multimodal.parse import (
17
+ AudioEmbeddingItems,
18
+ ImageEmbeddingItems,
19
+ MultiModalDataParser,
20
+ VideoEmbeddingItems,
21
+)
22
+
23
+
24
+class TestMultiModalParserShapeValidation:
25
+ """Test hidden_size validation in MultiModalDataParser."""
26
+
27
+ def test_image_embeddings_correct_hidden_size_accepted(self):
28
+ """Baseline: Image embeddings with correct hidden_size should work."""
29
+ expected_hidden_size = 768
30
+ parser = MultiModalDataParser(expected_hidden_size=expected_hidden_size)
31
+
32
+ valid_embeds = torch.randn(2, 100, expected_hidden_size)
33
+
34
+ result = parser.parse_mm_data({"image": valid_embeds})
35
+
36
+ assert "image" in result
37
+ assert isinstance(result["image"], ImageEmbeddingItems)
38
+ assert result["image"].get_count() == 2
39
+
40
+ def test_image_embeddings_wrong_hidden_size_rejected(self):
41
+ """Security: Image embeddings with wrong hidden_size should be rejected."""
42
+ expected_hidden_size = 768
43
+ wrong_hidden_size = 4096
44
+ parser = MultiModalDataParser(expected_hidden_size=expected_hidden_size)
45
+
46
+ invalid_embeds = torch.randn(2, 100, wrong_hidden_size)
47
+
48
+ with pytest.raises(ValueError) as exc_info:
49
+ parser.parse_mm_data({"image": invalid_embeds})
50
+
51
+ error_msg = str(exc_info.value).lower()
52
+ assert "image" in error_msg
53
+ assert "hidden dimension mismatch" in error_msg
54
+
55
+ def test_audio_embeddings_wrong_hidden_size_rejected(self):
56
+ """Security: Audio embeddings with wrong hidden_size should be rejected."""
57
+ expected_hidden_size = 768
58
+ wrong_hidden_size = 2048
59
+ parser = MultiModalDataParser(expected_hidden_size=expected_hidden_size)
60
+
61
+ invalid_embeds = torch.randn(2, 100, wrong_hidden_size)
62
+
63
+ with pytest.raises(ValueError) as exc_info:
64
+ parser.parse_mm_data({"audio": invalid_embeds})
65
+
66
+ error_msg = str(exc_info.value).lower()
67
+ assert "audio" in error_msg
68
+ assert "hidden dimension mismatch" in error_msg
69
+
70
+ def test_video_embeddings_wrong_hidden_size_rejected(self):
71
+ """Security: Video embeddings with wrong hidden_size should be rejected."""
72
+ expected_hidden_size = 768
73
+ wrong_hidden_size = 512
74
+ parser = MultiModalDataParser(expected_hidden_size=expected_hidden_size)
75
+
76
+ invalid_embeds = torch.randn(2, 100, wrong_hidden_size)
77
+
78
+ with pytest.raises(ValueError) as exc_info:
79
+ parser.parse_mm_data({"video": invalid_embeds})
80
+
81
+ error_msg = str(exc_info.value).lower()
82
+ assert "video" in error_msg
83
+ assert "hidden dimension mismatch" in error_msg
84
+
85
+ def test_list_of_embeddings_validates_each(self):
86
+ """Security: Each embedding in list should be validated."""
87
+ expected_hidden_size = 768
88
+ wrong_hidden_size = 1024
89
+ parser = MultiModalDataParser(expected_hidden_size=expected_hidden_size)
90
+
91
+ # List with second tensor having wrong hidden_size
92
+ invalid_embeds = [
93
+ torch.randn(100, expected_hidden_size),
94
+ torch.randn(100, wrong_hidden_size),
95
+ ]
96
+
97
+ with pytest.raises(ValueError) as exc_info:
98
+ parser.parse_mm_data({"image": invalid_embeds})
99
+
100
+ # Should identify which embedding failed
101
+ assert "[1]" in str(exc_info.value)
102
+
103
+ def test_validation_disabled_allows_any_size(self):
104
+ """When validation disabled (legacy), any hidden_size allowed."""
105
+ parser = MultiModalDataParser(expected_hidden_size=None)
106
+
107
+ any_hidden_size = 12345
108
+ embeds = torch.randn(2, 100, any_hidden_size)
109
+
110
+ # Should not raise
111
+ result = parser.parse_mm_data({"image": embeds})
112
+ assert "image" in result
113
+ assert isinstance(result["image"], ImageEmbeddingItems)
114
+
115
+
116
+class TestEmbeddingItemsDirectValidation:
117
+ """Direct tests for EmbeddingItems hidden_size validation."""
118
+
119
+ def test_image_embedding_items_validates_batched_tensor(self):
120
+ """Test validation for batched (3D) image embeddings."""
121
+ expected = 768
122
+ wrong = 1024
123
+
124
+ # Valid
125
+ valid = torch.randn(2, 100, expected)
126
+ items = ImageEmbeddingItems(valid, expected_hidden_size=expected)
127
+ assert items.get_count() == 2
128
+
129
+ # Invalid
130
+ invalid = torch.randn(2, 100, wrong)
131
+ with pytest.raises(ValueError) as exc_info:
132
+ ImageEmbeddingItems(invalid, expected_hidden_size=expected)
133
+
134
+ assert str(wrong) in str(exc_info.value)
135
+ assert str(expected) in str(exc_info.value)
136
+
137
+ def test_image_embedding_items_validates_list_of_tensors(self):
138
+ """Test validation for list of 2D image embeddings."""
139
+ expected = 768
140
+ wrong = 512
141
+
142
+ # Valid list
143
+ valid_list = [torch.randn(100, expected), torch.randn(50, expected)]
144
+ items = ImageEmbeddingItems(valid_list, expected_hidden_size=expected)
145
+ assert items.get_count() == 2
146
+
147
+ # Invalid list
148
+ invalid_list = [torch.randn(100, expected), torch.randn(50, wrong)]
149
+ with pytest.raises(ValueError) as exc_info:
150
+ ImageEmbeddingItems(invalid_list, expected_hidden_size=expected)
151
+
152
+ assert "[1]" in str(exc_info.value)
153
+
154
+ def test_audio_embedding_items_validates(self):
155
+ """Test validation for audio embeddings."""
156
+ expected = 768
157
+ wrong = 256
158
+
159
+ invalid = torch.randn(2, 100, wrong)
160
+ with pytest.raises(ValueError) as exc_info:
161
+ AudioEmbeddingItems(invalid, expected_hidden_size=expected)
162
+
163
+ assert "audio" in str(exc_info.value).lower()
164
+
165
+ def test_video_embedding_items_validates(self):
166
+ """Test validation for video embeddings."""
167
+ expected = 768
168
+ wrong = 384
169
+
170
+ invalid = torch.randn(2, 100, wrong)
171
+ with pytest.raises(ValueError) as exc_info:
172
+ VideoEmbeddingItems(invalid, expected_hidden_size=expected)
173
+
174
+ assert "video" in str(exc_info.value).lower()
175
+
176
+
177
+class TestShapeValidationIntegration:
178
+ """Integration tests verifying attack scenarios are blocked."""
179
+
180
+ def test_attack_scenario_multimodal_image(self):
181
+ """
182
+ Simulate attack through Chat API with image embeddings.
183
+
184
+ Verifies validation occurs in multimodal parser path.
185
+ """
186
+ expected_hidden_size = 768
187
+ wrong_hidden_size = 4096
188
+ parser = MultiModalDataParser(expected_hidden_size=expected_hidden_size)
189
+
190
+ attack_tensor = torch.randn(1, 100, wrong_hidden_size)
191
+
192
+ with pytest.raises(ValueError):
193
+ parser.parse_mm_data({"image": attack_tensor})
194
+
195
+ def test_attack_scenario_multimodal_audio(self):
196
+ """
197
+ Simulate attack through Chat API with audio embeddings.
198
+
199
+ Verifies validation occurs in multimodal parser path.
200
+ """
201
+ expected_hidden_size = 768
202
+ wrong_hidden_size = 2048
203
+ parser = MultiModalDataParser(expected_hidden_size=expected_hidden_size)
204
+
205
+ attack_tensor = torch.randn(1, 100, wrong_hidden_size)
206
+
207
+ with pytest.raises(ValueError):
208
+ parser.parse_mm_data({"audio": attack_tensor})
209
+
210
+ def test_attack_scenario_multimodal_video(self):
211
+ """
212
+ Simulate attack through Chat API with video embeddings.
213
+
214
+ Verifies validation occurs in multimodal parser path.
215
+ """
216
+ expected_hidden_size = 768
217
+ wrong_hidden_size = 1024
218
+ parser = MultiModalDataParser(expected_hidden_size=expected_hidden_size)
219
+
220
+ attack_tensor = torch.randn(1, 100, wrong_hidden_size)
221
+
222
+ with pytest.raises(ValueError):
223
+ parser.parse_mm_data({"video": attack_tensor})
tests/multimodal/test_embedding_shape_validation_unit.py
new
+249
@@ -0,0 +1,249 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Unit tests for embedding shape validation.
5
+
6
+Simple, fast unit tests that can run without server fixtures.
7
+Run with: pytest tests/multimodal/test_embedding_shape_validation_unit.py -v
8
+"""
9
+
10
+import pytest
11
+import torch
12
+
13
+from vllm.multimodal.parse import (
14
+ AudioEmbeddingItems,
15
+ ImageEmbeddingItems,
16
+)
17
+
18
+
19
+class TestImageEmbedBasicValidation:
20
+ """Test basic ndim validation in image embeddings via ImageEmbeddingItems."""
21
+
22
+ def test_valid_2d_tensor_accepted(self):
23
+ """Baseline: 2D tensors should be accepted."""
24
+ valid_tensor = torch.randn(10, 768, dtype=torch.float32)
25
+
26
+ # Should not raise - 2D is valid
27
+ items = ImageEmbeddingItems(valid_tensor)
28
+ assert items.get_count() == 10
29
+
30
+ def test_valid_3d_tensor_accepted(self):
31
+ """Baseline: 3D tensors should be accepted."""
32
+ valid_tensor = torch.randn(2, 10, 768, dtype=torch.float32)
33
+
34
+ # Should not raise - 3D is valid
35
+ items = ImageEmbeddingItems(valid_tensor)
36
+ assert items.get_count() == 2
37
+
38
+ def test_valid_list_of_2d_tensors_accepted(self):
39
+ """Baseline: List of 2D tensors should be accepted."""
40
+ tensors = [
41
+ torch.randn(10, 768, dtype=torch.float32),
42
+ torch.randn(15, 768, dtype=torch.float32),
43
+ ]
44
+
45
+ # Should not raise
46
+ items = ImageEmbeddingItems(tensors)
47
+ assert items.get_count() == 2
48
+
49
+ def test_1d_tensor_rejected(self):
50
+ """Security: 1D tensors should be rejected (invalid ndim)."""
51
+ invalid_tensor = torch.randn(768, dtype=torch.float32) # 1D
52
+
53
+ with pytest.raises(ValueError) as exc_info:
54
+ ImageEmbeddingItems(invalid_tensor)
55
+
56
+ assert "must be 2D" in str(exc_info.value) or "3D" in str(exc_info.value)
57
+
58
+ def test_4d_tensor_rejected(self):
59
+ """Security: 4D tensors should be rejected (invalid ndim)."""
60
+ invalid_tensor = torch.randn(1, 2, 10, 768, dtype=torch.float32) # 4D
61
+
62
+ with pytest.raises(ValueError) as exc_info:
63
+ ImageEmbeddingItems(invalid_tensor)
64
+
65
+ assert "must be 2D" in str(exc_info.value) or "3D" in str(exc_info.value)
66
+
67
+ def test_hidden_size_validation_correct_size(self):
68
+ """Embeddings with correct hidden size should be accepted."""
69
+ expected_hidden_size = 768
70
+ valid_tensor = torch.randn(10, expected_hidden_size, dtype=torch.float32)
71
+
72
+ # Should not raise
73
+ items = ImageEmbeddingItems(
74
+ valid_tensor, expected_hidden_size=expected_hidden_size
75
+ )
76
+ assert items.get_count() == 10
77
+
78
+ def test_hidden_size_validation_wrong_size_rejected(self):
79
+ """Embeddings with wrong hidden size should be rejected."""
80
+ expected_hidden_size = 768
81
+ wrong_hidden_size = 4096
82
+ invalid_tensor = torch.randn(10, wrong_hidden_size, dtype=torch.float32)
83
+
84
+ with pytest.raises(ValueError) as exc_info:
85
+ ImageEmbeddingItems(
86
+ invalid_tensor, expected_hidden_size=expected_hidden_size
87
+ )
88
+
89
+ error_msg = str(exc_info.value)
90
+ assert "hidden dimension mismatch" in error_msg.lower()
91
+ assert str(wrong_hidden_size) in error_msg
92
+ assert str(expected_hidden_size) in error_msg
93
+
94
+
95
+class TestAudioEmbedBasicValidation:
96
+ """Test basic ndim validation in audio embeddings via AudioEmbeddingItems."""
97
+
98
+ def test_valid_2d_tensor_accepted(self):
99
+ """Baseline: 2D tensors should be accepted."""
100
+ valid_tensor = torch.randn(10, 768, dtype=torch.float32)
101
+
102
+ # Should not raise - 2D is valid
103
+ items = AudioEmbeddingItems(valid_tensor)
104
+ assert items.get_count() == 10
105
+
106
+ def test_valid_3d_tensor_accepted(self):
107
+ """Baseline: 3D tensors should be accepted."""
108
+ valid_tensor = torch.randn(2, 10, 768, dtype=torch.float32)
109
+
110
+ # Should not raise - 3D is valid
111
+ items = AudioEmbeddingItems(valid_tensor)
112
+ assert items.get_count() == 2
113
+
114
+ def test_valid_list_of_2d_tensors_accepted(self):
115
+ """Baseline: List of 2D tensors should be accepted."""
116
+ tensors = [
117
+ torch.randn(10, 768, dtype=torch.float32),
118
+ torch.randn(15, 768, dtype=torch.float32),
119
+ ]
120
+
121
+ # Should not raise
122
+ items = AudioEmbeddingItems(tensors)
123
+ assert items.get_count() == 2
124
+
125
+ def test_1d_tensor_rejected(self):
126
+ """Security: 1D tensors should be rejected (invalid ndim)."""
127
+ invalid_tensor = torch.randn(768, dtype=torch.float32) # 1D
128
+
129
+ with pytest.raises(ValueError) as exc_info:
130
+ AudioEmbeddingItems(invalid_tensor)
131
+
132
+ assert "must be 2D" in str(exc_info.value) or "3D" in str(exc_info.value)
133
+
134
+ def test_scalar_rejected(self):
135
+ """Security: Scalar tensors should be rejected."""
136
+ invalid_tensor = torch.tensor(1.0) # 0D (scalar)
137
+
138
+ with pytest.raises(ValueError):
139
+ AudioEmbeddingItems(invalid_tensor)
140
+
141
+ def test_hidden_size_validation_correct_size(self):
142
+ """Embeddings with correct hidden size should be accepted."""
143
+ expected_hidden_size = 768
144
+ valid_tensor = torch.randn(10, expected_hidden_size, dtype=torch.float32)
145
+
146
+ # Should not raise
147
+ items = AudioEmbeddingItems(
148
+ valid_tensor, expected_hidden_size=expected_hidden_size
149
+ )
150
+ assert items.get_count() == 10
151
+
152
+ def test_hidden_size_validation_wrong_size_rejected(self):
153
+ """Embeddings with wrong hidden size should be rejected."""
154
+ expected_hidden_size = 768
155
+ wrong_hidden_size = 4096
156
+ invalid_tensor = torch.randn(10, wrong_hidden_size, dtype=torch.float32)
157
+
158
+ with pytest.raises(ValueError) as exc_info:
159
+ AudioEmbeddingItems(
160
+ invalid_tensor, expected_hidden_size=expected_hidden_size
161
+ )
162
+
163
+ error_msg = str(exc_info.value)
164
+ assert "hidden dimension mismatch" in error_msg.lower()
165
+ assert str(wrong_hidden_size) in error_msg
166
+ assert str(expected_hidden_size) in error_msg
167
+
168
+
169
+class TestShapeValidationDoSPrevention:
170
+ """
171
+ Tests for DoS prevention through shape validation.
172
+
173
+ Verifies that embeddings with incorrect shapes are rejected early,
174
+ preventing crashes during model inference.
175
+ """
176
+
177
+ def test_prevent_crash_from_wrong_shape_image_embeds(self):
178
+ """
179
+ Prevent crash scenario: wrong hidden size in image embeddings.
180
+
181
+ Without validation, this would pass initial checks but crash later
182
+ during model forward pass when dimensions don't match.
183
+ """
184
+ expected_hidden_size = 768 # Typical model hidden size
185
+ wrong_hidden_size = 4096 # Wrong size (e.g., Llama-sized)
186
+
187
+ wrong_embedding = torch.randn(100, wrong_hidden_size, dtype=torch.float32)
188
+
189
+ # Should be rejected at instantiation time, not during inference
190
+ with pytest.raises(ValueError) as exc_info:
191
+ ImageEmbeddingItems(
192
+ wrong_embedding, expected_hidden_size=expected_hidden_size
193
+ )
194
+
195
+ error_msg = str(exc_info.value)
196
+ assert "hidden dimension mismatch" in error_msg.lower()
197
+ assert str(expected_hidden_size) in error_msg # Expected
198
+ assert str(wrong_hidden_size) in error_msg # Received
199
+
200
+ def test_prevent_crash_from_wrong_shape_audio_embeds(self):
201
+ """
202
+ Prevent crash scenario: wrong hidden size in audio embeddings.
203
+ """
204
+ expected_hidden_size = 768
205
+ wrong_hidden_size = 4096
206
+
207
+ wrong_embedding = torch.randn(100, wrong_hidden_size, dtype=torch.float32)
208
+
209
+ with pytest.raises(ValueError) as exc_info:
210
+ AudioEmbeddingItems(
211
+ wrong_embedding, expected_hidden_size=expected_hidden_size
212
+ )
213
+
214
+ error_msg = str(exc_info.value)
215
+ assert "hidden dimension mismatch" in error_msg.lower()
216
+
217
+ def test_extremely_large_hidden_size_rejected(self):
218
+ """Security: Prevent DoS from extremely large embeddings."""
219
+ expected_hidden_size = 768
220
+ huge_hidden_size = 100000 # Large but not extreme to avoid test OOM
221
+
222
+ invalid_tensor = torch.randn(10, huge_hidden_size, dtype=torch.float32)
223
+
224
+ with pytest.raises(ValueError) as exc_info:
225
+ ImageEmbeddingItems(
226
+ invalid_tensor, expected_hidden_size=expected_hidden_size
227
+ )
228
+
229
+ assert "hidden dimension mismatch" in str(exc_info.value).lower()
230
+
231
+ def test_batch_with_mixed_hidden_sizes_rejected(self):
232
+ """All embeddings in a list must have the same hidden size."""
233
+ expected_hidden_size = 768
234
+
235
+ # One correct, one wrong
236
+ batch = [
237
+ torch.randn(10, expected_hidden_size, dtype=torch.float32),
238
+ torch.randn(10, expected_hidden_size + 100, dtype=torch.float32), # Wrong!
239
+ ]
240
+
241
+ # Should fail on the second one
242
+ with pytest.raises(ValueError) as exc_info:
243
+ ImageEmbeddingItems(batch, expected_hidden_size=expected_hidden_size)
244
+
245
+ assert "hidden dimension mismatch" in str(exc_info.value).lower()
246
+
247
+
248
+if __name__ == "__main__":
249
+ pytest.main([__file__, "-v", "--tb=short"])
vllm/multimodal/parse.py
+108
-9
@@ -126,6 +126,30 @@ class ProcessorBatchItems(ModalityDataItems[Sequence[_T], _T]):
126
return {}
127
128
129
+def validate_embedding_ndim(
130
+ tensor: torch.Tensor,
131
+ modality: str,
132
+ index: int | None = None,
133
+) -> None:
134
+ """Validate tensor ndim for multimodal embeddings.
135
+
136
+ Single embeddings should be 2D (seq_len, hidden_size).
137
+ Batched embeddings should be 3D (batch, seq_len, hidden_size).
138
+
139
+ Args:
140
+ tensor: The tensor to validate.
141
+ modality: The modality name for error messages (e.g., "image", "audio").
142
+ index: Optional index for list items, included in error messages.
143
+ """
144
+ if tensor.ndim < 2 or tensor.ndim > 3:
145
+ idx_str = f" [{index}]" if index is not None else ""
146
+ raise ValueError(
147
+ f"{modality.capitalize()} embedding{idx_str} must be 2D "
148
+ f"(seq_len, hidden_size) or 3D (batch, seq_len, hidden_size), "
149
+ f"got {tensor.ndim}D tensor with shape {tuple(tensor.shape)}"
150
+ )
151
+
152
+
153
class EmbeddingItems(
154
ModalityDataItems[torch.Tensor | list[torch.Tensor], torch.Tensor]
155
):
@@ -134,6 +158,63 @@ class EmbeddingItems(
158
or a list of embedding tensors (one per item).
159
"""
160
161
+ def __init__(
162
+ self,
163
+ data: torch.Tensor | list[torch.Tensor],
164
+ modality: str,
165
+ expected_hidden_size: int | None = None,
166
+ ) -> None:
167
+ super().__init__(data, modality)
168
+
169
+ # Validate ndim first (before hidden_size which depends on correct ndim)
170
+ self._validate_ndim()
171
+
172
+ # Validate hidden dimension if expected size is provided
173
+ if expected_hidden_size is not None:
174
+ self._validate_hidden_size(expected_hidden_size)
175
+
176
+ def _validate_ndim(self) -> None:
177
+ """Validate that embedding tensors have correct ndim (2D or 3D)."""
178
+ if isinstance(self.data, torch.Tensor):
179
+ validate_embedding_ndim(self.data, self.modality)
180
+ else:
181
+ # List of tensors: each should be 2D (seq_len, hidden_size)
182
+ for idx, tensor in enumerate(self.data):
183
+ if tensor.ndim != 2:
184
+ raise ValueError(
185
+ f"{self.modality.capitalize()} embedding [{idx}] must be "
186
+ f"2D (seq_len, hidden_size), got {tensor.ndim}D tensor "
187
+ f"with shape {tuple(tensor.shape)}"
188
+ )
189
+
190
+ def _validate_hidden_size(self, expected_hidden_size: int) -> None:
191
+ """Validate that embedding hidden dimension matches expected size.
192
+
193
+ This validates hidden dimensions to prevent vulnerabilities: Embeddings
194
+ with correct ndim but wrong hidden dimension could bypass initial
195
+ checks and cause crashes during model inference when dimensions don't match.
196
+ """
197
+ if isinstance(self.data, torch.Tensor):
198
+ # Batched tensor: shape is (batch, seq_len, hidden_size)
199
+ actual_hidden_size = self.data.shape[-1]
200
+ if actual_hidden_size != expected_hidden_size:
201
+ raise ValueError(
202
+ f"{self.modality.capitalize()} embedding hidden dimension "
203
+ f"mismatch: got {actual_hidden_size}, but model expects "
204
+ f"{expected_hidden_size}. Embedding shape: {tuple(self.data.shape)}"
205
+ )
206
+ else:
207
+ # List of tensors: each has shape (seq_len, hidden_size)
208
+ for idx, tensor in enumerate(self.data):
209
+ actual_hidden_size = tensor.shape[-1]
210
+ if actual_hidden_size != expected_hidden_size:
211
+ raise ValueError(
212
+ f"{self.modality.capitalize()} embedding [{idx}] hidden "
213
+ f"dimension mismatch: got {actual_hidden_size}, but model "
214
+ f"expects {expected_hidden_size}. "
215
+ f"Embedding shape: {tuple(tensor.shape)}"
216
+ )
217
+
218
def _unwrap(
219
self, item: torch.Tensor | MediaWithBytes[torch.Tensor]
220
) -> torch.Tensor:
@@ -228,8 +309,12 @@ class AudioProcessorItems(ProcessorBatchItems[HfAudioItem]):
309
310
311
class AudioEmbeddingItems(EmbeddingItems):
231
- def __init__(self, data: torch.Tensor | list[torch.Tensor]) -> None:
232
- super().__init__(data, "audio")
312
+ def __init__(
313
+ self,
314
+ data: torch.Tensor | list[torch.Tensor],
315
+ expected_hidden_size: int | None = None,
316
+ ) -> None:
317
+ super().__init__(data, "audio", expected_hidden_size)
318
319
320
class ImageSize(NamedTuple):
@@ -256,8 +341,12 @@ class ImageProcessorItems(ProcessorBatchItems[HfImageItem]):
341
342
343
class ImageEmbeddingItems(EmbeddingItems):
259
- def __init__(self, data: torch.Tensor | list[torch.Tensor]) -> None:
260
- super().__init__(data, "image")
344
+ def __init__(
345
+ self,
346
+ data: torch.Tensor | list[torch.Tensor],
347
+ expected_hidden_size: int | None = None,
348
+ ) -> None:
349
+ super().__init__(data, "image", expected_hidden_size)
350
351
352
class VideoProcessorItems(ProcessorBatchItems[HfVideoItem]):
@@ -287,8 +376,12 @@ class VideoProcessorItems(ProcessorBatchItems[HfVideoItem]):
376
377
378
class VideoEmbeddingItems(EmbeddingItems):
290
- def __init__(self, data: torch.Tensor | list[torch.Tensor]) -> None:
291
- super().__init__(data, "video")
379
+ def __init__(
380
+ self,
381
+ data: torch.Tensor | list[torch.Tensor],
382
+ expected_hidden_size: int | None = None,
383
+ ) -> None:
384
+ super().__init__(data, "video", expected_hidden_size)
385
386
387
_D = TypeVar("_D", bound=ModalityDataItems[Any, Any])
@@ -363,6 +456,10 @@ class MultiModalDataParser:
456
Args:
457
target_sr (float, optional): Enables automatic resampling of audio
458
items to the model's expected sampling rate.
459
+ expected_hidden_size (int, optional): Expected hidden dimension for
460
+ embedding inputs. If provided, validates that user-supplied
461
+ embeddings have the correct hidden size to prevent crashes
462
+ during model inference.
463
"""
464
465
def __init__(
@@ -371,6 +468,7 @@ class MultiModalDataParser:
468
target_sr: float | None = None,
469
audio_resample_method: Literal["librosa", "scipy"] = "librosa",
470
video_needs_metadata: bool = False,
471
+ expected_hidden_size: int | None = None,
472
) -> None:
473
super().__init__()
474
@@ -379,6 +477,7 @@ class MultiModalDataParser:
477
method=audio_resample_method,
478
)
479
self.video_needs_metadata = video_needs_metadata
480
+ self.expected_hidden_size = expected_hidden_size
481
482
@classmethod
483
def is_embeddings(
@@ -443,7 +542,7 @@ class MultiModalDataParser:
542
return None
543
544
if self.is_embeddings(data):
446
- return AudioEmbeddingItems(data)
545
+ return AudioEmbeddingItems(data, self.expected_hidden_size)
546
547
data_items: list[AudioItem]
548
if (
@@ -481,7 +580,7 @@ class MultiModalDataParser:
580
return None
581
582
if self.is_embeddings(data):
484
- return ImageEmbeddingItems(data)
583
+ return ImageEmbeddingItems(data, self.expected_hidden_size)
584
585
if (
586
isinstance(data, (PILImage.Image, MediaWithBytes))
@@ -507,7 +606,7 @@ class MultiModalDataParser:
606
return None
607
608
if self.is_embeddings(data):
510
- return VideoEmbeddingItems(data)
609
+ return VideoEmbeddingItems(data, self.expected_hidden_size)
610
611
data_items: list[VideoItem]
612
if (
vllm/multimodal/processing.py
+9
-1
@@ -1330,7 +1330,15 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
1330
of [`MultiModalDataParser`][vllm.multimodal.parse.MultiModalDataParser]
1331
that has additional subparsers.
1332
"""
1333
- return MultiModalDataParser()
1333
+ # Get expected hidden size for embedding validation if mm_embeds enabled
1334
+ # This validates hidden dimensions to prevent vulnerabilities: embeddings
1335
+ # with correct ndim but wrong shape could cause crashes at inference time
1336
+ mm_config = self.info.ctx.model_config.get_multimodal_config()
1337
+ expected_hidden_size = None
1338
+ if mm_config.enable_mm_embeds:
1339
+ expected_hidden_size = self.info.ctx.model_config.get_inputs_embeds_size()
1340
+
1341
+ return MultiModalDataParser(expected_hidden_size=expected_hidden_size)
1342
1343
def validate_num_items(
1344
self,