FlashAttention and online softmax

Before this lesson: the roofline model (lesson 2) and tiling and fusion (lesson 3). Attention with queries, keys and values (deep learning 9; the matrix form is in linear algebra 15). The log-sum-exp trick (subtract the max before exp; deep learning 2). Time: about 80 minutes, plus the notebook.

---

Why this matters #

Attention is the only transformer operation whose cost grows with $N^2$, where $N$ is the sequence length. At $N = 16{,}384$ tokens, one head's score matrix in bf16 is $16{,}384^2 \cdot 2$ bytes $= 512$ MiB, and there is one per head, per layer and per sequence. FlashAttention (Dao et al. 2022) computes exact attention tile by tile and never stores that matrix. PyTorch's scaled_dot_product_attention uses FlashAttention-2 kernels on recent NVIDIA GPUs, and LLM inference engines build on the same kernels.

Its key formula, the merge of two partial softmax results, appears again and again: in split-KV decoding (this lesson and lesson 12), in ring attention across GPUs (lesson 10), and in any softmax over a long axis. For science, $O(N)$ memory makes long token sequences trainable: global weather fields cut into patches, long seismic records, genomes.

1. Predict first #

2. The picture #

One query, one pass #

Start with the simplest case: one query, and a number (not a vector) as the value of each key. The query has a score $x_i$ for each key $i = 1, \dots, N$. We want the softmax-weighted average of the values:

Why subtract the max. In floating point, $e^{x}$ overflows quickly: $e^{89} > 3.4 \cdot 10^{38}$, the largest fp32 number. So code first subtracts $m = \max_i x_i$. Every exponent $x_i - m$ is then $\le 0$, every $e^{x_i - m}$ lies in $(0, 1]$, and nothing overflows. The factor $e^{-m}$ appears in the numerator and in the denominator, so it cancels:

The problem. $m$ is known only after the last score. So the simple "safe softmax" reads the scores three times: pass 1 finds $m$, pass 2 sums $e^{x_i - m}$, pass 3 forms the weighted sum. On a GPU every pass over a long row is a pass over HBM.

The online idea. Read the scores in blocks, and keep three numbers:

• $m$: the largest score seen so far;
• $\ell$: the sum $\sum_\text{seen} e^{x_i - m}$, relative to the current $m$;
• $\tilde o$: the sum $\sum_\text{seen} e^{x_i - m} v_i$, also relative to the current $m$.

When a new block raises the max from $m_\text{old}$ to $m_\text{new}$, the old sums are measured against the wrong max. Fix them with one multiplication. Each old term $e^{x_i - m_\text{old}}$ times $e^{m_\text{old} - m_\text{new}}$ is $e^{x_i - m_\text{new}}$, exactly the term we want. So multiply $\ell$ and $\tilde o$ by the rescale factor $e^{m_\text{old} - m_\text{new}} \le 1$, then add the new block's terms. At the end, $y = \tilde o / \ell$.

Small numbers. Scores $(1, 2)$ arrive first: $m = 2$ and $\ell = e^{1-2} + e^{2-2} = 1.368$. Then the score $4$ arrives: $m = 4$, the factor is $e^{2-4} = 0.135$, and the rescaled old sum is $1.368 \cdot 0.135 = 0.185 = e^{1-4} + e^{2-4}$. Adding the new term $e^{4-4} = 1$ gives $\ell = 1.185$, which is $\sum_i e^{x_i - 4}$ over all three scores, as if we had known the max from the start.

Interactive figure: loading…

The GPU picture: keep the tile on chip #

Now use the real sizes. $Q, K, V$ are $N \times d$ matrices: $N$ tokens, head dimension $d$. The output is $O = \operatorname{softmax}(\tau Q K^\top) V$ with $\tau = 1/\sqrt{d}$, the softmax taken row by row. Recall from lesson 1 that HBM is large and slow (80 GB, 3.35 TB/s on an H100 SXM), and that each SM has a small, fast on-chip memory (up to 228 KB of shared memory plus a 256 KB register file).

Standard attention runs three kernels, and each one goes through HBM:

Step	Reads from HBM	Writes to HBM
$S = \tau Q K^\top$	$Q, K$	$S$ ($N \times N$)
$P = \operatorname{softmax}(S)$	$S$	$P$ ($N \times N$)
$O = P V$	$P, V$	$O$

FlashAttention fuses the three steps into one kernel. A thread block owns $B_r$ query rows. It loads their $Q$ tile once into on-chip memory and keeps it there. Then it streams tiles of $B_c$ keys and values through on-chip memory. For each $K, V$ tile it computes the $B_r \times B_c$ score tile, updates the online state $(m, \ell, \tilde o)$ of each of its rows, and throws the score tile away. At the end it divides $\tilde o$ by $\ell$, writes its rows of $O$, and writes one extra number per row: the logsumexp $\mathrm{lse}_i = m_i + \log \ell_i$, which the backward pass needs.

Interactive figure: loading…

3. Definitions and main results #

The online softmax theorem #

