[kalectory]

Description

Restore the FlashAttention 2 head-dimension behavior from #2836. #2629 preserved the new padded Q/V head-dimension handling but accidentally reintroduced the exact compute-capability allowlist that #2836 removed.

On B300 (sm103), that allowlist rejects a padded head dimension of 256 and falls back to quadratic unfused attention.

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

• Preserve the padded Q/V head-dimension validation introduced by [PyTorch] Pad V when Q/V head dims differ (MLA) for THD #2629.
• Remove the stale architecture allowlist for FlashAttention 2 head dimensions above 192.
• Update the backend-selection diagnostic to match.

Testing

Status quo

• Layout: one 8xB300 node with four B300s assigned to the Megatron trainer (TP=4) and four to rollout inference. Qwen3.6-27B used BF16 THD attention, head dim 256, Transformer Engine 2.11.0, and FlashAttention 2.8.3.
• Our production workload requested FlashAttention, but TE's exact compute-capability allowlist omitted B300 (sm103) and silently selected UnfusedDotProductAttention instead.
• XID 1043331 replayed an unsliced 65,332-token datum with a 65,536-token microbatch budget. The quadratic unfused softmax tried to allocate 202.75 GiB and raised torch.OutOfMemoryError during forward.

After the fix

• Used the same image, package versions, TP4 trainer layout, model, BF16 THD configuration, head dim, 65,536-token budget, and exact 65,332-token datum. Only this selector gate was patched.
• The full 64k training update completed without OOM: forward_backward took 413.2s, the Adam optimizer step took 9.6s, and grad norm was finite at 1776.48. This validates FlashAttention 2 and 64k context on the four-B300 trainer configuration.
• Python compilation and diff checks pass.

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

[greptile-apps]

Greptile Summary

This PR restores FlashAttention 2 eligibility for padded Q/V head dimensions up to 256 on sm103 by removing a stale architecture allowlist.

• Retains validation that the padded head dimension is divisible by 8 and no greater than 256.
• Updates the backend-selection diagnostic to reflect the revised eligibility condition.

Confidence Score: 5/5

The PR appears safe to merge with no concrete blocking or independently actionable non-blocking issues identified.

The change narrowly removes an obsolete architecture filter while retaining FlashAttention 2’s padded head-dimension bounds, and the reviewed selection and dispatch paths reveal no established regression.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Removes the stale architecture-specific FA2 head-dimension restriction while preserving the general padded-dimension limits and aligning the diagnostic message.

_{Reviews (1): Last reviewed commit: "Restore FlashAttention 2 head dim suppor..." | Re-trigger Greptile}

[cyanguwa]

/te-ci L0

[0z5a]

@kalectory @cyanguwa This regression is still on main (eb3c860a). utils.py:1045 still allows FA2 for a padded head_dim > 192 only on ((8, 0), (9, 0), (10, 0), (12, 0)), so sm103 (B300 / GB300) still falls back to fused or unfused attention for head dims in (192, 256]. This PR was approved but closed without merging. Was that intentional, or could it be reopened?

One thing to consider before dropping the clause entirely. 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, in the branch commented "sm86 and sm89, max smem is 99 KB. Only works without dropout", it launches no kernel at all when dropout is enabled. The mha_bwd check that used to reject "head dim > 192 with dropout" on those GPUs was removed in Dao-AILab/flash-attention@bbb21d6, and neither v2.8.3 nor main has a replacement. Removing TE's allowlist would therefore send sm86/sm89 training with dropout down that path. This comes from reading the source; I have not run it on an sm86/sm89 GPU.

A narrower fix is to add (10, 3) to the tuple. sm103 has the same 227 KB of opt-in shared memory per block as sm90 and sm100, so it takes the same backward kernel as sm100. I have that change plus a regression test in tests/pytorch/attention/test_attention_backend_selection.py. The test mocks the compute capability, so it needs no GPU:

• sm103 with padded head dims 224 and 256, and MLA 192/256, select FA2;
• the 192, 264 and 252 boundaries;
• an sm100 control.

The three sm103 cases above 192 fail on current main and pass with the change. I'm happy to open it as a PR, or to leave it to you if you'd rather reopen this one.

[cyanguwa]

@0z5a, thanks for the comment. Looks like the author closed the PR himself, cc @kalectory.

Overall, I still think this change is necessary, at least for sm103. @0z5a, could you please submit a PR for the narrower fix you mentioned? I think it's safer to just enable sm103 for now and leave sm86/89 out. Thanks!

[kalectory]

Thanks @cyanguwa. I agree with the solution @0z5a suggested, will keep this PR closed and defer the changes to their fix for sm103.

[0z5a]

@cyanguwa Thanks! The narrower fix is up as #3714. It adds only (10, 3) to the allowlist and leaves sm86/sm89 unchanged. It also adds a backend-selection regression test: the three sm103 cases above 192 fail on current main and pass with the change. The PR is a draft for now; I'll mark it ready after a final review on my side.

The (12, 0) dropout case mentioned above is tracked separately in #3711, with the fix in #3715.