PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Junlin Chen, Daize Dong, Huanwei Di, Haolong Jia, Jiawei Wu, Haotian Xie, Mingkai Zheng, Yang Li, Leshang Chen, Huishu Wang, Eric P. Xing, Hongyi Wang

arXiv:2609.34272v1cs.LGcs.DCcs.NAmath.NA

TL;DR

Low-precision fused attention is widely trusted, but inaccurate backward gradients can arise silently during training. The paper introduces GProj to restore a zero-sum attention-gradient law broken by BF16 rounding. It brings query/key gradient errors to FP32-attention levels and reaches the same final loss as FP32 attention with 4.7% more step time.

  • Problem

    Attention kernels are commonly validated on random, well-conditioned inputs, while their numerical fidelity during training is rarely examined; backward errors can occur without a crash or NaN.

    Source: §1 · pp. 1–2

  • Method

    GProj restores the score gradient’s zero row sum after the BF16 cast using two rank-one corrections per row.

    Source: §1 · pp. 2–3

  • Results

    0.34%/0.37% median query/key gradient errors, down from 219%/13% after forward repair, with 4.7% more step time and the same final loss as FP32 attention.

    Source: §1 · p. 3, §5.1 · p. 7

  • Takeaways & Limitations

    Checking exact identities on the operands a kernel actually multiplies offers a practical way to audit and repair other low-precision kernels.

    Source: §7 · p. 10

    The evaluation covers BF16 FlashAttention-3 on Hopper GPUs with head dimension 64 and a 450M-parameter model; the 4.7% overhead is measured at 4096 tokens, and GProj leaves the value gradient to the native kernel.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper traces a silent late-training FlashAttention-3 failure to backward-pass error that persists after a known forward-softmax repair. It identifies how BF16 rounding breaks a zero-sum attention-gradient law and introduces GProj to restore it.

1 INTRODUCTION

  • A 450M-parameter model trained on 50B tokens with FA3 ran healthily for 25B tokens, then its gradient norm grew 1,000-fold and its loss ended 0.2 nats above FP32 attention, without NaNs.Recomputing the attention backward of two affected layers in FP32 almost eliminated the excess gradient while keeping forward activations identical.Source: §1 · pp. 1–2
  • Repairing the forward softmax removes the thousandfold growth, but query gradients remain wrong by more than their own size and attention scores still become extreme.Key smoothing delays the blow-up by about 3B tokens but does not prevent it.Source: §1 · p. 2
  • The BF16 cast leaves a nonzero row sum in the score gradient, leaking the mean key into the query gradient and breaking key-translation invariance.The leak can outweigh the true gradient when attention concentrates on one key and keys are large relative to their spread.Source: §1 · p. 2
  • GProj restores the zero row sum after casting, using the mass of the BF16 probabilities consumed by the kernel and two rank-one corrections per row.The correction is applied once within the existing backward pass and once in a second pass for the key gradient.Source: §1 · pp. 2–3

2 RELATED WORK

Prior work studies low-precision arithmetic and fused attention, but numerical assessments have focused mainly on forward outputs. This paper relates backward errors to other softmax, conservation-aware, and training-stability remedies.

2 RELATED WORK

  • Unlike prior work quantifying mainly forward-output deviations in FlashAttention, this paper finds a backward error that can dominate the gradient.Golden et al. (2024) quantify numerical deviation mainly through forward outputs.Source: §2 · p. 3
  • Dynamic softmax shifting and FP32 attention outputs address the saved-output error channel, not the score-gradient cast.Qiu & Yao (2026) motivate conditional dynamic shifting, while Kimi Team (2026, Section 2.1) keeps attention outputs in FP32 during training.Source: §2 · p. 3
  • SageBwd smooths keys before quantization, while Direct-P and MXAttention align normalization with quantized operands; GProj measures mass after the BF16 cast and projects contractions with ρ/m.The cited works are SageBwd (Zhang et al., 2026, Sections 4.2 and 6), Direct-P (Hu, 2026, Section 4.3), and MXAttention (Yu et al., 2026).Source: §2 · p. 3
  • Unlike interventions that change the model or optimizer to address growing logits, this paper shows kernel gradient error can drive that growth and GProj repairs it.The cited interventions include query–key normalization, entropy control, and direct query/key weight rescaling.Source: §2 · p. 3

3 WHERE THE GRADIENT ERROR COMES FROM

The excess model gradient is numerical error in two attention backward passes, and analysis traces it to violations of a zero-row-sum conservation law. One channel comes from the saved output; another arises when BF16 rounding breaks that law after the score gradient is formed.

