| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # Schema/aliasing tests for the AITER FP8 quantization custom ops. |
| 4 | # |
| 5 | # These use the shared opcheck helper, whose test_schema check catches custom |
| 6 | # ops whose implementation aliases an input that the registered schema declares |
| 7 | # as non-aliasing -- the failure mode behind the rocm_aiter_per_tensor_quant |
| 8 | # regression (a returned scale that aliased the input scale). |
| 9 | # |
| 10 | # Skipped if AITER is not installed or the platform is not ROCm. |
| 11 | |
| 12 | import importlib.util |
| 13 | |
| 14 | import pytest |
| 15 | import torch |
| 16 | |
| 17 | from tests.kernels.utils import opcheck |
| 18 | |
| 19 | # this import statement is needed to ensure the ops are registered |
| 20 | from vllm._aiter_ops import rocm_aiter_ops |
| 21 | from vllm.platforms import current_platform |
| 22 | |
| 23 | aiter_available = importlib.util.find_spec("aiter") is not None |
| 24 | |
| 25 | pytestmark = pytest.mark.skipif( |
| 26 | not (current_platform.is_rocm() and aiter_available), |
| 27 | reason="AITER ops are only available on ROCm with aiter package installed", |
| 28 | ) |
| 29 | |
| 30 | FP8_DTYPE = current_platform.fp8_dtype() |
| 31 | |
| 32 | |
| 33 | def _x(M=128, N=4096): |
| 34 | return torch.randn((M, N), dtype=torch.float16, device="cuda") |
| 35 | |
| 36 | |
| 37 | def test_per_tensor_quant_static_schema(): |
| 38 | """Static per-tensor: caller provides scale (the aliasing regression).""" |
| 39 | x = _x() |
| 40 | out = torch.empty_like(x, dtype=FP8_DTYPE) |
| 41 | scale = torch.ones(1, dtype=torch.float32, device="cuda") |
| 42 | opcheck( |
| 43 | torch.ops.vllm.rocm_aiter_per_tensor_quant, |
| 44 | (out, x, scale, False), |
| 45 | ) |
| 46 | |
| 47 | |
| 48 | def test_per_tensor_quant_dynamic_schema(): |
| 49 | """Dynamic per-tensor: op computes scale into the caller's buffer.""" |
| 50 | x = _x() |
| 51 | out = torch.empty_like(x, dtype=FP8_DTYPE) |
| 52 | scale = torch.empty(1, dtype=torch.float32, device="cuda") |
| 53 | opcheck( |
| 54 | torch.ops.vllm.rocm_aiter_per_tensor_quant, |
| 55 | (out, x, scale, True), |
| 56 | ) |
| 57 | |
| 58 | |
| 59 | def test_per_token_quant_dynamic_schema(): |
| 60 | """Dynamic per-token: op computes scale into a freshly allocated buffer.""" |
| 61 | x = _x() |
| 62 | opcheck( |
| 63 | torch.ops.vllm.rocm_aiter_per_token_quant, |
| 64 | (x, FP8_DTYPE, None), |
| 65 | ) |
| 66 | |
| 67 | |
| 68 | def test_group_fp8_quant_schema(): |
| 69 | """Dynamic per-token-group quant.""" |
| 70 | x = _x() |
| 71 | opcheck( |
| 72 | torch.ops.vllm.rocm_aiter_group_fp8_quant, |
| 73 | (x, 128), |
| 74 | ) |
| 75 | |
| 76 | |
| 77 | @pytest.mark.parametrize("dynamic", [True, False]) |
| 78 | def test_per_tensor_quant_matches_native(dynamic): |
| 79 | """Wrapper output matches the native scaled_fp8_quant reference.""" |
| 80 | from vllm import _custom_ops as ops |
| 81 | |
| 82 | torch.manual_seed(0) |
| 83 | x = _x() |
| 84 | if dynamic: |
| 85 | scale_in = None |
| 86 | else: |
| 87 | scale_in = torch.tensor([0.5], dtype=torch.float32, device="cuda") |
| 88 | |
| 89 | out, scale = rocm_aiter_ops.per_tensor_quant(x, FP8_DTYPE, scale_in) |
| 90 | ref_out, ref_scale = ops.scaled_fp8_quant(x, scale_in) |
| 91 | |
| 92 | assert out.shape == x.shape |
| 93 | assert out.dtype == FP8_DTYPE |
| 94 | assert scale.shape == ref_scale.shape |
| 95 | deq = out.to(torch.float32) * scale |
| 96 | if dynamic: |
| 97 | # Dynamic mode: AITER and native each compute their own scale, so their |
| 98 | # outputs differ and can't be compared. Just check that AITER's output |
| 99 | # dequantizes back to the input, within fp8 rounding error. |
| 100 | torch.testing.assert_close(deq, x.to(torch.float32), rtol=0.07, atol=5e-2) |
| 101 | else: |
| 102 | # Static mode: both use the caller's scale, so the outputs must match. |
| 103 | assert torch.equal(scale, scale_in) |
| 104 | ref_deq = ref_out.to(torch.float32) * ref_scale |
| 105 | torch.testing.assert_close(deq, ref_deq, rtol=2e-2, atol=2e-2) |
| 106 | |
| 107 | |
| 108 | def test_per_token_quant_matches_native(): |
| 109 | """Per-token quant output dequantizes back to the input within FP8 error.""" |
| 110 | torch.manual_seed(0) |
| 111 | x = _x() |
| 112 | |
| 113 | out, scale = rocm_aiter_ops.per_token_quant(x, FP8_DTYPE) |
| 114 | |
| 115 | assert out.shape == x.shape |
| 116 | assert out.dtype == FP8_DTYPE |
| 117 | assert scale.shape[0] == x.shape[0] |
| 118 | |
| 119 | # Dequantize and compare to original |
| 120 | deq = out.to(torch.float32) * scale.view(-1, 1) |
| 121 | torch.testing.assert_close(deq, x.to(torch.float32), rtol=0.07, atol=5e-2) |
| 122 | |
| 123 | |
| 124 | def test_rms_norm_determinism(): |
| 125 | """AITER RMSNorm produces bitwise-identical results across repeated calls.""" |
| 126 | # Import to ensure ops are registered |
| 127 | import vllm.kernels.aiter_ops # noqa: F401 |
| 128 | |
| 129 | torch.manual_seed(0) |
| 130 | M, N = 32, 512 |
| 131 | x = torch.randn(M, N, dtype=torch.bfloat16, device="cuda") |
| 132 | weight = torch.ones(N, dtype=torch.bfloat16, device="cuda") |
| 133 | eps = 1e-5 |
| 134 | |
| 135 | reference = torch.ops.vllm_aiter.rms_norm(x, weight, eps) |
| 136 | for i in range(3): |
| 137 | result = torch.ops.vllm_aiter.rms_norm(x, weight, eps) |
| 138 | assert torch.equal(reference, result), f"Run {i + 1} differs from reference" |