| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Embedding shape validation in multimodal APIs. |
| 4 | |
| 5 | Tests verify that embeddings with correct ndim but incorrect hidden_size |
| 6 | are rejected before they can cause crashes during model inference. |
| 7 | |
| 8 | Validation is performed by the parser (MultiModalDataParser) and EmbeddingItems |
| 9 | classes, not by MediaIO classes. |
| 10 | """ |
| 11 | |
| 12 | import pytest |
| 13 | import torch |
| 14 | |
| 15 | from vllm.multimodal.parse import ( |
| 16 | AudioEmbeddingItems, |
| 17 | ImageEmbeddingItems, |
| 18 | MultiModalDataParser, |
| 19 | VideoEmbeddingItems, |
| 20 | ) |
| 21 | |
| 22 | |
| 23 | class TestMultiModalParserShapeValidation: |
| 24 | """Test hidden_size validation in MultiModalDataParser.""" |
| 25 | |
| 26 | def test_image_embeddings_correct_hidden_size_accepted(self): |
| 27 | """Baseline: Image embeddings with correct hidden_size should work.""" |
| 28 | expected_hidden_size = 768 |
| 29 | parser = MultiModalDataParser(expected_hidden_size=expected_hidden_size) |
| 30 | |
| 31 | valid_embeds = torch.randn(2, 100, expected_hidden_size) |
| 32 | |
| 33 | result = parser.parse_mm_data({"image": valid_embeds}) |
| 34 | |
| 35 | assert "image" in result |
| 36 | assert isinstance(result["image"], ImageEmbeddingItems) |
| 37 | assert result["image"].get_count() == 2 |
| 38 | |
| 39 | def test_image_embeddings_wrong_hidden_size_rejected(self): |
| 40 | """Security: Image embeddings with wrong hidden_size should be rejected.""" |
| 41 | expected_hidden_size = 768 |
| 42 | wrong_hidden_size = 4096 |
| 43 | parser = MultiModalDataParser(expected_hidden_size=expected_hidden_size) |
| 44 | |
| 45 | invalid_embeds = torch.randn(2, 100, wrong_hidden_size) |
| 46 | |
| 47 | with pytest.raises(ValueError) as exc_info: |
| 48 | parser.parse_mm_data({"image": invalid_embeds}) |
| 49 | |
| 50 | error_msg = str(exc_info.value).lower() |
| 51 | assert "image" in error_msg |
| 52 | assert "hidden dimension mismatch" in error_msg |
| 53 | |
| 54 | def test_audio_embeddings_wrong_hidden_size_rejected(self): |
| 55 | """Security: Audio embeddings with wrong hidden_size should be rejected.""" |
| 56 | expected_hidden_size = 768 |
| 57 | wrong_hidden_size = 2048 |
| 58 | parser = MultiModalDataParser(expected_hidden_size=expected_hidden_size) |
| 59 | |
| 60 | invalid_embeds = torch.randn(2, 100, wrong_hidden_size) |
| 61 | |
| 62 | with pytest.raises(ValueError) as exc_info: |
| 63 | parser.parse_mm_data({"audio": invalid_embeds}) |
| 64 | |
| 65 | error_msg = str(exc_info.value).lower() |
| 66 | assert "audio" in error_msg |
| 67 | assert "hidden dimension mismatch" in error_msg |
| 68 | |
| 69 | def test_video_embeddings_wrong_hidden_size_rejected(self): |
| 70 | """Security: Video embeddings with wrong hidden_size should be rejected.""" |
| 71 | expected_hidden_size = 768 |
| 72 | wrong_hidden_size = 512 |
| 73 | parser = MultiModalDataParser(expected_hidden_size=expected_hidden_size) |
| 74 | |
| 75 | invalid_embeds = torch.randn(2, 100, wrong_hidden_size) |
| 76 | |
| 77 | with pytest.raises(ValueError) as exc_info: |
| 78 | parser.parse_mm_data({"video": invalid_embeds}) |
| 79 | |
| 80 | error_msg = str(exc_info.value).lower() |
| 81 | assert "video" in error_msg |
| 82 | assert "hidden dimension mismatch" in error_msg |
| 83 | |
| 84 | def test_list_of_embeddings_validates_each(self): |
| 85 | """Security: Each embedding in list should be validated.""" |
| 86 | expected_hidden_size = 768 |
| 87 | wrong_hidden_size = 1024 |
| 88 | parser = MultiModalDataParser(expected_hidden_size=expected_hidden_size) |
| 89 | |
| 90 | # List with second tensor having wrong hidden_size |
| 91 | invalid_embeds = [ |
| 92 | torch.randn(100, expected_hidden_size), |
| 93 | torch.randn(100, wrong_hidden_size), |
| 94 | ] |
| 95 | |
| 96 | with pytest.raises(ValueError) as exc_info: |
| 97 | parser.parse_mm_data({"image": invalid_embeds}) |
| 98 | |
| 99 | # Should identify which embedding failed |
| 100 | assert "[1]" in str(exc_info.value) |
| 101 | |
| 102 | def test_validation_disabled_allows_any_size(self): |
| 103 | """When validation disabled (legacy), any hidden_size allowed.""" |
| 104 | parser = MultiModalDataParser(expected_hidden_size=None) |
| 105 | |
| 106 | any_hidden_size = 12345 |
| 107 | embeds = torch.randn(2, 100, any_hidden_size) |
| 108 | |
| 109 | # Should not raise |
| 110 | result = parser.parse_mm_data({"image": embeds}) |
| 111 | assert "image" in result |
| 112 | assert isinstance(result["image"], ImageEmbeddingItems) |
| 113 | |
| 114 | |
| 115 | class TestEmbeddingItemsDirectValidation: |
| 116 | """Direct tests for EmbeddingItems hidden_size validation.""" |
| 117 | |
| 118 | def test_image_embedding_items_validates_batched_tensor(self): |
| 119 | """Test validation for batched (3D) image embeddings.""" |
| 120 | expected = 768 |
| 121 | wrong = 1024 |
| 122 | |
| 123 | # Valid |
| 124 | valid = torch.randn(2, 100, expected) |
| 125 | items = ImageEmbeddingItems(valid, expected_hidden_size=expected) |
| 126 | assert items.get_count() == 2 |
| 127 | |
| 128 | # Invalid |
| 129 | invalid = torch.randn(2, 100, wrong) |
| 130 | with pytest.raises(ValueError) as exc_info: |
| 131 | ImageEmbeddingItems(invalid, expected_hidden_size=expected) |
| 132 | |
| 133 | assert str(wrong) in str(exc_info.value) |
| 134 | assert str(expected) in str(exc_info.value) |
| 135 | |
| 136 | def test_image_embedding_items_validates_list_of_tensors(self): |
| 137 | """Test validation for list of 2D image embeddings.""" |
| 138 | expected = 768 |
| 139 | wrong = 512 |
| 140 | |
| 141 | # Valid list |
| 142 | valid_list = [torch.randn(100, expected), torch.randn(50, expected)] |
| 143 | items = ImageEmbeddingItems(valid_list, expected_hidden_size=expected) |
| 144 | assert items.get_count() == 2 |
| 145 | |
| 146 | # Invalid list |
| 147 | invalid_list = [torch.randn(100, expected), torch.randn(50, wrong)] |
| 148 | with pytest.raises(ValueError) as exc_info: |
| 149 | ImageEmbeddingItems(invalid_list, expected_hidden_size=expected) |
| 150 | |
| 151 | assert "[1]" in str(exc_info.value) |
| 152 | |
| 153 | def test_audio_embedding_items_validates(self): |
| 154 | """Test validation for audio embeddings.""" |
| 155 | expected = 768 |
| 156 | wrong = 256 |
| 157 | |
| 158 | invalid = torch.randn(2, 100, wrong) |
| 159 | with pytest.raises(ValueError) as exc_info: |
| 160 | AudioEmbeddingItems(invalid, expected_hidden_size=expected) |
| 161 | |
| 162 | assert "audio" in str(exc_info.value).lower() |
| 163 | |
| 164 | def test_video_embedding_items_validates(self): |
| 165 | """Test validation for video embeddings.""" |
| 166 | expected = 768 |
| 167 | wrong = 384 |
| 168 | |
| 169 | invalid = torch.randn(2, 100, wrong) |
| 170 | with pytest.raises(ValueError) as exc_info: |
| 171 | VideoEmbeddingItems(invalid, expected_hidden_size=expected) |
| 172 | |
| 173 | assert "video" in str(exc_info.value).lower() |
| 174 | |
| 175 | |
| 176 | class TestShapeValidationIntegration: |
| 177 | """Integration tests verifying attack scenarios are blocked.""" |
| 178 | |
| 179 | def test_attack_scenario_multimodal_image(self): |
| 180 | """Simulate attack through Chat API with image embeddings. |
| 181 | |
| 182 | Verifies validation occurs in multimodal parser path. |
| 183 | """ |
| 184 | expected_hidden_size = 768 |
| 185 | wrong_hidden_size = 4096 |
| 186 | parser = MultiModalDataParser(expected_hidden_size=expected_hidden_size) |
| 187 | |
| 188 | attack_tensor = torch.randn(1, 100, wrong_hidden_size) |
| 189 | |
| 190 | with pytest.raises(ValueError): |
| 191 | parser.parse_mm_data({"image": attack_tensor}) |
| 192 | |
| 193 | def test_attack_scenario_multimodal_audio(self): |
| 194 | """Simulate attack through Chat API with audio embeddings. |
| 195 | |
| 196 | Verifies validation occurs in multimodal parser path. |
| 197 | """ |
| 198 | expected_hidden_size = 768 |
| 199 | wrong_hidden_size = 2048 |
| 200 | parser = MultiModalDataParser(expected_hidden_size=expected_hidden_size) |
| 201 | |
| 202 | attack_tensor = torch.randn(1, 100, wrong_hidden_size) |
| 203 | |
| 204 | with pytest.raises(ValueError): |
| 205 | parser.parse_mm_data({"audio": attack_tensor}) |
| 206 | |
| 207 | def test_attack_scenario_multimodal_video(self): |
| 208 | """Simulate attack through Chat API with video embeddings. |
| 209 | |
| 210 | Verifies validation occurs in multimodal parser path. |
| 211 | """ |
| 212 | expected_hidden_size = 768 |
| 213 | wrong_hidden_size = 1024 |
| 214 | parser = MultiModalDataParser(expected_hidden_size=expected_hidden_size) |
| 215 | |
| 216 | attack_tensor = torch.randn(1, 100, wrong_hidden_size) |
| 217 | |
| 218 | with pytest.raises(ValueError): |
| 219 | parser.parse_mm_data({"video": attack_tensor}) |