Theorem (online softmax). Take one query row. Its scores $x \in \mathbb{R}^N$ arrive in blocks $x^{(1)}, \dots, x^{(T)}$, with value rows $v_i \in \mathbb{R}^d$. Start with $m_0 = -\infty$, $\ell_0 = 0$, $\tilde o_0 = 0$. For block $j = 1, \dots, T$:

Then $\tilde o_T / \ell_T = \sum_i \operatorname{softmax}(x)_i\, v_i$ exactly, and $m_T + \log \ell_T = \mathrm{lse}(x)$.

Proof (by induction on the blocks).

1. Claim after block $j$: $m_j$ is the max of all scores seen so far, $\ell_j = \sum_{\text{seen}} e^{x_i - m_j}$ and $\tilde o_j = \sum_{\text{seen}} e^{x_i - m_j} v_i$.
2. Base case $j = 1$: $m_1$ is the max of block 1. The factor $e^{m_0 - m_1} = e^{-\infty} = 0$ removes the empty start values, so $\ell_1$ and $\tilde o_1$ are the block-1 sums relative to $m_1$.
3. Step $j - 1 \to j$: the max of the old max and the new block's max is the max so far. By the claim for $j - 1$, $\ell_{j-1} = \sum_\text{old} e^{x_i - m_{j-1}}$. Multiplying by $e^{m_{j-1} - m_j}$ turns each old term into $e^{x_i - m_j}$. Adding the new block's terms gives the sum over all seen scores, relative to $m_j$. The same argument works for $\tilde o_j$, term by term, because each term of $\tilde o$ is a term of $\ell$ times a fixed vector $v_i$.
4. End: after block $T$ every score has been seen and $m_T = m = \max_i x_i$. So $\tilde o_T / \ell_T = \sum_i e^{x_i - m} v_i \big/ \sum_i e^{x_i - m}$. The factor $e^{-m}$ cancels, and this is $\sum_i \operatorname{softmax}(x)_i v_i$, with no approximation.
5. No overflow: $x_i \le m_j$ and $m_{j-1} \le m_j$, so every exponent is $\le 0$ and every exponential lies in $[0, 1]$.
6. Logsumexp: $m_T + \log \ell_T = \log\big(e^{m_T} \sum_i e^{x_i - m_T}\big) = \log \sum_i e^{x_i} = \mathrm{lse}(x)$. $\blacksquare$

Why the rescale is exact. The factor depends only on the two maxima, not on $i$. So it is the same for every old term, and one multiplication fixes the whole sum. This is the only property the proof needs (see the quiz at the end of this section for a variant where it fails).

Keep $\tilde o$ unnormalised. FlashAttention-2 keeps $\tilde o$ as a sum and divides by $\ell$ once, at the end. The first FlashAttention kept a normalised output and rescaled it at every block, which costs extra non-matmul FLOPs.

The tiled forward pass. For one thread block with query rows $Q_i$ ($B_r \times d$):

1. Load $Q_i$ into on-chip memory. Set $m = -\infty$, $\ell = 0$, $\tilde o = 0$ for each row.
2. For each key tile $j$: load $K_j, V_j$ ($B_c \times d$ each). Compute the score tile $S_{ij} = \tau Q_i K_j^\top$ ($B_r \times B_c$) on chip.
3. Update each row with the theorem: new max, rescale factor, $\ell$ and $\tilde o$ ($\tilde o \leftarrow e^{m_\text{old} - m}\tilde o + e^{S_{ij} - m} V_j$, a small matmul).
4. After the last tile: write $O_i = \tilde o / \ell$ and $\mathrm{lse}_i = m + \log \ell$ to HBM.

Steps 2 and 3 are two small matmuls ($Q_i K_j^\top$ and $\tilde P_{ij} V_j$) plus elementwise work, so the tensor cores do most of the arithmetic.

Merging two partial results #

Split the keys of one query into two disjoint sets $A$ and $B$. Suppose a thread block (or a GPU) has computed attention over $A$ alone: the normalised output and the logsumexp

and the same for $B$. The merge formula gives attention over all keys:

Proof.

1. $e^{\mathrm{lse}_A} + e^{\mathrm{lse}_B} = \sum_{i \in A} e^{x_i} + \sum_{i \in B} e^{x_i} = \sum_{i \in A \cup B} e^{x_i}$, because $A$ and $B$ are disjoint. Take the log: $\mathrm{lse}$ is the logsumexp over all keys.
2. $e^{\mathrm{lse}_A - \mathrm{lse}}\, o_A = \sum_{i \in A} e^{\mathrm{lse}_A - \mathrm{lse}}\, e^{x_i - \mathrm{lse}_A} v_i = \sum_{i \in A} e^{x_i - \mathrm{lse}}\, v_i$. The same holds for $B$.
3. Add the two: $o = \sum_{i \in A \cup B} e^{x_i - \mathrm{lse}}\, v_i = \sum_i \operatorname{softmax}(x)_i\, v_i$. $\blacksquare$

