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

arXiv: https://arxiv.org/abs/2609.34272 · HTML: https://mlsanity.com/paper/2609.34272/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 22 of its 31 references that are in the map.
- First posted on arXiv 28 Sep 2026.

## Builds on (most cited first, 22 of 22)

- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [Muon is Scalable for LLM Training](https://mlsanity.com/paper/2502.16982/) (2025)
- [Deep Learning with Limited Numerical Precision](https://mlsanity.com/paper/1502.02551/) (2015)
- [Massive Activations in Large Language Models](https://mlsanity.com/paper/2402.17762/) (2024)
- [Query-Key Normalization for Transformers](https://mlsanity.com/paper/2010.04245/) (2020)
- [A Study of BFLOAT16 for Deep Learning Training](https://mlsanity.com/paper/1905.12322/) (2019)
- [When Attention Sink Emerges in Language Models: An Empirical View](https://mlsanity.com/paper/2410.10781/) (2024)
- [Small-scale proxies for large-scale Transformer training instabilities](https://mlsanity.com/paper/2309.14322/) (2023)
- [Stabilizing Transformer Training by Preventing Attention Entropy Collapse](https://mlsanity.com/paper/2303.06296/) (2023)
- [SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization](https://mlsanity.com/paper/2411.10958/) (2024)
- [Scaling FP8 training to trillion-token LLMs](https://mlsanity.com/paper/2409.12517/) (2024)
- [Revisiting BFloat16 Training](https://mlsanity.com/paper/2010.06192/) (2020)
- [Why Low-Precision Transformer Training Fails: An Analysis on Flash Attention](https://mlsanity.com/paper/2510.04212/) (2025)
- [Is Flash Attention Stable?](https://mlsanity.com/paper/2405.02803/) (2024)
- [MXAttention: Data-Free Optimal Scaling and Pre-Normalization Quantization for MXFP4 Attention](https://mlsanity.com/paper/2607.24377/) (2026)
- [Hardware-Aware FP4 FlashAttention-4](https://mlsanity.com/paper/2609.04105/) (2026)
- [Online Pseudo-average Shifting Attention(PASA) for Robust Low-precision LLM Inference: Algorithms and Numerical Analysis](https://mlsanity.com/paper/2503.01873/) (2025)

## Abstract (arXiv)

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.