Broken Symmetry in BF16 Attention: Why FlashAttention Gradients Blow Up Late in Training

TL;DR: FlashAttention-3 in BF16 — the kernel essentially everyone trusts — can return attention gradients that are wrong by multiple times their own magnitude, silently, with no NaN and no crash, late in training. The authors trace the failure to a broken conservation law: the exact softmax score gradient sums to zero along every row, which is what makes the query gradient invariant to translating all keys; rounding that gradient to BF16 leaves a tiny nonzero row sum that multiplies the mean key and can dwarf the true signal once attention gets sharp and keys get large. Their fix, GProj, restores the zero sum after the cast with two rank-one corrections, costs 4.7% per step, and recovers FP32-attention training exactly. The error, they show, is shared by cuDNN and PyTorch's fused SDPA backends too — it's a property of fused BF16 attention backwards in general.

A failure with no alarm

The setup is mundane: a 450M-parameter transformer, pretrained on 50B tokens with FlashAttention-3 (FA3) in BF16. For 25B tokens everything looks healthy. Then the pre-clipping gradient norm grows a thousandfold, loss drifts up 0.2 nats above an otherwise bitwise-matched FP32-attention run, and — crucially — nothing overflows and nothing is NaN. In the real world this run would be blamed on data, learning rate, or bad luck with a seed.

The localization experiment is the first thing worth internalizing. The authors freeze the model at 33.6B tokens, keep every forward activation bitwise identical, and recompute only the attention backward of the two worst layers (5 and 11) in FP32. The full-model gradient norm drops from ~5354 to ~18 (Table 11). Almost the entire observed gradient is numerical error from two backward passes of a kernel whose forward outputs are accurate. That answers the paper's framing question — how can a kernel with correct outputs return gradients this wrong? — by showing it's not a contradiction: gradient correctness is never audited in the regime where it fails.

The conservation law hiding in softmax's gradient

Write the attention backward for one query. With scores s_j = \alpha q^\top k_j (scale \alpha, keys k_j), probabilities p_j = \mathrm{softmax}(s)_j, output o = \sum_j p_j v_j, and incoming derivative u = \partial\mathcal{L}/\partial o, define a_j = u^\top v_j and \mu_a = \sum_j p_j a_j. The score gradient is

Notice that \sum_j g_j = \sum_j p_j a_j - \mu_a \sum_j p_j = 0, identically. This zero row sum is not a numerical nicety; it's a symmetry. It means dq can be rewritten as a covariance, dq = \alpha\sum_j p_j (k_j - \mu_k)(a_j - \mu_a) with \mu_k = \sum_j p_j k_j the attention-weighted mean key — so the query gradient depends only on how keys differ from one another, never on where they sit as a group. Translate every key by the same vector b and dq is unchanged, exactly because the g_j sum to zero and the common part cancels. The paper calls this gauge invariance, and it's the right word: there's a direction in input space (the common key offset) to which the exact gradient is blind by construction.

FA3's backward computes g, rounds it to a BF16 operand t, and contracts t with the keys. Rounding does not preserve a zero sum. With error e = t - g and residual row mass \rho = \sum_j t_j, the computed contraction decomposes exactly as

The first term is bounded by the spread of the keys — benign. The second term injects the mean key itself, scaled by a rounding-sized scalar. Whether that matters depends entirely on where the keys live.

Why late training detonates it

Here's the vicious regime. Late in training, two layers of this model drift toward near-one-hot attention with huge keys (an attention-sink-like pattern: logit scales of 10^4–5\times10^4, p_{\max}>0.999 on 99% of rows, key RMS above 200). When attention is nearly one-hot, the true gradient — a covariance over a nearly degenerate distribution — is close to zero. Meanwhile \|\mu_k\| is close to the norm of the dominant key, which is enormous. A single BF16 rounding step of row mass then beats the signal by orders of magnitude.

The two-key bound makes the nastiest property explicit. With attention weights 1-\varepsilon and \varepsilon, both the true gradient and the maximal row mass scale as \varepsilon(1-\varepsilon), so the sharpness cancels: the worst-case relative leak is 2u_b\|\mu_k\|/\|k_1-k_2\|, where u_b = 2^{-8} is BF16's unit roundoff. Concentrating attention gives no protection; only key norm relative to key spread matters — and that's exactly what late training inflates.

And it's a feedback loop, which the ablations nail down cleanly: FA3-SBS (forward repaired, backward untouched) and GProj differ only in the backward, and only GProj's run keeps all 24 layers out of the large-logit regime. The gradient error drives logit growth, which grows the keys, which amplifies the error. The instability isn't the model's fault being exposed by the kernel; the kernel is creating the regime.

The four-key witness (Section 3.4 / Appendix A) is a beautiful minimal case: every input and the saved output exactly BF16-representable, uniform attention, true query-gradient coordinate exactly 0. Casting g=(1+\tfrac{3}{1024}, -1, -\tfrac{3}{1024}, 0) to BF16 rounds 1+\tfrac{3}{1024} to 1, leaving \rho = -3/1024; with keys sharing first coordinate 65536, the computed gradient coordinate is -24 instead of 0. No forward repair can touch this — the forward was already exact.

Why existing fixes stop halfway

