flash-attn v2.0.3 fixed this issue

facebookresearch · Aug 3, 2023 · 698532d · 698532d
1 parent 8d1a8e5
commit 698532d
Showing 1 changed file with 0 additions and 12 deletions.
diff --git a/xformers/ops/fmha/flash.py b/xformers/ops/fmha/flash.py
@@ -302,18 +302,6 @@ class BwOp(AttentionBwOpBase):
 
     MAX_HEADDIM_SM8x = 192
 
-    @classmethod
-    def shape_not_supported_reasons(
-        cls, Mq: int, Mkv: int, K: int, Kv: int
-    ) -> List[str]:
-        reasons = super().shape_not_supported_reasons(Mq, Mkv, K, Kv)
-        if (Mq % 128) or (Mkv % 128):
-            reasons.append(
-                "flashv2 beta: BW is incorrect when seqlen is not aligned on 128 "
-                "(https://github.com/Dao-AILab/flash-attention/issues/334)"
-            )
-        return reasons
-
     @classmethod
     def not_supported_reasons(cls, d: Inputs) -> List[str]:
         reasons = super(BwOp, cls).not_supported_reasons(d)