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

Abstract: BF16 is now standard in large-scale pretraining, including in fused attention kernels such as FlashAttention, and these kernels are widely trusted. When we used FlashAttention-3 to pretrain a 450M-parameter transformer on 50B tokens, however, we ran into a problem: training was healthy for 25B tokens, then the gradient norm grew a thousandfold and the loss ended 0.2 nats above FP32 attention, without a single NaN. Recomputing the attention backward of just two layers in FP32 removes almost all of the excess gradient. Part of the cause is known: a fused multiply-add in the forward softmax, so far treated as an extreme-input NaN case and never fixed in FlashAttention-3. Repairing it stops the blow-up, but the query gradient is still wrong by more than its own size, and training still drives attention logits to thousands of times their size under accurate gradients. The remaining error comes from a broken conservation law. The softmax score gradient sums to zero along every row, which makes the query gradient blind to where the keys sit as a group; rounding it to BF16 leaves a small nonzero sum that leaks the mean key into the gradient, and the leak grows exactly as late training makes keys large and attention sharp. We introduce GProj (gauge projection), which restores the zero sum after the cast with two rank-one corrections per row. It cuts the remaining median query/key gradient errors from 219%/13% to 0.34%/0.37%, on par with FP32 attention, for 4.7% more time per training step. In matched from-scratch runs it trains to the same loss as FP32 attention, while FlashAttention-3 and key smoothing both destabilize.

• The paper identifies two numerical issues in BF16 fused attention backward passes: a saved-output error due to fused softmax arithmetic and a conservation-law violation after BF16 casting, both leading to large gradient errors.
• The authors propose a corrected version of FlashAttention-3 (FA3), called GProj, which reduces the median gradient error to approximately 0.34% and 0.37% for query and key gradients, respectively, preserving computational efficiency and reducing large-logit regime issues.
• Experimental validation shows that GProj matches the performance of FP32 attention in stability and loss, resolving drift and incorrect gradient growth in both synthetic and practical training contexts, ensuring reliable training without NaNs or overflows.

Problem setting and empirical failure

The paper investigates a silent numerical failure in BF16 fused attention backward passes. The motivating experiment is a 450M-parameter transformer trained for 50B tokens with FlashAttention-3 (FA3). Training appears healthy for approximately 25B tokens, after which the pre-clipping gradient norm increases by roughly three orders of magnitude and the loss diverges from an otherwise matched FP32-attention run by approximately 0.2 nats. No NaNs or infinities occur. The failure is therefore not detected by conventional overflow monitoring and could plausibly be misattributed to optimization, data, or model-scale effects.

The authors localize the excess gradient to the attention backward. At 33.6B tokens, replacing only the backward computations of layers 5 and 11 with FP32 arithmetic, while holding all forward activations bitwise fixed, reduces the full-model gradient norm from a median of 5353.856 to 18.185. Replacing the saved attention output alone reduces it to 19.746. This intervention establishes that the dominant pathology is a numerical backward error in two attention layers rather than a forward loss anomaly or a general optimizer failure.

The training trajectory exhibits a pronounced delayed-onset behavior. FA3 remains competitive through the early and middle stages, then enters a regime in which layers 5 and 11 develop exceptionally large query and key norms, nearly one-hot attention distributions, and increasingly inaccurate gradients. The delayed onset is important: random well-conditioned validation inputs and early-training comparisons do not expose the defect.

Figure 1: Matched from-scratch training shows late gradient growth and loss drift for FA3, whereas GProj remains aligned with FP32 attention throughout 50B tokens.

The paper identifies two distinct numerical channels. The first is a known forward-softmax issue: FA3 fuses scaling and maximum subtraction in an FMA, so the maximum score need not map exactly to zero. The resulting error contaminates the BF16 attention output saved for the backward pass. The second channel is the paper’s principal contribution: even with an exact saved output, casting the score gradient to BF16 violates an exact zero-sum invariant and produces a spurious gradient component proportional to the mean key.

The saved-output channel

For a query row, keys $k_j$, values $v_j$, attention probabilities $p_j$, and incoming output derivative $u$, the exact score-gradient entries are

$g_j = p_j(a_j-\mu_a),$

