flash-attn-4 4.0.0b11: DotProductAttention batched backward returns wrong dQ/dK/dV at head_dim=256 on Blackwell (sm_100)

[XiaohanZhangCMU]

Summary

On a B200 (sm_100), transformer_engine.pytorch.DotProductAttention with the FlashAttention-4 backend (flash-attn-4==4.0.0b11) computes a correct forward but an incorrect backward for head_dim=256 when the batch contains more than one sequence. The returned dQ is ~5x larger than the true gradient and points in nearly the opposite direction (cosine ≈ -0.11). The UnfusedDotProductAttention backend and a plain fp32 PyTorch SDPA reference agree with each other to <1% on the same inputs. This silently corrupts training gradients (grad_norm explodes ~1000x, loss diverges) while the forward/loss looks normal.

Environment

• GPU: NVIDIA B200, compute capability sm_100 (10, 0)
• flash-attn-4==4.0.0b11
• TransformerEngine 2.16.0
• torch 2.11.0+cu128, CUDA 12.8
• nvidia-cudnn-cu12 9.19.0.56 (as pinned by flash-attn-4)

Conditions

• DotProductAttention, qkv_format="sbhd", attn_mask_type="causal", num_attention_heads=16, num_gqa_groups=2, head_dim=256, dtype bfloat16.
• Input-dependent: random Gaussian Q/K/V do not trigger it (even with a large incoming gradient). It reproduces with real transformer activations (attached).
• Requires batch >= 2 sequences: a batch=1 slice of the same data is correct under FA4; batch>=2 is wrong.
• Forward is correct (matches fp32 to ~1e-3). Only the backward (dQ/dK/dV) is wrong.
• cuDNN-fused attention is unavailable for head_dim=256 ("No dot product attention backend is available"), so on Blackwell the usable backends for this shape are FA4 (buggy) and unfused (correct, slow).

Reproduce

Attached: fa4_repro_data.pt (real captured q,k,v,d_out + meta; sbhd shape [384, 2, 16, 256], bf16 — the minimal 2-sequence batch) and the script below.

# FA4 (buggy):
NVTE_FLASH_ATTN=1 NVTE_FUSED_ATTN=0 NVTE_UNFUSED_ATTN=0 python fa4_bug_repro.py fa4_repro_data.pt
# Unfused (correct, sanity):
NVTE_FLASH_ATTN=0 NVTE_FUSED_ATTN=0 NVTE_UNFUSED_ATTN=1 python fa4_bug_repro.py fa4_repro_data.pt

fa4_bug_repro.py

import math, sys, torch
import transformer_engine, transformer_engine.pytorch as te

def fp32_reference(q, k, v, causal):  # sbhd q[s,b,hq,d], k/v[s,b,hkv,d] -> [s,b,hq,d]
    s, b, hq, d = q.shape; hkv = k.shape[2]
    qt = q.permute(1,2,0,3).float(); kt = k.permute(1,2,0,3).float(); vt = v.permute(1,2,0,3).float()
    rep = hq // hkv
    if rep > 1: kt = kt.repeat_interleave(rep, 1); vt = vt.repeat_interleave(rep, 1)
    sc = torch.matmul(qt, kt.transpose(-1,-2)) / math.sqrt(d)
    if causal: sc = sc.masked_fill(torch.triu(torch.ones(s,s,device=q.device,dtype=torch.bool),1), float("-inf"))
    return torch.matmul(torch.softmax(sc,-1), vt).permute(2,0,1,3).contiguous()

cap = torch.load(sys.argv[1], map_location="cpu"); m = cap["meta"]; dev = torch.device("cuda")
hq, hkv, hd = m["num_attention_heads"], m["num_gqa_groups"], m["kv_channels"] or cap["q"].shape[-1]
causal = "causal" in str(m["attn_mask_type"]).lower()
def leaf(n, dt): x = cap[n].to(dev).to(dt).detach().clone(); x.requires_grad_(True); return x
rq, rk, rv = leaf("q",torch.float32), leaf("k",torch.float32), leaf("v",torch.float32)
rout = fp32_reference(rq, rk, rv, causal).reshape(cap["d_out"].shape); rout.backward(cap["d_out"].to(dev).float())
q, k, v = leaf("q",torch.bfloat16), leaf("k",torch.bfloat16), leaf("v",torch.bfloat16)
dpa = te.DotProductAttention(num_attention_heads=hq, kv_channels=hd, num_gqa_groups=hkv,
    attention_dropout=0.0, qkv_format="sbhd", attn_mask_type=m["attn_mask_type"] or "causal").to(dev)
out = dpa(q, k, v); out.backward(cap["d_out"].to(dev).to(torch.bfloat16))
print(f"GPU={torch.cuda.get_device_name(0)} sm={torch.cuda.get_device_capability(0)} TE={transformer_engine.__version__}")
for nm, a, b in [("out",rout,out),("dQ",rq.grad,q.grad),("dK",rk.grad,k.grad),("dV",rv.grad,v.grad)]:
    a, b = a.float().flatten(), b.float().flatten()
    rel = (a-b).norm().item()/(a.norm().item() or 1.0); cos = torch.nn.functional.cosine_similarity(a,b,0).item()
    print(f"  {nm}: rel_l2={rel:.3e} cos={cos:+.4f} ref_norm={a.norm():.1f} te_norm={b.norm():.1f}"
          + ("  <-- WRONG" if (rel>5e-2 or cos<0.99) else ""))

Observed (vs fp32 reference, identical input)

fa4_repro_data.pt ([384, 2, 16, 256]), FA4 backend:

tensor	rel_l2 vs fp32 ref	cos	ref norm	FA4 norm
out (fwd)	1.7e-3	+1.000	577.9	577.9
dQ	4.56	-0.11	5,441	23,588
dK	3.48	+0.12	2,531	8,747
dV	1.75	+0.07	8,275	12,499

Same input through UnfusedDotProductAttention matches the fp32 reference (dQ rel_l2 6.7e-3, cos +1.000). A batch=1 slice of the same tensors is correct under FA4; an 8-sequence batch behaves the same as batch=2 (FA4 dQ 85,710 vs ref 17,487). Since batch=1 is correct and batch>=2 is not, this looks like a batch-indexing/accumulation bug in the FA4 backward at head_dim=256, not a precision issue.