From 34c19a32c13af63f61dbf18c9bd0c646fbb1c345 Mon Sep 17 00:00:00 2001 From: strgrb Date: Wed, 4 Mar 2026 22:47:29 +0800 Subject: [PATCH] fix flaky test for test_kda_kernels (#19864) --- test/registered/attention/test_kda_kernels.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/registered/attention/test_kda_kernels.py b/test/registered/attention/test_kda_kernels.py index e9c712db7..57d3962eb 100644 --- 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))