[Cute,Fwd,Sm120] Fix three forward-pass correctness bugs (use_tma_O, smem-store atom, pack_gqa)

[jganbar]

Summary

Three independent bugs prevent FA4's forward kernel from producing correct output on consumer Blackwell (SM_120) GPUs as of flash-attn-4==4.0.0b12. All three reproduce on main at the current HEAD.

• Compile crash in the epilogue. use_tma_O was True on SM_120 but the SM80 base class never wires a TMA atom for O — it passes None to self.epilogue, and the epilogue crashes inside quack.copy_utils.tma_get_copy_fn.
• Silent numerical wrongness. The rmem→smem store atom in the epilogue was the SM90 stmatrix (chosen by get_smem_store_atom for any arch >= 90), but the SM80 base class's mma.sync.aligned.m16n8k16 output register layout is not compatible with it. The kernel completed, returned the correct shape, and the output norm matched SDPA — but element-wise values were wrong by 0.5–3.9 (bf16/fp16, vs SDPA-bf16's own ~0.003).
• Compile crash on every default-policy GQA/MQA shape. pack_gqa.compute_ptr trips cute.crd2idx because cuTeDSL 4.4.2 collapses the composite (qhead_per_kvhead, seqlen_q) mode in mO[None, 0] to a rank-1 layout. Pragmatic mitigation here: default pack_gqa=False on SM_120 until the deeper cuTeDSL issue is fixed; explicit pack_gqa=True is still honoured.

After this PR, the SM_120 forward kernel matches PyTorch SDPA-bf16's own numerical accuracy across a 240-shape grid (240/240 pass, max abs diffs bit-identical to SDPA's on 219/240 shapes and within 3.5e-5 on the rest) and is ~1.15× faster on average (range 0.74×–1.95×).

Bugs and fixes at a glance

Each commit has a focused message with the full root-cause writeup; this is the summary.

#	File	Symptom	Fix
1	flash_fwd.py:652	AttributeError: 'NoneType' object has no attribute '_trait' at trace time	self.use_tma_O = False unconditionally in the SM80 base class
2	flash_fwd.py:347	Kernel runs to completion, returns plausible output, but max abs diff vs fp32 reference is 0.5–3.9 (vs ~0.003 for SDPA-bf16)	Force get_smem_store_atom(80, ...) so the SM80 base class always uses the universal copy, not the SM90 stmatrix
3	interface.py:430	ValueError: Operation creation failed (unable to compute crd2idx) on every default-policy GQA/MQA shape	Default pack_gqa=False on SM_120

Bug 2 (the stmatrix mismatch) is the headline fix — silent numerical wrongness is the worst kind of bug to ship and this one would propagate into training loss / eval quality for any model using FA4 on consumer Blackwell.

How the headline bug was found

The error profile was very specific and ruled out most hypotheses up-front:

• Bitwise-deterministic across reruns → not a race / scheduler issue.
• Uniform across all M-tiles (not concentrated in m_block=0) → not a kernel prologue / initialisation issue.
• Scales linearly with input magnitude → precision / permutation signature, not a logic bug.
• Output norm preserved to ~3 significant digits → values are reshuffled with their same-magnitude neighbours, not zeroed.

That pattern points at a memory transfer using the wrong layout. Grepping self.arch for arch-conditional branches in the SM80 base class surfaces exactly one candidate: get_smem_store_atom in the epilogue. The branch picks the SM90 stmatrix path for any arch ≥ 90, but stmatrix is hardware-paired with WGMMA's output register layout, not SM80 MMA's. Patch is one literal arg change.

Test plan

Hardware: NVIDIA GeForce RTX 5090, capability (12, 0), driver 590.48.01, CUDA 12.8.

Software: Python 3.11.15, torch==2.10.0+cu128, nvidia-cutlass-dsl==4.4.2, this branch installed editable.

Correctness (240 shapes × 3 backends = 720 evaluations, all pass)

dtypes    : fp16, bf16
causal    : True, False
B         : 1, 2
S         : 128, 512, 1024, 2048, 4096
heads     : MHA(Hq=Hkv=16), GQA(Hq=16,Hkv=8), MQA(Hq=16,Hkv=1)
D         : 64, 128
backends  : sdpa_bshd (reference), fa4 (default pack_gqa), fa4_no_packgqa
reference : PyTorch SDPA in fp32 on the same inputs

status by backend:
  sdpa_bshd       240 / 240 pass
  fa4             240 / 240 pass
  fa4_no_packgqa  240 / 240 pass

FA4's per-shape max_abs_diff vs the fp32 SDPA reference is bit-identical to SDPA-bf16's own on 219/240 shapes; the remaining 21 differ by ≤ 3.5e-5 (FP reordering noise).

Before this PR the analogous run was 240 SDPA pass, 80 FA4 "pass" (MHA only — actually producing element-wise diffs of 0.5–3.9 against fp32 reference, flagged fail by the harness; the kernel had reported success), 160 FA4 error (GQA/MQA, the pack_gqa cute crash).

Performance

Same harness, 10 warmup + 50 measured iters per cell, torch.cuda.Event timing, hard per-process memory cap so we don't disturb other GPU tenants.

backend	mean TFLOPs/s	peak
sdpa_bshd (PyTorch eager + repeat_interleave for GQA)	87.9	193.7
fa4 (this PR)	94.8	196.2

FA4 speedup vs sdpa_bshd across 240 shapes: mean 1.15×, median 1.07×, max 1.95×, min 0.74×. Biggest FA4 wins are GQA at small S (SDPA pays for repeat_interleave); the few shapes where FA4 still loses are S=512 D=128 causal MHA, which is a tile-size tuning opportunity orthogonal to this PR.

Minimal reproducer for the silent-wrongness bug

import torch
import torch.nn.functional as F
from flash_attn.cute.interface import flash_attn_func

torch.manual_seed(0)
B, S, H, D = 1, 256, 4, 64
dt = torch.bfloat16
q = torch.randn(B, S, H, D, device="cuda", dtype=dt)
k = torch.randn(B, S, H, D, device="cuda", dtype=dt)
v = torch.randn(B, S, H, D, device="cuda", dtype=dt)

out = flash_attn_func(q, k, v, causal=False)
if isinstance(out, tuple):
    out = out[0]

ref = F.scaled_dot_product_attention(
    q.float().transpose(1, 2),
    k.float().transpose(1, 2),
    v.float().transpose(1, 2),
).transpose(1, 2).contiguous()

print(f"max abs diff: {(out.float() - ref).abs().max().item():.4f}")
# before this PR: ~0.61   (bf16 elementwise tolerance should be ~0.003)
# after this PR:  ~0.002

Limitations / out of scope

• Backward pass. This PR is forward-only. flash_bwd_postprocess.py:537 calls get_smem_store_atom(self.arch.major*10+self.arch.minor, ...) and almost certainly needs the analogous fix when going through FlashAttentionBackwardSm120. A backward-correctness check (autograd + grad.allclose(grad_ref)) should follow as a separate PR.
• The cuTeDSL composite-mode flattening (deeper root cause behind commit 3). Commit 3 only flips the default to pack_gqa=False on SM_120; it does not fix pack_gqa.compute_ptr or the cuTeDSL behaviour itself. A proper fix lives in either pack_gqa.compute_ptr (compute the flat index directly, without relying on the slice preserving compositeness) or in cuTeDSL.
• No new tile-size tuning for SM_120. The few shapes where FA4 loses to SDPA after this PR are tile-tuning candidates but out of scope.

🤖 Generated with Claude Code