From ac2819c81fcd15ac272dba8873784be21f34da46 Mon Sep 17 00:00:00 2001 From: Rahul Vijayaraghavan Date: Wed, 4 Mar 2026 03:13:58 +0530 Subject: [PATCH] Fix assertion tolerance for bf16 precision in triton attention UT (#17461) Signed-off-by: Rahul Vijayaraghavan --- test/registered/attention/test_triton_attention_kernels.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/registered/attention/test_triton_attention_kernels.py b/test/registered/attention/test_triton_attention_kernels.py index b38e673b4..d9520172b 100644 --- a/test/registered/attention/test_triton_attention_kernels.py +++ b/test/registered/attention/test_triton_attention_kernels.py @@ -304,8 +304,10 @@ class TestTritonAttention(CustomTestCase): max_len_in_batch, ) - self.assertTrue(torch.allclose(o_extend, o_redundant, rtol=1e-2)) - self.assertTrue(torch.allclose(o_extend_mask, o_redundant, rtol=1e-2)) + self.assertTrue(torch.allclose(o_extend, o_redundant, rtol=1e-2, atol=1e-3)) + self.assertTrue( + torch.allclose(o_extend_mask, o_redundant, rtol=1e-2, atol=1e-3) + ) def test_extend_attention(self):