diff --git a/wanvideo/modules/attention.py b/wanvideo/modules/attention.py
index 0735663..1cc45ee 100644
--- a/wanvideo/modules/attention.py
+++ b/wanvideo/modules/attention.py
@@ -151,11 +151,12 @@ def attention(
     deterministic=False,
     dtype=torch.bfloat16,
     attention_mode='sdpa',
-):
-    if attention_mode == 'flash_attention_2':
-        fa_version = 2
-    elif attention_mode == 'flash_attention_3':
-        fa_version = 3
+):  
+    if "flash" in attention_mode:
+        if attention_mode == 'flash_attention_2':
+            fa_version = 2
+        elif attention_mode == 'flash_attention_3':
+            fa_version = 3
         return flash_attention(
             q=q,
             k=k,