[ROCm] Fix AITER FP8 quantization schema tests (#46414)

Signed-off-by: Djordje Ramic <djoramic@amd.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

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.