[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
• Added automated selector coverage for (192, 128) and (64, 512) in both training and inference (4 cases total), asserting training fallback while retaining FA3 for inference

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]

Confidence Score: 5/5

The PR appears safe to merge because the unsupported FA3 backward configuration is filtered during selection while supported inference behavior remains available.

Findings

1.  Mode Behavior Lacks Coverage ▶

Summary

• Preserves FA3 selection for supported forward-only inference.
• Adds parameterized selector tests for both unequal-head-dimension support branches and both runtime modes.
• Resolves the previous coverage concern by asserting training fallback and inference retention.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Attention request] --> B{FA3 otherwise supported?}
  B -- No --> C[Evaluate other backends]
  B -- Yes --> D{Training with QK dim != V dim?}
  D -- Yes --> C
  D -- No --> E[FA3 remains eligible]
  C --> F[Select viable fused or unfused fallback]

Loading

_{Reviews (2) · Last reviewed commit: "[PyT] Test FA3 mismatched-dimension mode..."}