Skip to content

[bugfix] fix flash attention 2 unavailable error on Ascend NPU #39166

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ArthurZucker merged 2 commits into from
Jul 7, 2025
Merged

[bugfix] fix flash attention 2 unavailable error on Ascend NPU #39166

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
df6b916
[bugfix] fix flash attention 2 error on Ascend NPU
FightingZhen Jul 2, 2025
19a6b5c
Merge branch 'main' into bugfix_fa2_npu
FightingZhen Jul 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 10 additions & 16 deletions src/transformers/modeling_flash_attention_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,16 @@ def _fa3_pad_input(hidden_states, indices, batch, seqlen):
from flash_attn.bert_padding import unpad_input as unpad_input_fa2
from flash_attn.layers.rotary import apply_rotary_emb

HAS_FA2 = True
FA_VERSION = 2
elif is_torch_npu_available():
# patch functions in package `flash-attn` when using flash-attention on Ascend NPU.
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_2_func
from .integrations.npu_flash_attention import npu_flash_attn_varlen_func as flash_attn_2_varlen_func
from .integrations.npu_flash_attention import pad_input as pad_input_fa2
from .integrations.npu_flash_attention import unpad_input as unpad_input_fa2

HAS_FA2 = True
FA_VERSION = 2
else:
Expand Down Expand Up @@ -136,22 +146,6 @@ def _fa3_pad_input(hidden_states, indices, batch, seqlen):
unpad_input = globals()[f"unpad_input_fa{FA_VERSION}"]
pad_input = globals()[f"pad_input_fa{FA_VERSION}"]

# patch functions in package `flash-attn` when using flash-attention on Ascend NPU.
if is_torch_npu_available():
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

Expand Down