[ROCm] Relax MLA rope+cache test tolerances for bf16 (#51083)
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Rohan Potdar committed
Aug 5, 2026 at 17:50 UTC
38ebd97bcab958cbca63eeca87eb1fb57116b772
1 file changed
+17
-5
tests/kernels/core/test_rotary_embedding_mla_cache_fused.py
+17
-5
@@ -181,10 +181,14 @@ def test_concat_and_cache_mla_rope_fused(
181
kv_cache_scale,
182
)
183
184
- # ROCm neox-style Triton FMA diverges slightly from the fused kernel, so
185
- # relax the affected tolerance: rtol for fp8 (one e4m3 ULP ~12.5%) and atol
186
- # otherwise (bounded ~6e-4). Other paths use the CUDA defaults.
184
+ # On ROCm the AITER Triton rope diverges by ~1 ULP from the fused kernel,
185
+ # which the tight CUDA tolerances don't cover for the low-precision paths:
186
+ # - bf16: up to ~1 bf16 ULP (0.0156-0.03125 at values ~2-3), so relax the
187
+ # kv-cache atol (bounded ~0.02) and the query atol (~0.04).
188
+ # - neox-style fp8: one e4m3 ULP (~12.5%), so relax rtol.
189
+ # Other paths use the CUDA defaults.
190
rocm_neox = current_platform.is_rocm() and is_neox_style
191
+ rocm_bf16 = current_platform.is_rocm() and dtype == torch.bfloat16
192
if kv_cache_dtype == "fp8":
193
result_temp = torch.empty_like(kv_cache, dtype=torch.float16)
194
ops.convert_fp8(
@@ -198,13 +202,21 @@ def test_concat_and_cache_mla_rope_fused(
202
expected_temp, ref_kv_cache, kv_cache_scale.item(), kv_dtype=kv_cache_dtype
203
)
204
torch.testing.assert_close(
201
- result_temp, expected_temp, atol=0.001, rtol=0.15 if rocm_neox else 0.1
205
+ result_temp,
206
+ expected_temp,
207
+ atol=0.004 if rocm_bf16 else 0.001,
208
+ rtol=0.15 if rocm_neox or rocm_bf16 else 0.1,
209
)
210
+ elif rocm_bf16:
211
+ torch.testing.assert_close(kv_cache, ref_kv_cache, atol=0.02, rtol=1e-3)
212
elif rocm_neox:
213
torch.testing.assert_close(kv_cache, ref_kv_cache, atol=1e-3, rtol=1e-3)
214
else:
215
torch.testing.assert_close(kv_cache, ref_kv_cache)
216
217
torch.testing.assert_close(
209
- query, ref_q_pe, atol=get_default_atol(query), rtol=get_default_rtol(query)
218
+ query,
219
+ ref_q_pe,
220
+ atol=0.04 if rocm_bf16 else get_default_atol(query),
221
+ rtol=get_default_rtol(query),
222
)