where $a_j=u^\top v_j$ and $\mu_a=\sum_jp_ja_j$. Consequently,

$\sum_j g_j=0.$

The query gradient can be written as

$G_Q=\alpha\sum_j g_j k_j,$

with $\alpha$ denoting the attention scale. FA3 estimates the scalar reduction $v_j$0 through the BF16 output saved during the forward pass. If the saved output induces an error $v_j$1, the resulting query-gradient error is

$v_j$2

where $v_j$3 is the attention-weighted mean key. Thus, even a small saved-output error is amplified when the mean key is large.

The source of the saved-output error is traced to the fused softmax arithmetic. The implementation effectively computes a fused scaled-and-shifted exponent argument rather than explicitly subtracting the row maximum before scaling. Because the FMA does not necessarily round the intermediate product separately, the maximum element may receive a nonzero exponent argument. In a one-key row, where exact attention requires $v_j$4 and $v_j$5, this can produce $v_j$6 and a nonzero query gradient.

The paper verifies this mechanism using 96 prescribed configurations whose dot products are analytically controlled. A scalar model of the FMA, BF16 exponentiation operand, and FP32 denominator reproduces the native outputs bitwise, including all six cases in which the one-key identity fails. The proposed subtract-before-scale modification, FA3-SBS, restores the identity in all prescribed cases and substantially reduces the observed gradient error.

The saved-output channel accounts for most of the initial failure. Across a confirmation set of 127 documents, the predicted saved-output error has near-perfect alignment with the observed native query-gradient error in the most affected layers. At checkpoint 8000 in layer 11, replacing only the saved output reduces the absolute query-gradient error by a median factor of 72.02. However, this repair does not eliminate the remaining backward defect.

Figure 2: Across 127 documents, the saved-output prediction explains most of the native query-gradient error, with high error cosine and a small unexplained residual.

The distinction between saved output and saved log-sum-exp is established by state-cross experiments. Substituting the corrected output while retaining the native log-sum-exp reproduces the improvement, whereas substituting only the log-sum-exp does not. The implication is direct: the forward repair must target the saved attention output, not merely normalization statistics.

The broken conservation law

The central theoretical observation is that exact attention backward propagation possesses a translation symmetry. Adding a common vector $v_j$7 to every key in a row changes every logit by the same constant, leaves the softmax probabilities unchanged, and therefore leaves the exact local VJP unchanged. Algebraically, this invariance follows from the zero row sum of the score gradient:

$v_j$8

The query gradient can equivalently be expressed in centered form,

$v_j$9

Only key differences matter. A numerical backward that responds to the absolute common offset of the keys violates this symmetry.

FA3 casts the score gradient $p_j$0 to BF16 before contracting it with the keys. Let $p_j$1 be the cast operand and define the rounding error $p_j$2 and row mass $p_j$3. Since the exact score gradient has zero mass,

$p_j$4

The first term is an ordinary centered rounding error. The second is the leakage term. It is generated solely by the nonzero row sum left by BF16 rounding and is multiplied by the mean key. A common key translation $p_j$5 leaves the exact gradient unchanged but changes the numerical contraction by $p_j$6.

This mechanism becomes severe in the late-training regime. When attention is nearly one-hot, the exact query gradient is a covariance over a distribution with very small effective spread and can be close to zero. The mean key, however, can remain large. A single BF16 rounding unit in the score gradient can therefore dominate the true signal. The paper’s explicit four-key witness makes the effect non-asymptotic: with an exact forward and exactly representable inputs, BF16 casting changes an exact zero query-gradient coordinate into a value of $p_j$7.

Figure 3: BF16 casting leaves nonzero score-gradient row mass, producing a mean-key leak whose magnitude grows with key offset; GProj removes the leak while retaining BF16-level error.

This result is stronger than a diagnosis of inaccurate softmax outputs. It shows that a perfect forward pass cannot remove the principal residual error, because the defect is introduced after the forward state has been computed. It also explains why FP32 accumulation alone is insufficient: the incorrect BF16 multiplicands have already broken the invariant before the matrix products are accumulated.

