[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>
djramic committed
Jun 24, 2026 at 16:29 UTC
61ee183d28ff3dbd169b61e0205e693a2314337f
1 file changed
+29
-59
tests/rocm/aiter/test_quant_op_schema.py
+29
-59
@@ -2,9 +2,9 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
# Schema/aliasing tests for the AITER FP8 quantization custom ops.
4
#
5
-# These use torch.library.opcheck, whose test_schema check catches custom ops
6
-# whose implementation aliases an input that the registered schema declares as
7
-# non-aliasing -- the failure mode behind the rocm_aiter_per_tensor_quant
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.
@@ -14,6 +14,8 @@ import importlib.util
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
@@ -32,26 +34,14 @@ def _x(M=128, N=4096):
34
return torch.randn((M, N), dtype=torch.float16, device="cuda")
35
36
35
-# The in-place per-tensor op takes the fp8 output buffer as an input, which
36
-# opcheck's test_schema cannot exercise ("mul_cuda" is unimplemented for fp8),
37
-# so restrict to the utils that run on fp8 inputs. The aliasing contract for
38
-# this op is instead covered by test_per_tensor_quant_torch_compile below.
39
-_INPLACE_OPCHECK_UTILS = (
40
- "test_faketensor",
41
- "test_aot_dispatch_dynamic",
42
- "test_autograd_registration",
43
-)
44
-
45
-
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")
51
- torch.library.opcheck(
42
+ opcheck(
43
torch.ops.vllm.rocm_aiter_per_tensor_quant,
44
(out, x, scale, False),
54
- test_utils=_INPLACE_OPCHECK_UTILS,
45
)
46
47
@@ -60,17 +50,16 @@ def test_per_tensor_quant_dynamic_schema():
50
x = _x()
51
out = torch.empty_like(x, dtype=FP8_DTYPE)
52
scale = torch.empty(1, dtype=torch.float32, device="cuda")
63
- torch.library.opcheck(
53
+ opcheck(
54
torch.ops.vllm.rocm_aiter_per_tensor_quant,
55
(out, x, scale, True),
66
- test_utils=_INPLACE_OPCHECK_UTILS,
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()
73
- torch.library.opcheck(
62
+ opcheck(
63
torch.ops.vllm.rocm_aiter_per_token_quant,
64
(x, FP8_DTYPE, None),
65
)
@@ -79,7 +68,7 @@ def test_per_token_quant_dynamic_schema():
68
def test_group_fp8_quant_schema():
69
"""Dynamic per-token-group quant."""
70
x = _x()
82
- torch.library.opcheck(
71
+ opcheck(
72
torch.ops.vllm.rocm_aiter_group_fp8_quant,
73
(x, 128),
74
)
@@ -103,43 +92,24 @@ def test_per_tensor_quant_matches_native(dynamic):
92
assert out.shape == x.shape
93
assert out.dtype == FP8_DTYPE
94
assert scale.shape == ref_scale.shape
106
- if not dynamic:
107
- # static scale is passed through unchanged
108
- assert torch.equal(scale, scale_in)
109
- # Compare dequantized values to be robust to 1-ULP fp8 boundary flips.
95
deq = out.to(torch.float32) * scale
111
- ref_deq = ref_out.to(torch.float32) * ref_scale
112
- torch.testing.assert_close(deq, ref_deq, rtol=2e-2, atol=2e-2)
113
-
114
-
115
-@pytest.mark.parametrize("dynamic", [True, False])
116
-def test_per_tensor_quant_torch_compile(monkeypatch, dynamic):
117
- """per_tensor_quant compiles under inductor without an aliasing error.
118
-
119
- Forces the custom-op aliasing check to error (it is otherwise only a
120
- warning outside CI), so a regression that returns an input-aliasing
121
- scale fails here regardless of the CI env var.
122
- """
123
- aliasing_cfg = pytest.importorskip("torch._functorch.config")
124
- monkeypatch.setattr(
125
- aliasing_cfg, "error_on_custom_op_aliasing", True, raising=False
126
- )
127
-
128
- x = _x()
129
- scale = None if dynamic else torch.tensor([0.5], dtype=torch.float32, device="cuda")
130
-
131
- def fn(x, s):
132
- return rocm_aiter_ops.per_tensor_quant(x, FP8_DTYPE, s)
133
-
134
- compiled = torch.compile(fn, fullgraph=True, backend="inductor", dynamic=False)
135
-
136
- out_eager, scale_eager = fn(x, scale)
137
- out_compiled, scale_compiled = compiled(x, scale)
138
-
139
- assert out_compiled.shape == out_eager.shape
140
- torch.testing.assert_close(
141
- out_compiled.to(torch.float32) * scale_compiled,
142
- out_eager.to(torch.float32) * scale_eager,
143
- rtol=2e-2,
144
- atol=2e-2,
145
- )
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
+# A test_per_tensor_quant_torch_compile test previously lived here to validate
109
+# the per-tensor aliasing contract. It existed because opcheck's test_schema
110
+# could not check this op directly: test_schema compares the op's outputs with
111
+# torch.allclose, but on fp8 outputs that comparison runs arithmetic fp8 does not
112
+# support and raises "mul_cuda" is unimplemented for fp8. The fp8-safe opcheck
113
+# helper fixes that by casting to double before the comparison, so the per-tensor
114
+# schema tests above can now run test_schema directly. That makes this test
115
+# redundant, so it has been removed.