diff --git a/tests/kernels/core/test_fused_quant_layernorm.py b/tests/kernels/core/test_fused_quant_layernorm.py index 02bdc90bb61..06df13d48c4 100644 --- a/tests/kernels/core/test_fused_quant_layernorm.py +++ b/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: diff --git a/tests/kernels/core/test_layernorm.py b/tests/kernels/core/test_layernorm.py index 1a6c893ad6b..a1ca0c09cf6 100644 --- a/tests/kernels/core/test_layernorm.py +++ b/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),