The theoretical distinction between forward repair and projection is formalized through a key-translation argument. A backward pass that directly contracts the cast score gradient has an error component proportional to $p_j$8, and its worst-case relative error is unbounded under common key translations. A projected backward pass removes this dependence and leaves an error bounded by the attention-weighted key spread. The result identifies the relevant numerical quantity not as the absolute key magnitude but as the key offset relative to within-row variation.

Gauge projection

GProj restores the zero row sum after the score-gradient cast. The implementation uses the BF16 probability operand $p_j$9 that the kernel already forms and its actual mass

$u$0

Given $u$1, GProj subtracts $u$2:

$u$3

This correction satisfies $u$4 even when the BF16 probability mass is not exactly one. Subtracting $u$5 without dividing by $u$6 would leave a residual mass $u$7, which is why the correction must use the mass of the actual operand consumed by the kernel.

The corrected query contraction is implemented without materializing $u$8:

$u$9

The key gradient requires a second pass because $g_j = p_j(a_j-\mu_a),$0 is known only after a query row has processed all key tiles, whereas the native key accumulator is organized by key tile. GProj therefore adds a second contraction to correct $g_j = p_j(a_j-\mu_a),$1 before its final BF16 cast. The value gradient remains native.

The projection is exact on any zero-sum score gradient and acts only on the cast error. In geometric terms, it removes the component of the numerical error that lies in the row-mass direction while preserving the other BF16 rounding components. The corrected query contraction depends on centered keys and is invariant to common key translation. Its remaining error is consequently controlled by key spread rather than key offset.

Gradient accuracy and training behavior

The principal accuracy evaluation uses eight captured attention inputs from layers 5 and 11 at 23.1B and 33.6B tokens. All methods receive identical BF16 inputs, and gradients are compared with an FP64 centered VJP. FA3-SBS and GProj have byte-identical forward outputs, so their comparison isolates the backward.

Method	$g_j = p_j(a_j-\mu_a),$2 error	$g_j = p_j(a_j-\mu_a),$3 error	$g_j = p_j(a_j-\mu_a),$4 error
FA3	773%	85.6%	0.369%
FA3-SBS	219%	13.3%	0.334%
GProj-Q	0.342%	13.3%	0.334%
GProj	0.342%	0.371%	0.334%
Fused FP32 attention	0.385%	0.371%	0.165%

The forward repair reduces the median $g_j = p_j(a_j-\mu_a),$5 error from 773% to 219%, but leaves the query error larger than the gradient itself. The query projection reduces it further to 0.342%, essentially matching fused FP32 attention. The additional key pass reduces $g_j = p_j(a_j-\mu_a),$6 error from 13.3% to 0.371%. These results separate the roles of the two corrections: GProj-Q repairs the query contraction, while the second pass repairs the key contraction.

The synthetic suite contains 58 stress cases covering common-key translations, singleton rows, packed layouts, unequal lengths, tile boundaries, sharp attention, tiny upstream derivatives, and long sequences. GProj passes all cases and does not increase $g_j = p_j(a_j-\mu_a),$7 or $g_j = p_j(a_j-\mu_a),$8 error relative to FA3-SBS. Its largest reported synthetic relative errors are 2.49% for $g_j = p_j(a_j-\mu_a),$9 and 2.76% for $a_j=u^\top v_j$0, both occurring in numerically small-gradient cases. For exact-zero targets, residuals remain near the arithmetic floor; for example, singleton-row $a_j=u^\top v_j$1 is exactly zero and the corresponding $a_j=u^\top v_j$2 residual is approximately $a_j=u^\top v_j$3.

The training study uses matched from-scratch runs with the same seed, data order, architecture, and optimization configuration. GProj and FP32 attention remain stable through 50.0B tokens and finish with essentially identical loss: 1.665. FA3 finishes at 1.866 over the final 0.84B tokens, approximately 0.20 nats worse, with a median gradient norm of 30.705. FA3-SBS remains numerically stable in the narrow sense that its gradient norm does not exceed 10, but its final loss is 1.679 and it drives several layers into the same extreme-logit regime. Key smoothing delays the failure by approximately 3B tokens but ultimately produces a final loss of 1.903 and a median gradient norm of 233.384.

