FlashAttention 1, 2, and 3 Explained: Fast IO-Aware Exact Attention

In deep learning hardware, the most precious resource is not raw arithmetic capacity, but memory bandwidth. While GPU compute power (FLOPS) has scaled exponentially over the past decade, memory bandwidth between High Bandwidth Memory (HBM) and on-chip Static RAM (SRAM) has grown at a much slower rate.

In 2022, Tri Dao et al. introduced FlashAttention, fundamentally changing how transformers are trained and served. By restructuring the self-attention algorithm to be IO-Aware (aware of the memory hierarchy speeds between HBM and SRAM), FlashAttention computes exact attention with $2\times -5\times$ speedups while cutting memory footprint from quadratic $\mathcal{O}(N^2)$ to linear $\mathcal{O}(N)$.

---

1. The Standard Attention Memory Wall

In traditional self-attention:

$\mathbf{S}=\mathbf{Q}{\mathbf{K}}^T\in {\mathbb{R}}^{N\times N}$

$\mathbf{P}=\text{softmax}(\mathbf{S})\in {\mathbb{R}}^{N\times N}$

$\mathbf{O}=\mathbf{PV}\in {\mathbb{R}}^{N\times d}$

The IO Roundtrip Penalty

For a sequence length of $N=16,384$:

• The attention matrix $\mathbf{S}$ contains $16384\times 16384=268,435,456$ elements ($536.8\text{ MB}$ per attention head).
• Standard implementations write $\mathbf{S}$ to HBM, read $\mathbf{S}$ from HBM to compute Softmax, write $\mathbf{P}$ back to HBM, and read $\mathbf{P}$ again to multiply by $\mathbf{V}$.
• The GPU spends over 80% of its execution time waiting on memory bus read/write transfers, rather than performing useful matrix multiplications on Tensor Cores.

---

2. FlashAttention-1: Tiling and Online Softmax

FlashAttention achieves its breakthrough through two core mathematical and engineering innovations: Tiling and Online Softmax Reduction.

The Mathematics of Online Softmax

Standard Softmax requires a full pass across all elements in a row to find the maximum value $m=\max (x_1,\dots ,x_N)$ for numerical stability, followed by a second pass to compute the denominator sum $\ell =\sum e^{x_i-m}$.

The Online Softmax algorithm allows computing Softmax incrementally across blocks without seeing the entire row upfront:

1. When processing block 1, compute local max $m^{(1)}$ and local sum ${\ell }^{(1)}$.
2. When processing block 2, compute local max $m^{(2)}$. If $m^{(2)}>m^{(1)}$, the new global max is $m^{(\text{new})}=\max (m^{(1)},m^{(2)})$.
3. Rescale the previous accumulated output block by $e^{m^{(1)}-m^{(\text{new})}}$ and update the running sum:

${\ell }^{(\text{new})}=e^{m^{(1)}-m^{(\text{new})}}{\ell }^{(1)}+e^{m^{(2)}-m^{(\text{new})}}{\ell }^{(2)}$

${\mathbf{O}}^{(\text{new})}=\text{diag}\left(e^{m^{(1)}-m^{(\text{new})}}\right){\mathbf{O}}^{(1)}+{\mathbf{P}}^{(2)}{\mathbf{V}}^{(2)}$

This allows the entire attention operation to be computed in a single fused CUDA kernel pass, reading $\mathbf{Q},\mathbf{K},\mathbf{V}$ once from HBM and writing $\mathbf{O}$ once to HBM!

---

3. The Generational Evolution: FlashAttention 1 vs 2 vs 3

Feature Comparison Matrix

---

4. FlashAttention-3: Harnessing NVIDIA Hopper Silicon

FlashAttention-3 exploits the hardware features of NVIDIA Hopper (H100/H200) architecture:

1. Tensor Memory Accelerator (TMA): A dedicated hardware copy engine that transfers 2D/3D tensor tiles between global HBM and shared memory without using register files or stalling SM execution pipelines.
2. Asynchronous Ping-Pong Pipelining: While Tensor Cores multiply Block $i$, the TMA concurrently loads Block $i+1$ into alternate SRAM buffers, hiding memory latency completely.
3. WGMMA (Warpgroup Matrix Multiply and Accumulate): Groups of 4 warps (128 threads) execute unified $64\times 128\times 16$ matrix multiplications in lockstep.

---

5. Summary & Key Takeaways

1. IO-Aware Design is the Future of AI Systems: Designing algorithms to minimize slow HBM memory roundtrips yields orders-of-magnitude higher speedups than adding raw FLOPs.
2. Tiling and Online Softmax Eliminate Quadratic Memory: Computing running Softmax statistics allows full attention to execute inside small on-chip SRAM scratchpads.
3. FlashAttention-2 Achieves 70%+ GPU Utilization: Streamlining non-matrix math and optimizing warp parallelism made FA-2 the default standard in open-source inference engines.
4. FlashAttention-3 Unlocks Hopper Potential: Leveraging TMA and asynchronous ping-pong execution on H100 GPUs delivers up to 1.4 PetaFLOPS of FP8 attention throughput.