Relax the atol for test_forward and test_dropout due to the using of …

…packed fp16_2_fp32 conversion in ck_tile
facebookresearch · Jun 20, 2024 · 8e9dc32 · 8e9dc32
1 parent 2655be6
commit 8e9dc32
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/tests/test_mem_eff_attention.py b/tests/test_mem_eff_attention.py
@@ -910,7 +910,7 @@ def test_dropout(op, q_len, kv_len, batch_size, k_len, p, seed, attn_bias):
     if dtype is torch.float:
         assert_allclose(out, ref, atol=2e-4), f"{(out - ref).abs().max()}"
     else:
-        assert_allclose(out.float(), ref, atol=2.2e-2), f"{(out - ref).abs().max()}"
+        assert_allclose(out.float(), ref, atol=2.8e-2), f"{(out - ref).abs().max()}"
 
     num_trials = 1000
     p_val_tol = 1e-6

diff --git a/xformers/ops/fmha/ck.py b/xformers/ops/fmha/ck.py
@@ -183,7 +183,7 @@ class FwOp(AttentionFwOpBase):
 
     ERROR_ATOL: Mapping[torch.dtype, float] = {
         torch.float: 3e-4,
-        torch.half: 4e-3,
+        torch.half: 6e-3,
         torch.bfloat16: 2.8e-2,
     }
     ERROR_RTOL: Mapping[torch.dtype, float] = {