The two weights $w_A = e^{\mathrm{lse}_A - \mathrm{lse}}$ and $w_B = e^{\mathrm{lse}_B - \mathrm{lse}}$ are positive and add up to 1: $w_A$ is the share of the total softmax mass that lies in $A$. So the merge is a weighted average, not a plain average and not a sum.

Compute it stably. Write $\mathrm{lse} = \max(\mathrm{lse}_A, \mathrm{lse}_B) + \log\big(1 + e^{-|\mathrm{lse}_A - \mathrm{lse}_B|}\big)$ (NumPy: np.logaddexp). An empty or fully masked set has $\mathrm{lse} = -\infty$ and weight 0. If both sets are empty, return $o = 0$ and $\mathrm{lse} = -\infty$. Never compute $-\infty - (-\infty)$: it is nan, and one nan spreads through the whole output.

Interactive figure: loading…

HBM traffic #

Theorem (Dao et al. 2022, Theorem 2). Let $d \le M \le Nd$. Standard attention needs $\Theta(Nd + N^2)$ HBM accesses. FlashAttention needs $\Theta(N^2 d^2 / M)$ HBM accesses.

Standard attention, counted. Per head, from the table in section 2:

• $S = \tau QK^\top$: read $Q$ and $K$ ($2Nd$), write $S$ ($N^2$);
• $P = \operatorname{softmax}(S)$: read $S$, write $P$ ($2N^2$);
• $O = PV$: read $P$ and $V$ ($N^2 + Nd$), write $O$ ($Nd$).

Total: $4N^2 + 4Nd = \Theta(Nd + N^2)$. A mask or dropout adds more passes over $N \times N$ matrices.

FlashAttention, counted. Use the FlashAttention-2 loop order: the outer loop runs over the $N / B_r$ query tiles. Each query tile reads its $Q$ rows once, reads all of $K$ and $V$ ($2Nd$ elements), and writes its $O$ rows once. Over all tiles:

plus $N$ for the lse values, which we ignore. The query tile is the reuse factor: each $K, V$ element read from HBM is used by $B_r$ query rows. This is the same idea as the tile size $T$ of the tiled matmul in lesson 3.

How large can $B_r$ be? The tiles must fit on chip. Up to a constant, they are $B_r \times d$ matrices, so the largest $B_r$ that fits is a constant times $M/d$. Put $B_r = \Theta(M/d)$ into the count: $2N^2 d / B_r = \Theta(N^2 d^2 / M)$. That is where $d^2/M$ comes from.

The on-chip budget, and the original loop order

On chip, one thread block holds $Q_i$ and $\tilde o_i$ ($B_r \times d$ each), $K_j$ and $V_j$ ($B_c \times d$ each), the score tile ($B_r \times B_c$) and $m, \ell$ ($B_r$ each). Choose $B_r = \lfloor M/(6d) \rfloor$ (at least 1) and $B_c = \min(B_r, d)$. Then, for $M \ge 6d$ and $d \ge 8$:

1. $Q_i$ and $\tilde o_i$ take $2B_r d \le M/3$ elements;
2. $K_j$ and $V_j$ take $2B_c d \le 2B_r d \le M/3$;
3. the score tile takes $B_r B_c \le B_r d \le M/6$;
4. $m$ and $\ell$ take $2B_r \le M/(3d) \le M/24$.

The sum is at most $0.875\,M$, so everything fits. This choice is $B_r = \Theta(M/d)$, and the traffic is $2Nd + 2N^2 d/B_r = \Theta(N^2 d^2/M)$. The $2Nd$ term is the smaller one, because $M \le Nd$ gives $N^2 d^2 / M \ge Nd$. Real kernels choose the tile sizes by measurement: FlashAttention-2 takes $B_r$ and $B_c$ from $\{64, 128\}$, depending on $d$ and the shared memory of the GPU.

The first FlashAttention paper loops the other way: the outer loop runs over $K, V$ tiles of size $B_c = \lceil M/(4d) \rceil$, the inner loop over query tiles of size $B_r = \min(\lceil M/(4d)\rceil, d)$. Then $K$ and $V$ are read once, but $Q$, $O$, $m$ and $\ell$ are read and written once per $K, V$ tile, that is $N/B_c = \Theta(Nd/M)$ times. Each pass moves $\Theta(Nd)$ elements, so the total is again $\Theta(N^2 d^2 / M)$.

Dao et al. also prove a lower bound (their Proposition 3): no exact attention algorithm needs only $o(N^2 d^2 / M)$ HBM accesses for all $M$ in $[d, Nd]$. So, up to constants, the tiled order is optimal.

Compare the two. For $N \gg d$ the ratio standard / FlashAttention is about $4N^2 / (2N^2 d / B_r) = 2B_r / d$. With $B_r = 128$: 4 for $d = 64$, and 2 for $d = 128$.

The roofline view (lesson 2). Both versions do the same forward FLOPs: $QK^\top$ and $PV$ cost $2N^2 d$ each, $4N^2 d$ in total. In bf16 (2 bytes per element) the arithmetic intensity is about

