[PyTorch] Restore FlashAttention 2 head dim support on sm103

[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