Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

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

Aperta
  • Qualcosa non va?
#3,711 2 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 2 giorni

@0z5a ci sta già lavorando.

Dal 20/10/2026.

  • #3714 di @0z5a — aperta
  • #3715 di @0z5a — aperta

Valutazione

Difficoltà
3/5
Tempo stimato
1-2 giorni
Idoneità per principianti
69/100
Tipo di issue
Bug
Chiarezza
Specificata chiaramente
Stato di attività
Attiva
Stack tecnologico
python, pytorch
Ambito
backend, machine-learning

Direzione di ricerca

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.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

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).
Lingua principale
Python
Stelle
3.6k
Fork
872
Merge medio
4g 11h
PR unite (30g)
59

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di NVIDIA/TransformerEngine

Tutte le issue di NVIDIA/TransformerEngine

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.