Skip to content

[PyTorch] Fix FA4 selection when FA3 is unavailable. #2909

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
ptrendx merged 1 commit into from
Apr 23, 2026
+8 −3
Merged

[PyTorch] Fix FA4 selection when FA3 is unavailable. #2909

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
11 changes: 8 additions & 3 deletions transformer_engine/pytorch/attention/dot_product_attention/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -464,9 +464,14 @@ def get_attention_backend(
# On SM90, prefer FA3 over FA4 when FA3 is available.
# FA3 is more mature on Hopper; FA4's SM90 backward has limitations
# (MLA, non-standard head dims, SplitKV).
if use_flash_attention_4 and use_flash_attention_3 and device_compute_capability == (9, 0):
if FlashAttentionUtils.v4_is_installed:
logger.debug("Disabling FlashAttention 4 to prefer FlashAttention 3 on SM90")
if (
device_compute_capability == (9, 0)
and use_flash_attention_3
and FlashAttentionUtils.v3_is_installed
and use_flash_attention_4
and FlashAttentionUtils.v4_is_installed
):
logger.debug("Disabling FlashAttention 4 to prefer FlashAttention 3 on SM90")
use_flash_attention_4 = False

# Filter: Data type
Expand Down
Loading