[Test][ROCm] Account for gfx950 FP8 RMSNorm rounding (#49839)

Signed-off-by: Andreas Karatzas <[email protected]>
Signed-off-by: Andreas Karatzas <[email protected]>
Co-authored-by: OpenAI Codex <[email protected]>

tests/kernels/core/test_fused_quant_layernorm.py

@@ -8,7 +8,7 @@ import pytest
 import torch
 
 import vllm._custom_ops as ops
-from tests.kernels.utils import fp8_ulp_distance, opcheck
+from tests.kernels.utils import fp8_allclose, fp8_ulp_distance, opcheck
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
     per_token_group_quant_fp8,
@@ -19,6 +19,12 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import (
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
 
+ON_GFX950 = False
+if current_platform.is_rocm():
+    from vllm.platforms.rocm import on_gfx950
+
+    ON_GFX950 = on_gfx950()
+
 DTYPES = [torch.bfloat16, torch.float]
 QUANT_DTYPES = [torch.int8, current_platform.fp8_dtype()]
 VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029]
@@ -317,6 +323,13 @@ def test_rms_norm(
         and dtype == torch.bfloat16
         and current_platform.is_rocm()
     )
+    use_gfx950_fp8_allclose = (
+        current_platform.is_rocm()
+        and ON_GFX950
+        and group_size is None
+        and dtype == torch.bfloat16
+        and quant_dtype == current_platform.fp8_dtype()
+    )
 
     def scales_close(rtol: float, atol: float) -> bool:
         if torch.allclose(ref_scales, ops_scales, rtol=rtol, atol=atol):
@@ -341,6 +354,10 @@ def test_rms_norm(
                 ulp = fp8_ulp_distance(ref_out, ops_out)
                 max_outliers = ulp.numel() // 100_000 + 8
                 ok = int((ulp > 0).sum().item()) <= max_outliers
+            elif use_gfx950_fp8_allclose:
+                # Valid gfx950 reduction trees can straddle an E4M3 boundary.
+                ok = fp8_allclose(ops_out, ref_out, rtol=0.125, atol=2e-3)
+                ok = ok and int(fp8_ulp_distance(ops_out, ref_out).max()) <= 1
             else:
                 # CUDA (& non-bf16): compare dequantized values with relaxed tolerance.
                 if group_size is None:

tests/kernels/core/test_layernorm.py

@@ -5,7 +5,7 @@ import pytest
 import torch
 
 from tests.kernels.quant_utils import FP8_DTYPE
-from tests.kernels.utils import fp8_ulp_distance, opcheck
+from tests.kernels.utils import fp8_allclose, fp8_ulp_distance, opcheck
 from vllm import ir
 from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
 from vllm.platforms import current_platform
@@ -207,14 +207,22 @@ def test_fused_rms_norm_quant(
         )
 
     if current_platform.is_rocm():
-        # Fused and unfused FP8 paths can land on opposite sides of an E4M3 tie;
-        # tolerate a tiny number of isolated fp8 outliers on ROCm.
-        ulp = fp8_ulp_distance(out_quant, out_quant_fused)
-        max_outliers = ulp.numel() // 100_000 + 8
-        num_outliers = int((ulp > 0).sum().item())
-        assert num_outliers <= max_outliers, (
-            f"FP8 quant mismatch: {num_outliers} fp8 outliers (allowed {max_outliers})"
-        )
+        from vllm.platforms.rocm import on_gfx950
+
+        if on_gfx950() and dtype == torch.float16 and not add_residual:
+            # Fusion may round normalized FP16 across an E4M3 boundary on gfx950.
+            assert fp8_allclose(out_quant_fused, out_quant, rtol=0.125, atol=2e-3)
+            assert int(fp8_ulp_distance(out_quant_fused, out_quant).max()) <= 1
+        else:
+            # Fused and unfused FP8 paths can land on opposite sides of an E4M3
+            # tie; tolerate a tiny number of isolated fp8 outliers on ROCm.
+            ulp = fp8_ulp_distance(out_quant, out_quant_fused)
+            max_outliers = ulp.numel() // 100_000 + 8
+            num_outliers = int((ulp > 0).sum().item())
+            assert num_outliers <= max_outliers, (
+                f"FP8 quant mismatch: {num_outliers} fp8 outliers "
+                f"(allowed {max_outliers})"
+            )
     else:
         torch.testing.assert_close(
             out_quant.to(dtype=torch.float32),