[0z5a]

Description

get_attention_backend disables FlashAttention 2 for padded head dims in (192, 256] unless the device is in an exact allowlist, ((8, 0), (9, 0), (10, 0), (12, 0)). sm103 (B300 / GB300) is not in it. So on sm103, every config with max(head_dim_qk, head_dim_v) in (192, 256] falls back to FusedAttention or UnfusedDotProductAttention, even when flash-attn is installed and requested. Examples are head_dim 256, or MLA with head_dim_v 256.

This is a regression:

• [PyTorch] Fix FlashAttention 2 head_dim > 192 on sm103 and other architectures #2836 (merged 2026-06-04) removed the allowlist.
• [PyTorch] Pad V when Q/V head dims differ (MLA) for THD #2629 (merged 2026-06-05) restored it in the hunk that added padded Q/V head dims.

#3360 reported the same problem on B300 and was closed by its author. As discussed there, this PR is the narrower fix: it adds only (10, 3) and leaves sm86/sm89 out.

The fallback also defeats an existing preference. For THD with dropout on (10, 0) and (10, 3), is_slow_fused_thd_dropout disables FusedAttention so that FA2 is used. On sm103 with head_dim > 192, FA2 was already disabled by the allowlist, so the preference never applies.

Why a list entry rather than dropping the clause: flash-attn's hdim256 backward picks its kernel by opt-in shared memory per block. Below 144 KB (sm86/sm89/sm120 have 99 KB) it launches no kernel at all when dropout is enabled (Dao-AILab/flash-attention#3053). sm103 has 227 KB, the same as sm90 and sm100, so it takes the same backward kernel as sm100. The existing (12, 0) entry has that problem with dropout. It is tracked in #3711 and fixed separately.

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

• Add (10, 3) to the FA2 head_dim > 192 allowlist. Update the debug message from "sm80/90/100+" (which wrongly suggested every sm100+ device) to "sm80/90/100/103/120".
• Add test_fa2_large_head_dim_backend_selection to tests/pytorch/attention/test_attention_backend_selection.py. It mocks the compute capability and the flash-attn version, like the FA4 cases in that file. It covers:
  • sm103 with padded head dim 224, 256, and MLA 192/256: FA2 is selected;
  • the 192 boundary, 264 (> 256) and 252 (not a multiple of 8): FA2 is not selected;
  • an sm100 control.

Testing

• New test against unpatched utils.py: sm103-d256, sm103-d224 and sm103-mla-padded-d256 fail; the other 4 pass.
• New test with the patch: 7/7 pass. Whole test_attention_backend_selection.py: 35 passed, 2 skipped (the existing SM12x-GPU tests).
• Together with the [PyTorch] FA2 is selected for head_dim > 192 with dropout on sm120, where flash-attn's backward launches no kernel #3711 fix, 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.
• I have not run FA2 at head_dim 256 on sm103 hardware. This PR changes backend selection only.

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

Allows FlashAttention 2 on sm103 when the padded head size is greater than 192 and at most 256.

• sm103 attention calls can use FlashAttention 2 with large head dimensions.

No actionable issues were found. This review checked code paths, not GPU execution; tests were not run.

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