maybe fix flash attention

2025-02-26 12:15:35 +02:00
parent 7606f28122
commit 0a020a6f18
1 changed files with 6 additions and 5 deletions
@@ -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,