[0z5a]

Description

Fixes #3711.

flash-attn's head_dim 256 backward (run_mha_bwd_hdim256) picks its kernel by opt-in shared memory per block. Below 144 KB it launches no kernel at all when dropout is enabled, and returns without an error, so dq/dk/dv are never written. sm86, sm89, sm120 and sm121 all have 99 KB. Padded head dims 224 and 256 both use this kernel. The flash-attn side is Dao-AILab/flash-attention#3053; the check that used to reject this case was removed in Dao-AILab/flash-attention@bbb21d6.

TE's FA2 head_dim > 192 allowlist includes (12, 0). Training with dropout at a padded head dim in (192, 256] on sm120 therefore selects FA2 and gets unwritten gradients. With flash-attn 2.8.3, a backward with dq/dk/dv pre-filled with NaN leaves them 100% NaN:

• on an RTX 5090 (sm_120, 101376 B opt-in shared memory per block), for head_dim 224 and 256 with dropout_p=0.1;
• not without dropout, and not at head_dim ≤ 192;
• not on B200 (232448 B), where every case computes.

This PR disables FA2 in exactly that case:

• padded head dim > 192;
• attention_dropout > 0;
• training mode;
• backward required;
• one of the 99 KB GPUs, (8, 6), (8, 9), (12, 0), (12, 1).

Inference, dropout-free training and forward-only calls keep FA2. sm86/sm89/sm121 are already excluded by the allowlist for head_dim > 192. Listing them keeps the gate correct if that list is relaxed later.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• get_attention_backend: a separate check right after the existing FA2 head-dim check. It turns FA2 off for the case above and logs the reason at debug level.
• test_fa2_large_head_dim_dropout_backend_selection in tests/pytorch/attention/test_attention_backend_selection.py, which mocks the compute capability:
  • sm120 with dropout at padded head dim 256, 224, and MLA 192/256: FA2 is not selected;
  • controls where FA2 stays selected: sm120 without dropout, in eval, forward only, and at head dim 192; sm100 with dropout.

Testing

• New test against unpatched utils.py: the three sm120 dropout cases fail; the 5 controls pass.
• New test with the patch: 8/8 pass. Whole test_attention_backend_selection.py: 36 passed, 2 skipped (the existing SM12x-GPU tests).
• Together with [PyTorch] Allow FlashAttention 2 head_dim > 192 on sm103 #3714 (sm103 allowlist), the whole file gives 43 passed, 2 skipped. The two changes touch separate hunks.
• How the tests ran: these tests need no GPU, so they ran in a process with no visible GPU. That process used a source build of TE at 9d9bd9e5 with this branch's Python sources on top. The attention selector and its tests are unchanged between 9d9bd9e5 and this branch's base a7b6c0cb.
• black 24.4.2 with the repository's pre-commit arguments reports no changes.
• The hardware numbers above are from flash-attn directly; see [FA2] Backward with dropout and head_dim > 192 launches no kernel on GPUs with 99 KB of shared memory per block (sm86/sm89/sm120): dq/dk/dv are left unwritten Dao-AILab/flash-attention#3053 for the reproducer. I have not run TE end to end on sm120.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

🤖 Generated with Claude Code

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no actionable issues were found.

Summary

This PR prevents FA2 from running an unsafe dropout backward on sm120 when the padded head dimension exceeds 192.

• Large-dimension dropout backward calls skip FlashAttention 2 on affected GPUs.

_{Reviews (1) · Last reviewed commit: "[PyTorch] Keep FlashAttention 2 off for ..." · Reviewed by Greptile}