3.1 THE EXCESS GRADIENT IS A NUMERICAL ERROR

  • At 33.6B tokens, recomputing only the attention backward of these two layers in FP32, while keeping every attention output and the loss bitwise identical, reduced the model-gradient norm by more than two orders of magnitude.Source: §3.1 · p. 4
  • The failure grows in layers 5 and 11, where queries and keys become large and attention nearly one-hot; in matched runs, the backward error drives this regime, and larger keys amplify the error.The other 22 of 24 layers remain at moderate score magnitudes.Source: §3.1 · p. 4

3.2 A CONSERVATION LAW OF THE ATTENTION BACKWARD

  • The exact score gradient sums to zero on each row, making attention gradients insensitive to common translations of the keys; numerical error can violate this before or after the score gradient is formed.The paper calls these error channels one and two, respectively.Source: §3.2 · p. 4

3.3 CHANNEL ONE: THE SAVED OUTPUT

  • An error in the saved-output reduction shifts the score gradient and creates query-gradient error along the attention-weighted mean key; a corrected saved output removes most of this channel.FA3-SBS, which subtracts the unscaled row maximum before scaling, removes the saved-output error.Source: §3.3 · pp. 4–5

3.4 CHANNEL TWO: THE CAST BREAKS THE CONSERVATION LAW

  • BF16 rounding can break the score-gradient row sum; its residual mass multiplies the mean key, producing a leak that can dominate as keys grow and attention nears one-hot.The same rounded operand also affects the key gradient.Source: §3.4 · p. 5
  • In a four-key example, BF16 casting changed a true zero first-coordinate query gradient to −24 through the nonzero score-gradient row mass.The example uses keys sharing a first coordinate and a score-gradient entry that rounds from 1 + η to 1.Source: §3.4 · pp. 5–6

4 GPROJ

GProj restores the score gradient’s zero row sum after BF16 casting by projecting along the probabilities the kernel actually uses. This removes offset-dependent query-gradient error and bounds the remaining error by key spread under the stated assumptions.

4 GPROJ

  • GProj subtracts λr from the cast score gradient t, choosing λ = ρ/m so the resulting t⊥ has zero row sum even when BF16 probabilities have mass m ≠ 1.Here ρ is the row mass left by casting and m is the actual mass of the BF16 probabilities used by the kernel.Source: §4 · p. 6
  • The projection need not be materialized because both gradient contractions are linear in t⊥; GProj adds products, rank-one terms, and row sums to the backward pass.The key correction requires a second pass because λ is known only after each query row has seen every key.Source: §4 · p. 6
  • Because the exact score gradient already sums to zero, projection leaves it unchanged and removes only cast error, canceling common key offsets in the query contraction.The projected query contraction centers keys using the r-weighted center.Source: §4 · p. 6
  • The stated error bound concerns an idealized comparison, while the implemented kernel’s saved-output and reconstructed-probability operands add row-mass error of order ub.GProj removes that row mass regardless of its source.Source: §4 · p. 7
  • With r = p, direct contraction has unbounded worst-case relative query error under key translation, while GProj’s error is bounded by |α| ub σkσa independently of the keys’ mean.The comparison assumes exact probabilities and center before casting, with exact arithmetic after the cast.Source: §4 · p. 7

5 RESULTS

Across matched training runs, GProj stays stable and ends at the FP32-attention loss, while FA3 and key smoothing finish with larger gradients and higher loss. GProj also brings gradient errors close to FP32 attention and preserves translation invariance across tested cases.

5.1 ACCURACY WITH AN IDENTICAL FORWARD

  • 0.34%/0.37% median query/key gradient error for GProj, versus 219%/13% after forward repair, on eight captured inputs.FA3-SBS and GProj share byte-identical forward outputs, isolating the backward comparison; the query correction improves dQ and the key pass improves dK, while dV remains at FA3-SBS error.Source: §5.1 · pp. 7–8, Table 1 · p. 8
  • The Figure 3 inset shows GProj coinciding with FP32 loss, while FA3-SBS drifts 0.01–0.02 nats above it.The plotted gradient norm is pre-clipping and on a log scale; the dotted line marks the approximate 25.2B-token onset of FA3’s gradient growth.Source: Figure 3 · p. 8

5.2 TRANSLATION INVARIANCE AND ROBUSTNESS

  • GProj returns essentially zero error in the translation test’s zero-target query coordinate, while FA3 and FA3-SBS return errors of order one.It passes all 60 synthetic cases; under random common-key translations, its zero-target-coordinate error is four orders of magnitude below FA3-SBS, and dK/dV error never rises relative to FA3-SBS.Source: §5.2 · pp. 8–9

