Hacktoberfest 2026：维护者为十月标记出来的 issue，仍然开放、适合新手。 浏览 Hacktoberfest issue

[PyTorch] FA2 is selected for head_dim > 192 with dropout on sm120, where flash-attn's backward launches no kernel

未关闭
  • 哪里不对？
#3,711 1 条评论 0 个 reaction 已指派 0 人 在 GitHub 查看

维护者通常 2 天内回复

还没有人认领这个 Issue。

评估

难度
3/5
预计耗时
1-2 天
新手友好度
69/100
Issue 类型
缺陷
描述清晰度
描述清楚
活跃度
活跃
技术栈
python, pytorch
领域
backend, machine-learning

调研方向

Start at transformer_engine/pytorch/attention/dot_product_attention/utils.py:1045 and read get_attention_backend, then inspect the flash-attn backward behavior described in the issue. Reproduce with head dimensions 224 or 256, nonzero dropout, training mode and an sm120 GPU; compare against dropout-free training or inference. Done means FA2 is not selected for the affected training configuration while supported cases remain available.

由索引模型根据 Issue 内容生成。

描述

Describe the bug

get_attention_backend allows FlashAttention 2 for a padded head dim in (192, 256] on ((8, 0), (9, 0), (10, 0), (12, 0)) (transformer_engine/pytorch/attention/dot_product_attention/utils.py:1045 on main). On (12, 0), training with attention_dropout > 0 at such a head dim gets gradients that flash-attn never wrote:

  • flash-attn's hdim256 backward picks its kernel by opt-in shared memory per block (run_mha_bwd_hdim256 in csrc/flash_attn/src/flash_bwd_launch_template.h).
  • Below 144 KB (sm86/sm89/sm120/sm121 have 99 KB), the dropout instantiation of that branch launches nothing.
  • flash-attn no longer rejects the case; the old head dim > 192 with dropout requires A100/A800 or H100/H800 check was removed in Dao-AILab/flash-attention@bbb21d6.

Measured with flash-attn 2.8.3 (official wheel, torch 2.10.0+cu130), calling _flash_attn_backward with dq/dk/dv pre-filled with NaN:

head_dim dropout_p RTX 5090 (sm_120, 101376 B smem/block opt-in) B200 (sm_100, 232448 B)
≤ 192 0.0 / 0.1 computed computed
224, 256 0.0 computed computed
224, 256 0.1 100% of dq/dk/dv still NaN, no error computed

TE's FA2 path allocates the gradient buffers with torch.empty_like, so in practice the gradients are whatever memory was there.

I have reported the flash-attn side as Dao-AILab/flash-attention#3053.

Steps to reproduce

The reproducer is in the flash-attn issue above. Through TE, any DotProductAttention with head_dim (or max(head_dim_qk, head_dim_v)) in (192, 256], attention_dropout > 0 and training mode on an sm120 GPU reaches it, whenever FA2 is selected. That is the case with fused attention disabled or unsupported for the config, and with the FA2 preference for THD + dropout.

Expected behavior

FA2 should not be selected for padded head dim > 192 with dropout in training on GPUs with less than 144 KB of opt-in shared memory per block. Inference and dropout-free training on those GPUs work, so the gate can stay narrow:

low_smem_fa2_bwd = device_compute_capability in ((8, 6), (8, 9), (12, 0), (12, 1))
if fa2_padded_head_dim > 192 and attention_dropout != 0.0 and is_training and low_smem_fa2_bwd:
    use_flash_attention_2 = False

Related: #3360 / #2836 discuss the same > 192 clause for sm103. The comment I left there argued against removing the clause entirely for this reason. sm103 itself has 227 KB and is not affected.

Environment

  • RTX 5090: driver 580.105.08.
  • B200: driver 580.126.09.
  • Both: torch 2.10.0+cu130, flash-attn 2.8.3.
  • TE main at eb3c860a (code reference only; the measurement is at the flash-attn level).
主要语言
Python
星标
3.6k
派生
862
平均合并
4 天 1 小时
30 天内合并 PR
57

环境准备

从这里开始

  1. 先读完整个 Issue，再读项目的贡献指南。
  2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
  3. Fork 仓库，在一个分支上完成修改。
  4. 提交 Pull Request，并在描述里引用这个 Issue 编号。

NVIDIA/TransformerEngine 的其他 Issue

查看 NVIDIA/TransformerEngine 的全部 Issue

相似的 Issue

更多 Python Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。