Skip to content

[PyT] Disable FA3 for training when head_dim_qk != head_dim_v #3490

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
yuweih205 wants to merge 1 commit into
base: main
Choose a base branch
from
+11 −0
Open

[PyT] Disable FA3 for training when head_dim_qk != head_dim_v #3490

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: 11 additions & 0 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 @@ -983,6 +983,17 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
)
use_flash_attention_3 = False

if use_flash_attention_3 and is_training and head_dim_qk != head_dim_v:
logger.debug(
"Disabling FlashAttention 3 for training with head_dim_qk != head_dim_v, "
"as its backward pass does not support it "
"(Dao-AILab/flash-attention#1487). "
"Found: head_dim_qk = %s, head_dim_v = %s.",
head_dim_qk,
head_dim_v,
)
use_flash_attention_3 = False
Comment on lines +986 to +995

@greptile-apps greptile-apps Bot Sep 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

P2 Mode Behavior Lacks Coverage

This new training-only FA3 restriction has no automated regression test. Existing attention tests do not assert that unequal QK/V head dimensions reject FA3 during training or confirm that inference retains FA3. A future selector change could therefore restore the backward crash or unnecessarily disable the supported forward-only path. Please add coverage for both modes, including representative configurations from each unequal-dimension support branch.

Knowledge Base Used: PyTorch attention execution

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!


if (
use_flash_attention_4
and FlashAttentionUtils.v4_is_installed
Expand Down