diff --git a/tests/rocm/aiter/test_quant_op_schema.py b/tests/rocm/aiter/test_quant_op_schema.py index 9b2fac6e017..d5668f95a6d 100644 --- a/tests/rocm/aiter/test_quant_op_schema.py +++ b/tests/rocm/aiter/test_quant_op_schema.py @@ -2,9 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Schema/aliasing tests for the AITER FP8 quantization custom ops. # -# These use torch.library.opcheck, whose test_schema check catches custom ops -# whose implementation aliases an input that the registered schema declares as -# non-aliasing -- the failure mode behind the rocm_aiter_per_tensor_quant +# These use the shared opcheck helper, whose test_schema check catches custom +# ops whose implementation aliases an input that the registered schema declares +# as non-aliasing -- the failure mode behind the rocm_aiter_per_tensor_quant # regression (a returned scale that aliased the input scale). # # Skipped if AITER is not installed or the platform is not ROCm. @@ -14,6 +14,8 @@ import importlib.util import pytest import torch +from tests.kernels.utils import opcheck + # this import statement is needed to ensure the ops are registered from vllm._aiter_ops import rocm_aiter_ops from vllm.platforms import current_platform @@ -32,26 +34,14 @@ def _x(M=128, N=4096): return torch.randn((M, N), dtype=torch.float16, device="cuda") -# The in-place per-tensor op takes the fp8 output buffer as an input, which -# opcheck's test_schema cannot exercise ("mul_cuda" is unimplemented for fp8), -# so restrict to the utils that run on fp8 inputs. The aliasing contract for -# this op is instead covered by test_per_tensor_quant_torch_compile below. -_INPLACE_OPCHECK_UTILS = ( - "test_faketensor", - "test_aot_dispatch_dynamic", - "test_autograd_registration", -) - - def test_per_tensor_quant_static_schema(): """Static per-tensor: caller provides scale (the aliasing regression).""" x = _x() out = torch.empty_like(x, dtype=FP8_DTYPE) scale = torch.ones(1, dtype=torch.float32, device="cuda") - torch.library.opcheck( + opcheck( torch.ops.vllm.rocm_aiter_per_tensor_quant, (out, x, scale, False), - test_utils=_INPLACE_OPCHECK_UTILS, ) @@ -60,17 +50,16 @@ def test_per_tensor_quant_dynamic_schema(): x = _x() out = torch.empty_like(x, dtype=FP8_DTYPE) scale = torch.empty(1, dtype=torch.float32, device="cuda") - torch.library.opcheck( + opcheck( torch.ops.vllm.rocm_aiter_per_tensor_quant, (out, x, scale, True), - test_utils=_INPLACE_OPCHECK_UTILS, ) def test_per_token_quant_dynamic_schema(): """Dynamic per-token: op computes scale into a freshly allocated buffer.""" x = _x() - torch.library.opcheck( + opcheck( torch.ops.vllm.rocm_aiter_per_token_quant, (x, FP8_DTYPE, None), ) @@ -79,7 +68,7 @@ def test_per_token_quant_dynamic_schema(): def test_group_fp8_quant_schema(): """Dynamic per-token-group quant.""" x = _x() - torch.library.opcheck( + opcheck( torch.ops.vllm.rocm_aiter_group_fp8_quant, (x, 128), ) @@ -103,43 +92,24 @@ def test_per_tensor_quant_matches_native(dynamic): assert out.shape == x.shape assert out.dtype == FP8_DTYPE assert scale.shape == ref_scale.shape - if not dynamic: - # static scale is passed through unchanged - assert torch.equal(scale, scale_in) - # Compare dequantized values to be robust to 1-ULP fp8 boundary flips. deq = out.to(torch.float32) * scale - ref_deq = ref_out.to(torch.float32) * ref_scale - torch.testing.assert_close(deq, ref_deq, rtol=2e-2, atol=2e-2) + if dynamic: + # Dynamic mode: AITER and native each compute their own scale, so their + # outputs differ and can't be compared. Just check that AITER's output + # dequantizes back to the input, within fp8 rounding error. + torch.testing.assert_close(deq, x.to(torch.float32), rtol=0.07, atol=5e-2) + else: + # Static mode: both use the caller's scale, so the outputs must match. + assert torch.equal(scale, scale_in) + ref_deq = ref_out.to(torch.float32) * ref_scale + torch.testing.assert_close(deq, ref_deq, rtol=2e-2, atol=2e-2) -@pytest.mark.parametrize("dynamic", [True, False]) -def test_per_tensor_quant_torch_compile(monkeypatch, dynamic): - """per_tensor_quant compiles under inductor without an aliasing error. - - Forces the custom-op aliasing check to error (it is otherwise only a - warning outside CI), so a regression that returns an input-aliasing - scale fails here regardless of the CI env var. - """ - aliasing_cfg = pytest.importorskip("torch._functorch.config") - monkeypatch.setattr( - aliasing_cfg, "error_on_custom_op_aliasing", True, raising=False - ) - - x = _x() - scale = None if dynamic else torch.tensor([0.5], dtype=torch.float32, device="cuda") - - def fn(x, s): - return rocm_aiter_ops.per_tensor_quant(x, FP8_DTYPE, s) - - compiled = torch.compile(fn, fullgraph=True, backend="inductor", dynamic=False) - - out_eager, scale_eager = fn(x, scale) - out_compiled, scale_compiled = compiled(x, scale) - - assert out_compiled.shape == out_eager.shape - torch.testing.assert_close( - out_compiled.to(torch.float32) * scale_compiled, - out_eager.to(torch.float32) * scale_eager, - rtol=2e-2, - atol=2e-2, - ) +# A test_per_tensor_quant_torch_compile test previously lived here to validate +# the per-tensor aliasing contract. It existed because opcheck's test_schema +# could not check this op directly: test_schema compares the op's outputs with +# torch.allclose, but on fp8 outputs that comparison runs arithmetic fp8 does not +# support and raises "mul_cuda" is unimplemented for fp8. The fp8-safe opcheck +# helper fixes that by casting to double before the comparison, so the per-tensor +# schema tests above can now run test_schema directly. That makes this test +# redundant, so it has been removed.