Repository: huggingface/transformers
Pull request: [bugfix] fix flash_attention_2 unavailable error on Ascend NPU

# What does this PR do?
After PR #39474 merged, flash_attention_2 is unavailable on Ascend NPU, because package `flash-attn` can not be installed on Ascend NPU :(

This PR is committed for solving this problem.

**Detail modification:**
1. Follow newest flash attention preparation logics in `src/transformers/modeling_flash_attention_utils.py` module, use `get_npu_flash_attn_funcs` return flash attention related functions for Ascend NPU.
2. Ignore logics related with package `flash-attn` in func `_flash_attn_2_can_dispatch` from `src/transformers/modeling_utils.py` module, it is not required for Ascend NPU.

Fixes # (issue)
Not related.

## Before submitting
- [x] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


Changed files:
filename: src/transformers/integrations/npu_flash_attention.py
patch: @@ -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

filename: src/transformers/modeling_flash_attention_utils.py
patch: @@ -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,7 +299,12 @@ 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_3" or (impl is None and is_fa3):
+    elif 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()
+    elif impl == "flash_attention_3" or (impl is None and is_fa3):
         from flash_attn_interface import flash_attn_func, flash_attn_varlen_func
 
         pad_input, unpad_input = _fa3_pad_input, _fa3_unpad_input

filename: src/transformers/modeling_utils.py
patch: @@ -2483,8 +2483,12 @@ def _flash_attn_2_can_dispatch(self, is_init_check: bool = False) -> bool:
             preface = "FlashAttention2 has been toggled on, but it cannot be used due to the following error:"
             install_message = "Please refer to the documentation of https://huggingface.co/docs/transformers/perf_infer_gpu_one#flashattention-2 to install Flash Attention 2."
 
-            # package `flash-attn` can not be installed on Ascend NPU, ignore related validation logi
-            if importlib.util.find_spec("flash_attn") is None and not is_torch_npu_available():
+            # package `flash-attn` can not be installed on Ascend NPU, following validation logics can be ignored.
+            if is_torch_npu_available():
+                logger.info("Detect using FlashAttention2 on Ascend NPU.")
+                return True
+
+            if importlib.util.find_spec("flash_attn") is None:
                 raise ImportError(f"{preface} the package flash_attn seems to be not installed. {install_message}")
             else:
                 # Check FA2 installed version compatibility