# Is Flash Attention Stable?

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

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 11 of its 11 references that are in the map.
- First posted on arXiv 5 May 2024.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Understanding the Difficulty of Training Transformers](https://mlsanity.com/paper/2004.08249/) (2020)
- [Statistical Aspects of Wasserstein Distances](https://mlsanity.com/paper/1806.05500/) (2018)
- [A Theory on Adam Instability in Large-Scale Machine Learning](https://mlsanity.com/paper/2304.09871/) (2023)
- [A Loss Curvature Perspective on Training Instability in Deep Learning](https://mlsanity.com/paper/2110.04369/) (2021)
- [Improving Deep Neural Network Random Initialization Through Neuronal Rewiring](https://mlsanity.com/paper/2207.08148/) (2022)
- [Accuracy Booster: Enabling 4-bit Fixed-point Arithmetic for DNN Training](https://mlsanity.com/paper/2211.10737/) (2022)
- [A PDE-based Explanation of Extreme Numerical Sensitivities and Edge of Stability in Training Neural Networks](https://mlsanity.com/paper/2206.02001/) (2022)

## Built on by (most cited first, 10 of 10)

- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [The Open Source Advantage in Large Language Models (LLMs)](https://mlsanity.com/paper/2412.12004/) (2024)
- [Gated KalmaNet: A Fading Memory Layer Through Test-Time Ridge Regression](https://mlsanity.com/paper/2511.21016/) (2025)
- [Characterizing and Efficiently Accelerating Multimodal Generation Model Inference](https://mlsanity.com/paper/2410.00215/) (2024)
- [Fingerprinting Inference Systems of Large Language Models](https://mlsanity.com/paper/2605.29979/) (2026)
- [UniForm: A Reuse Attention Mechanism Optimized for Efficient Vision Transformers on Edge Devices](https://mlsanity.com/paper/2412.02344/) (2024)
- [An Efficient vLLM-Based Inference Pipeline for Unified Audio Understanding and Generation](https://mlsanity.com/paper/2607.02119/) (2026)
- [FastTPS: An Optimized Method for LLM Token Phase for AI accelerators](https://mlsanity.com/paper/2607.11211/) (2026)
- [Broken Symmetry in BF16 Attention: Why FlashAttention Gradients Blow Up Late in Training](https://mlsanity.com/paper/2609.34272/) (2026)
- [Backward-State Policy Is Part of the Learning Algorithm](https://mlsanity.com/paper/2609.39813/) (2026)

## Abstract (arXiv)

Training large-scale machine learning models poses distinct system challenges, given both the size and complexity of today's workloads. Recently, many organizations training state-of-the-art Generative AI models have reported cases of instability during training, often taking the form of loss spikes. Numeric deviation has emerged as a potential cause of this training instability, although quantifying this is especially challenging given the costly nature of training runs. In this work, we develop a principled approach to understanding the effects of numeric deviation, and construct proxies to put observations into context when downstream effects are difficult to quantify. As a case study, we apply this framework to analyze the widely-adopted Flash Attention optimization. We find that Flash Attention sees roughly an order of magnitude more numeric deviation as compared to Baseline Attention at BF16 when measured during an isolated forward pass. We then use a data-driven analysis based on the Wasserstein Distance to provide upper bounds on how this numeric deviation impacts model weights during training, finding that the numerical deviation present in Flash Attention is 2-5 times less significant than low-precision training.