5.3 FROM-SCRATCH TRAINING AND COST

  • GProj and FP32 attention finish 50.0B tokens with small gradients at the same loss (1.665), while FA3 and key smoothing end with higher loss and large gradients.FA3-SBS finishes stably, but its backward error still drives several layers into the large-logit regime; key smoothing delays blow-up by about 3B tokens without preventing it.Source: Figure 3 · p. 8, §5.3 · p. 9
  • 4.7% additional complete-step time for GProj, compared with 33.4% for fused FP32 attention.GProj saves no extra forward state.Source: §5.3 · p. 9

6 LIMITATIONS

The study evaluates BF16 FlashAttention-3 on Hopper GPUs with head dimension 64 and a 450M-parameter model.

6 LIMITATIONS

  • The study covers BF16 FlashAttention-3 on Hopper GPUs with head dimension 64 and a 450M-parameter model.The 4.7% overhead is measured at 4096 tokens and grows with attention’s share of the step; GProj leaves the value gradient to the native kernel.Source: §6 · p. 10

7 CONCLUSION

The conclusion identifies BF16 casting of the softmax score gradient as breaking a zero-sum conservation law, and presents GProj as a two-correction repair. The appendices provide translation analysis, implementation details, and proofs of what the projection removes and what error can remain.

7 CONCLUSION

  • GProj restores the zero-sum law after the cast with two rank-one corrections per row, bringing query/key gradient errors to FP32-attention levels and training stably where FA3 degrades.The conclusion proposes checking exact identities on the operands a kernel actually multiplies as a practical audit-and-repair approach for other low-precision kernels.Source: §7 · p. 10

Appendix roadmap.

  • The appendix roadmap assigns cast witnesses and proofs to Appendices A–C, method and protocol details to D–E, synthetic tests to F, and training runs to G.The remaining appendices cover model configuration, FA3 failure localization, and a plain-Transformer control.Source: Appendix roadmap · p. 13

A.1 EXACT TRANSLATIONS AND NUMERICAL ROW-MASS LEAKAGE

  • A common key translation leaves exact attention probabilities unchanged, but a nonzero rounded score-gradient row sum makes the query-gradient contraction respond to that translation.Key smoothing changes the key operand but leaves the post-formation row sum intact; GProj corrects the cast score operand itself.Source: §A.1 · p. 13
  • In the four-key witness, an exact BF16 forward still yields a query-gradient error of −24 after rounding the score-gradient operand.The true first query-gradient coordinate is zero; the example embeds in dimension 64 by zero padding.Source: §A.1 · p. 13

B IMPLEMENTED ARITHMETIC AND ROUNDING

  • The implemented GProj kernel uses BF16 matrix operands and returned gradients, with FP32 accumulators, row statistics, and scalar corrections.It runs a native backward followed by a separate dK correction pass; it stores no quadratic matrix and does not project dV.Source: B IMPLEMENTED ARITHMETIC AND ROUNDING · p. 14

C.1 PROJECTION AND THE ERROR THAT REMAINS

  • The projection preserves zero-row-sum components and removes only the row mass of the error; zero-sum error remains.Its correction uses the actual mass of the BF16 probabilities consumed by the kernel; centering error can still leave query error because projection cannot restore value differences lost in the cast.Source: §A.1 · p. 13, §C.1 · p. 15

C.2 PROOF OF THEOREM 4.2

The proof shows that projecting the rounded score gradient restores translation-invariant query gradients, while implementation rounding leaves additional error terms in both correction passes.

C.2 PROOF OF THEOREM 4.2

  • Under a common key translation, the unprojected relative query-gradient error can grow without bound, while projection removes dependence on the mean key.The projected error bound depends on key variation, not the common translation.Source: §C.2 · pp. 15–16
  • In 256-row, 1024-key synthetic sweeps, translation leaves p, g, and GQ unchanged, while the GProj floor ranges from 0.12% to 0.17% as logit standard deviation varies from 1 to 6.The sweep varies the common key offset and evaluates the forward-repair and GProj corrections.Source: §C.2 · p. 16

C.3 ROUNDING IN THE IMPLEMENTED QUERY CORRECTION

  • The implemented query-correction error comprises surviving operand error plus accumulation, ratio, and final-rounding errors.The total error norm is bounded by the sum of the component-error norms.Source: §C.3 · p. 16

