Skip to content

[PyTorch] Fix FlashAttention 2 head_dim > 192 on sm103 and other architectures #2836

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

Open
pedramr wants to merge 1 commit into
base: main
Choose a base branch
from
+2 −5
Open

[PyTorch] Fix FlashAttention 2 head_dim > 192 on sm103 and other architectures #2836

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
7 changes: 2 additions & 5 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 @@ -631,16 +631,13 @@ def get_attention_backend(
if use_flash_attention_2 and (
head_dim_qk > 256
or head_dim_qk % 8 != 0
or (
head_dim_qk > 192
and device_compute_capability not in ((8, 0), (9, 0), (10, 0), (12, 0))
)
or (head_dim_qk > 192 and device_compute_capability < (8, 0))
Copy link
Copy Markdown
Contributor

@greptile-apps greptile-apps bot Apr 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Dead code: condition is always False at this point

Lines 428–431 unconditionally set use_flash_attention_2 = False whenever device_compute_capability < (8, 0). By the time execution reaches line 634, use_flash_attention_2 can only be True if device_compute_capability >= (8, 0), so the sub-expression device_compute_capability < (8, 0) is never true and the entire third or branch is unreachable. The bug-fix intent is correct (no longer blocking head_dim > 192 on sm86/sm89/sm103), but the residual condition could be confusing to future readers who might believe it provides a meaningful guard.

Consider removing the dead branch entirely:

Suggested change
or (head_dim_qk > 192 and device_compute_capability < (8, 0))
or head_dim_qk % 8 != 0

):
if FlashAttentionUtils.is_installed:
logger.debug(
"Disabling FlashAttention 2 due to unsupported head_dim_qk and head_dim_v. "
"Supported: head_dim_qk = head_dim_v, head_dim_qk %%8 = 0, "
"head_dim_qk <= 256 (>192 requires sm80/90/100+). "
"head_dim_qk <= 256 (>192 requires sm80+). "
"Found: head_dim_qk = %s, head_dim_v = %s, on sm%s.",
head_dim_qk,
head_dim_v,
Expand Down