[bugfix] fix flash_attention_2 unavailable error on Ascend NPU

Author: FightingZhen

src/transformers/integrations/npu_flash_attention.py

@@ -267,3 +267,8 @@ def npu_apply_rotary_emb(x, cos, sin, **kwargs):
         sin = sin.unsqueeze(0).unsqueeze(2)
 
     return npu_rotary_mul(x, cos, sin)
+
+
+def get_npu_flash_attn_funcs():
+    # return flash attention related functions used for Ascend NPU in order
+    return npu_flash_attn_func, npu_flash_attn_varlen_func, pad_input, unpad_input, False

src/transformers/modeling_flash_attention_utils.py

@@ -261,7 +261,7 @@ def fa_peft_integration_check(q, k, v, target_dtype: Optional[torch.dtype] = Non
 
 def _lazy_imports(impl: Optional[str]):
     # returns funcs and pad/unpad based on impl
-    is_fa2 = is_flash_attn_2_available() or is_torch_npu_available()
+    is_fa2 = is_flash_attn_2_available()
     is_fa3 = is_flash_attn_3_available()
     if impl == "flash_attention_2" or (impl is None and is_fa2 and not is_fa3):
         try:
@@ -299,6 +299,10 @@ def _lazy_imports(impl: Optional[str]):
                 raise ImportError(
                     "Failed to import flash attention 2, please install it or use another implementation."
                 ) from e
+    if impl == "flash_attention_2" and is_torch_npu_available():
+        # get flash attention related functions from `.integrations.npu_flash_attention` module for Ascend NPU
+        from .integrations.npu_flash_attention import get_npu_flash_attn_funcs
+        return get_npu_flash_attn_funcs()
     if impl == "flash_attention_3" or (impl is None and is_fa3):
         from flash_attn_interface import flash_attn_func, flash_attn_varlen_func