Part of the story was already known: a fused multiply-add in the forward softmax means the row maximum doesn't map exactly to zero, corrupting the saved output \hat o and hence the backward's reduction \hat\delta = u^\top\hat o. This was treated as an extreme-input NaN corner case and never fixed in FA3. Repairing it ("FA3-SBS": subtract the max before scaling) kills the thousandfold blow-up and the run finishes stably — but the query gradient is still wrong by 219% (median), training still drives logits to \sim 4\times10^4, and the final loss sits 0.01–0.02 nats above FP32. Stable training is not evidence of correct gradients — maybe the paper's most quotable lesson.

Key smoothing (subtracting the mean key before the kernel, as in SageBwd) shrinks \mu_k but leaves \rho intact; when attention is one-hot, the mean that matters is the dominant key itself, which smoothing doesn't remove. In the runs it delays the blow-up by ~3B tokens, then the run ends worse than vanilla FA3 (1.903 vs 1.866). FP32 saved outputs (Kimi K3's remedy) and dynamic shifting also act only on the forward channel; Theorem 4.2 shows any method that contracts t directly has worst-case relative error that is unbounded under key translation, while the projected version is bounded by |\alpha|u_b\sigma_k\sigma_a — key spread times value-score spread, independent of where the keys sit.

GProj: two rank-one corrections

The fix follows directly from the diagnosis: restore the law on the operand the kernel actually multiplies. FA3's backward already forms the BF16 probabilities r = \mathrm{BF16}(P) for the dV product. GProj measures the row masses \rho = \sum_j t_j and m = \sum_j r_j, and projects t^\perp = t - \lambda r with \lambda = \rho/m — dividing by m rather than assuming the BF16 probabilities sum to 1, which they don't. Since \mathbf{1}^\top g = 0, the projection leaves the exact gradient untouched and removes only the error's row mass. And t^\perp never needs to be materialized, because both contractions are linear:

The dQ correction rides inside the existing backward (one extra rK product and two row sums). The dK correction needs a second pass over the tiles, because \lambda for a query row is only known after that row has seen every key, while dK accumulates per key tile. Everything stays BF16 products with FP32 accumulators; no extra forward state is saved. Equivalently, the query correction replaces every key with k_j - \mu_r (the r-weighted center), so a common key offset cancels exactly — the projection restores the gauge symmetry by construction, and the authors prove it's the smallest zero-sum correction in the r-weighted norm.

The evidence

On eight attention inputs captured from the failing layers:

The query correction alone fixes dQ (0.342%) while leaving dK at 13.3%; the second pass fixes dK to 0.371%. GProj passes a 58-case synthetic suite (packed, causal boundaries, singleton rows, long sequences) and, under random common-key translations, its error in the mathematically-zero coordinate is four orders of magnitude below FA3-SBS's.

In matched from-scratch 50B-token runs (same seed, same data order, only the attention computation varies):

Cost: 238.1 ms/step vs FA3's 227.4 ms — 4.73% overhead — against +33.4% for an optimized fused FP32 attention and +216% for eager FP32, with identical peak memory (Table 4).

One more result deserves emphasis: on fixed captured operands, FA2 reproduces FA3's errors bit-for-near-bit, and cuDNN and PyTorch's flash/efficient SDPA backends behave like FA3-SBS — forward fixed, backward error of ~220% intact (Table 13). This is not an FA3 bug. It's a shared property of every fused BF16 attention backward tested.

What to make of it

If this holds, the practical takeaways are substantial. First, a trusted, universally deployed kernel family has a silent gradient pathology in a regime (sharp attention, large keys) that real training visits — and the pathology causes that regime rather than merely suffering from it. Second, the methodological template generalizes: audit low-precision kernels by checking the exact identities (conservation laws, symmetries) the math guarantees, on the operands the kernel actually multiplies, rather than by comparing against a reference on well-conditioned random inputs. The authors say this explicitly in the conclusion, and it feels like the start of a genre.

Reasons for calibrated skepticism. Everything is one model (450M, a nonstandard "Full-AttnRes" architecture with learned residual aggregation, Muon optimizer), one seed, head dimension 64, Hopper. Notably, QK normalization and softcapping were disabled — many production recipes use exactly these to suppress logit growth, which would likely keep training out of the regime where the leak dominates; a 102M plain-transformer control in Appendix K indeed stays near the FP32 reference (0.26–2.4% gradient discrepancy) even under a 4× learning-rate stress, because its keys never get large. So the honest claim is conditional: fused BF16 attention gradients are unreliable when attention sinks / massive activations push keys large, and that condition arises naturally in at least some setups. Whether it bites at frontier scale, with modern stabilizers, is the open question — though the asymmetry is appealing: GProj costs 4.7% (growing with attention's share of the step), needs no hyperparameters, and provably restores an exact invariance, whereas QK-norm and friends change the model. Also note GProj leaves dV uncorrected, and the overhead was measured at 4K context.

If you read one thing: Section 3.4 (channel two, with the four-key witness) for the aha, then Section 4 / Theorem 4.2 for why forward repairs are fundamentally insufficient — the worst-case error of any direct contraction is unbounded under a symmetry the true gradient doesn't even see. Appendix J.4's backend sweep is the short table that should worry kernel maintainers everywhere.