• standard: $\dfrac{4N^2 d}{2 \cdot 4N^2} = \dfrac{d}{2}$ FLOP/byte (32 for $d = 64$);
• FlashAttention: $\dfrac{4N^2 d}{2 \cdot 2N^2 d / B_r} = B_r$ FLOP/byte (128 for $B_r = 128$).

The H100 ridge point is $989 / 3.35 \approx 295$ FLOP/byte, so standard attention is deep in the memory-bound region. The count for FlashAttention is pessimistic: thread blocks that run at the same time usually work on the same head, so many $K, V$ re-reads hit the 50 MB L2 cache instead of HBM. The real HBM traffic is lower, and well-tuned kernels are close to compute-bound.

Interactive figure: loading…

Backward pass: recompute instead of store #

The backward pass needs $P$. Standard autograd stores it: $N^2$ numbers per head. FlashAttention stores only $O$ and the per-row $\mathrm{lse}_i$, and recomputes $P$ tile by tile. With $dO$ the gradient of the loss with respect to $O$, for each tile:

1. recompute $S_{ij} = \tau Q_i K_j^\top$, then $P_{ij} = e^{S_{ij} - \mathrm{lse}_i}$ (no max and no sum needed: the lse already contains both, since $e^{s - m}/\ell = e^{s - (m + \log \ell)}$);
2. $dV_j \mathrel{+}= P_{ij}^\top dO_i$ and $dP_{ij} = dO_i V_j^\top$;
3. $dS_{ij} = P_{ij} \circ (dP_{ij} - D_i)$, with $D_i = \sum_k (dO \circ O)_{ik}$ (one number per row). This is the softmax backward (linear algebra 16); the trick $D_i = \sum_j P_{ij}\, dP_{ij} = dO_i \cdot O_i$ avoids a sum over the whole row;
4. $dQ_i \mathrel{+}= \tau\, dS_{ij} K_j$ and $dK_j \mathrel{+}= \tau\, dS_{ij}^\top Q_i$.

FLOPs. The standard backward has four matmuls ($dV$, $dP$, $dQ$, $dK$), $2N^2 d$ each, so $8N^2 d$. FlashAttention adds the recomputation of $S$: $10N^2 d$, which is 25% more.

Time. It is still much faster, because attention is memory-bound and FlashAttention removes most of the memory traffic. Dao et al. (2022, Figure 2) measured forward plus backward for GPT-2 medium (sequence length 1024, head dimension 64, 16 heads, batch 64) on an A100:

	Standard	FlashAttention
GFLOPs	66.6	75.2
HBM reads and writes	40.3 GB	4.4 GB
run time	41.7 ms	7.3 ms

13% more FLOPs, 9 times fewer bytes, 5.7 times faster: the run time follows the bytes.

What FlashAttention-2 changed #

FlashAttention-2 (Dao 2023) keeps the algorithm and changes how the work is split:

1. Fewer non-matmul FLOPs. It keeps $\tilde o$ unnormalised, divides once at the end, and stores only the lse. This matters because on an A100 a bf16 matmul runs at 312 TFLOP/s but non-matmul fp32 work at 19.5 TFLOP/s: each non-matmul FLOP costs 16 times more.
2. Parallelism over the sequence. FlashAttention-1 used one thread block per (sequence, head). With long sequences the batch is small, and batch × heads can be smaller than the 108 SMs of an A100. FlashAttention-2 gives each query tile its own thread block, so long sequences also fill the GPU.
3. Work split across warps. Inside a thread block, FlashAttention-1 split $K$ and $V$ across the 4 warps, so the warps had to add their partial outputs through shared memory. FlashAttention-2 splits the query rows across warps. All warps read the same $K, V$ tiles and need no communication.

Result: about 2 times faster than FlashAttention-1, at 50–73% of the A100 peak (up to 73% in the forward pass, up to 63% in the backward pass).

Split-KV for decoding (flash-decoding) #

During decoding (lesson 12) each sequence has one new query per step. Parallelism over query tiles then gives only batch × heads thread blocks. With batch 1 and 32 heads, that is 32 blocks for the 132 SMs of an H100: at most 32 SMs work. Decode attention is memory-bound (it reads the whole KV cache, lesson 2), and one SM cannot pull a large share of the HBM bandwidth. So idle SMs mean wasted bandwidth.

Split-KV cuts each sequence's KV cache into $k$ chunks and gives each chunk its own thread block. Each block returns a partial $(o, \mathrm{lse})$ for its chunk. A small second kernel merges the $k$ partials of each row with the merge formula. The extra work is tiny ($k$ vectors per row), and now batch × heads × $k$ blocks fill the GPU. Dao, Haziza, Massa and Sizov (2023) report up to 8 times faster decoding for very long sequences.

Interactive figure: loading…

Answers to section 1. (1) Yes. Keep the running max $m$, the running sum $\ell$ and the running weighted sum $\tilde o$; rescale both sums by $e^{m_\text{old} - m_\text{new}}$ when the max grows, and output $\tilde o / \ell$ at the end. Three numbers per query (with vector values: $m$, $\ell$ and a $d$-vector). (2) Attention is memory-bound, and FlashAttention removes the $N^2$ traffic; the run time follows bytes, not FLOPs. It stores $O$ and one logsumexp per query row.

