fix flaky test for test_kda_kernels (#19864)

2026-03-04 22:47:29 +08:00
parent 738ebfd330
commit 34c19a32c1
1 changed files with 3 additions and 1 deletions
--- a/test/registered/attention/test_kda_kernels.py
+++ b/test/registered/attention/test_kda_kernels.py
@@ -115,7 +115,9 @@ class TestKDAFusedSigmoidGatingRecurrent(unittest.TestCase):
        abs_diff_out = (core_attn_out - core_attn_out_ref).abs().max()
        abs_diff_state = (last_state - last_state_ref).abs().max()
        print(f"{abs_diff_out=}, {abs_diff_state=}")
-        self.assertTrue(torch.allclose(core_attn_out, core_attn_out_ref))
+        self.assertTrue(
+            torch.allclose(core_attn_out, core_attn_out_ref, rtol=1e-3, atol=1e-4)
+        )
        self.assertTrue(torch.allclose(last_state, last_state_ref))