Figure 4: Holding the forward fixed localizes the excess full-model gradient to the backward of layers 5 and 11; correcting the saved output removes most of the error, while FP32 removes both numerical channels.

The regime audit clarifies why stable loss alone is insufficient. At 33.6B tokens, FA3 has logit scales of 12,843 and 51,057 in layers 5 and 11, respectively; 99.4% of layer-11 rows satisfy $a_j=u^\top v_j$4. GProj has maximum logit scale 26, layer-11 query/key RMS values of 2 and 3, and no rows with $a_j=u^\top v_j$5. Thus, the inaccurate backward does not merely perturb an otherwise healthy trajectory: it drives the model toward a large-key, nearly one-hot regime that amplifies the very error responsible for the instability.

The released Qiu–Yao dynamic-shift baseline performs worse in this setup. Its shift rule can move positive near-tied tile maxima by $a_j=u^\top v_j$6, making all tile weights exponentially small when logits are large. Tile sums then fall below the implementation’s $a_j=u^\top v_j$7 clamp. The run spikes at approximately 6.3B tokens and is stopped at 13.2B tokens. This result does not refute dynamic shifting generally, but it shows that the particular released arithmetic and clamp policy do not provide a reliable substitute for conservation-aware backward correction in this evaluation.

Computational cost

GProj preserves the BF16 matrix products and does not require additional forward state or quadratic attention storage. On an H200 with batch one and sequence length 4096, complete training-step time is 238.120 ms for GProj versus 227.357 ms for FA3, a 4.73% increase. Eager FP32 attention requires 719.420 ms, or 216.43% additional time. An optimized fused FP32 implementation still costs 308.276 ms, 33.43% above FA3.

Peak allocated memory is effectively unchanged: 9.191 GiB for both FA3 and GProj in the reported comparison. The overhead is therefore principally computational, arising from the additional $a_j=u^\top v_j$8 contraction, row-statistic handling, and second key-correction pass. The reported 4.7% figure is specific to the complete training step and the tested 4096-token configuration; it should not be interpreted as a sequence-length-independent kernel overhead.

Limitations and open questions

The empirical scope is narrow. The main experiments use BF16 FA3 on Hopper GPUs, head dimension 64, grouped-query attention, sequence length 4096, and a single 450M-parameter architecture. Although the fixed-operand backend sweep shows related backward errors in FA2, cuDNN, and PyTorch fused SDPA, the paper does not establish equivalent late-training behavior across GPU generations, head dimensions, precision formats, model families, or training scales.

GProj does not correct $a_j=u^\top v_j$9, whose error remains that of FA3-SBS. Its projection also assumes that the relevant conservation law is the appropriate invariant for the implemented attention formulation and that the BF16 probability mass remains finite and nonzero. The synthetic evaluation satisfies these conditions, but more adversarial kernels or alternate sparsity and quantization schemes may require different projections.

Finally, the training evidence consists of matched runs with one principal seed and a fixed model/data configuration. The agreement between GProj and FP32 attention is strong within that setting, but the extent to which the same mechanism controls instability in larger models, longer contexts, FP8 kernels, or other fused backward implementations remains an open empirical question.

Conclusion

The paper shows that BF16 fused attention can produce severely incorrect gradients without NaNs or obvious forward failure. The defect has two components: a saved-output error caused by fused softmax arithmetic and a more fundamental violation of the exact zero-sum score-gradient identity after BF16 casting. The latter creates a mean-key leakage term that grows with key offset and becomes dominant when training produces large keys and nearly one-hot attention.

FA3-SBS repairs the forward channel but cannot repair the cast-induced conservation-law violation. GProj restores the zero row sum using a probability-weighted projection, reducing median $\mu_a=\sum_jp_ja_j$0 and $\mu_a=\sum_jp_ja_j$1 errors to approximately 0.34% and 0.37%, matching FP32 attention at a reported 4.7% training-step overhead. In the matched 50B-token study, it reaches the same final loss as FP32 attention while avoiding the large-logit regime induced by uncorrected BF16 backward arithmetic. The paper’s broader methodological result is that low-precision kernels should be audited not only against aggregate numerical error, but also against exact structural identities—particularly those that encode invariance and conservation in the underlying derivative.