Causal masks and FlashAttention-3

Causal masking. With a causal mask, query $i$ sees only keys $j \le i$. Cut $S$ into $T \times T$ tiles with $T = N/B$ (square tiles, $B_r = B_c = B$). Tiles fully above the diagonal contain only masked scores: the kernel skips them and does not even load their $K, V$. Tiles on the diagonal are computed with the mask; tiles below it need no mask. So $T(T+1)/2$ of the $T^2$ tiles are computed, about half for large $T$ (for $N = 16{,}384$ and $B = 128$: 8,256 of 16,384 tiles). Causal attention therefore costs about half the FLOPs and half the $K, V$ traffic.

FlashAttention-3 (Shah et al. 2024) targets Hopper GPUs. FlashAttention-2 reaches only about 35% of the H100 peak. FlashAttention-3 uses the asynchronous Hopper hardware (the TMA copy engine and warp specialization, so that loads, matmuls and the softmax overlap) and fp8 tensor cores. It reports up to 740 TFLOP/s in fp16 (75% of the H100 peak) and close to 1.2 PFLOP/s in fp8. One reason the overlap matters: an H100 SXM5 does 989 TFLOP/s of fp16 matmul but only 3.9 TFLOP/s of special functions such as the exponential, so the exp of the softmax is not free.

4. Worked example #

(a) Online softmax by hand #

Scores $x = (\ln 2,\ 0,\ \ln 8,\ \ln 4)$, so $e^{x} = (2, 1, 8, 4)$: whole numbers you can follow. Values $v = (1, 2, 3, 4)$. Two blocks of two scores.

Block 1, scores $(\ln 2, 0)$:

1. $m_1 = \max(-\infty, \ln 2, 0) = \ln 2$. The factor $e^{m_0 - m_1} = e^{-\infty} = 0$ clears the start values.
2. $\ell_1 = e^{\ln 2 - \ln 2} + e^{0 - \ln 2} = 1 + \tfrac12 = 1.5$.
3. $\tilde o_1 = 1 \cdot 1 + \tfrac12 \cdot 2 = 2$.

Block 2, scores $(\ln 8, \ln 4)$:

1. $m_2 = \max(\ln 2, \ln 8, \ln 4) = \ln 8$. The rescale factor is $e^{\ln 2 - \ln 8} = 2/8 = \tfrac14$.
2. $\ell_2 = \tfrac14 \cdot 1.5 + e^{\ln 8 - \ln 8} + e^{\ln 4 - \ln 8} = 0.375 + 1 + 0.5 = 1.875$.
3. $\tilde o_2 = \tfrac14 \cdot 2 + 1 \cdot 3 + \tfrac12 \cdot 4 = 0.5 + 3 + 2 = 5.5$.

step	$m$	factor $e^{m_\text{old} - m}$	$\ell$	$\tilde o$
block 1: $(\ln 2, 0)$	$\ln 2$	0 (start)	1.5	2
block 2: $(\ln 8, \ln 4)$	$\ln 8$	$\tfrac14$	1.875	5.5

• Output: $\tilde o_2 / \ell_2 = 5.5 / 1.875 = 44/15 \approx 2.933$.
• Direct check: $(2 \cdot 1 + 1 \cdot 2 + 8 \cdot 3 + 4 \cdot 4) / (2 + 1 + 8 + 4) = 44/15$. ✓
• Logsumexp: $m_2 + \log \ell_2 = \ln 8 + \ln 1.875 = \ln 15$, and $\log(2 + 1 + 8 + 4) = \ln 15$. ✓

(b) The same result by a merge #

Treat the two blocks as two key sets, computed separately.

• Set $A$ (keys 1, 2): $\mathrm{lse}_A = \ln(2 + 1) = \ln 3$ and $o_A = (2 \cdot 1 + 1 \cdot 2)/3 = 4/3$.
• Set $B$ (keys 3, 4): $\mathrm{lse}_B = \ln(8 + 4) = \ln 12$ and $o_B = (8 \cdot 3 + 4 \cdot 4)/12 = 10/3$.
• Merge: $\mathrm{lse} = \ln(3 + 12) = \ln 15$, weights $w_A = 3/15 = 0.2$ and $w_B = 12/15 = 0.8$.
• $o = 0.2 \cdot \tfrac43 + 0.8 \cdot \tfrac{10}{3} = \tfrac{4}{15} + \tfrac{40}{15} = 44/15$. ✓

The plain average $(4/3 + 10/3)/2 = 7/3 \approx 2.333$ would be wrong.

(c) At scale #

$N = 16{,}384$ tokens (for example a 2-D field cut into $128 \times 128$ patches), $d = 64$, 16 heads, bf16, one H100 SXM (989 TFLOP/s dense bf16, 3.35 TB/s).

