Hacktoberfest 2026: los issues que los mantenedores marcaron para octubre, abiertos y aptos para principiantes. Explorar issues de Hacktoberfest

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

Abierto
  • ¿Algo no está bien?
#3,711 1 comentario 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 2 días

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
3/5
Tiempo estimado
1-2 días
Aptitud para principiantes
69/100
Tipo de issue
Error
Claridad
Bien especificado
Estado de actividad
Activo
Stack tecnológico
python, pytorch
Área
backend, machine-learning

Línea de trabajo

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.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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).
Lenguaje dominante
Python
Estrellas
3.6k
Forks
862
Merge medio
4 d 1 h
PR fusionados (30 d)
57

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de NVIDIA/TransformerEngine

Todos los issues de NVIDIA/TransformerEngine

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.