From 7bd16d8891a37c0875fbdb962b6e96f24b4329f9 Mon Sep 17 00:00:00 2001 From: hyw <109567717+yuweih205@users.noreply.github.com> Date: Mon, 7 Sep 2026 12:05:41 +0800 Subject: [PATCH] [PyT] Disable FA3 for training when head_dim_qk != head_dim_v Signed-off-by: hyw <109567717+yuweih205@users.noreply.github.com> --- .../pytorch/attention/dot_product_attention/utils.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 33d612b4f4..b0d4def95a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -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 + if ( use_flash_attention_4 and FlashAttentionUtils.v4_is_installed