Fix assertion tolerance for bf16 precision in triton attention UT (#17461)

Signed-off-by: Rahul Vijayaraghavan <rahul.vijayaraghavan@intel.com>
This commit is contained in:
Rahul Vijayaraghavan
2026-03-04 03:13:58 +05:30
committed by GitHub
parent 85ab6a7f54
commit ac2819c81f

View File

@@ -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):