[Bugfix][ROCm] Fix rocm_aiter_per_tensor_quant custom op aliasing (#45747)
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Rohan Potdar committed
Jun 18, 2026 at 11:30 UTC
509947463375cc27e2a60d05ce5463f6dd059171
2 files changed
+166
-13
tests/rocm/aiter/test_quant_op_schema.py
new
+145
@@ -0,0 +1,145 @@
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 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
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
+# this import statement is needed to ensure the ops are registered
18
+from vllm._aiter_ops import rocm_aiter_ops
19
+from vllm.platforms import current_platform
20
+
21
+aiter_available = importlib.util.find_spec("aiter") is not None
22
+
23
+pytestmark = pytest.mark.skipif(
24
+ not (current_platform.is_rocm() and aiter_available),
25
+ reason="AITER ops are only available on ROCm with aiter package installed",
26
+)
27
+
28
+FP8_DTYPE = current_platform.fp8_dtype()
29
+
30
+
31
+def _x(M=128, N=4096):
32
+ return torch.randn((M, N), dtype=torch.float16, device="cuda")
33
+
34
+
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
+
46
+def test_per_tensor_quant_static_schema():
47
+ """Static per-tensor: caller provides scale (the aliasing regression)."""
48
+ x = _x()
49
+ out = torch.empty_like(x, dtype=FP8_DTYPE)
50
+ scale = torch.ones(1, dtype=torch.float32, device="cuda")
51
+ torch.library.opcheck(
52
+ torch.ops.vllm.rocm_aiter_per_tensor_quant,
53
+ (out, x, scale, False),
54
+ test_utils=_INPLACE_OPCHECK_UTILS,
55
+ )
56
+
57
+
58
+def test_per_tensor_quant_dynamic_schema():
59
+ """Dynamic per-tensor: op computes scale into the caller's buffer."""
60
+ x = _x()
61
+ out = torch.empty_like(x, dtype=FP8_DTYPE)
62
+ scale = torch.empty(1, dtype=torch.float32, device="cuda")
63
+ torch.library.opcheck(
64
+ torch.ops.vllm.rocm_aiter_per_tensor_quant,
65
+ (out, x, scale, True),
66
+ test_utils=_INPLACE_OPCHECK_UTILS,
67
+ )
68
+
69
+
70
+def test_per_token_quant_dynamic_schema():
71
+ """Dynamic per-token: op computes scale into a freshly allocated buffer."""
72
+ x = _x()
73
+ torch.library.opcheck(
74
+ torch.ops.vllm.rocm_aiter_per_token_quant,
75
+ (x, FP8_DTYPE, None),
76
+ )
77
+
78
+
79
+def test_group_fp8_quant_schema():
80
+ """Dynamic per-token-group quant."""
81
+ x = _x()
82
+ torch.library.opcheck(
83
+ torch.ops.vllm.rocm_aiter_group_fp8_quant,
84
+ (x, 128),
85
+ )
86
+
87
+
88
+@pytest.mark.parametrize("dynamic", [True, False])
89
+def test_per_tensor_quant_matches_native(dynamic):
90
+ """Wrapper output matches the native scaled_fp8_quant reference."""
91
+ from vllm import _custom_ops as ops
92
+
93
+ torch.manual_seed(0)
94
+ x = _x()
95
+ if dynamic:
96
+ scale_in = None
97
+ else:
98
+ scale_in = torch.tensor([0.5], dtype=torch.float32, device="cuda")
99
+
100
+ out, scale = rocm_aiter_ops.per_tensor_quant(x, FP8_DTYPE, scale_in)
101
+ ref_out, ref_scale = ops.scaled_fp8_quant(x, scale_in)
102
+
103
+ assert out.shape == x.shape
104
+ assert out.dtype == FP8_DTYPE
105
+ 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.
110
+ 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
+ )
vllm/_aiter_ops.py
+21
-13
@@ -1019,23 +1019,26 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_fake(
1019
1020
1021
def _rocm_aiter_per_tensor_quant_impl(
1022
+ out: torch.Tensor,
1023
x: torch.Tensor,
1023
- quant_dtype: torch.dtype,
1024
- scale: torch.Tensor | None = None,
1025
-) -> tuple[torch.Tensor, torch.Tensor]:
1026
- from aiter.ops.quant import per_tensor_quant_hip
1024
+ scale: torch.Tensor,
1025
+ is_dynamic: bool,
1026
+) -> None:
1027
+ from aiter.ops.quant import dynamic_per_tensor_quant, static_per_tensor_quant
1028
1028
- return per_tensor_quant_hip(x, scale, quant_dtype)
1029
+ if is_dynamic:
1030
+ dynamic_per_tensor_quant(out, x, scale)
1031
+ else:
1032
+ static_per_tensor_quant(out, x, scale)
1033
1034
1035
def _rocm_aiter_per_tensor_quant_fake(
1036
+ out: torch.Tensor,
1037
x: torch.Tensor,
1033
- quant_dtype: torch.dtype,
1034
- scale: torch.Tensor | None = None,
1035
-) -> tuple[torch.Tensor, torch.Tensor]:
1036
- return torch.empty_like(x, dtype=quant_dtype), torch.empty(
1037
- 1, dtype=torch.float32, device=x.device
1038
- )
1038
+ scale: torch.Tensor,
1039
+ is_dynamic: bool,
1040
+) -> None:
1041
+ pass
1042
1043
1044
def _rocm_aiter_per_token_quant_impl(
@@ -1979,7 +1982,7 @@ class rocm_aiter_ops:
1982
direct_register_custom_op(
1983
op_name="rocm_aiter_per_tensor_quant",
1984
op_func=_rocm_aiter_per_tensor_quant_impl,
1982
- mutates_args=[],
1985
+ mutates_args=["out", "scale"],
1986
fake_impl=_rocm_aiter_per_tensor_quant_fake,
1987
dispatch_key=current_platform.dispatch_key,
1988
)
@@ -2392,7 +2395,12 @@ class rocm_aiter_ops:
2395
quant_dtype: torch.dtype,
2396
scale: torch.Tensor | None = None,
2397
) -> tuple[torch.Tensor, torch.Tensor]:
2395
- return torch.ops.vllm.rocm_aiter_per_tensor_quant(x, quant_dtype, scale)
2398
+ out = torch.empty_like(x, dtype=quant_dtype)
2399
+ is_dynamic = scale is None
2400
+ if is_dynamic:
2401
+ scale = torch.empty(1, dtype=torch.float32, device=x.device)
2402
+ torch.ops.vllm.rocm_aiter_per_tensor_quant(out, x, scale, is_dynamic)
2403
+ return out, scale
2404
2405
@staticmethod
2406
def per_token_quant(