| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for embedding shape validation. |
| 4 | |
| 5 | Simple, fast unit tests that can run without server fixtures. |
| 6 | Run with: pytest tests/multimodal/test_embedding_shape_validation_unit.py -v |
| 7 | """ |
| 8 | |
| 9 | import pytest |
| 10 | import torch |
| 11 | |
| 12 | from vllm.multimodal.parse import ( |
| 13 | AudioEmbeddingItems, |
| 14 | ImageEmbeddingItems, |
| 15 | ) |
| 16 | |
| 17 | |
| 18 | class TestImageEmbedBasicValidation: |
| 19 | """Test basic ndim validation in image embeddings via ImageEmbeddingItems.""" |
| 20 | |
| 21 | def test_valid_2d_tensor_accepted(self): |
| 22 | """Baseline: 2D tensors should be accepted.""" |
| 23 | valid_tensor = torch.randn(10, 768, dtype=torch.float32) |
| 24 | |
| 25 | # Should not raise - 2D is valid |
| 26 | items = ImageEmbeddingItems(valid_tensor) |
| 27 | assert items.get_count() == 10 |
| 28 | |
| 29 | def test_valid_3d_tensor_accepted(self): |
| 30 | """Baseline: 3D tensors should be accepted.""" |
| 31 | valid_tensor = torch.randn(2, 10, 768, dtype=torch.float32) |
| 32 | |
| 33 | # Should not raise - 3D is valid |
| 34 | items = ImageEmbeddingItems(valid_tensor) |
| 35 | assert items.get_count() == 2 |
| 36 | |
| 37 | def test_valid_list_of_2d_tensors_accepted(self): |
| 38 | """Baseline: List of 2D tensors should be accepted.""" |
| 39 | tensors = [ |
| 40 | torch.randn(10, 768, dtype=torch.float32), |
| 41 | torch.randn(15, 768, dtype=torch.float32), |
| 42 | ] |
| 43 | |
| 44 | # Should not raise |
| 45 | items = ImageEmbeddingItems(tensors) |
| 46 | assert items.get_count() == 2 |
| 47 | |
| 48 | def test_1d_tensor_rejected(self): |
| 49 | """Security: 1D tensors should be rejected (invalid ndim).""" |
| 50 | invalid_tensor = torch.randn(768, dtype=torch.float32) # 1D |
| 51 | |
| 52 | with pytest.raises(ValueError) as exc_info: |
| 53 | ImageEmbeddingItems(invalid_tensor) |
| 54 | |
| 55 | assert "must be 2D" in str(exc_info.value) or "3D" in str(exc_info.value) |
| 56 | |
| 57 | def test_4d_tensor_rejected(self): |
| 58 | """Security: 4D tensors should be rejected (invalid ndim).""" |
| 59 | invalid_tensor = torch.randn(1, 2, 10, 768, dtype=torch.float32) # 4D |
| 60 | |
| 61 | with pytest.raises(ValueError) as exc_info: |
| 62 | ImageEmbeddingItems(invalid_tensor) |
| 63 | |
| 64 | assert "must be 2D" in str(exc_info.value) or "3D" in str(exc_info.value) |
| 65 | |
| 66 | def test_hidden_size_validation_correct_size(self): |
| 67 | """Embeddings with correct hidden size should be accepted.""" |
| 68 | expected_hidden_size = 768 |
| 69 | valid_tensor = torch.randn(10, expected_hidden_size, dtype=torch.float32) |
| 70 | |
| 71 | # Should not raise |
| 72 | items = ImageEmbeddingItems( |
| 73 | valid_tensor, expected_hidden_size=expected_hidden_size |
| 74 | ) |
| 75 | assert items.get_count() == 10 |
| 76 | |
| 77 | def test_hidden_size_validation_wrong_size_rejected(self): |
| 78 | """Embeddings with wrong hidden size should be rejected.""" |
| 79 | expected_hidden_size = 768 |
| 80 | wrong_hidden_size = 4096 |
| 81 | invalid_tensor = torch.randn(10, wrong_hidden_size, dtype=torch.float32) |
| 82 | |
| 83 | with pytest.raises(ValueError) as exc_info: |
| 84 | ImageEmbeddingItems( |
| 85 | invalid_tensor, expected_hidden_size=expected_hidden_size |
| 86 | ) |
| 87 | |
| 88 | error_msg = str(exc_info.value) |
| 89 | assert "hidden dimension mismatch" in error_msg.lower() |
| 90 | assert str(wrong_hidden_size) in error_msg |
| 91 | assert str(expected_hidden_size) in error_msg |
| 92 | |
| 93 | |
| 94 | class TestAudioEmbedBasicValidation: |
| 95 | """Test basic ndim validation in audio embeddings via AudioEmbeddingItems.""" |
| 96 | |
| 97 | def test_valid_2d_tensor_accepted(self): |
| 98 | """Baseline: 2D tensors should be accepted.""" |
| 99 | valid_tensor = torch.randn(10, 768, dtype=torch.float32) |
| 100 | |
| 101 | # Should not raise - 2D is valid |
| 102 | items = AudioEmbeddingItems(valid_tensor) |
| 103 | assert items.get_count() == 10 |
| 104 | |
| 105 | def test_valid_3d_tensor_accepted(self): |
| 106 | """Baseline: 3D tensors should be accepted.""" |
| 107 | valid_tensor = torch.randn(2, 10, 768, dtype=torch.float32) |
| 108 | |
| 109 | # Should not raise - 3D is valid |
| 110 | items = AudioEmbeddingItems(valid_tensor) |
| 111 | assert items.get_count() == 2 |
| 112 | |
| 113 | def test_valid_list_of_2d_tensors_accepted(self): |
| 114 | """Baseline: List of 2D tensors should be accepted.""" |
| 115 | tensors = [ |
| 116 | torch.randn(10, 768, dtype=torch.float32), |
| 117 | torch.randn(15, 768, dtype=torch.float32), |
| 118 | ] |
| 119 | |
| 120 | # Should not raise |
| 121 | items = AudioEmbeddingItems(tensors) |
| 122 | assert items.get_count() == 2 |
| 123 | |
| 124 | def test_1d_tensor_rejected(self): |
| 125 | """Security: 1D tensors should be rejected (invalid ndim).""" |
| 126 | invalid_tensor = torch.randn(768, dtype=torch.float32) # 1D |
| 127 | |
| 128 | with pytest.raises(ValueError) as exc_info: |
| 129 | AudioEmbeddingItems(invalid_tensor) |
| 130 | |
| 131 | assert "must be 2D" in str(exc_info.value) or "3D" in str(exc_info.value) |
| 132 | |
| 133 | def test_scalar_rejected(self): |
| 134 | """Security: Scalar tensors should be rejected.""" |
| 135 | invalid_tensor = torch.tensor(1.0) # 0D (scalar) |
| 136 | |
| 137 | with pytest.raises(ValueError): |
| 138 | AudioEmbeddingItems(invalid_tensor) |
| 139 | |
| 140 | def test_hidden_size_validation_correct_size(self): |
| 141 | """Embeddings with correct hidden size should be accepted.""" |
| 142 | expected_hidden_size = 768 |
| 143 | valid_tensor = torch.randn(10, expected_hidden_size, dtype=torch.float32) |
| 144 | |
| 145 | # Should not raise |
| 146 | items = AudioEmbeddingItems( |
| 147 | valid_tensor, expected_hidden_size=expected_hidden_size |
| 148 | ) |
| 149 | assert items.get_count() == 10 |
| 150 | |
| 151 | def test_hidden_size_validation_wrong_size_rejected(self): |
| 152 | """Embeddings with wrong hidden size should be rejected.""" |
| 153 | expected_hidden_size = 768 |
| 154 | wrong_hidden_size = 4096 |
| 155 | invalid_tensor = torch.randn(10, wrong_hidden_size, dtype=torch.float32) |
| 156 | |
| 157 | with pytest.raises(ValueError) as exc_info: |
| 158 | AudioEmbeddingItems( |
| 159 | invalid_tensor, expected_hidden_size=expected_hidden_size |
| 160 | ) |
| 161 | |
| 162 | error_msg = str(exc_info.value) |
| 163 | assert "hidden dimension mismatch" in error_msg.lower() |
| 164 | assert str(wrong_hidden_size) in error_msg |
| 165 | assert str(expected_hidden_size) in error_msg |
| 166 | |
| 167 | |
| 168 | class TestShapeValidationDoSPrevention: |
| 169 | """Tests for DoS prevention through shape validation. |
| 170 | |
| 171 | Verifies that embeddings with incorrect shapes are rejected early, |
| 172 | preventing crashes during model inference. |
| 173 | """ |
| 174 | |
| 175 | def test_prevent_crash_from_wrong_shape_image_embeds(self): |
| 176 | """Prevent crash scenario: wrong hidden size in image embeddings. |
| 177 | |
| 178 | Without validation, this would pass initial checks but crash later |
| 179 | during model forward pass when dimensions don't match. |
| 180 | """ |
| 181 | expected_hidden_size = 768 # Typical model hidden size |
| 182 | wrong_hidden_size = 4096 # Wrong size (e.g., Llama-sized) |
| 183 | |
| 184 | wrong_embedding = torch.randn(100, wrong_hidden_size, dtype=torch.float32) |
| 185 | |
| 186 | # Should be rejected at instantiation time, not during inference |
| 187 | with pytest.raises(ValueError) as exc_info: |
| 188 | ImageEmbeddingItems( |
| 189 | wrong_embedding, expected_hidden_size=expected_hidden_size |
| 190 | ) |
| 191 | |
| 192 | error_msg = str(exc_info.value) |
| 193 | assert "hidden dimension mismatch" in error_msg.lower() |
| 194 | assert str(expected_hidden_size) in error_msg # Expected |
| 195 | assert str(wrong_hidden_size) in error_msg # Received |
| 196 | |
| 197 | def test_prevent_crash_from_wrong_shape_audio_embeds(self): |
| 198 | """Prevent crash scenario: wrong hidden size in audio embeddings.""" |
| 199 | expected_hidden_size = 768 |
| 200 | wrong_hidden_size = 4096 |
| 201 | |
| 202 | wrong_embedding = torch.randn(100, wrong_hidden_size, dtype=torch.float32) |
| 203 | |
| 204 | with pytest.raises(ValueError) as exc_info: |
| 205 | AudioEmbeddingItems( |
| 206 | wrong_embedding, expected_hidden_size=expected_hidden_size |
| 207 | ) |
| 208 | |
| 209 | error_msg = str(exc_info.value) |
| 210 | assert "hidden dimension mismatch" in error_msg.lower() |
| 211 | |
| 212 | def test_extremely_large_hidden_size_rejected(self): |
| 213 | """Security: Prevent DoS from extremely large embeddings.""" |
| 214 | expected_hidden_size = 768 |
| 215 | huge_hidden_size = 100000 # Large but not extreme to avoid test OOM |
| 216 | |
| 217 | invalid_tensor = torch.randn(10, huge_hidden_size, dtype=torch.float32) |
| 218 | |
| 219 | with pytest.raises(ValueError) as exc_info: |
| 220 | ImageEmbeddingItems( |
| 221 | invalid_tensor, expected_hidden_size=expected_hidden_size |
| 222 | ) |
| 223 | |
| 224 | assert "hidden dimension mismatch" in str(exc_info.value).lower() |
| 225 | |
| 226 | def test_batch_with_mixed_hidden_sizes_rejected(self): |
| 227 | """All embeddings in a list must have the same hidden size.""" |
| 228 | expected_hidden_size = 768 |
| 229 | |
| 230 | # One correct, one wrong |
| 231 | batch = [ |
| 232 | torch.randn(10, expected_hidden_size, dtype=torch.float32), |
| 233 | torch.randn(10, expected_hidden_size + 100, dtype=torch.float32), # Wrong! |
| 234 | ] |
| 235 | |
| 236 | # Should fail on the second one |
| 237 | with pytest.raises(ValueError) as exc_info: |
| 238 | ImageEmbeddingItems(batch, expected_hidden_size=expected_hidden_size) |
| 239 | |
| 240 | assert "hidden dimension mismatch" in str(exc_info.value).lower() |
| 241 | |
| 242 | |
| 243 | if __name__ == "__main__": |
| 244 | pytest.main([__file__, "-v", "--tb=short"]) |