| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Tests for FusedMMInputNorm and the fused affine transform kernel.""" |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | |
| 9 | from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config |
| 10 | from vllm.config.compilation import CompilationMode |
| 11 | from vllm.model_executor.layers.fusion.mm_input_norm import ( |
| 12 | FusedMMInputNorm, |
| 13 | IdentityInputNorm, |
| 14 | NormParams, |
| 15 | fused_mm_input_norm_triton, |
| 16 | ) |
| 17 | from vllm.platforms import current_platform |
| 18 | from vllm.triton_utils import HAS_TRITON |
| 19 | from vllm.utils.torch_utils import set_random_seed |
| 20 | |
| 21 | # Module-level tests run on any platform: the CPU platform exercises |
| 22 | # forward_native, accelerators exercise the fused kernel. Only the raw |
| 23 | # kernel tests require an accelerator. |
| 24 | _DEVICE = torch.device(current_platform.device_type) |
| 25 | |
| 26 | requires_accelerator = pytest.mark.skipif( |
| 27 | current_platform.is_cpu(), |
| 28 | reason="fused Triton kernel requires a CUDA/XPU accelerator", |
| 29 | ) |
| 30 | requires_triton = pytest.mark.skipif(not HAS_TRITON, reason="requires Triton") |
| 31 | requires_vllm_config = pytest.mark.usefixtures("default_vllm_config") |
| 32 | |
| 33 | |
| 34 | def _reference_input_norm( |
| 35 | pixel_values: torch.Tensor, |
| 36 | image_mean: list[float], |
| 37 | image_std: list[float], |
| 38 | rescale_factor: float, |
| 39 | channel: int, |
| 40 | out_dtype: torch.dtype, |
| 41 | ) -> torch.Tensor: |
| 42 | """Straightforward per-channel affine: (x * rescale - mean) / std.""" |
| 43 | patches, size = pixel_values.shape |
| 44 | patch_size = size // channel |
| 45 | mean = torch.tensor( |
| 46 | image_mean, dtype=torch.float32, device=pixel_values.device |
| 47 | ).view(1, channel, 1) |
| 48 | std = torch.tensor(image_std, dtype=torch.float32, device=pixel_values.device).view( |
| 49 | 1, channel, 1 |
| 50 | ) |
| 51 | x = pixel_values.to(torch.float32).reshape(patches, channel, patch_size) |
| 52 | x = (x * rescale_factor - mean) / std |
| 53 | return x.view(patches, size).to(out_dtype) |
| 54 | |
| 55 | |
| 56 | # CLIP-style RGB normalization constants reused across tests. |
| 57 | _RGB_MEAN = [0.48145466, 0.4578275, 0.40821073] |
| 58 | _RGB_STD = [0.26862954, 0.26130258, 0.27577711] |
| 59 | _RGB_RESCALE = 1.0 / 255.0 |
| 60 | |
| 61 | |
| 62 | def _check_against_reference( |
| 63 | pixel_values: torch.Tensor, |
| 64 | out_dtype: torch.dtype, |
| 65 | channel: int = 3, |
| 66 | image_mean: list[float] | None = None, |
| 67 | image_std: list[float] | None = None, |
| 68 | rescale_factor: float = _RGB_RESCALE, |
| 69 | **close_kwargs, |
| 70 | ): |
| 71 | """Run ``FusedMMInputNorm`` on ``pixel_values`` and compare against the |
| 72 | eager per-channel affine reference. Uses the RGB constants by default.""" |
| 73 | if image_mean is None: |
| 74 | image_mean = list(_RGB_MEAN) |
| 75 | if image_std is None: |
| 76 | image_std = list(_RGB_STD) |
| 77 | norm = FusedMMInputNorm( |
| 78 | image_mean=image_mean, |
| 79 | image_std=image_std, |
| 80 | rescale_factor=rescale_factor, |
| 81 | channel=channel, |
| 82 | ).to(_DEVICE) |
| 83 | expected = _reference_input_norm( |
| 84 | pixel_values, image_mean, image_std, rescale_factor, channel, out_dtype |
| 85 | ) |
| 86 | torch.testing.assert_close( |
| 87 | norm(pixel_values, visual_dtype=out_dtype), expected, **close_kwargs |
| 88 | ) |
| 89 | |
| 90 | |
| 91 | # =========================================================================== |
| 92 | # Module-level behavior |
| 93 | # =========================================================================== |
| 94 | @requires_vllm_config |
| 95 | class TestFusedMMInputNormModule: |
| 96 | def test_matches_reference_above_cudnn_grid_limit(self): |
| 97 | """num_patches above the old cuDNN batch-norm grid limit (~65535) |
| 98 | used to raise CUDNN_STATUS_INTERNAL_ERROR.""" |
| 99 | set_random_seed(0) |
| 100 | pixel_values = torch.randint( |
| 101 | 0, 256, (70000, 3 * 196), dtype=torch.float32, device=_DEVICE |
| 102 | ) |
| 103 | _check_against_reference(pixel_values, torch.float32) |
| 104 | |
| 105 | def test_identity_passthrough(self): |
| 106 | norm = IdentityInputNorm() |
| 107 | |
| 108 | pixel_values = torch.randn(8, 3 * 196, dtype=torch.float32, device=_DEVICE) |
| 109 | out = norm(pixel_values, visual_dtype=torch.bfloat16) |
| 110 | torch.testing.assert_close(out, pixel_values.to(torch.bfloat16)) |
| 111 | |
| 112 | |
| 113 | # =========================================================================== |
| 114 | # dtype coverage |
| 115 | # =========================================================================== |
| 116 | @requires_vllm_config |
| 117 | class TestFusedMMInputNormDtypes: |
| 118 | @pytest.mark.parametrize( |
| 119 | "in_dtype,out_dtype", |
| 120 | [ |
| 121 | (torch.float32, torch.float32), |
| 122 | (torch.float32, torch.bfloat16), |
| 123 | (torch.float32, torch.float16), |
| 124 | (torch.bfloat16, torch.bfloat16), |
| 125 | (torch.bfloat16, torch.float32), |
| 126 | (torch.float16, torch.float16), |
| 127 | (torch.uint8, torch.bfloat16), |
| 128 | (torch.uint8, torch.float32), |
| 129 | ], |
| 130 | ) |
| 131 | def test_dtype_combinations(self, in_dtype: torch.dtype, out_dtype: torch.dtype): |
| 132 | set_random_seed(0) |
| 133 | if in_dtype == torch.uint8: |
| 134 | pixel_values = torch.randint( |
| 135 | 0, 256, (32, 3 * 16), dtype=torch.uint8, device=_DEVICE |
| 136 | ) |
| 137 | else: |
| 138 | pixel_values = torch.rand( |
| 139 | (32, 3 * 16), dtype=torch.float32, device=_DEVICE |
| 140 | ).to(in_dtype) |
| 141 | |
| 142 | _check_against_reference( |
| 143 | pixel_values, |
| 144 | out_dtype, |
| 145 | image_mean=[0.5, 0.5, 0.5], |
| 146 | image_std=[0.5, 0.5, 0.5], |
| 147 | atol=1e-2, |
| 148 | rtol=1e-2, |
| 149 | ) |
| 150 | |
| 151 | |
| 152 | # =========================================================================== |
| 153 | # Shape / channel coverage |
| 154 | # =========================================================================== |
| 155 | @requires_vllm_config |
| 156 | class TestFusedMMInputNormShapes: |
| 157 | @pytest.mark.parametrize("channel", [1, 3, 4]) |
| 158 | def test_channel_variants(self, channel: int): |
| 159 | set_random_seed(0) |
| 160 | pixel_values = torch.randint( |
| 161 | 0, 256, (16, channel * 64), dtype=torch.float32, device=_DEVICE |
| 162 | ) |
| 163 | _check_against_reference( |
| 164 | pixel_values, |
| 165 | torch.float32, |
| 166 | channel=channel, |
| 167 | image_mean=[0.5] * channel, |
| 168 | image_std=[0.25] * channel, |
| 169 | ) |
| 170 | |
| 171 | @pytest.mark.parametrize("patch_size", [1, 2048]) |
| 172 | def test_block_boundaries(self, patch_size: int): |
| 173 | """Masked-tail (12 elements) and exact-multiple (24 blocks) paths.""" |
| 174 | set_random_seed(0) |
| 175 | pixel_values = torch.randint( |
| 176 | 0, 256, (4, 3 * patch_size), dtype=torch.float32, device=_DEVICE |
| 177 | ) |
| 178 | _check_against_reference(pixel_values, torch.float32) |
| 179 | |
| 180 | |
| 181 | # =========================================================================== |
| 182 | # Input handling: non-contiguous inputs |
| 183 | # =========================================================================== |
| 184 | @requires_vllm_config |
| 185 | class TestFusedMMInputNormInputHandling: |
| 186 | def test_non_contiguous_input_matches_reference(self): |
| 187 | set_random_seed(0) |
| 188 | base = torch.randint( |
| 189 | 0, 256, (8, 3 * 32, 2), dtype=torch.float32, device=_DEVICE |
| 190 | ) |
| 191 | non_contig = base[..., 0] |
| 192 | assert not non_contig.is_contiguous() |
| 193 | |
| 194 | _check_against_reference(non_contig, torch.float32) |
| 195 | |
| 196 | |
| 197 | # =========================================================================== |
| 198 | # Raw kernel entry point |
| 199 | # =========================================================================== |
| 200 | @requires_accelerator |
| 201 | @requires_triton |
| 202 | class TestFusedMMInputNormKernel: |
| 203 | @pytest.mark.parametrize( |
| 204 | "N, C, L", |
| 205 | [ |
| 206 | (1, 1, 1), |
| 207 | (1, 3, 1), |
| 208 | (3, 3, 7), # tile straddles channel boundaries |
| 209 | (2, 8, 63), |
| 210 | (4, 3, 1000), |
| 211 | ], |
| 212 | ) |
| 213 | def test_channel_boundary_crossing(self, N: int, C: int, L: int): |
| 214 | """The 1D kernel recovers ``c = (offs // L) % C`` per lane, so a tile |
| 215 | can straddle a channel boundary and gather weight/bias per lane. |
| 216 | These shapes force that path by keeping ``L`` well below the block.""" |
| 217 | set_random_seed(0) |
| 218 | x = torch.randn(N, C, L, dtype=torch.float32, device=_DEVICE) |
| 219 | w = torch.randn(C, dtype=torch.float32, device=_DEVICE) |
| 220 | b = torch.randn(C, dtype=torch.float32, device=_DEVICE) |
| 221 | |
| 222 | out = torch.empty_like(x) |
| 223 | fused_mm_input_norm_triton(x, out, w, b) |
| 224 | |
| 225 | expected = x * w.view(1, C, 1) + b.view(1, C, 1) |
| 226 | torch.testing.assert_close(out, expected) |
| 227 | |
| 228 | def test_non_contiguous_inputs_materialized(self): |
| 229 | """Read-only tensors are made contiguous internally.""" |
| 230 | N, C, L = 3, 3, 100 |
| 231 | set_random_seed(0) |
| 232 | base = torch.randn(N, C, L, 2, dtype=torch.float32, device=_DEVICE) |
| 233 | x = base[..., 0] |
| 234 | assert not x.is_contiguous() |
| 235 | w = torch.randn(C, dtype=torch.float32, device=_DEVICE) |
| 236 | b = torch.randn(C, dtype=torch.float32, device=_DEVICE) |
| 237 | |
| 238 | out = torch.empty(N, C, L, dtype=torch.float32, device=_DEVICE) |
| 239 | fused_mm_input_norm_triton(x, out, w, b) |
| 240 | |
| 241 | expected = x * w.view(1, C, 1) + b.view(1, C, 1) |
| 242 | torch.testing.assert_close(out, expected) |
| 243 | |
| 244 | def test_rejects_mismatched_output_shape(self): |
| 245 | """``outputs`` must be shaped exactly like ``inputs``.""" |
| 246 | N, C, L = 3, 3, 100 |
| 247 | x = torch.randn(N, C, L, dtype=torch.float32, device=_DEVICE) |
| 248 | w = torch.randn(C, dtype=torch.float32, device=_DEVICE) |
| 249 | b = torch.randn(C, dtype=torch.float32, device=_DEVICE) |
| 250 | |
| 251 | for out_shape in ((N, C + 1, L), (N, C, L + 1), (N + 2, C, L)): |
| 252 | with pytest.raises(AssertionError): |
| 253 | fused_mm_input_norm_triton( |
| 254 | x, |
| 255 | torch.empty(*out_shape, dtype=torch.float32, device=_DEVICE), |
| 256 | w, |
| 257 | b, |
| 258 | ) |
| 259 | |
| 260 | |
| 261 | # =========================================================================== |
| 262 | # Construction / configuration |
| 263 | # =========================================================================== |
| 264 | @requires_vllm_config |
| 265 | class TestFusedMMInputNormConstruction: |
| 266 | """Weight/bias buffer semantics at init time.""" |
| 267 | |
| 268 | def test_input_norm_is_always_enabled(self): |
| 269 | config = VllmConfig( |
| 270 | compilation_config=CompilationConfig( |
| 271 | mode=CompilationMode.VLLM_COMPILE, |
| 272 | compile_mm_encoder=True, |
| 273 | custom_ops=["none"], |
| 274 | ) |
| 275 | ) |
| 276 | with set_current_vllm_config(config): |
| 277 | FusedMMInputNorm(_RGB_MEAN, _RGB_STD, _RGB_RESCALE) |
| 278 | assert ( |
| 279 | config.compilation_config.enabled_custom_ops["fused_mm_input_norm"] == 1 |
| 280 | ) |
| 281 | assert not config.compilation_config.disabled_custom_ops[ |
| 282 | "fused_mm_input_norm" |
| 283 | ] |
| 284 | |
| 285 | def test_identity_config_buffers(self): |
| 286 | """Numerically identity parameters still build plain weight/bias |
| 287 | buffers; the identity shortcut lives in ``from_model_config``.""" |
| 288 | norm = FusedMMInputNorm( |
| 289 | image_mean=[0.0, 0.0, 0.0], |
| 290 | image_std=[1.0, 1.0, 1.0], |
| 291 | rescale_factor=1.0, |
| 292 | ) |
| 293 | torch.testing.assert_close(norm.weight, torch.ones(3)) |
| 294 | torch.testing.assert_close(norm.bias, torch.zeros(3)) |
| 295 | |
| 296 | def test_non_identity_buffers(self): |
| 297 | channel = 3 |
| 298 | norm = FusedMMInputNorm( |
| 299 | image_mean=_RGB_MEAN, |
| 300 | image_std=_RGB_STD, |
| 301 | rescale_factor=_RGB_RESCALE, |
| 302 | channel=channel, |
| 303 | ) |
| 304 | assert norm.weight.shape == (channel,) |
| 305 | assert norm.bias.shape == (channel,) |
| 306 | |
| 307 | mean = torch.tensor(_RGB_MEAN, dtype=torch.float32) |
| 308 | std = torch.tensor(_RGB_STD, dtype=torch.float32) |
| 309 | torch.testing.assert_close(norm.weight, _RGB_RESCALE / std) |
| 310 | torch.testing.assert_close(norm.bias, -mean / std) |
| 311 | |
| 312 | def test_is_identity_params(self): |
| 313 | """The numeric-identity check mirrors the old allclose detection.""" |
| 314 | assert NormParams([0.0, 0.0, 0.0], [1.0, 1.0, 1.0], 1.0).is_identity |
| 315 | assert not NormParams(_RGB_MEAN, _RGB_STD, _RGB_RESCALE).is_identity |
| 316 | |
| 317 | @pytest.mark.parametrize( |
| 318 | ("image_mean", "image_std", "rescale_factor"), |
| 319 | [ |
| 320 | ([0.0, 0.0, 0.0], [1.0, 1.0, 1.0], 1.0), |
| 321 | ([0.5, 0.5, 0.5], [0.25, 0.25, 0.25], 1 / 255), |
| 322 | ], |
| 323 | ) |
| 324 | def test_buffers_initialized_on_default_device( |
| 325 | self, |
| 326 | image_mean: list[float], |
| 327 | image_std: list[float], |
| 328 | rescale_factor: float, |
| 329 | ): |
| 330 | """Buffers are built on CPU first, then moved to the caller's |
| 331 | default device, so init never touches the accelerator.""" |
| 332 | # The meta device gives the CPU-only test shard the same non-CPU |
| 333 | # default-device semantics without requiring a CUDA build. |
| 334 | default_device = "cuda" if torch.cuda.is_available() else "meta" |
| 335 | with torch.device(default_device): |
| 336 | norm = FusedMMInputNorm(image_mean, image_std, rescale_factor) |
| 337 | |
| 338 | assert norm.weight.device.type == default_device |
| 339 | assert norm.bias.device.type == default_device |