Source
Qiu et al. (2025), LIT-198 — ARXIV-2510.04212.
Industrial practice puts memory-bound operations like flash attention in BF16 while pushing compute-bound ones like the FFN to FP8. That configuration sometimes explodes, and the explosion had been treated as a random instability to work around.
It is not random, and it takes two causes — neither sufficient alone:
- Attention produces low-rank representations that repeat, across training steps and across tokens. The same direction is updated again and again.
- Rounding in low-precision addition is biased, not symmetric. BF16 addition loses precision at exponent alignment and again at round-to-nearest, and the residual does not average out.
Put together, the biased errors act as coefficients on the repeated low-rank update. Instead of cancelling over steps they compound into a systematic gradient bias, driving the spectral norm of weights and activations up until training derails.
That account is LIT-198’s, and it is the first mechanistic one: the paper isolates the failure to biased BF16 rounding in flash attention and tests it causally, by de-biasing the rounding and watching training stabilise. The recommendation to fix the attention output’s precision starts from that diagnosis; the FP32 form below is the one Kimi K3 adopted citing it.
What to do
Keep the attention output in FP32 during training. Kimi K3 does exactly this at 2.8T and redesigns its kernel to absorb the doubled on-chip footprint — which is the cost, and it is a real one.
The paper’s own test is the narrower version: a minimal modification to flash attention that removes the bias in the rounding error lets the low-rank updates cancel again and stabilises training. If you control the kernel, that is the cheaper fix; if you do not, the FP32 output is the one available.
Why the isolation matters more than the remedy
The paper works down: tiling is not the source; the failure originates in a single layer; it is linked to one specific computation; the numerical error there is the source; and it is localised to particular attention heads. That sequence turns “low precision is unstable” into a mechanism, and the mechanism is what tells you which of several remedies is the cheap one.
This is the second instance in the record of an instability with a named cause and a targeted fix rather than a hyperparameter workaround — SOTA-131’s QK-Clip, for attention-logit growth, is the first. The pattern worth carrying is that a loss spike is not a flake.
What it qualifies
SOTA-085 says to use flash attention wherever the hardware supports it, because the result is exact. The one numerical change it names is the order of accumulation, and it points here for why that matters at scale. This practice extends it with that part, which turned out to be the part that matters.
Known implementations
- Kimi K3, which keeps the attention output in FP32 during training.