C.4 ROUNDING IN THE IMPLEMENTED KEY CORRECTION

  • The key pass reconstructs probabilities and casts scaled queries, so query and key gradients need not use identical effective score operands; the value gradient is not projected.The correction is applied to the FP32 key accumulator before the final BF16 cast.Source: §C.4 · pp. 16–17

D COMPARISON METHODS AND ARITHMETIC CONTRACTS

The section defines the attention-kernel variants and their arithmetic contracts, then distinguishes the backward ablations and the dynamic-shift implementation.

D COMPARISON METHODS AND ARITHMETIC CONTRACTS

  • GProj adds query and key corrections to FA3-SBS while retaining BF16 matrix products with FP32 accumulation; all methods expose BF16 outputs and gradients to the model.FA3-SBS changes only the forward softmax, subtracting the row maximum before scaling.Source: D COMPARISON METHODS AND ARITHMETIC… · p. 17
  • QY-shift uses a tile-dependent shift with 10^-3 epsilon, 512 × 512 tiles, and BF16 intermediate state and gradient accumulation.Its previous numerator and denominator are rescaled before the online update.Source: D COMPARISON METHODS AND ARITHMETIC… · p. 17
  • QY-shift rounds the elementwise product before row reduction and updates BF16 gradient buffers per tile, whereas GProj retains FP32 gradient accumulators and applies α after query correction.The C++/CUDA port matches the authors’ code byte for byte.Source: D COMPARISON METHODS AND ARITHMETIC… · p. 17
  • The ablations progress from FA3-SBS to query-corrected GProj-Q and then full GProj with the key pass, isolating backward differences because all share one forward.GProj denotes the complete method.Source: D COMPARISON METHODS AND ARITHMETIC… · p. 18

E ACCURACY AND COST: SETUP AND ADDITIONAL RESULTS

The evaluation compares attention-gradient accuracy across captured and synthetic inputs, reports the key-pass ablation, and specifies the timing conditions.

E ACCURACY AND COST: SETUP AND ADDITIONAL RESULTS

  • Eight captured inputs cover layers 5 and 11 at 23.1B and 33.6B tokens; relative L2 gradient errors are measured against FP64 references and summarized by medians and maxima.Each capture is a causal 4096-token sequence with 16 query heads, four KV heads, and head dimension 64.Source: E ACCURACY AND COST: SETUP AND… · p. 18
  • 13.2651% to 0.3712%: adding the key pass lowers median dK error, while leaving O, dQ, and dV unchanged on the same execution.The query correction alone leaves the median dK error at 13.2651%.Source: E ACCURACY AND COST: SETUP AND… · p. 18
  • Peak allocated memory is 9.191 GiB for both FA3 and GProj and 9.182 GiB for eager FP32 under the specified one-H200 training-step timing setup.The reported step-time statistic is the median across workers of their mean times after warmup.Source: E ACCURACY AND COST: SETUP AND… · p. 19

F SYNTHETIC NUMERICAL SUITE

  • The synthetic suite evaluates O, dQ, dK, and dV across 58 stress cases and two analytic four-key calibrations in BF16 with head dimension 64 and GQA.Cases include key and value shifts, sharp attention, empty causal rows, and tile-boundary rows.Source: F SYNTHETIC NUMERICAL SUITE · p. 19

G TRAINING STUDY

The training study reports that Qiu–Yao shifting stalls and that its tile-shift rule increasingly clamps rows as training proceeds; matched runs and checkpoint evaluations define the study’s scope.

G TRAINING STUDY

  • Matched runs share the same 16-H200 recipe and differ only in attention computation; checkpoint regime statistics are evaluated on two fixed documents.The runs use seed 42 and 11,930 updates totaling 50.0B tokens; Qiu–Yao stopped early.Source: G TRAINING STUDY · p. 20
  • 0.31 nats at 8.4–8.9B tokens: Qiu–Yao trails FA3 after stalling near loss 2.28–2.30.The run was stopped at 13.2B tokens.Source: G TRAINING STUDY · pp. 20–21
  • For positive row maxima with near-ties, the shift rule uses 2r; large logits can push weights below the FP32 normal range and leave tile sums below the 10^-10 clamp.The resulting row output is dominated by rounding.Source: G TRAINING STUDY · p. 21
  • At 12.6B tokens, about half the rows in layers 15–23 are fully clamped, and attention output is uncorrelated with FP32 attention (median relative error ≈1.0).At step 1500, row maxima reach approximately 5,300.Source: G TRAINING STUDY · p. 21

H MODEL AND TRAINING CONFIGURATION

