[Bug] Backend selection picks FA3 for training with head_dim_qk=192 / v_head_dim=128, but FA3 backward cannot run it

Describe the bug

With an MLA-style geometry (head_dim_qk=192, v_head_dim=128, GQA), when FusedAttention is unavailable (e.g. NVTE_FUSED_ATTN=0, which Megatron-LM sets when launched with --attention-backend flash), TE selects FlashAttention 3 in training mode. The forward pass succeeds; the backward crashes deep in the kernel:

[DEBUG | DotProductAttention]: Disabling FusedAttention due to NVTE_FUSED_ATTN=0
[DEBUG | DotProductAttention]: Disabling FlashAttention 2 as it does not support MLA.
[DEBUG | DotProductAttention]: Selected backend = FlashAttention (3.0.0b1)
...forward OK...
File "flash_attn_3/flash_attn_interface.py", line 123, in _flash_attn_backward
RuntimeError: out must have shape (batch_size, seqlen_q, num_heads, head_size)

Both plain causal and sliding-window (window_size=(127, 0)) hit the same crash.

Root cause (as far as I can tell)

_is_fa3_supported() in dot_product_attention/utils.py allows head_dim_qk != head_dim_v when 128 < qk <= 192 and 96 < v <= 128 — this matches FA3's forward support, but the function never consults is_training. FA3's backward for hdimQK=192/hdimV=128 is not implemented (open feature request: Dao-AILab/flash-attention#1487). So the unsupported-backward config passes selection and only fails mid-backward with an opaque shape error.

To Reproduce

# NVTE_FLASH_ATTN=1 NVTE_FUSED_ATTN=0 NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=2 python repro.py
import torch
from transformer_engine.pytorch import DotProductAttention

HQ, KV, DQK, DV, S = 16, 1, 192, 128, 4096
q = torch.randn(S, 1, HQ, DQK, dtype=torch.bfloat16, device="cuda", requires_grad=True)
k = torch.randn(S, 1, KV, DQK, dtype=torch.bfloat16, device="cuda", requires_grad=True)
v = torch.randn(S, 1, KV, DV,  dtype=torch.bfloat16, device="cuda", requires_grad=True)
dpa = DotProductAttention(num_attention_heads=HQ, kv_channels=(DQK, DV), num_gqa_groups=KV,
                          attention_dropout=0.0, qkv_format="sbhd", attn_mask_type="causal").cuda().train()
out = dpa(q, k, v)       # forward OK
out.sum().backward()     # RuntimeError: out must have shape (batch_size, seqlen_q, num_heads, head_size)

Expected behavior

During training, FA3 should be filtered out for geometries whose backward it cannot run — same as other capability filters — so selection either falls back to a viable backend or fails fast with a clear "no viable backend" error, instead of crashing inside flash_attn_3_cuda.bwd after a successful forward.

Environment

TE 2.10.0+769ed778 · torch 2.9.1+cu130 · flash-attn 2.7.4.post1 · flash_attn_3 3.0.0b1 · H200 (sm90) · CUDA 13.0

Start in dot_product_attention/utils.py at _is_fa3_supported() and compare its capability checks with the training case described in the issue. Run repro.py with the listed environment settings to confirm backend selection and the backward failure. Done means training avoids the unsupported FA3 geometry and either selects a viable backend or reports that none is available.