Fix ROCm AITER FP8 KV test tolerances. (#53590)
Signed-off-by: Aarushi Jain <aarushjain29@users.noreply.github.com> Co-authored-by: Aarushi Jain <aarushjain29@users.noreply.github.com>
Aarushi Jain committed
Sep 9, 2026 at 12:22 UTC
1454b7172796c00955ae32e3fe865b98b262855a
2 files changed
+20
-12
tests/kernels/attention/test_rocm_aiter_fa.py
+9
-8
@@ -805,14 +805,14 @@ def test_aiter_mha_decode_single_token(dtype):
805
def test_aiter_mha_varlen_fp8_kv(dtype):
806
"""AITER flash attention with FP8 KV cache matches reference on BF16-cast KV.
807
808
- cp_mha_gather_cache is called with dequant=True to dequantize FP8 to dtype
809
- before passing to flash_attn_varlen_func. We compare to ref_paged_attn on
810
- the dtype-cast KV cache with FP8 quantization tolerance.
811
-
812
- Exercises: VLLM_ROCM_USE_AITER, VLLM_ROCM_USE_AITER_MHA, FP8 KV path.
808
+ The reference dequantizes the same FP8 KV tensors the kernel reads, so
809
+ e4m3 rounding cancels on both sides; the residual is bf16/fp16 kernel
810
+ rounding (~few millidecimals), not fp8-scale quantization error.
811
"""
814
- atol = 0.15
815
- rtol = 0.05
812
+ if dtype == torch.bfloat16:
813
+ atol, rtol = 6e-3, 1e-2
814
+ else:
815
+ atol, rtol = 7e-4, 1e-3
816
_assert_aiter_supported()
817
if not current_platform.supports_fp8():
818
pytest.skip("FP8 not supported on this hardware")
@@ -900,7 +900,8 @@ def test_aiter_mha_varlen_fp8_kv(dtype):
900
scale=scale,
901
)
902
903
- # FP8 quantization + dequantization introduces noise
903
+ # FP8 quantization + dequantization cancels on the reference side; compare
904
+ # kernel rounding at the working precision.
905
_print_close_stats(
906
f"varlen_fp8_kv dtype={dtype} query_len={query_len} kv_len={kv_len}",
907
output,
tests/kernels/attention/test_rocm_aiter_unified_attn.py
+11
-4
@@ -49,14 +49,20 @@ PREFILL_SEQ_LENS = [
49
[(256, 1024), (128, 2048)],
50
]
51
52
+Fp8Variant = Literal["fp8_kv", "fp8_query", "fp8_query_kv"]
53
+
54
DEFAULT_ATOL, DEFAULT_RTOL = 1.5e-2, 1e-2
53
-FP8_ATOL, FP8_RTOL = 1.5e-1, 1.5e-1
55
+# fp8_kv: reference dequantizes the same KV, so error is bf16-scale (~2e-3).
56
+# fp8_query / fp8_query_kv: query and/or KV e4m3 rounding dominates (~0.1-0.4).
57
+FP8_TOL_BY_VARIANT: dict[Fp8Variant, tuple[float, float]] = {
58
+ "fp8_kv": (2.5e-3, 1e-2),
59
+ "fp8_query": (1.5e-1, 1.5e-1),
60
+ "fp8_query_kv": (1.5e-1, 1.5e-1),
61
+}
62
# Non-unity scale so q_descale handling is exercised explicitly.
63
Q_SCALE = 0.75
64
K_SCALE, V_SCALE = 0.5, 0.25
65
58
-Fp8Variant = Literal["fp8_kv", "fp8_query", "fp8_query_kv"]
59
-
66
FP8_VARIANTS = [
67
pytest.param("fp8_kv", id="fp8_kv"),
68
pytest.param("fp8_query", id="fp8_query"),
@@ -336,4 +342,5 @@ def test_aiter_unified_attn_fp8(
342
block_size=block_size,
343
variant=variant,
344
)
339
- _assert_matches_reference(case, atol=FP8_ATOL, rtol=FP8_RTOL)
345
+ atol, rtol = FP8_TOL_BY_VARIANT[variant]
346
+ _assert_matches_reference(case, atol=atol, rtol=rtol)