Update attention.py

2025-02-26 12:17:52 +02:00
parent 0a020a6f18
commit 3c70da4192
1 changed files with 2 additions and 2 deletions
@@ -153,9 +153,9 @@ def attention(
    attention_mode='sdpa',
 ):  
    if "flash" in attention_mode:
-        if attention_mode == 'flash_attention_2':
+        if attention_mode == 'flash_attn_2':
            fa_version = 2
-        elif attention_mode == 'flash_attention_3':
+        elif attention_mode == 'flash_attn_3':
            fa_version = 3
        return flash_attention(
            q=q,