CroissAttn — fixing FlashAttention backward's zero-sum bug for end-to-end FP4 training

TL;DR

Alexander Johansen (Stanford CS PhD) announces CroissAttn, an ICLR’27 submission that extends the zero-sum-bug framing Jon Barron just fixed in softmax (A Fix for Softmax and Cross-Entropy) to the FlashAttention backward pass under FP4 stochastic rounding. The standard FA backward computes its correction term δ from the stored high-precision output; under FP4 stochastic rounding, dP is a random draw and the standard formula is wrong on every draw. The fix is one line: compute δ from the same rounded dP row, not from the stored output — same mean, zero row sum on every draw. The regime where this matters is attention sinks: the sink key’s true gradient is near zero, so the rounding error of O dominates, making the shortcut unbounded on sink rows (1× on diffuse rows, 2–3× on peaky rows, unbounded on a sink). On 150 draws on one sink row, the standard shortcut lands inside 1× of the true gradient 1 of 150 times; CroissAttn hits 150 of 150.

Key claims

The standard FA backward’s δ computation implicitly assumes dP was computed in high precision; under FP4 stochastic rounding dP is a random draw and the identity breaks on every draw [Thread 1/, App. A.3].
The fix is to compute δ from the same rounded dP row, not from the stored output — matched mean, zero row sum per draw. In the fused kernel this is two sweeps over the row and one product fewer than the standard formulation [Thread 3/, App. A.3].
The pathology concentrates on attention-sink rows: the sink key’s true gradient is near zero, so the rounding error of O dominates. Qiu & Yao (ICLR’26, arxiv:2609.34272 in the embedded thread) show this blows up even BF16 training [Thread 4/].
Error scaling vs row sharpness: ~1× on diffuse rows, 2–3× on peaky rows, unbounded on a sink [Thread 7/, App. C.1].
Empirical: 150 FP4 stochastic-rounding draws on one sink row. The shortcut lands within 1× of the true gradient 1 of 150 times; CroissAttn lands within 1× all 150 times [Thread 7/].
Credit lineage spelled out in-thread: Prieto ‘25, Petrenko ‘26, Liu ‘26 saw the output-layer version; Barron et al. (A Fix for Softmax and Cross-Entropy) fixed it in softmax; Qiu & Yao did BF16 attention. CroissAttn is the stochastic-rounding edition (Sept 2026) [Thread 9/].
RL cousin: BF16 log-probs bias the importance ratio and collapse entropy (Petrenko §4.1). FP4 rollouts make it ~100× worse in a toy test — exact replay may be useful for the fix [Thread 10/].
The authors’ original submitted kernels used the shortcut via a high-precision auxiliary product; the per-draw rule deletes that product. Whether this is a net wall-clock speedup is still being profiled [Thread 8/].

Method

Johansen frames the FA backward’s correction term δ as the attention analogue of the softmax zero-sum identity: in exact arithmetic, the row of dS = P ⊙ (dP − δ·1) must sum to zero because P sums to one and δ is a row-wise scalar matching ⟨P, dP⟩. The standard implementation substitutes δ = ⟨dO, O⟩ using the stored high-precision O, which is correct in expectation but wrong per draw once dP is produced by stochastic FP4 rounding. The one-line fix is to compute δ from the actually-rounded dP row rather than from the stored O: same mean, exact zero row sum on every single draw. In the fused Blackwell kernel this becomes two sweeps over the dP row and one fewer high-precision product than the FA3/FA4 shortcut. The failure mode is sharpened on attention-sink rows because the sink key’s true gradient is tiny while the FP4 rounding error of O is not — the shortcut’s error becomes unbounded as the row sharpens. The author identifies three strictly compositional families where this matters: (i) the FA3/FA4 BF16-forward + FP4-stochastic backward, (ii) training with FP4-attention end-to-end, and (iii) RL post-training where BF16 or FP4 log-probs feed importance ratios that then compound into entropy collapse.

Results

Headline numbers from the thread:

Error distribution over 150 FP4 stochastic-rounding draws on a single sink row: standard shortcut inside 1× of the true gradient 1/150; CroissAttn 150/150 [Thread 7/].
Sharpness-dependent error scaling: diffuse ≈ 1×, peaky 2–3×, sink unbounded [Thread 7/, App. C.1].
RL toy: FP4 rollouts amplify the BF16 Petrenko §4.1 entropy-collapse pathology by ~100× [Thread 10/].
No end-to-end training throughput or final-loss numbers in-thread yet; a second round of experiments (“per-draw rule deletes the auxiliary product”) is in progress.

Why it’s interesting

CroissAttn lands directly in the ongoing wiki argument between loss-implementation stability (A Fix for Softmax and Cross-Entropy) and training-side low-precision kernel co-design (Attn-QAT: 4-Bit Attention With Quantization-Aware Training). Johansen explicitly inherits Barron’s framing — a shortcut that is correct in expectation but wrong per draw once a lower-precision upstream input is introduced — and ports it to the attention backward under stochastic rounding. The failure locus (attention sinks) is the same one A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Training already names as load-bearing for training stability, and the pathology (BF16 training spikes from the shortcut) is the same one Qiu & Yao reported for FP32-vs-BF16 FA3 — this time weaponized by FP4 stochastic rounding. The RL cousin extends the thread to the training–inference mismatch cluster alongside Defeating the Training-Inference Mismatch via FP16 and The 4-bitter Lesson: Balancing Stability and Performance in NVFP4 RL: BF16/FP4 log-probs bias the importance ratio, and exact-replay-style corrections may compose with the per-draw δ rule. The main thing to watch is whether the kernel-level implementation of the per-draw rule is competitive with the shortcut’s auxiliary-product recipe (the submitted paper’s kernels used the latter) — if so, this is the first filed end-to-end FP4 FA backward that is numerically consistent on sink rows without carrying a BF16 auxiliary output of the shape Attn-QAT uses.