[PyT] Disable FA3 for training when head_dim_qk != head_dim_v

[yuweih205]

Description

Fixes #3481

_is_fa3_supported() matches FA3's forward support matrix, so mismatched-head-dim configs (e.g. head_dim_qk=192 / head_dim_v=128) pass backend selection in training mode — but FA3 has no backward for mismatched head dims (Dao-AILab/flash-attention#1487), so the forward succeeds and the backward crashes with an opaque RuntimeError: out must have shape ... from flash_attn_3_cuda.bwd.

Add a selection-time filter disabling FA3 for training when head_dim_qk != head_dim_v, following the existing filter style.

Verified on H200 (TE 2.10, torch 2.9.1+cu130, flash_attn_3 3.0.0b1), with NVTE_FUSED_ATTN=0 to force the FA3 path:

• training 64/128, 64/512, 192/128 (both mismatch branches of _is_fa3_supported): previously crashed in backward; with the fix, selection falls back to a viable backend and forward+backward complete
• training, default env: FusedAttention still selected, no behavior change
• inference 192/128: FA3 still selectable, forward-only use unaffected

Type of change

• Bug fix (non-breaking change which fixes an issue)

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

[greptile-apps]

Greptile Summary

This PR prevents FlashAttention 3 from being selected for training when query/key and value head dimensions differ, while preserving its supported inference path.

• Adds a training-aware FA3 capability filter before final backend selection.
• Emits a diagnostic message containing the incompatible dimensions.
• Allows selection to continue to another eligible attention backend.

Confidence Score: 4/5

The implementation appears safe to merge, with non-blocking automated test coverage recommended for the new training-versus-inference selector behavior.

The guard uses the selector’s existing inputs, propagates correctly through backend aggregation, and preserves supported fallback behavior; the remaining concern is that its mode-sensitive contract is not protected by a regression test.

Files Needing Attention: transformer_engine/pytorch/attention/dot_product_attention/utils.py

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Adds the correct training-only FA3 guard for unequal head dimensions, but does not add automated coverage for the new mode-sensitive selection behavior.

_{Reviews (1): Last reviewed commit: "[PyT] Disable FA3 for training when head..." | Re-trigger Greptile}

[greptile-apps]

Mode Behavior Lacks Coverage

This new training-only FA3 restriction has no automated regression test. Existing attention tests do not assert that unequal QK/V head dimensions reject FA3 during training or confirm that inference retains FA3. A future selector change could therefore restore the backward crash or unnecessarily disable the supported forward-only path. Please add coverage for both modes, including representative configurations from each unequal-dimension support branch.

Knowledge Base Used: PyTorch attention execution

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!