For five years following the introduction of the Transformer, self-attention was treated purely as an algebraic computation: $text{Softmax}(mathbf{Q}mathbf{K}^T / sqrt{d})mathbf{V}$. In practice, however, standard attention implementations scaled catastrophically not because of arithmetic complexity ($O(N^2)$ FLOPs), but because of Memory IO. In 2022, Tri Dao and Christopher Ré introduced FlashAttention, restructuring the attention algorithm from first principles to be hardware-aware, achieving 2x to 4x training acceleration and up to 10x memory reduction with zero approximation error.
1. The GPU Memory Hierarchy: SRAM vs HBM
A modern GPU (such as NVIDIA A100 or H100) comprises two primary memory tiers:
- High Bandwidth Memory (HBM / DRAM): Vast capacity (80 GB), but relatively slow bandwidth (1.5–3.35 TB/s) and high latency (~hundreds of clock cycles).
- Static RAM (SRAM / Shared Memory): Ultra-fast bandwidth (~19 TB/s) and tiny latency (a few clock cycles), but extremely scarce (only ~192 KB per Streaming Multiprocessor, or ~20 MB total across the entire chip).
In standard PyTorch attention, intermediate tensors are materialized in global HBM at every step:
- Load $mathbf{Q}, mathbf{K} in mathbb{R}^{N times d}$ from HBM $rightarrow$ compute $mathbf{S} = mathbf{Q}mathbf{K}^T in mathbb{R}^{N times N}$ $rightarrow$ write $mathbf{S}$ back to HBM.
- Read $mathbf{S}$ from HBM $rightarrow$ compute $mathbf{P} = text{Softmax}(mathbf{S}) in mathbb{R}^{N times N}$ $rightarrow$ write $mathbf{P}$ back to HBM.
- Read $mathbf{P}$ and $mathbf{V}$ from HBM $rightarrow$ compute $mathbf{O} = mathbf{P}mathbf{V} in mathbb{R}^{N times d}$ $rightarrow$ write $mathbf{O}$ to HBM.
For sequence length $N = 32,768$, materializing the $N times N$ attention matrix requires storing a single $32text{K} times 32text{K}$ FP16 tensor consuming 2 Gigabytes per attention head per batch item. Global memory traffic suffocates GPU performance.
2. The Online Softmax Mathematical Trick
The core algorithmic obstacle to computing attention in small SRAM tiles is the Softmax normalization denominator: $sum_{j=1}^N exp(S_{ij})$. Standard softmax requires a complete pass over the entire row to find the maximum value $m$ (for numerical stability) and the sum of exponentials $l$ before any output can be computed.
FlashAttention overcomes this using the Online Softmax algorithm (Milakov & Gimelshein, 2018). Given two blocks of unnormalized scores $x^{(1)}, x^{(2)} in mathbb{R}^B$, their local maxima are $m_1 = max(x^{(1)})$ and $m_2 = max(x^{(2)})$. The combined maximum is:
$m_{text{new}} = max(m_1, m_2)$
The local exponential sums $l_1 = sum exp(x_i^{(1)} – m_1)$ and $l_2 = sum exp(x_i^{(2)} – m_2)$ are updated seamlessly via rescale factors:
$l_{text{new}} = l_1 e^{m_1 – m_{text{new}}} + l_2 e^{m_2 – m_{text{new}}}$
Similarly, the intermediate accumulated output vector $mathbf{O}$ is rescaled on the fly:
$mathbf{O}_{text{new}} = mathbf{O}_1 left(frac{l_1 e^{m_1 – m_{text{new}}}}{l_{text{new}}}right) + mathbf{O}_2 left(frac{e^{m_2 – m_{text{new}}}}{l_{text{new}}}right)$
3. FlashAttention Tiling Algorithm
Equipped with online softmax, FlashAttention partitions the Query, Key, and Value matrices into discrete blocks that fit entirely within fast on-chip SRAM (block size $B_c times d$ and $B_r times d$):
- Load a block $mathbf{K}_j, mathbf{V}_j$ into fast SRAM.
- Load a block $mathbf{Q}_i$ into SRAM.
- Compute the local sub-matrix $mathbf{S}_{ij} = mathbf{Q}_i mathbf{K}_j^T$ entirely within SRAM registers.
- Update the running statistics $m_i, l_i$ and the running output accumulator $mathbf{O}_i$ using online softmax updates.
- Write back only the final output $mathbf{O}$ to slow global HBM once all Key-Value blocks are exhausted.
The massive $N times N$ attention matrix is never materialized in global HBM. HBM accesses drop from $O(N^2)$ to $O(N d)$, transforming attention from an IO-bound nightmare into a compute-bound operation executed at peak Tensor Core speed.