1. Memory. $S$ is $N^2 \cdot 2$ B $= 512$ MiB per head, 8 GiB for 16 heads, for each layer and sequence. FlashAttention keeps $16 \cdot N \cdot 4$ B $= 1$ MiB of fp32 lse values.
2. Compute per head. $4N^2 d = 6.87 \cdot 10^{10}$ FLOP, so at least $6.87 \cdot 10^{10} / 989 \cdot 10^{12} = 69$ µs.
3. Standard traffic. $4N^2 + 4Nd = 1.08 \cdot 10^9$ elements $= 2.16$ GB, so at least $2.16 \cdot 10^9 / 3.35 \cdot 10^{12} = 644$ µs. Memory-bound: 9 times the compute time.
4. FlashAttention traffic, $B_r = 128$, no cache. $2Nd + 2N^2 d / B_r = 2.71 \cdot 10^8$ elements $= 541$ MB, so at least 162 µs. That is 4 times less ($= 2B_r/d$).
5. With L2. $K$ and $V$ of one head are $2 \cdot N d \cdot 2$ B $= 4$ MiB, so they fit in the 50 MB L2. If L2 serves every re-read, the HBM traffic falls to about $4Nd$ elements (read $Q, K, V$, write $O$): 8.4 MB, 2.5 µs. The kernel is then compute-bound.

5. In code #

NumPy checks the hand computation and a full tiled forward pass, with and without a causal mask. The notebook has these cells and the exercises.

import numpy as np
from scipy.special import logsumexp

# (a) the online softmax of the worked example, two blocks of two
x, v = np.log([2., 1., 8., 4.]), np.array([1., 2., 3., 4.])
m, l, o = -np.inf, 0.0, 0.0
for blk in (slice(0, 2), slice(2, 4)):
    m_new = max(m, x[blk].max())
    f = np.exp(m - m_new)                                  # rescale factor (0 at the start)
    l = f * l + np.exp(x[blk] - m_new).sum()
    o = f * o + (np.exp(x[blk] - m_new) * v[blk]).sum()
    m = m_new
    print(round(f, 4), l, o)                               # 0.0 1.5 2.0, then 0.25 1.875 5.5
print(o / l, 44 / 15, np.isclose(m + np.log(l), np.log(15)))   # 2.9333... 2.9333... True

# (b) a tiled forward pass: Q tile on chip, K/V tiles streamed, O(N) extra memory
def flash_forward(Q, K, V, Br=16, Bc=16, causal=False):
    N, d = Q.shape
    O, lse = np.zeros_like(Q), np.empty(N)
    for i in range(0, N, Br):
        q = Q[i:i + Br]; r = len(q)
        m, l, acc = np.full(r, -np.inf), np.zeros(r), np.zeros((r, d))
        for j in range(0, N, Bc):
            if causal and j > i + r - 1:
                break                                      # tile above the diagonal: skip
            s = q @ K[j:j + Bc].T / np.sqrt(d)
            if causal:
                rows, cols = np.arange(i, i + r)[:, None], np.arange(j, j + s.shape[1])
                s = np.where(cols <= rows, s, -np.inf)
            m_new = np.maximum(m, s.max(axis=1))
            p, f = np.exp(s - m_new[:, None]), np.exp(m - m_new)
            l = f * l + p.sum(axis=1)
            acc = f[:, None] * acc + p @ V[j:j + Bc]
            m = m_new
        O[i:i + r], lse[i:i + r] = acc / l[:, None], m + np.log(l)
    return O, lse

def naive(Q, K, V, causal=False):
    S = Q @ K.T / np.sqrt(Q.shape[1])                      # the full N x N matrix
    if causal:
        S = np.where(np.tri(len(Q), dtype=bool), S, -np.inf)
    lse = logsumexp(S, axis=1)
    return np.exp(S - lse[:, None]) @ V, lse

rng = np.random.default_rng(0)
Q, K, V = (rng.normal(size=(100, 32)) for _ in range(3))   # N = 100 is not a multiple of 16
for causal in (False, True):
    (O, lse), (O2, lse2) = flash_forward(Q, K, V, causal=causal), naive(Q, K, V, causal)
    print(causal, np.abs(O - O2).max() < 1e-12, np.abs(lse - lse2).max() < 1e-12)
# False True True
# True True True

PyTorch has a fused attention kernel. A short check (CPU, float64) shows the two facts you rely on: the result is exact, and autograd keeps much less memory for the backward pass.

import math
import torch
import torch.nn.functional as F

torch.manual_seed(0)
N, d = 1000, 64
# shapes (batch, heads, N, d)
q, k, v = (torch.randn(1, 1, N, d, dtype=torch.float64) for _ in range(3))
O_std = torch.softmax(q @ k.transpose(-1, -2) / math.sqrt(d), dim=-1) @ v  # builds S, P
print(torch.allclose(F.scaled_dot_product_attention(q, k, v), O_std))      # True

def saved_bytes(fn, *args):
    """Bytes of the distinct tensors autograd keeps for the backward pass."""
    seen = {}
    def pack(t):
        seen[t.untyped_storage().data_ptr()] = t.untyped_storage().nbytes()
        return t
    with torch.autograd.graph.saved_tensors_hooks(pack, lambda t: t):
        fn(*args)
    return sum(seen.values())

