[ROCm][CI] Add aiter per-token FP8 quant roundtrip and RMSNorm determinism tests (#50905)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Divakar Verma committed
Aug 4, 2026 at 22:32 UTC
bb543efdfd1ed340147f51a31141988e68d8990a
1 file changed
+31
-8
tests/rocm/aiter/test_quant_op_schema.py
+31
-8
@@ -105,11 +105,34 @@ def test_per_tensor_quant_matches_native(dynamic):
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.
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"