The paper’s training runs use a 450M-parameter Full-AttnRes transformer with shared-read configuration, and the analyses use checkpoints from a from-scratch FA3 run.

H MODEL AND TRAINING CONFIGURATION

  • The 450M-parameter Full-AttnRes transformer uses learned softmax aggregation over earlier sublayer outputs and a shared residual-depth read for token-attention queries, keys, and values.All training runs, captured operands, and frozen-state diagnostics use this configuration and are trained from scratch.Source: H MODEL AND TRAINING CONFIGURATION · p. 21, Table 9 · p. 21

I LOCALIZING THE NUMERICAL FAILURE

  • All analyses use checkpoints from the from-scratch FA3 run.Source: I LOCALIZING THE NUMERICAL FAILURE · p. 21
  • Table 10 reports per-document medians across 127 documents per checkpoint and layer, using 128 sampled query rows and all 4096 keys.Its residual fraction and error cosine compare the saved-output error prediction with observed native error.Source: Table 10 · p. 22

I.1 THE SAVED-OUTPUT CHANNEL PREDICTS THE LOCAL ERROR

Saved-output error explains much of the late query-gradient failure, and correcting it nearly eliminates the excess full-model gradient. A forward-softmax repair fixes one defect, but a remaining backward error grows with the common key component and appears across fused BF16 backends.

I.1 THE SAVED-OUTPUT CHANNEL PREDICTS THE LOCAL ERROR

  • At checkpoint 8000, layer 11, replacing saved O reduces median absolute dQ error 72.02×, while the predicted error leaves a 0.0197 median residual fraction across 127 documents.The local analysis also found a 0.9999 cosine between predicted and observed error for one document; the measurements sample query rows against all 4096 keys.Source: §I.1 · p. 22, §J.1 · p. 24

I.2 FIXED-FORWARD INTERVENTIONS REACH THE FULL MODEL

  • At checkpoint 8000, median full-model gradient norm falls from 5353.856 natively to 19.746 with saved-O correction and 18.185 with FP32 VJP; at checkpoint 6500, 36.653 falls to about 2.3–2.5.The attention outputs and loss remain bitwise identical across modes.Source: §I.2 · p. 22, Table 11 · p. 23

I.3 THE SOURCE OF THE SAVED-OUTPUT ERROR

  • A fused multiply-add can keep the row maximum from mapping exactly to zero; explicitly rounding the subtraction before scaling repairs the one-key identities and lowers late-capture dQ and dK errors.The FA3-SBS forward repair reduces absolute dQ error by 10.8–43.0× and dK error by 12.0–58.6× on four checkpoint-8000 captures.Source: §I.3 · p. 23

I.4 WHAT THE FORWARD REPAIR LEAVES BEHIND

  • On late captures, changing saved O alone recovers the forward repair’s improvement, whereas changing only LSE slightly worsens error; even with an exact saved output, query-gradient error grows linearly with a common key coordinate.Source: §I.4 · p. 23, §J.3 · p. 26

J.2 PRESCRIBED FMA CASES

  • The scalar arithmetic model reproduces all 96 prescribed one-key outputs bitwise, including six violations, and the explicit rounded-subtraction repair restores both exact identities in all 96 cases.Source: §J.2 · pp. 24–25

J.3 SAVED-STATE CROSSES AND THE REMAINING GAUGE DEFECT

  • At score gap 64, the common-key error slopes are 0.0101494 for FA3 and 0.0101553 for FA3-SBS; the measured error scales linearly with the common key value.The prescribed gauge construction confirms unchanged output and LSE as that value varies.Source: §J.3 · p. 26

J.4 BACKEND COMPARISON AT IDENTICAL CAPTURED OPERANDS

  • FA2 reproduces FA3’s errors, while cuDNN and fused SDPA retain median relative query-gradient error of about 2.2, versus 0.003–0.004 for materialized math paths.The backend sweep holds captured Q, K, V, and incoming derivatives fixed.Source: §J.4 · pp. 26–27, Table 13 · p. 27

K A PLAIN-TRANSFORMER CONTROL

  • In the plain-transformer control, LR4 attention-backward discrepancies remain small at early checks, but the runs become unstable as query RMS stays far below the FA3 run’s late value.At 500 updates, discrepancy is 1.5–2.4% at LR4; maximum query RMS is 42.5–44.5 versus above 200 in FA3 layer 11 at 33.6B tokens.Source: K A PLAIN-TRANSFORMER CONTROL · p. 27, Table 14 · p. 28, Figure 10 · p. 28

Related papers

Loading 2609.34272v1…