Someone reads you a long list of numbers, one at a time, and asks for the softmax of the whole list at the end. You may keep only a couple of numbers in your head. What do you keep?

The obvious answer is the running sum of $e^{s_j}$, and in exact arithmetic it works. In floating point it overflows as soon as one score passes about 88 (in float32) or 11 (in float16), and the standard cure from §4.3, subtracting the maximum score first, seems to need the maximum before you start summing, which you only know at the end. So the textbook "safe" softmax reads the list three times, once for the maximum, once for the sum, and once to divide.

This section shows that two numbers are enough, a running maximum and a running sum that you rescale whenever the maximum changes. We derive the recurrence, prove it correct, extend it to the weighted sum of values that attention actually needs, and then find the property that makes FlashAttention possible. Partial softmax results computed on separate pieces of a row can be merged, exactly, in any order. That turns the softmax from a whole-row operation into one you can compute tile by tile, which is what §12.1 showed we need.

Three passes, and why that is a problem

Fix one row of scores $s_1, \dots, s_n$ (for attention, $n = T_k$, the number of keys, and $s_j$ is one query's score against key $j$). The safe softmax is

Subtracting $m$ changes nothing mathematically, since the factor $e^{-m}$ cancels between numerator and denominator, but it makes the largest exponent $e^0 = 1$, so nothing overflows and $\ell \ge 1$. Each of the three quantities is a pass over the row. If the row lives in slow memory, that is three reads of every score plus one write of the output. Milakov and Gimelshein (2018) counted these as memory accesses per element. The unsafe version (no maximum) needs 3, two loads and a store. The safe version needs 4.

For one row of a vocabulary softmax that is a 33% tax. For attention it is worse than a tax. The scores are themselves expensive to produce (each is a $d_k$-dimensional dot product), and FlashAttention never wants to store them at all. Three passes would mean computing $\mathbf{Q}\mathbf{K}^\top$ three times.

Two numbers that can be updated

Here is the idea, on the four scores $2, 5, 1, 7$. Keep a running maximum $m_j$ and a running sum $\ell_j$ of exponentials measured relative to the current maximum.

• After $s_1 = 2$, the maximum is 2 and the sum is $e^{2-2} = 1$.
• Then $s_2 = 5$ arrives and the maximum jumps to 5. The old sum was measured relative to 2. Relative to 5, the same term is $e^{2-5} = e^{2-2}\cdot e^{2-5}$, so we multiply the old sum by $e^{m_{\text{old}} - m_{\text{new}}} = e^{-3} \approx 0.0498$ and add the new term $e^{5-5} = 1$, giving $1.0498$.
• $s_3 = 1$ does not change the maximum, so we just add $e^{1-5} \approx 0.0183$.
• $s_4 = 7$ raises the maximum again. Multiply by $e^{5-7}$ and add 1.

The cell runs exactly this and checks, at every step, that the running sum equals the sum it is supposed to equal.

The two columns on the right agree at every step, and at the end $(m, \ell) = (7, 1.1446)$ gives the softmax directly. Step through it yourself in the widget, including rows that you make up, before reading the proof.

The recurrence, and a proof that it is right

Start from $m_0 = -\infty$ and $\ell_0 = 0$, and for $j = 1, \dots, n$ set

In plain words, whenever the maximum goes up, shrink everything you have summed so far by the factor the maximum went up by, then add the new term.

Theorem (Milakov and Gimelshein 2018). For every $j \ge 1$,

Proof, by induction on $j$. For $j = 1$, $m_1 = \max(-\infty, s_1) = s_1$, and $\ell_1 = 0\cdot e^{-\infty} + e^{0} = 1 = e^{s_1 - m_1}$. (Here $e^{-\infty} = 0$, and the product is $0\cdot 0$, not an indeterminate form, because $m_1$ is finite.) Now suppose both statements hold for $j - 1$. The first statement for $j$ is immediate, since the maximum of a set with one element added is the larger of the old maximum and the new element. For the second,

where the middle step uses $e^{a}e^{b} = e^{a+b}$ to combine the exponents. So the statement holds for every $j$, and in particular after the last score, $\operatorname{softmax}(\mathbf{s})_i = e^{s_i - m_n}/\ell_n$. $\blacksquare$

Two bounds come for free, and they are why this is numerically safe, not only correct. Every term $e^{s_i - m_j}$ is at most 1 (because $m_j$ is the maximum) and at least one equals 1 (the maximum itself), so

In plain words, the running sum never overflows (it is at most the number of scores), never underflows to zero (it is at least 1), and the rescaling factor only ever shrinks what is already there. (At $j = 1$ the factor is $e^{-\infty} = 0$, which discards the empty initial sum.) Each rescaling multiplies the sum by a factor computed with a relative rounding error of about $\varepsilon_{\text{mach}}$, so after $k$ rescalings the relative error grows at most like $k\,\varepsilon_{\text{mach}}$, the same order as the error of summing $n$ floating-point numbers in the first place.

The result is a softmax with two passes and 3 memory accesses per element instead of 4, a 1.33 times reduction. On a V100, Milakov and Gimelshein measured up to 1.3 times faster softmax than the safe version at large batch sizes. The cell compares the pass counts directly by wrapping the row in an object that counts how many times it is iterated.

import numpy as np

def online_normalizer(scores):
    scores = list(scores)
    m = max(scores)
    return m, float(np.sum(np.exp(np.array(scores) - m)))   # two passes, not one

From weights to outputs

Attention does not need the softmax weights themselves. It needs their weighted sum of value vectors, $\mathbf{o} = \sum_j y_j \mathbf{v}_j$, and that removes the second pass too.

Carry a third quantity, the unnormalized output $\tilde{\mathbf{o}}_j = \sum_{i\le j} e^{s_i - m_j}\mathbf{v}_i$, a vector of width $d_v$. It is measured relative to the same running maximum as $\ell_j$, so it obeys the same recurrence,

The proof is the same induction with $\mathbf{v}_i$ riding along in every term. In plain words, the output is accumulated in one pass, the running sum is accumulated beside it, and one division at the end makes it a weighted average.

There is a second way to write it that keeps the output normalized at every step, $\mathbf{o}_j = \tilde{\mathbf{o}}_j / \ell_j$. Dividing the recurrence by $\ell_j$ gives

This is the form in the original FlashAttention (line 12 of Algorithm 1 in Dao et al. 2022). It does a division at every step, which is wasted work, and FlashAttention-2 switched to the unnormalized form for exactly that reason (§12.4). Both are exact.

The last thing worth keeping is the log-sum-exp of the scores, $\operatorname{lse}(\mathbf{s}) = \log\sum_j e^{s_j} = m_n + \log \ell_n$. It is one number per query, it can be computed from the two we already have, and it is all the backward pass will need to rebuild the weights (§12.3), since $y_j = e^{s_j - \operatorname{lse}(\mathbf{s})}$.

One query, one pass over the keys and values, memory of size $d_v + 2$ no matter how many keys there are. This is the observation behind Rabe and Staats (2021), who noted that attention for a single query therefore needs only constant memory in the sequence length.

Merging partial results

The recurrence processes one score at a time. FlashAttention processes a block of keys at a time, and different GPU thread blocks, or different GPUs, may process different blocks. What we need is a way to combine summaries of two pieces of a row.

Split the keys into two disjoint sets $\mathcal{A}$ and $\mathcal{B}$, and summarize each with a triple measured relative to its own maximum,

and the same for $\mathcal{B}$. Define the merge $\oplus$ by

In plain words, bring both summaries to the common reference $m$, the larger of the two maxima, then add.

Claim. The merge of the summaries of $\mathcal{A}$ and $\mathcal{B}$ is the summary of $\mathcal{A}\cup\mathcal{B}$. The maximum part is clear. For the sum, $\ell_{\mathcal{A}}e^{m_{\mathcal{A}} - m} = \sum_{i\in\mathcal{A}} e^{s_i - m_{\mathcal{A}}}e^{m_{\mathcal{A}} - m} = \sum_{i\in\mathcal{A}} e^{s_i - m}$, and likewise for $\mathcal{B}$, so the merged sum is $\sum_{i\in\mathcal{A}\cup\mathcal{B}} e^{s_i - m}$. The output is identical with $\mathbf{v}_i$ in every term. $\blacksquare$

Check it by hand on the row $2, 5, 1, 7$ from before, split into halves. The first half $\{2, 5\}$ has summary $(5,\; e^{-3} + 1) = (5, 1.0498)$ and the second $\{1, 7\}$ has $(7,\; e^{-6} + 1) = (7, 1.0025)$. The common reference is 7, so the merged sum is $1.0498\,e^{5-7} + 1.0025\,e^{7-7} = 0.1421 + 1.0025 = 1.1446$, exactly the $\ell$ that the one-score-at-a-time recurrence produced. The order in which the scores were seen made no difference.

Three properties follow, and each one is used later in the chapter.

1. Associative and commutative. Merging the summaries of $\mathcal{A}$, $\mathcal{B}$, $\mathcal{C}$ in any grouping or order gives the summary of $\mathcal{A}\cup\mathcal{B}\cup\mathcal{C}$, because union is associative and commutative. Milakov and Gimelshein state this for the $(m, \ell)$ pair. So the pieces of a row can be combined left to right, in a balanced tree of depth $\log_2$ of the number of pieces, or as they arrive from different processors. This is the structure of a parallel reduction (§11.4).
2. An identity element. The empty set has summary $(-\infty, 0, \mathbf{0})$, and merging with it changes nothing. In code the case where both maxima are $-\infty$ needs care, because $e^{-\infty - (-\infty)}$ is NaN. That case is a fully masked block, and it happens constantly with causal masks.
3. The one-step recurrence is a special case. Processing $s_j$ is merging the summary of the prefix with the summary of the single score, $(s_j, 1, \mathbf{v}_j)$.

The same merge, in log-sum-exp form

Kernels usually return the normalized output $\mathbf{o}_{\mathcal{A}} = \tilde{\mathbf{o}}_{\mathcal{A}}/\ell_{\mathcal{A}}$ together with $\operatorname{lse}_{\mathcal{A}} = m_{\mathcal{A}} + \log\ell_{\mathcal{A}} = \log\sum_{i\in\mathcal{A}} e^{s_i}$, two things with a direct meaning. The merge is then

The derivation is short. The weight that the union's softmax puts on all of $\mathcal{A}$ together is $\sum_{i\in\mathcal{A}} e^{s_i} / \sum_{i\in\mathcal{A}\cup\mathcal{B}} e^{s_i} = e^{\operatorname{lse}_{\mathcal{A}} - \operatorname{lse}}$, and within $\mathcal{A}$ the weights are exactly those of $\mathbf{o}_{\mathcal{A}}$. In plain words, the merged output is a convex combination of the two partial outputs, weighted by how much softmax mass each piece holds, and the two weights sum to 1. The first log is computed stably as $\max + \log(e^{\operatorname{lse}_{\mathcal{A}} - \max} + e^{\operatorname{lse}_{\mathcal{B}} - \max})$, NumPy's logaddexp.

This is the form used when the pieces are computed far apart. Flash-Decoding splits one query's keys across many thread blocks (§12.4), and ring attention splits them across GPUs (§12.7). Each piece reports $(\mathbf{o}, \operatorname{lse})$, and one merge combines them.

import numpy as np

def merge(a, b):
    (ma, la, oa), (mb, lb, ob) = a, b
    return max(ma, mb), la + lb, oa + ob          # forgets to rescale

def finalize(p):
    m, l, acc = p

Blocks, the bridge to FlashAttention

Combine the two ideas and you have the inner loop of FlashAttention. Hold a block of $B_r$ queries. Stream the keys and values past them in blocks of $B_c$ at a time. For each key block, compute the block of scores (a small matrix product), summarize it per query row, and merge the summary into the running one. In matrix form, with $\mathbf{m}, \boldsymbol{\ell}$ vectors over the rows and $\tilde{\mathbf{O}}$ a matrix, one step is

where the subtraction $\mathbf{S}_j - \mathbf{m}^{\text{new}}$ broadcasts each row's maximum across that row. In plain words, this is the scalar recurrence, run for all query rows at once, a block of keys at a time, with the rescaling applied once per block rather than once per key. The expensive parts, $\mathbf{S}_j = \mathbf{Q}\mathbf{K}_j^\top/\sqrt{d_k}$ and $\tilde{\mathbf{P}}_j\mathbf{V}_j$, are matrix products that run on tensor cores. The rescaling is a few elementwise operations on a $B_r\times B_c$ tile that never leaves on-chip memory. You will write this loop in the next exercise, and §12.3 adds the outer loop over query blocks, the causal mask, and the backward pass.

import numpy as np

def blocked_attention(Q, K, V, block):
    raise NotImplementedError

Precision, and what can still go wrong

The online softmax was designed for stability, so it is worth seeing it rescue a computation that the naive formula ruins. In float16 the largest representable number is 65,504, so $e^{12} \approx 163{,}000$ overflows (§4.3).

The naive version returns zeros and NaN, and the online version gets every weight to float16 precision, with a running sum of 1.79 that never came near the limit. In practice kernels keep $\mathbf{m}$, $\boldsymbol{\ell}$, and $\tilde{\mathbf{O}}$ in float32 even when $\mathbf{Q}$, $\mathbf{K}$, and $\mathbf{V}$ are bf16 or FP8, because they are a few numbers per row and their accuracy determines the output's (§8.5).

Four details still matter in real code.

• Fully masked blocks and rows. With a causal mask most key blocks above the diagonal are entirely masked for some rows. If a row's running maximum is still $-\infty$ after such a block, $e^{m - m^{\text{new}}}$ is $e^{-\infty - (-\infty)}$, NaN. Kernels either skip blocks that are fully masked (§12.3) or substitute 0 for an infinite maximum before subtracting. A row that is masked everywhere ends with $\ell = 0$ and must be handled explicitly, as in §4.4.
• The maximum is a convenience, not a requirement. Any number at least as large as every score would prevent overflow. If you knew a bound in advance, for instance because queries and keys are normalized (QK-norm, §5.7), you could subtract it once and never rescale. The running maximum is used because it is the tightest bound available without looking ahead, and a loose bound risks underflowing every term to zero.
• Kernels use base 2. GPUs compute $2^x$ in one special-function instruction, and $e^x = 2^{x\log_2 e}$. So production kernels fold the constant $\log_2 e$ into the score scale $1/\sqrt{d_k}$ once, keep $m$ in the scaled units, and call exp2 everywhere, which saves a multiplication per score and changes nothing else, because every identity above holds for any base. The stored log-sum-exp is converted back to natural log at the end if the caller expects it.
• Rescaling costs exponentials. Each block needs $B_r$ extra exponentials for $e^{\mathbf{m} - \mathbf{m}^{\text{new}}}$ on top of the $B_rB_c$ for the scores, a fraction $1/B_c$ of extra work. That is negligible for arithmetic, but exponentials run on a slow special-function unit, which is why FlashAttention-2 and 3 work hard to keep that unit busy in parallel with the tensor cores (§12.4).

import numpy as np

def softmax_topk(scores, k):
    x = np.array(list(scores), float)
    idx = np.argsort(-x)[:k]
    return idx, np.exp(x[idx]) / np.exp(x).sum()     # unstable, and not one pass in spirit

In PyTorch

PyTorch does not expose the online softmax as a function, because it lives inside fused kernels. You can still write the block recurrence in a few lines of torch and check it against torch.softmax and against scaled_dot_product_attention, which runs the same idea in C++ on CPU and in FlashAttention on GPU.

Read more

• Milakov & Gimelshein (2018). Online normalizer calculation for softmax. The recurrence, its proof, the parallel merge, and the fused softmax plus top-k, in eight pages.
• Rabe & Staats (2021). Self-attention Does Not Need O(n^2) Memory. Uses the merge to process keys in chunks, with constant memory per query.
• Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Section 3.1 applies the online softmax to tiles, in the normalized form.
• Ye (2023). From Online Softmax to FlashAttention. Lecture notes that derive FlashAttention from the online softmax in a few clean steps.
• Chan, Golub & LeVeque (1983). Algorithms for Computing the Sample Variance: Analysis and Recommendations. The parallel merge for mean and variance, the same pattern for another statistic.

Watch next

▶

▶