Hacktoberfest 2026: những issue maintainer đã đánh dấu cho tháng Mười, đang mở và phù hợp người mới. Xem issue Hacktoberfest

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

Đang mở
  • Có gì không ổn?
#3,711 1 bình luận 0 reaction 0 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 2 ngày

Chưa có ai nhận issue này.

Đánh giá

Độ khó
3/5
Thời gian dự kiến
1-2 ngày
Mức phù hợp với người mới
69/100
Loại issue
Lỗi
Độ rõ ràng
Đặc tả rõ ràng
Mức độ hoạt động
Sôi nổi
Công nghệ
python, pytorch
Lĩnh vực
backend, machine-learning

Hướng nghiên cứu

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.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

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).
Ngôn ngữ chính
Python
Star
3.6k
Fork
862
Merge trung bình
4 ngày 1 giờ
Pull request đã merge (30 ngày)
57

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của NVIDIA/TransformerEngine

Tất cả issue của NVIDIA/TransformerEngine

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.