Hacktoberfest 2026：メンテナが10月に向けて印を付けた、オープンで初心者向けの 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 コメント 2 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 2 日以内に返信

@0z5a がすでに取り組んでいます。

2026年10月20日 から。

  • #3714 @0z5a による — オープン
  • #3715 @0z5a による — オープン

評価

難易度
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日 9時間
マージ済み PR（30日）
61

環境構築

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

NVIDIA/TransformerEngine のほかの issue

NVIDIA/TransformerEngine の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。