Softmax and FlashAttention

Online softmax keeps a running maximum and rescales; FlashAttention uses it to never write the score matrix to HBM.

Loading the animation…

Concept

Attention computes $S = QK^\top$ (an $N \times N$ matrix of scores for a sequence of $N$ tokens), takes a softmax along each row to get $P$ , and returns $O = PV$ . Done as three separate kernels, standard attention writes $S$ to HBM, reads it back for the softmax, writes $P$ , and reads it again for the product: four passes over $N^2$ numbers. At $N = 4096$ and head dimension 64 that is 136 MB for one head, almost all of it the score matrix.

FlashAttention (Dao et al., 2022) never stores $S$ . It takes a block of $B_r$ queries, streams the keys and values past in blocks of $B_c$ , and for each pair computes the $B_r \times B_c$ score tile on chip, folds it into a running output, and discards it. The animation is FlashAttention-2's loop order (Dao, 2023): one row of tiles per query block. What goes to and from HBM is $Q$ and $O$ once, and $K$ and $V$ once per query block, between 2.1 MB (if every re-read of $K$ and $V$ hits L2, the "K and V from L2" bar) and 34.6 MB (if none does) at the same size.

The obstacle is the softmax: each output row needs $\sum_j e^{s_j}$ over the whole row before any of it can be normalised, and FlashAttention sees the row a tile at a time. The trick that makes it possible is online softmax.

Loading the animation…

Concept

A safe softmax subtracts the row's maximum before exponentiating, so nothing overflows: $p_i = e^{x_i - m} / \sum_j e^{x_j - m}$ . That seems to need two passes (find $m$ , then sum), and a third to normalise. Milakov and Gimelshein (2018) showed that one pass is enough: keep a running maximum $m$ and a running sum $\ell$ of $e^{x - m}$ . When a new block brings a larger maximum $m'$ , every term already in $\ell$ was computed against the old $m$ and is too large by $e^{m' - m}$ , the same factor for all of them, so multiplying $\ell$ by $e^{m - m'}$ fixes them all at once. Step through the animation and watch the bars below the scores: they all shrink together when the maximum rises.

FlashAttention keeps exactly these two numbers per query row, and rescales its partial output row $O_i$ by the same factor, which is why it can stream $K$ and $V$ through without ever seeing a whole row of scores.

Maths

Online softmax is exact. Let $m_k$ and $\ell_k$ be the maximum and the sum after $k$ blocks. Suppose $\ell_k = \sum_{j \le k\text{-th block}} e^{x_j - m_k}$ . After the next block, with $m_{k+1} = \max(m_k, \max_{\text{block}} x_j)$ ,

\ell_k\, e^{m_k - m_{k+1}} + \sum_{\text{block}} e^{x_j - m_{k+1}} = \sum_{\text{earlier}} e^{x_j - m_k + m_k - m_{k+1}} + \sum_{\text{block}} e^{x_j - m_{k+1}} = \sum_{\text{all so far}} e^{x_j - m_{k+1}} ,

so by induction $\ell$ ends as $\sum_j e^{x_j - m}$ with $m$ the true maximum, and $p_i = e^{x_i - m}/\ell$ is the ordinary softmax, in exact arithmetic. In floating point the two differ only by the rounding of the different order of additions: the animation's last frame shows the largest difference, of order $10^{-16}$ .

HBM traffic. With $e$ bytes per value, standard attention moves $e\,(2Nd + 4N^2 + 2Nd)$ : read $Q$ and $K$ , write and read $S$ , write and read $P$ , read $V$ , write $O$ . FlashAttention-2 reads $Q$ once ( $eNd$ ), writes $O$ once ( $eNd$ ) and reads $K$ and $V$ ( $2eNd$ ) once for each of the $N/B_r$ query blocks:

B_{\text{flash}} = e\,\Bigl(2Nd + \frac{2N^2 d}{B_r}\Bigr) .

The $N^2$ term survives, but divided by $B_r/d$ and halved, and with L2 reuse it disappears: the floor is $4eNd$ , every value read or written once. This is the paper's $\Theta(N^2 d^2 / M)$ result with $B_r \propto M/d$ for on-chip memory $M$ .

Code

The online softmax step, cut from src/lib/gpu/model.ts:

const mNew = m > bmax ? m : bmax;
const scale = Math.exp(m - mNew);
let s = 0.0;
for (const v of blk) s += Math.exp(v - mNew);
const lPrev = l;
l = l * scale + s;

And the bytes FlashAttention-2 moves per tile (no L2 reuse in moved, every re-read an L2 hit in once):

if (j === 0) {
  moved += eb * br * d;
  once += eb * br * d;
}
moved += 2 * eb * bc * d;
if (i === 0) once += 2 * eb * bc * d;

One head, forward pass (model, BF16, blocks of 128)

Sequence × head dim	Standard attention	FlashAttention, no L2 reuse	Every value once (floor)
4096 × 64	136 MB	34.6 MB	2.1 MB
4096 × 128	138 MB	69.2 MB	4.19 MB
8192 × 128	545 MB	273 MB	8.39 MB