[Test][ROCm] Account for gfx950 FP8 RMSNorm rounding (#49839)
Signed-off-by: Andreas Karatzas <akaratza@amd.com> Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Andreas Karatzas committed
Jul 30, 2026 at 05:46 UTC
0c64be8873307c8db5901178f00945ce41866f2d
2 files changed
+35
-10
tests/kernels/core/test_fused_quant_layernorm.py
+18
-1
@@ -8,7 +8,7 @@ import pytest
8
import torch
9
10
import vllm._custom_ops as ops
11
-from tests.kernels.utils import fp8_ulp_distance, opcheck
11
+from tests.kernels.utils import fp8_allclose, fp8_ulp_distance, opcheck
12
from vllm.model_executor.layers.layernorm import RMSNorm
13
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
14
per_token_group_quant_fp8,
@@ -19,6 +19,12 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import (
19
from vllm.platforms import current_platform
20
from vllm.utils.torch_utils import set_random_seed
21
22
+ON_GFX950 = False
23
+if current_platform.is_rocm():
24
+ from vllm.platforms.rocm import on_gfx950
25
+
26
+ ON_GFX950 = on_gfx950()
27
+
28
DTYPES = [torch.bfloat16, torch.float]
29
QUANT_DTYPES = [torch.int8, current_platform.fp8_dtype()]
30
VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029]
@@ -317,6 +323,13 @@ def test_rms_norm(
323
and dtype == torch.bfloat16
324
and current_platform.is_rocm()
325
)
326
+ use_gfx950_fp8_allclose = (
327
+ current_platform.is_rocm()
328
+ and ON_GFX950
329
+ and group_size is None
330
+ and dtype == torch.bfloat16
331
+ and quant_dtype == current_platform.fp8_dtype()
332
+ )
333
334
def scales_close(rtol: float, atol: float) -> bool:
335
if torch.allclose(ref_scales, ops_scales, rtol=rtol, atol=atol):
@@ -341,6 +354,10 @@ def test_rms_norm(
354
ulp = fp8_ulp_distance(ref_out, ops_out)
355
max_outliers = ulp.numel() // 100_000 + 8
356
ok = int((ulp > 0).sum().item()) <= max_outliers
357
+ elif use_gfx950_fp8_allclose:
358
+ # Valid gfx950 reduction trees can straddle an E4M3 boundary.
359
+ ok = fp8_allclose(ops_out, ref_out, rtol=0.125, atol=2e-3)
360
+ ok = ok and int(fp8_ulp_distance(ops_out, ref_out).max()) <= 1
361
else:
362
# CUDA (& non-bf16): compare dequantized values with relaxed tolerance.
363
if group_size is None:
tests/kernels/core/test_layernorm.py
+17
-9
@@ -5,7 +5,7 @@ import pytest
5
import torch
6
7
from tests.kernels.quant_utils import FP8_DTYPE
8
-from tests.kernels.utils import fp8_ulp_distance, opcheck
8
+from tests.kernels.utils import fp8_allclose, fp8_ulp_distance, opcheck
9
from vllm import ir
10
from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
11
from vllm.platforms import current_platform
@@ -207,14 +207,22 @@ def test_fused_rms_norm_quant(
207
)
208
209
if current_platform.is_rocm():
210
- # Fused and unfused FP8 paths can land on opposite sides of an E4M3 tie;
211
- # tolerate a tiny number of isolated fp8 outliers on ROCm.
212
- ulp = fp8_ulp_distance(out_quant, out_quant_fused)
213
- max_outliers = ulp.numel() // 100_000 + 8
214
- num_outliers = int((ulp > 0).sum().item())
215
- assert num_outliers <= max_outliers, (
216
- f"FP8 quant mismatch: {num_outliers} fp8 outliers (allowed {max_outliers})"
217
- )
210
+ from vllm.platforms.rocm import on_gfx950
211
+
212
+ if on_gfx950() and dtype == torch.float16 and not add_residual:
213
+ # Fusion may round normalized FP16 across an E4M3 boundary on gfx950.
214
+ assert fp8_allclose(out_quant_fused, out_quant, rtol=0.125, atol=2e-3)
215
+ assert int(fp8_ulp_distance(out_quant_fused, out_quant).max()) <= 1
216
+ else:
217
+ # Fused and unfused FP8 paths can land on opposite sides of an E4M3
218
+ # tie; tolerate a tiny number of isolated fp8 outliers on ROCm.
219
+ ulp = fp8_ulp_distance(out_quant, out_quant_fused)
220
+ max_outliers = ulp.numel() // 100_000 + 8
221
+ num_outliers = int((ulp > 0).sum().item())
222
+ assert num_outliers <= max_outliers, (
223
+ f"FP8 quant mismatch: {num_outliers} fp8 outliers "
224
+ f"(allowed {max_outliers})"
225
+ )
226
else:
227
torch.testing.assert_close(
228
out_quant.to(dtype=torch.float32),