| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Test batch-invariant RMS normalization against a PyTorch reference.""" |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | from utils import skip_if_not_cuda, skip_unsupported |
| 8 | |
| 9 | from vllm.model_executor.determinism.batch_invariant import ( |
| 10 | rms_norm_batch_invariant, |
| 11 | ) |
| 12 | from vllm.platforms import current_platform |
| 13 | |
| 14 | DEVICE_TYPE = current_platform.device_type |
| 15 | |
| 16 | |
| 17 | def _rms_norm_reference( |
| 18 | input_tensor: torch.Tensor, |
| 19 | weight: torch.Tensor, |
| 20 | eps: float, |
| 21 | ) -> torch.Tensor: |
| 22 | """Compute RMSNorm independently using PyTorch operations.""" |
| 23 | input_fp32 = input_tensor.float() |
| 24 | output = input_fp32 * torch.rsqrt( |
| 25 | input_fp32.square().mean(dim=-1, keepdim=True) + eps |
| 26 | ) |
| 27 | return (output * weight.float()).to(input_tensor.dtype) |
| 28 | |
| 29 | |
| 30 | @skip_if_not_cuda |
| 31 | @pytest.mark.parametrize("batch_size", [1, 4, 64, 300]) |
| 32 | @pytest.mark.parametrize("hidden_size", [512, 2048, 4096, 8192]) |
| 33 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) |
| 34 | @pytest.mark.parametrize("eps", [1e-6, 1e-5]) |
| 35 | @pytest.mark.parametrize("seed", list(range(4))) |
| 36 | def test_rms_norm_batch_invariant_vs_reference( |
| 37 | default_vllm_config, |
| 38 | batch_size: int, |
| 39 | hidden_size: int, |
| 40 | dtype: torch.dtype, |
| 41 | eps: float, |
| 42 | seed: int, |
| 43 | ): |
| 44 | """Compare batch-invariant Triton RMS norm against a PyTorch reference. |
| 45 | |
| 46 | Tests that the Triton-based batch-invariant RMS norm produces numerically |
| 47 | equivalent results to an independent implementation across configurations. |
| 48 | """ |
| 49 | device = torch.device(DEVICE_TYPE) |
| 50 | |
| 51 | # Create test input and weight |
| 52 | torch.manual_seed(seed) |
| 53 | input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device) |
| 54 | weight = torch.randn(hidden_size, dtype=dtype, device=device) |
| 55 | |
| 56 | reference_output = _rms_norm_reference(input_tensor, weight, eps) |
| 57 | |
| 58 | # Batch-invariant implementation (Triton) |
| 59 | triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps) |
| 60 | |
| 61 | # Compare outputs |
| 62 | # Use looser tolerance for bfloat16 due to its lower precision |
| 63 | if dtype == torch.bfloat16: |
| 64 | rtol, atol = 1e-1, 1e-1 # 10% relative tolerance for bfloat16 |
| 65 | else: |
| 66 | rtol, atol = 1e-2, 1e-2 # 1% for float16/float32 |
| 67 | |
| 68 | torch.testing.assert_close( |
| 69 | triton_output, |
| 70 | reference_output, |
| 71 | rtol=rtol, |
| 72 | atol=atol, |
| 73 | msg=f"RMS norm mismatch for batch_size={batch_size}, " |
| 74 | f"hidden_size={hidden_size}, " |
| 75 | f"dtype={dtype}, eps={eps}, seed={seed}", |
| 76 | ) |
| 77 | |
| 78 | |
| 79 | @skip_if_not_cuda |
| 80 | @pytest.mark.parametrize("hidden_size", [512, 4096]) |
| 81 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) |
| 82 | @pytest.mark.parametrize("eps", [1e-6]) |
| 83 | @pytest.mark.parametrize("n_extra", [3, 299]) |
| 84 | @pytest.mark.parametrize("seed", list(range(16))) |
| 85 | def test_fused_add_rms_norm_batch_invariant_residual_path( |
| 86 | hidden_size: int, |
| 87 | dtype: torch.dtype, |
| 88 | eps: float, |
| 89 | n_extra: int, |
| 90 | seed: int, |
| 91 | ): |
| 92 | """Test the batch-invariant fused residual-add + RMSNorm helper directly.""" |
| 93 | device = torch.device(DEVICE_TYPE) |
| 94 | |
| 95 | torch.manual_seed(seed) |
| 96 | x_single = torch.randn(1, hidden_size, dtype=dtype, device=device) |
| 97 | residual_single = torch.randn(1, hidden_size, dtype=dtype, device=device) |
| 98 | weight = torch.randn(hidden_size, dtype=dtype, device=device) |
| 99 | |
| 100 | x_batch = torch.cat( |
| 101 | [ |
| 102 | x_single, |
| 103 | torch.randn(n_extra, hidden_size, dtype=dtype, device=device), |
| 104 | ], |
| 105 | dim=0, |
| 106 | ) |
| 107 | residual_batch = torch.cat( |
| 108 | [ |
| 109 | residual_single, |
| 110 | torch.randn(n_extra, hidden_size, dtype=dtype, device=device), |
| 111 | ], |
| 112 | dim=0, |
| 113 | ) |
| 114 | |
| 115 | def fused_add_rms_norm(x, residual, w, e) -> tuple[torch.Tensor, torch.Tensor]: |
| 116 | import vllm._custom_ops as ops |
| 117 | |
| 118 | ops.fused_add_rms_norm(x, residual, w, e) |
| 119 | return x, residual |
| 120 | |
| 121 | out_single, residual_out_single = fused_add_rms_norm( |
| 122 | x_single.clone(), |
| 123 | residual_single.clone(), |
| 124 | weight, |
| 125 | eps, |
| 126 | ) |
| 127 | out_batch, residual_out_batch = fused_add_rms_norm( |
| 128 | x_batch.clone(), |
| 129 | residual_batch.clone(), |
| 130 | weight, |
| 131 | eps, |
| 132 | ) |
| 133 | |
| 134 | merged_single = x_single + residual_single |
| 135 | ref_out = _rms_norm_reference(merged_single, weight, eps) |
| 136 | |
| 137 | torch.testing.assert_close( |
| 138 | residual_out_single, |
| 139 | merged_single, |
| 140 | rtol=0.0, |
| 141 | atol=0.0, |
| 142 | msg="Residual output should equal x + residual exactly", |
| 143 | ) |
| 144 | torch.testing.assert_close( |
| 145 | residual_out_batch[:1], |
| 146 | merged_single, |
| 147 | rtol=0.0, |
| 148 | atol=0.0, |
| 149 | msg="Residual output should be batch invariant", |
| 150 | ) |
| 151 | torch.testing.assert_close( |
| 152 | out_single, |
| 153 | out_batch[:1], |
| 154 | rtol=0.0, |
| 155 | atol=0.0, |
| 156 | msg="Fused add RMSNorm output should be batch invariant", |
| 157 | ) |
| 158 | |
| 159 | if dtype == torch.bfloat16: |
| 160 | rtol, atol = 1e-1, 1e-1 |
| 161 | else: |
| 162 | rtol, atol = 1e-2, 1e-2 |
| 163 | |
| 164 | torch.testing.assert_close( |
| 165 | out_single, |
| 166 | ref_out, |
| 167 | rtol=rtol, |
| 168 | atol=atol, |
| 169 | msg="Fused add RMSNorm output should stay numerically close to the " |
| 170 | "PyTorch RMSNorm reference", |
| 171 | ) |
| 172 | |
| 173 | |
| 174 | FP8_DTYPE = current_platform.fp8_dtype() |
| 175 | |
| 176 | # The large launch (num_tokens=300 >= 256) drops an un-pinned kernel to block |
| 177 | # 256, while the small launch (255 rows) stays under the threshold and keeps the |
| 178 | # larger block (1024, or 512 for per-block quant). Under the pin the two launches |
| 179 | # use the same block, so the shared first 255 rows must match bit-for-bit; 255 is |
| 180 | # the most rows a single small launch can hold (< 256, and <= 256 for per-block). |
| 181 | _LARGE_TOKENS = 300 |
| 182 | _SMALL_TOKENS = 255 |
| 183 | |
| 184 | |
| 185 | def _assert_rows_bit_identical(small, large, msg): |
| 186 | if small.dtype == FP8_DTYPE: |
| 187 | assert torch.equal(small.view(torch.uint8), large.view(torch.uint8)), msg |
| 188 | else: |
| 189 | torch.testing.assert_close(small, large, rtol=0.0, atol=0.0, msg=msg) |
| 190 | |
| 191 | |
| 192 | @skip_if_not_cuda |
| 193 | @pytest.mark.parametrize("hidden_size", [512, 4096]) |
| 194 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) |
| 195 | @pytest.mark.parametrize("seed", list(range(4))) |
| 196 | def test_rms_norm_batch_invariant_nonresidual_kernel( |
| 197 | hidden_size: int, dtype: torch.dtype, seed: int |
| 198 | ): |
| 199 | """C++ ``rms_norm`` (no residual) must be batch invariant across the block |
| 200 | threshold. Reached in compiled mode with ``ir_op_priority.rms_norm=["vllm_c"]`` |
| 201 | (default priority is ``native``/inductor codegen when compiling). |
| 202 | """ |
| 203 | import vllm._custom_ops as ops |
| 204 | |
| 205 | device = torch.device(DEVICE_TYPE) |
| 206 | torch.manual_seed(seed) |
| 207 | rows = torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device) |
| 208 | weight = torch.randn(hidden_size, dtype=dtype, device=device) |
| 209 | |
| 210 | def rms_norm(x): |
| 211 | out = torch.empty_like(x) |
| 212 | ops.rms_norm(out, x, weight, 1e-6) |
| 213 | return out |
| 214 | |
| 215 | large = rms_norm(rows.clone()) |
| 216 | small = rms_norm(rows[:_SMALL_TOKENS].clone()) |
| 217 | _assert_rows_bit_identical( |
| 218 | small, |
| 219 | large[:_SMALL_TOKENS], |
| 220 | "rms_norm output depends on num_tokens (block size)", |
| 221 | ) |
| 222 | |
| 223 | |
| 224 | @skip_if_not_cuda |
| 225 | @pytest.mark.parametrize("hidden_size", [512, 4096]) |
| 226 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) |
| 227 | @pytest.mark.parametrize("seed", list(range(4))) |
| 228 | @pytest.mark.parametrize("add_residual", [False, True]) |
| 229 | def test_rms_norm_static_fp8_quant_batch_invariant( |
| 230 | hidden_size: int, dtype: torch.dtype, seed: int, add_residual: bool |
| 231 | ): |
| 232 | """C++ static per-tensor fp8-quant RMSNorm must be batch invariant across |
| 233 | the block threshold. Covers ``rms_norm_static_fp8_quant`` and, with |
| 234 | ``add_residual``, ``fused_add_rms_norm_static_fp8_quant`` (the compiled fp8 |
| 235 | path where ``RMSNormQuantFusionPass`` rewrites norm + quant into them). |
| 236 | """ |
| 237 | device = torch.device(DEVICE_TYPE) |
| 238 | torch.manual_seed(seed) |
| 239 | rows = torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device) |
| 240 | residual = ( |
| 241 | torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device) |
| 242 | if add_residual |
| 243 | else None |
| 244 | ) |
| 245 | weight = torch.randn(hidden_size, dtype=dtype, device=device) |
| 246 | quant_scale = torch.tensor(1.0, dtype=torch.float32, device=device) |
| 247 | |
| 248 | def quant(x, res): |
| 249 | out = torch.empty_like(x, dtype=FP8_DTYPE) |
| 250 | if add_residual: |
| 251 | torch.ops._C.fused_add_rms_norm_static_fp8_quant( |
| 252 | out, x, res, weight, quant_scale, 1e-6 |
| 253 | ) |
| 254 | else: |
| 255 | torch.ops._C.rms_norm_static_fp8_quant(out, x, weight, quant_scale, 1e-6) |
| 256 | return out |
| 257 | |
| 258 | large = quant(rows.clone(), residual.clone() if residual is not None else None) |
| 259 | small = quant( |
| 260 | rows[:_SMALL_TOKENS].clone(), |
| 261 | residual[:_SMALL_TOKENS].clone() if residual is not None else None, |
| 262 | ) |
| 263 | _assert_rows_bit_identical( |
| 264 | small, |
| 265 | large[:_SMALL_TOKENS], |
| 266 | "static-fp8-quant RMSNorm output depends on num_tokens (block size)", |
| 267 | ) |
| 268 | |
| 269 | |
| 270 | @skip_if_not_cuda |
| 271 | @pytest.mark.parametrize("hidden_size", [512, 4096]) |
| 272 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) |
| 273 | @pytest.mark.parametrize("seed", list(range(4))) |
| 274 | def test_rms_norm_per_block_quant_batch_invariant( |
| 275 | hidden_size: int, dtype: torch.dtype, seed: int |
| 276 | ): |
| 277 | """C++ ``rms_norm_per_block_quant`` must be batch invariant across the |
| 278 | block threshold (compiled fp8 block-quant path; block pinned to 512).""" |
| 279 | import vllm._custom_ops as ops |
| 280 | |
| 281 | device = torch.device(DEVICE_TYPE) |
| 282 | torch.manual_seed(seed) |
| 283 | rows = torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device) |
| 284 | weight = torch.randn(hidden_size, dtype=dtype, device=device) |
| 285 | group_size = [1, 128] |
| 286 | |
| 287 | def per_block_quant(x): |
| 288 | return ops.rms_norm_per_block_quant(x, weight, 1e-6, FP8_DTYPE, group_size) |
| 289 | |
| 290 | out_large, scale_large = per_block_quant(rows.clone()) |
| 291 | out_small, scale_small = per_block_quant(rows[:_SMALL_TOKENS].clone()) |
| 292 | _assert_rows_bit_identical( |
| 293 | out_small, |
| 294 | out_large[:_SMALL_TOKENS], |
| 295 | "rms_norm_per_block_quant output depends on num_tokens (block size)", |
| 296 | ) |
| 297 | torch.testing.assert_close( |
| 298 | scale_small, |
| 299 | scale_large[:_SMALL_TOKENS], |
| 300 | rtol=0.0, |
| 301 | atol=0.0, |
| 302 | msg="rms_norm_per_block_quant scales depend on num_tokens (block size)", |
| 303 | ) |
| 304 | |
| 305 | |
| 306 | @skip_if_not_cuda |
| 307 | @pytest.mark.parametrize("batch_size", [1, 16, 128]) |
| 308 | @pytest.mark.parametrize("seq_len", [1, 32, 512]) |
| 309 | @pytest.mark.parametrize("hidden_size", [2048, 4096]) |
| 310 | def test_rms_norm_3d_input( |
| 311 | default_vllm_config, batch_size: int, seq_len: int, hidden_size: int |
| 312 | ): |
| 313 | """Test RMS norm with 3D input tensors (batch, seq_len, hidden_size). |
| 314 | |
| 315 | Ensures that the batch-invariant RMS norm correctly handles multi-dimensional |
| 316 | inputs that are common in transformer models. |
| 317 | """ |
| 318 | device = torch.device(DEVICE_TYPE) |
| 319 | dtype = torch.bfloat16 |
| 320 | eps = 1e-6 |
| 321 | |
| 322 | torch.manual_seed(42) |
| 323 | input_tensor = torch.randn( |
| 324 | batch_size, seq_len, hidden_size, dtype=dtype, device=device |
| 325 | ) |
| 326 | weight = torch.randn(hidden_size, dtype=dtype, device=device) |
| 327 | |
| 328 | reference_output = _rms_norm_reference(input_tensor, weight, eps) |
| 329 | |
| 330 | # Batch-invariant implementation |
| 331 | triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps) |
| 332 | |
| 333 | # Use looser tolerance for bfloat16 |
| 334 | rtol, atol = 1e-1, 1e-1 # 10% tolerance for bfloat16 |
| 335 | |
| 336 | torch.testing.assert_close( |
| 337 | triton_output, |
| 338 | reference_output, |
| 339 | rtol=rtol, |
| 340 | atol=atol, |
| 341 | msg=f"RMS norm mismatch for 3D input with batch_size={batch_size}, " |
| 342 | f"seq_len={seq_len}, hidden_size={hidden_size}", |
| 343 | ) |
| 344 | |
| 345 | |
| 346 | @skip_if_not_cuda |
| 347 | def test_rms_norm_numerical_stability(default_vllm_config): |
| 348 | """Test RMS norm numerical stability with extreme values. |
| 349 | |
| 350 | Ensures that both implementations handle edge cases like very small or large |
| 351 | values without producing NaN or Inf. |
| 352 | """ |
| 353 | device = torch.device(DEVICE_TYPE) |
| 354 | dtype = torch.float16 |
| 355 | eps = 1e-6 |
| 356 | hidden_size = 2048 |
| 357 | |
| 358 | # Test cases with extreme values |
| 359 | test_cases = [ |
| 360 | # Very small values |
| 361 | torch.ones(4, hidden_size, dtype=dtype, device=device) * 1e-5, |
| 362 | # Very large values |
| 363 | torch.ones(4, hidden_size, dtype=dtype, device=device) * 1e4, |
| 364 | # Mixed small and large |
| 365 | torch.randn(4, hidden_size, dtype=dtype, device=device) * 100, |
| 366 | # Values near zero |
| 367 | torch.randn(4, hidden_size, dtype=dtype, device=device) * 1e-6, |
| 368 | ] |
| 369 | |
| 370 | weight = torch.ones(hidden_size, dtype=dtype, device=device) |
| 371 | |
| 372 | for idx, input_tensor in enumerate(test_cases): |
| 373 | reference_output = _rms_norm_reference(input_tensor, weight, eps) |
| 374 | |
| 375 | # Batch-invariant implementation |
| 376 | triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps) |
| 377 | |
| 378 | # Check for NaN or Inf |
| 379 | assert not torch.isnan(reference_output).any(), ( |
| 380 | f"Reference RMS norm produced NaN for test case {idx}" |
| 381 | ) |
| 382 | assert not torch.isinf(reference_output).any(), ( |
| 383 | f"Reference RMS norm produced Inf for test case {idx}" |
| 384 | ) |
| 385 | assert not torch.isnan(triton_output).any(), ( |
| 386 | f"Triton RMS norm produced NaN for test case {idx}" |
| 387 | ) |
| 388 | assert not torch.isinf(triton_output).any(), ( |
| 389 | f"Triton RMS norm produced Inf for test case {idx}" |
| 390 | ) |
| 391 | |
| 392 | # Compare outputs - very lenient for extreme values with float16 |
| 393 | torch.testing.assert_close( |
| 394 | triton_output, |
| 395 | reference_output, |
| 396 | rtol=2e-1, # 20% tolerance for extreme values |
| 397 | atol=2e-1, |
| 398 | msg=f"RMS norm mismatch for extreme value test case {idx}", |
| 399 | ) |
| 400 | |
| 401 | |
| 402 | @skip_unsupported |
| 403 | def test_rms_norm_formula(default_vllm_config): |
| 404 | """Test that RMS norm follows the correct mathematical formula. |
| 405 | |
| 406 | Verifies: output = input / sqrt(mean(input^2) + eps) * weight |
| 407 | """ |
| 408 | device = torch.device(DEVICE_TYPE) |
| 409 | dtype = torch.float32 # Use float32 for higher precision in formula check |
| 410 | eps = 1e-6 |
| 411 | hidden_size = 1024 |
| 412 | |
| 413 | torch.manual_seed(42) |
| 414 | input_tensor = torch.randn(8, hidden_size, dtype=dtype, device=device) |
| 415 | weight = torch.randn(hidden_size, dtype=dtype, device=device) |
| 416 | |
| 417 | # Compute expected output using the formula |
| 418 | variance = (input_tensor.pow(2).mean(dim=-1, keepdim=True)).to(dtype) |
| 419 | expected_output = input_tensor * torch.rsqrt(variance + eps) * weight |
| 420 | |
| 421 | # Batch-invariant implementation |
| 422 | triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps) |
| 423 | |
| 424 | # Compare against formula |
| 425 | torch.testing.assert_close( |
| 426 | triton_output, |
| 427 | expected_output, |
| 428 | rtol=1e-4, |
| 429 | atol=1e-4, |
| 430 | msg="Triton RMS norm doesn't match expected formula", |
| 431 | ) |
| 432 | |
| 433 | |
| 434 | @skip_if_not_cuda |
| 435 | @pytest.mark.parametrize("hidden_size", [128, 1024, 4096, 16384]) |
| 436 | def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int): |
| 437 | """Test RMS norm with various hidden sizes to ensure block size handling. |
| 438 | |
| 439 | The Triton kernel uses a fixed BLOCK_SIZE=1024, so this tests that it |
| 440 | correctly handles hidden sizes both smaller and larger than the block size. |
| 441 | """ |
| 442 | device = torch.device(DEVICE_TYPE) |
| 443 | dtype = torch.bfloat16 |
| 444 | eps = 1e-6 |
| 445 | batch_size = 16 |
| 446 | |
| 447 | torch.manual_seed(42) |
| 448 | input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device) |
| 449 | weight = torch.randn(hidden_size, dtype=dtype, device=device) |
| 450 | |
| 451 | reference_output = _rms_norm_reference(input_tensor, weight, eps) |
| 452 | |
| 453 | # Batch-invariant implementation |
| 454 | triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps) |
| 455 | |
| 456 | # Use looser tolerance for bfloat16 |
| 457 | rtol, atol = 1e-1, 1e-1 # 10% tolerance for bfloat16 |
| 458 | |
| 459 | torch.testing.assert_close( |
| 460 | triton_output, |
| 461 | reference_output, |
| 462 | rtol=rtol, |
| 463 | atol=atol, |
| 464 | msg=f"RMS norm mismatch for hidden_size={hidden_size}", |
| 465 | ) |
| 466 | |
| 467 | |
| 468 | @skip_unsupported |
| 469 | def test_rms_norm_determinism(default_vllm_config): |
| 470 | """Test that batch-invariant RMS norm produces deterministic results. |
| 471 | |
| 472 | Runs the same input through the kernel multiple times and verifies |
| 473 | identical outputs. |
| 474 | """ |
| 475 | device = torch.device(DEVICE_TYPE) |
| 476 | dtype = torch.bfloat16 |
| 477 | eps = 1e-6 |
| 478 | hidden_size = 4096 |
| 479 | batch_size = 32 |
| 480 | |
| 481 | torch.manual_seed(42) |
| 482 | input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device) |
| 483 | weight = torch.randn(hidden_size, dtype=dtype, device=device) |
| 484 | |
| 485 | # Run multiple times |
| 486 | outputs = [] |
| 487 | for _ in range(5): |
| 488 | output = rms_norm_batch_invariant(input_tensor.clone(), weight, eps=eps) |
| 489 | outputs.append(output) |
| 490 | |
| 491 | # All outputs should be identical |
| 492 | reference = outputs[0] |
| 493 | for idx, output in enumerate(outputs[1:], start=1): |
| 494 | torch.testing.assert_close( |
| 495 | output, |
| 496 | reference, |
| 497 | rtol=0.0, |
| 498 | atol=0.0, |
| 499 | msg=f"RMS norm not deterministic: run {idx} differs from reference", |
| 500 | ) |
| 501 | |
| 502 | |
| 503 | @skip_unsupported |
| 504 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) |
| 505 | def test_rms_norm_batch_invariance(dtype): |
| 506 | """Same row gives identical rms_norm result regardless of batch neighbors. |
| 507 | |
| 508 | This verifies that the output for a given row is independent of what other |
| 509 | rows are present in the batch — the core batch-invariance property. |
| 510 | """ |
| 511 | device = torch.device(DEVICE_TYPE) |
| 512 | torch.manual_seed(42) |
| 513 | hidden_size = 2048 |
| 514 | eps = 1e-6 |
| 515 | |
| 516 | weight = torch.randn(hidden_size, dtype=dtype, device=device) |
| 517 | row = torch.randn(1, hidden_size, dtype=dtype, device=device) |
| 518 | |
| 519 | # Compute rms_norm on the single row alone |
| 520 | out_single = rms_norm_batch_invariant(row, weight, eps=eps) |
| 521 | |
| 522 | # Embed the same row in a larger batch with random neighbors |
| 523 | batch = torch.randn(8, hidden_size, dtype=dtype, device=device) |
| 524 | batch[4] = row[0] |
| 525 | out_batch = rms_norm_batch_invariant(batch, weight, eps=eps) |
| 526 | |
| 527 | assert isinstance(out_single, torch.Tensor) and isinstance(out_batch, torch.Tensor) |
| 528 | assert torch.equal(out_single[0], out_batch[4]), ( |
| 529 | "rms_norm output for a row differs when batch context changes" |
| 530 | ) |
| 531 | |
| 532 | |
| 533 | if __name__ == "__main__": |
| 534 | # Run a quick smoke test |
| 535 | print("Running quick smoke test of RMS norm implementations...") |
| 536 | |
| 537 | device = torch.device(DEVICE_TYPE) |
| 538 | batch_size = 8 |
| 539 | hidden_size = 4096 |
| 540 | dtype = torch.bfloat16 |
| 541 | eps = 1e-6 |
| 542 | |
| 543 | torch.manual_seed(42) |
| 544 | input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device) |
| 545 | weight = torch.randn(hidden_size, dtype=dtype, device=device) |
| 546 | |
| 547 | reference_output = _rms_norm_reference(input_tensor, weight, eps) |
| 548 | |
| 549 | # Batch-invariant implementation |
| 550 | triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps) |
| 551 | assert isinstance(triton_output, torch.Tensor) |
| 552 | |
| 553 | # Compare |
| 554 | max_diff = (triton_output - reference_output).abs().max().item() |
| 555 | mean_diff = (triton_output - reference_output).abs().mean().item() |
| 556 | |
| 557 | print(f"Max difference: {max_diff:.6e}") |
| 558 | print(f"Mean difference: {mean_diff:.6e}") |
| 559 | print(f"Reference output sample: {reference_output[0, :5].tolist()}") |
| 560 | print(f"Triton output sample: {triton_output[0, :5].tolist()}") |
| 561 | |
| 562 | if max_diff < 1e-3: |
| 563 | print("✓ Smoke test passed!") |
| 564 | else: |
| 565 | print("✗ Smoke test failed - differences too large") |