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

[yuweih205]

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

[modelpath-dev]

I will take this issue. Please assign it to me.

The problem seems to be in the _is_fa3_supported() function in dot_product_attention/utils.py. It allows the selection of FlashAttention 3 for training even when the backward pass is unsupported for the given dimensions. I would first modify this function to check if the mode is training and ensure it only selects FA3 when both forward and backward passes are supported. I will test the fix by reproducing the issue and verifying that the correct backend is selected or a clear error is shown.

[yuweih205]

Thanks for offering. I filed this after reproducing and tracing the failure, and I plan to work on the fix myself. I'll send a PR shortly, so please leave it unassigned for now.

[modelpath-dev]

Got it, thanks for letting me know. I'll step back and let you handle it. If you need any help or a review later, feel free to reach out.

[yuweih205]

Follow-up with more data (same setup: H200, TE 2.10, flash_attn_3 3.0.0b1, training, NVTE_FUSED_ATTN=0): tested 8 mismatched-dim shapes across both branches of _is_fa3_supported()'s mismatch condition. Every shape the filter lets through to FA3 crashes in backward — 64/128 and 64/512 (the qk<=64 branch) the same way as 160/128 and 192/128 (the (128,192]×(96,128] branch). Shapes like 128/64 or 256/128 only "work" because the existing filter already rejects them and they fall back to another backend. So FA3's backward supports no mismatched head dims at all here, and the fix disables FA3 for training whenever head_dim_qk != head_dim_v.