ci fix

Author: EduardDurech

src/transformers/modeling_flash_attention_utils.py

@@ -30,6 +30,7 @@
 
 
 logger = logging.get_logger(__name__)
+flash_attn_func = None
 
 
 def _index_first_axis(tensor, indices):
@@ -92,6 +93,7 @@ def _fa3_pad_input(hidden_states, indices, batch, seqlen):
     output[indices] = hidden_states
     return output.view(batch, seqlen, *dim)
 
+
 FA_VERSION = None
 if is_flash_attn_2_available():
     from flash_attn import flash_attn_func as flash_attn_2_func
@@ -135,10 +137,19 @@ def _fa3_pad_input(hidden_states, indices, batch, seqlen):
 
 # patch functions in package `flash-attn` when using flash-attention on Ascend NPU.
 if is_torch_npu_available():
-    from .integrations.npu_flash_attention import pad_input, unpad_input
-    from .integrations.npu_flash_attention import npu_apply_rotary_emb as apply_rotary_emb  # noqa
-    from .integrations.npu_flash_attention import npu_flash_attn_func as flash_attn_func
-    from .integrations.npu_flash_attention import npu_flash_attn_varlen_func as flash_attn_varlen_func
+    from .integrations.npu_flash_attention import (
+        npu_apply_rotary_emb as apply_rotary_emb,  # noqa: F401
+    )
+    from .integrations.npu_flash_attention import (
+        npu_flash_attn_func as flash_attn_func,
+    )
+    from .integrations.npu_flash_attention import (
+        npu_flash_attn_varlen_func as flash_attn_varlen_func,
+    )
+    from .integrations.npu_flash_attention import (
+        pad_input,
+        unpad_input,
+    )
 
 
 _flash_supports_window_size = False
@@ -279,9 +290,7 @@ def _upad_input(
     else:
         # The -q_len: slice assumes left padding.
         attention_mask = attention_mask[:, -query_length:]
-        query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q, *_ = unpad_input_func(
-            query_layer, attention_mask
-        )
+        query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q, *_ = unpad_input_func(query_layer, attention_mask)
 
     return (
         query_layer,

src/transformers/modeling_utils.py

@@ -2549,7 +2549,7 @@ def _check_and_enable_flash_attn_3(
             else:
                 raise ValueError(
                     f"{preface} Flash Attention 3 is not available on CPU. Please make sure torch can access a CUDA device."
-                    )
+                )
 
         _is_bettertransformer = getattr(cls, "use_bettertransformer", False)
 
@@ -2570,9 +2570,7 @@ def _check_and_enable_flash_attn_3(
             )
 
         if getattr(config, "alibi", False) or getattr(config, "use_alibi", False):
-            raise ValueError(
-                "Model is configured to use ALiBi, which is not supported by Flash Attention 3."
-            )
+            raise ValueError("Model is configured to use ALiBi, which is not supported by Flash Attention 3.")
 
         # Check for attention dropout, which is incompatible with FA3
         if hasattr(config, "attention_dropout") and config.attention_dropout > 0:

src/transformers/models/aria/modeling_aria.py

@@ -667,6 +667,7 @@ class AriaPreTrainedModel(PreTrainedModel):
     supports_gradient_checkpointing = True
     _no_split_modules = ["AriaDecoderLayer"]
     _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn_3 = True
     _supports_flash_attn_2 = True
     _supports_sdpa = True
     _supports_flex_attn = True

src/transformers/models/bitnet/modeling_bitnet.py

@@ -318,6 +318,7 @@ class BitNetPreTrainedModel(PreTrainedModel):
     supports_gradient_checkpointing = True
     _no_split_modules = ["BitNetDecoderLayer"]
     _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn_3 = True
     _supports_flash_attn_2 = True
     _supports_sdpa = True
     _supports_flex_attn = True

src/transformers/models/cohere/modeling_cohere.py

@@ -355,6 +355,7 @@ class CoherePreTrainedModel(PreTrainedModel):
     supports_gradient_checkpointing = True
     _no_split_modules = ["CohereDecoderLayer"]
     _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn_3 = True
     _supports_flash_attn_2 = True
     _supports_sdpa = True
     _supports_flex_attn = True

src/transformers/models/cohere2/modeling_cohere2.py

@@ -334,6 +334,7 @@ class Cohere2PreTrainedModel(PreTrainedModel):
     supports_gradient_checkpointing = True
     _no_split_modules = ["Cohere2DecoderLayer"]
     _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn_3 = True
     _supports_flash_attn_2 = True
     _supports_sdpa = True
     _supports_flex_attn = True

src/transformers/models/deepseek_v3/modeling_deepseek_v3.py

@@ -504,6 +504,7 @@ class DeepseekV3PreTrainedModel(PreTrainedModel):
     supports_gradient_checkpointing = True
     _no_split_modules = ["DeepseekV3DecoderLayer"]
     _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn_3 = True
     _supports_flash_attn_2 = True
     _supports_sdpa = True
     _supports_flex_attn = True

src/transformers/models/diffllama/modeling_diffllama.py

@@ -556,6 +556,7 @@ class DiffLlamaPreTrainedModel(PreTrainedModel):
     supports_gradient_checkpointing = True
     _no_split_modules = ["DiffLlamaDecoderLayer"]
     _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn_3 = True
     _supports_flash_attn_2 = True
     _supports_sdpa = True
     _supports_flex_attn = False

src/transformers/models/gemma/modeling_gemma.py

@@ -318,6 +318,7 @@ class GemmaPreTrainedModel(PreTrainedModel):
     supports_gradient_checkpointing = True
     _no_split_modules = ["GemmaDecoderLayer"]
     _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn_3 = True
     _supports_flash_attn_2 = True
     _supports_sdpa = True
     _supports_flex_attn = True

src/transformers/models/gemma2/modeling_gemma2.py

@@ -339,6 +339,7 @@ class Gemma2PreTrainedModel(PreTrainedModel):
     supports_gradient_checkpointing = True
     _no_split_modules = ["Gemma2DecoderLayer"]
     _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn_3 = True
     _supports_flash_attn_2 = True
     _supports_sdpa = True
     _supports_flex_attn = True