x = [t.float().requires_grad_() for t in (q, k, v)]       # fp32: 4 bytes per element
naive_t = lambda q, k, v: torch.softmax(q @ k.transpose(-1, -2) / math.sqrt(d), dim=-1) @ v
print(saved_bytes(naive_t, *x), saved_bytes(F.scaled_dot_product_attention, *x))
# 4768000 1028000

The naive version keeps $Q, K, V$ (768,000 bytes) and the $1000 \times 1000$ matrix $P$ (4,000,000 bytes). The fused kernel keeps $Q, K, V$, $O$ (256,000 bytes) and the 1,000 lse values (4,000 bytes). The gap grows as $N^2$. (Checked with PyTorch 2.x on a CPU; which kernel PyTorch picks depends on the device and the inputs.)

6. Where it shows up #

• Training. torch.nn.functional.scaled_dot_product_attention dispatches to FlashAttention-2 kernels on recent NVIDIA GPUs (since PyTorch 2.2). Long-context training uses causal and sliding-window masks, both with tile skipping. The activation memory of attention drops from $O(N^2)$ to $O(N)$ (lesson 6).
• Inference. Prefill uses the same kernel. Decoding uses split-KV and paged KV caches (lesson 12).
• Across GPUs. Ring attention (lesson 10) sends $K, V$ blocks around a ring of GPUs; each GPU merges the partial results with exactly the merge formula of this lesson.
• ML for science. A global 0.25° grid has $721 \times 1440 \approx 10^6$ points per level. Even after patching, weather and climate transformers see tens of thousands of tokens; $O(N)$ memory makes this trainable. The same holds for long seismic waveforms and genome sequences.
• Beyond attention. Any softmax over a long axis can be streamed this way, for example the cross-entropy loss over a large vocabulary: process the vocabulary in chunks and keep $(m, \ell)$.

7. Common mistakes #

• Rescaling $\ell$ but not $\tilde o$ when the max grows. Both sums are relative to $m$; both need the factor $e^{m_\text{old} - m_\text{new}}$.
• Normalising each tile and adding (or averaging) the tiles. Partial results need the weights $e^{\mathrm{lse}_A - \mathrm{lse}}$; a plain sum or average is wrong.
• nan from fully masked rows or tiles. $e^{-\infty - (-\infty)}$ is nan. Guard the case $m = -\infty$ (and $\mathrm{lse} = -\infty$ in a merge).
• "FlashAttention approximates attention." It is exact. Only the order of the floating-point operations changes, so results agree up to rounding.
• "FlashAttention reduces FLOPs." The forward FLOPs are the same, and the backward does more. It reduces memory traffic and memory use.
• Keeping $m$, $\ell$, $\mathrm{lse}$ or $\tilde o$ in bf16. Keep them in fp32. bf16 has only about 3 significant decimal digits (lesson 5), and the sums collect many terms.
• Expecting the training speed-up at decode. With one query there is no $N \times N$ matrix to avoid; decoding needs split-KV to use all SMs.

8. Check yourself #

1. The running state of one query is $m = 3$, $\ell = 2$, $\tilde o = 5$ (scalar values). A block with scores $(5, 3)$ and values $(1, 0)$ arrives. Give the new $m$, $\ell$ and $\tilde o$.
   Answer

   $m = 5$ and the factor is $e^{3-5} = e^{-2} \approx 0.135$. $\ell = 2e^{-2} + e^{5-5} + e^{3-5} = 1 + 3e^{-2} \approx 1.406$. $\tilde o = 5e^{-2} + 1 \cdot 1 + e^{-2} \cdot 0 = 1 + 5e^{-2} \approx 1.677$. The output so far is $\tilde o/\ell \approx 1.193$.
2. Merge $(o_A, \mathrm{lse}_A) = (1, \ln 2)$ and $(o_B, \mathrm{lse}_B) = (4, \ln 6)$.
   Answer

   $\mathrm{lse} = \ln(2 + 6) = \ln 8$. The weights are $2/8$ and $6/8$, so $o = \tfrac28 \cdot 1 + \tfrac68 \cdot 4 = 3.25$.
3. $N = 32{,}768$. What does FlashAttention store per head for the backward pass, and how does it compare with storing $P$ in bf16?
   Answer

   $O$ (needed anyway, as the output) and $N$ fp32 lse values: $32{,}768 \cdot 4$ B $= 128$ KiB. $P$ would be $N^2 \cdot 2$ B $= 2$ GiB, which is 16,384 times more.
4. With $d = 128$ and $B_r = 128$, the no-cache traffic ratio standard / FlashAttention is only 2. Give two reasons why FlashAttention still wins clearly.
   Answer

   (1) The L2 cache serves many $K, V$ re-reads, so the real HBM traffic of FlashAttention is far below the no-cache count. (2) The $N \times N$ matrices disappear: memory is $O(N)$, so longer sequences and larger batches fit. Also, masks and dropout add no extra $N \times N$ passes, and one kernel replaces several.
