hotfix pipefusion using flash_attn (#411)

xdit-project · Dec 26, 2024 · 81700db · 81700db
1 parent 92187b8
commit 81700db
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 21 deletions.
diff --git a/setup.py b/setup.py
@@ -43,7 +43,7 @@ def get_cuda_version():
         ],
         extras_require={
             "diffusers": [
-                "diffusers>=0.32.0",  # NOTE: diffusers>=0.32.0.dev is necessary for CogVideoX and Flux
+                "diffusers>=0.31.0",  # NOTE: diffusers>=0.32.0.dev is necessary for CogVideoX and Flux
                 "flash_attn>=2.6.3",
             ]
         },

diff --git a/xfuser/core/long_ctx_attention/ring/ring_flash_attn.py b/xfuser/core/long_ctx_attention/ring/ring_flash_attn.py
@@ -11,6 +11,7 @@
 except ImportError:
     flash_attn = None
     _flash_attn_forward = None
+    from yunchang.kernels.attention import pytorch_attn_forward
 
 def xdit_ring_flash_attn_forward(
     process_group,
@@ -85,34 +86,43 @@ def xdit_ring_flash_attn_forward(
             key, value = k, v
 
         if not causal or step <= comm.rank:
-            assert flash_attn is not None, f"FlashAttention is not available, please install flash_attn"
-            if flash_attn.__version__ <= "2.6.3":
-                block_out, _, _, _, _, block_lse, _, _ = _flash_attn_forward(
+            if flash_attn is None:
+                block_out, block_lse = pytorch_attn_forward(
                     q,
                     key,
                     value,
                     dropout_p,
                     softmax_scale,
                     causal=causal and step == 0,
-                    window_size=window_size,
-                    softcap=0.0,
-                    alibi_slopes=alibi_slopes,
-                    return_softmax=True and dropout_p > 0,
                 )
             else:
-                block_out, block_lse, _, _ = _flash_attn_forward(
-                    q,
-                    key,
-                    value,
-                    dropout_p,
-                    softmax_scale,
-                    causal=causal and step == 0,
-                    window_size_left=window_size[0],
-                    window_size_right=window_size[1],
-                    softcap=0.0,
-                    alibi_slopes=alibi_slopes,
-                    return_softmax=True and dropout_p > 0,
-                )
+                if flash_attn.__version__ <= "2.6.3":
+                    block_out, _, _, _, _, block_lse, _, _ = _flash_attn_forward(
+                        q,
+                        key,
+                        value,
+                        dropout_p,
+                        softmax_scale,
+                        causal=causal and step == 0,
+                        window_size=window_size,
+                        softcap=0.0,
+                        alibi_slopes=alibi_slopes,
+                        return_softmax=True and dropout_p > 0,
+                    )
+                else:
+                    block_out, block_lse, _, _ = _flash_attn_forward(
+                        q,
+                        key,
+                        value,
+                        dropout_p,
+                        softmax_scale,
+                        causal=causal and step == 0,
+                        window_size_left=window_size[0],
+                        window_size_right=window_size[1],
+                        softcap=0.0,
+                        alibi_slopes=alibi_slopes,
+                        return_softmax=True and dropout_p > 0,
+                    )
             out, lse = update_out_and_lse(out, lse, block_out, block_lse)
 
         if step + 1 != comm.world_size: