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:
@@ -304,8 +304,10 @@ class TestTritonAttention(CustomTestCase):
|
|||||||
max_len_in_batch,
|
max_len_in_batch,
|
||||||
)
|
)
|
||||||
|
|
||||||
self.assertTrue(torch.allclose(o_extend, 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))
|
self.assertTrue(
|
||||||
|
torch.allclose(o_extend_mask, o_redundant, rtol=1e-2, atol=1e-3)
|
||||||
|
)
|
||||||
|
|
||||||
def test_extend_attention(self):
|
def test_extend_attention(self):
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user