5. Decoding on an H100 (132 SMs): batch 4, 16 heads, multi-head attention. How many thread blocks run without split-KV, and with the KV cache split into 8 chunks?
   Answer

   Without split: $4 \cdot 16 = 64$ blocks, so at most 64 of 132 SMs are busy. With 8 chunks: $64 \cdot 8 = 512$ blocks, enough to fill all SMs (about 4 waves). A small merge kernel then combines 8 partials per row.
6. Why does the backward pass not need $m$ and $\ell$ separately?
   Answer

   $P_{ij} = e^{S_{ij} - m_i} / \ell_i = e^{S_{ij} - (m_i + \log \ell_i)} = e^{S_{ij} - \mathrm{lse}_i}$. One number per row is enough, and every exponent is $\le 0$, so nothing overflows.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. Write the online softmax update for $\ell$ when a new block arrives.

   Show answer

2. Why is multiplying the old sums by $e^{m_\text{old} - m_\text{new}}$ exact?

   Show answer

3. Why does the safe softmax subtract the max?

   Show answer

4. Give the formula that merges $(o_A, \mathrm{lse}_A)$ and $(o_B, \mathrm{lse}_B)$.

   Show answer

5. Merge $(o_A, \mathrm{lse}_A) = (2, \ln 1)$ with $(o_B, \mathrm{lse}_B) = (5, \ln 3)$.

   Show answer

6. HBM traffic of standard attention and of FlashAttention, per head?

   Show answer

7. What does FlashAttention store for the backward pass, and what does it recompute?

   Show answer

8. Why is the FlashAttention backward faster despite 25% more matmul FLOPs?

   Show answer

9. What does split-KV (flash-decoding) do?

   Show answer

10. How large is one head's score matrix at $N = 16{,}384$ in bf16?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Write online softmax, or a tiled attention forward pass, in NumPy." Keep $(m, \ell, \tilde o)$ per row, rescale by $e^{m_\text{old} - m_\text{new}}$, divide once at the end, return lse. Guard fully masked rows.
• "Why is FlashAttention faster if it does the same FLOPs?" Attention is memory-bound. Fusion keeps $S$ and $P$ on chip; traffic falls from $\Theta(N^2)$ to $\Theta(N^2 d^2 / M)$, and in practice L2 serves most re-reads.
• "What is its memory cost?" $O(N)$ extra per head (the lse values) instead of $O(N^2)$.
• "How does it handle a causal mask?" It skips tiles above the diagonal and masks only the diagonal tiles; about half the work.
• "Does it help decoding?" Not by itself: there is one query and no $N \times N$ matrix. Split-KV adds parallelism over the KV length and merges the partials with lse.
• "How do you combine attention computed on different GPUs?" With the lse merge: it is associative, so partials can be merged in any order (ring attention, split-KV).

10. Practice #

The notebook for this lesson runs in Colab, with NumPy only. It has five exercises, each with a check cell and a folded solution:

1. online_softmax_stats(x, block): the running $(m, \ell)$ of a stream of scores, checked against NumPy and SciPy's logsumexp, also with a huge score and with $-\infty$ entries.
2. merge_partials(o1, lse1, o2, lse2) for many rows at once, with empty (fully masked) sets.
3. flash_attention_forward(Q, K, V, Br, Bc, causal): the tiled forward pass, checked against naive attention to $10^{-10}$, with and without a causal mask.
4. Count HBM traffic with a CountingMemory: write the formulas for standard and tiled attention, check them against the counts, and plot the ratio over $N$ and the tile size.
5. Split-KV decoding: split a KV cache into $k$ chunks, compute the partials, merge them, and check.

11. Watch and read #

How FlashAttention Accelerates Generative AI RevolutionJia-Bin Huang · 12 min

A short visual tour of tiling and the online softmax; a good first look before the papers.

FlashAttention - Tri Dao | Stanford MLSys #67Tri Dao · Stanford MLSys Seminars · 59 min

The author explains IO-awareness, tiling and recomputation, and the measured speed-ups.

Lecture 12: Flash AttentionGPU MODE · 72 min

How the algorithm maps onto a real GPU kernel, step by step.

Flash Attention derived and coded from first principles with Triton (Python)Umar Jamil · 7 h 38 min

A long, complete derivation of the forward and backward passes and a Triton implementation; watch the parts you need.

• Dao, Fu, Ermon, Rudra, Ré (2022), FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness: §3.1 (Algorithm 1, tiling and recomputation), §3.2 (Theorem 2, Proposition 3), Appendix B (backward pass), Appendix C (proofs).
• Dao (2023), FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning: §3.1 (algorithm, causal masking), §3.2 (parallelism), §3.3 (work split between warps).
• Shah et al. (2024), FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision.
• Dao, Haziza, Massa, Sizov (2023), Flash-Decoding for long-context inference.
• Milakov and Gimelshein (2018), Online normalizer calculation for softmax: the online normaliser itself.
• Triton tutorial Fused Attention: a short, readable forward and backward kernel.