We now have both halves of an idea. §12.1 showed that standard attention is slow because it writes the $T\times T$ matrices $\mathbf{S}$ and $\mathbf{P}$ to HBM and reads them back, and that a kernel which read $\mathbf{Q}$, $\mathbf{K}$, $\mathbf{V}$ once and wrote $\mathbf{O}$ once would be compute-bound. §12.2 showed that the softmax, the step that seemed to need a whole row at once, can be computed a block at a time with exact corrections.

Put them together and you get FlashAttention (Dao et al. 2022). This section derives it in full. We build the forward pass tile by tile, with the causal mask and the block skipping it allows. We work out what the forward pass must leave behind for the backward pass (one number per query) and derive the backward pass that recomputes everything else. Then we prove the IO complexity theorem, $\Theta(T^2d_k^2/M)$ HBM accesses for SRAM of size $M$, and its matching lower bound. Along the way we write the whole algorithm in NumPy with a counter on every byte that crosses between HBM and SRAM, and check it against naive attention to machine precision.

The picture, a grid of tiles

Fix one head. The queries $\mathbf{Q} \in \R^{T_q\times d_k}$ are cut into row blocks $\mathbf{Q}_1, \dots, \mathbf{Q}_{n_r}$ of $B_r$ rows each, and the keys and values into column blocks $\mathbf{K}_1, \dots, \mathbf{K}_{n_c}$ and $\mathbf{V}_1, \dots, \mathbf{V}_{n_c}$ of $B_c$ rows each, so $n_r = \lceil T_q/B_r\rceil$ and $n_c = \lceil T_k/B_c\rceil$. (Here a subscript on a bold capital is a block of rows, not a single row.) The score matrix becomes a grid of $n_r\times n_c$ tiles,

FlashAttention visits the tiles one at a time. For each tile it loads the blocks it needs into SRAM, computes $\mathbf{S}_{ij}$ there, uses it to update a running result, and throws it away. The tile never travels to HBM. Choose the block sizes so that a query block, a key block, a value block, the tile, and the running output all fit in one SM's SRAM, which on an A100 means blocks of 64 or 128 rows for head widths of 64 or 128.

Two orders of the loops are possible, and both are used. The original FlashAttention put the key and value blocks in the outer loop. FlashAttention-2 swapped them, putting the query blocks outside, which is simpler and parallelizes better (§12.4). We use the FlashAttention-2 order for the forward pass and come back to the original when we prove the theorem.

The forward pass

With the query block $\mathbf{Q}_i$ held in SRAM, the inner loop streams key and value blocks past it and runs the block form of the online softmax from §12.2. Keep, for the $B_r$ rows of the block, a running maximum $\mathbf{m}\in\R^{B_r}$, a running sum $\boldsymbol{\ell}\in\R^{B_r}$, and an unnormalized output $\tilde{\mathbf{O}}\in\R^{B_r\times d_v}$, initialized to $-\infty$, $0$, and $\mathbf{0}$.

1. Load $\mathbf{Q}_i$ from HBM into SRAM.
2. For each key block, $j = 1, \dots, n_c$, do the following.
   1. Load $\mathbf{K}_j$ and $\mathbf{V}_j$ into SRAM.
   2. Compute the tile $\mathbf{S}_{ij} = \mathbf{Q}_i\mathbf{K}_j^\top/\sqrt{d_k}$ (a matrix product on tensor cores).
   3. Update the maximum, $\mathbf{m}^{\text{new}} = \max(\mathbf{m}, \operatorname{rowmax}(\mathbf{S}_{ij}))$.
   4. Exponentiate, $\tilde{\mathbf{P}}_{ij} = \exp(\mathbf{S}_{ij} - \mathbf{m}^{\text{new}})$, with each row's maximum subtracted from that row.
   5. Rescale and accumulate, $\boldsymbol{\ell} \leftarrow e^{\mathbf{m}-\mathbf{m}^{\text{new}}}\odot\boldsymbol{\ell} + \operatorname{rowsum}(\tilde{\mathbf{P}}_{ij})$ and $\tilde{\mathbf{O}} \leftarrow \operatorname{diag}(e^{\mathbf{m}-\mathbf{m}^{\text{new}}})\,\tilde{\mathbf{O}} + \tilde{\mathbf{P}}_{ij}\mathbf{V}_j$, then set $\mathbf{m} \leftarrow \mathbf{m}^{\text{new}}$.
3. Write $\mathbf{O}_i = \operatorname{diag}(\boldsymbol{\ell})^{-1}\tilde{\mathbf{O}}$ and $\operatorname{lse}_i = \mathbf{m} + \log\boldsymbol{\ell}$ to HBM.

In plain words, each query block makes one sweep over all the keys and values, keeping a few numbers per row that are corrected each time a larger score shows up, and writes its finished output exactly once.

Why is the result exactly attention? Step 2 is the merge of §12.2 applied, for each row, to the summary of keys $1, \dots, (j-1)B_c$ and the summary of key block $j$. By the merge claim proved there, after the last block each row holds the summary of all its keys, $\ell_r = \sum_l e^{s_{rl} - m_r}$ and $\tilde{\mathbf{o}}_r = \sum_l e^{s_{rl} - m_r}\mathbf{v}_l$ with $m_r$ the row maximum. Dividing gives $\sum_l \operatorname{softmax}(\mathbf{s}_r)_l\,\mathbf{v}_l$, which is row $r$ of $\mathbf{P}\mathbf{V}$. Nothing is approximated. FlashAttention is exact attention, computed in a different order, and its output differs from the naive one only by floating-point rounding.

Here it is in NumPy, with the HBM counter from §12.1 on every load and store. The outer loop is over query blocks, the inner loop over key blocks, and the comments match the steps above.

The outputs match naive attention to about $10^{-15}$, and so do the log-sum-exps. The traffic tells the story. Each of the $n_r = 8$ query blocks reads all of $\mathbf{K}$ and $\mathbf{V}$ once, so the count is

1.18 million values here against 4.46 million for standard attention. And nothing of size $T\times T$ was ever stored.

The causal mask, and the tiles you never visit

A causal mask (§4.4) lets query $r$ see keys $l \le r$. In the tile grid that splits the tiles into three kinds.

• Below the diagonal, every key in the tile is at or before every query, so the tile is computed with no mask at all.
• On the diagonal, the tile crosses $r = l$, so it is computed and then masked with $-\infty$ above the diagonal.
• Above the diagonal, every key is after every query, so every weight would be zero. The tile is skipped entirely, with no load, no product, and no exponential.

That last line is the block skipping that makes causal FlashAttention cost about half of non-causal. In the run above, 72 of the 128 tiles were visited. With $n$ square blocks per side the visited fraction is $n(n+1)/2$ out of $n^2$, approaching one half. A naive implementation, by contrast, computes the full $T\times T$ matrix and throws half of it away (§4.7).

There is one subtlety. The skip test compares the first key of the block with the last query of the block, j0 > rows[-1], because a tile can be skipped only if every one of its entries is masked. Comparing first with first, an off-by-one that is easy to make when $B_r \ne B_c$, silently drops keys that some rows should see.

import numpy as np

def flash_attention_forward(Q, K, V, Br, Bc, causal=False):
    S = Q @ K.T / np.sqrt(Q.shape[1])                    # materializes T_q x T_k: not allowed
    if causal:
        S = np.where(np.tril(np.ones(S.shape, bool)), S, -np.inf)
    P = np.exp(S - S.max(1, keepdims=True))
    return P @ V / P.sum(1, keepdims=True), np.log(P.sum(1)), 0

What the forward pass leaves behind

Standard attention keeps $\mathbf{P}$ for the backward pass, $T_qT_k$ numbers per head. FlashAttention keeps the output $\mathbf{O}$, which the layer returns anyway, and one extra number per query row, the log-sum-exp

That is enough, because every attention weight can be rebuilt from it and its score,

In plain words, the log-sum-exp is the softmax's normalizer in log space. Subtract it from a score and exponentiate, and you have the weight, with no maximum and no sum needed. And since $s_{rl} \le \max_l s_{rl} \le \operatorname{lse}_r$, the exponent is never positive, so the reconstruction cannot overflow.

The original FlashAttention stored $\mathbf{m}$ and $\boldsymbol{\ell}$ separately. FlashAttention-2 stores only their combination $\operatorname{lse}$ (Dao 2023), half the statistics and one fewer operation per weight in the backward pass. Either way the extra memory is $O(T)$ per head instead of $O(T^2)$, which is Theorem 1 of Dao et al. (2022), $O(T^2d_k)$ FLOPs and $O(T)$ memory beyond the inputs and output.

For one layer of Llama 3 8B at 8,192 tokens, the difference is between storing $\mathbf{P}$, $32\times8192^2$ numbers or 4 GiB in bf16, and storing the log-sum-exps, $32\times8192$ numbers or 1 MiB in float32. That factor of $T/2 \approx 4{,}000$ is why FlashAttention made long-context training possible on its own, even setting its speed aside.

The backward pass, derived

Recall the backward pass of attention from §4.3. With $c = 1/\sqrt{d_k}$, $\mathbf{S} = c\,\mathbf{Q}\mathbf{K}^\top$, $\mathbf{P} = \operatorname{softmax}(\mathbf{S})$, $\mathbf{O} = \mathbf{P}\mathbf{V}$, and upstream gradient $\mathrm{d}\mathbf{O}$,

with $D_r = \mathrm{d}\mathbf{O}_r\cdot\mathbf{O}_r$, the dot product of row $r$ of $\mathrm{d}\mathbf{O}$ with row $r$ of $\mathbf{O}$. Each of these is written in terms of $T\times T$ matrices. The job is to rewrite every one in terms of tiles, so that the backward pass, like the forward pass, never holds more than one tile at a time. Three observations do it.

Observation 1, the weights can be recomputed tile by tile. By the reconstruction above, the tile of weights is

with block $i$ of the stored log-sum-exps subtracted from each row. It needs only $\mathbf{Q}_i$, $\mathbf{K}_j$, and $\operatorname{lse}_i$, and it equals the forward pass's weights exactly (up to rounding), not approximately.

Observation 2, the row terms need no row of $\mathbf{P}$. The softmax backward needs $D_r = \sum_l P_{rl}\,\mathrm{d}P_{rl}$, a sum over a whole row. §4.3 showed that this equals $\mathrm{d}\mathbf{O}_r\cdot\mathbf{O}_r$, because $\sum_l P_{rl}\,\mathrm{d}P_{rl} = \sum_l P_{rl}\,(\mathrm{d}\mathbf{O}_r\cdot\mathbf{v}_l) = \mathrm{d}\mathbf{O}_r\cdot\sum_l P_{rl}\mathbf{v}_l = \mathrm{d}\mathbf{O}_r\cdot\mathbf{O}_r$. So $\mathbf{D}$ is computed before the main loop, in one pass over $\mathrm{d}\mathbf{O}$ and $\mathbf{O}$, and each tile then needs only its block $\mathbf{D}_i$.

Observation 3, every product splits into a sum over tiles. With the tile of score gradients

block matrix multiplication gives

To check the tile formula for $\mathrm{d}\mathbf{S}$, look at one entry. $\mathrm{d}S_{rl} = P_{rl}(\mathrm{d}P_{rl} - D_r)$ involves only entry $(r, l)$ of $\mathbf{P}$ and $\mathrm{d}\mathbf{P}$ and the row term $D_r$, and $\mathrm{d}P_{rl} = \mathrm{d}\mathbf{O}_r\cdot\mathbf{v}_l$, so restricting $r$ to block $i$ and $l$ to block $j$ gives exactly the tile formula. The three sums are the definition of a matrix product, written block by block. In plain words, the backward pass is the forward pass's tile loop run again, recomputing each tile of weights from the stored log-sum-exps, and adding each tile's contribution into the three gradients.

Which loop goes outside? $\mathrm{d}\mathbf{K}_j$ and $\mathrm{d}\mathbf{V}_j$ are sums over query blocks, and $\mathrm{d}\mathbf{Q}_i$ a sum over key blocks. FlashAttention-2 puts the key blocks outside, so $\mathrm{d}\mathbf{K}_j$ and $\mathrm{d}\mathbf{V}_j$ accumulate in SRAM and are written once, and lets the contributions to $\mathrm{d}\mathbf{Q}_i$ go to HBM as additions. On a GPU, different thread blocks handle different key blocks at the same time, so their additions to the same $\mathrm{d}\mathbf{Q}_i$ use atomic adds (Dao 2023). In our sequential NumPy version that is a read, an add, and a write. Atomic additions land in whatever order the thread blocks finish, and floating-point addition is not associative, so $\mathrm{d}\mathbf{Q}$ can differ in its last bits from one run to the next. The flash-attn package offers a deterministic backward pass (deterministic=True) that fixes the order at some cost in speed, which is worth turning on when you need bitwise reproducible training or are hunting a numerical bug.

The gradients agree with the naive formulas to about $10^{-15}$, with and without the causal mask, and the backward pass's traffic is a few dozen times $Td_k$, with no $T^2$ term. Causal block skipping works in the backward pass exactly as in the forward pass.

The price of recomputation

The backward pass now does five tile products instead of four. Besides $\mathrm{d}\mathbf{O}_i\mathbf{V}_j^\top$, $\mathbf{P}_{ij}^\top\mathrm{d}\mathbf{O}_i$, $\mathrm{d}\mathbf{S}_{ij}\mathbf{K}_j$, and $\mathrm{d}\mathbf{S}_{ij}^\top\mathbf{Q}_i$, it recomputes $\mathbf{Q}_i\mathbf{K}_j^\top$. Counting $2T^2d_k$ FLOPs per product, a training step costs $4T^2d_k$ forward plus $10T^2d_k$ backward instead of $4 + 8$, about 17% more arithmetic. Dao et al. measured 13% more FLOPs (75.2 against 66.6 GFLOPs) for GPT-2 medium, and a 5.7 times shorter runtime, because the arithmetic was never what took the time.

This is selective recomputation, the same trade as activation checkpointing (§10.1), applied to the one activation that is quadratic in length. Here it is a pure win, because recomputing a tile on chip is cheaper than reading it from HBM.

import numpy as np

def flash_attention_backward(Q, K, V, O, dO, lse, Br, Bc, causal=False):
    raise NotImplementedError

The IO complexity theorem

Now we can say precisely how much FlashAttention saves. Use the external memory model of §12.1, with SRAM holding $M$ numbers, and write $d = d_k = d_v$ for the head width and $T = T_q = T_k$, to match the paper's statement.

Theorem (Dao et al. 2022, Theorem 2). For $d \le M \le Td$, standard attention requires $\Theta(Td + T^2)$ HBM accesses, and FlashAttention requires $\Theta(T^2d^2/M)$.

The proof is a count. We give it for the loop order of the paper's Algorithm 1, which puts the key and value blocks outside, and then note why the other order gives the same answer.

Block sizes. Algorithm 1 sets $B_c = \lceil M/4d\rceil$ and $B_r = \min(\lceil M/4d\rceil, d)$. Then everything that must be on chip at once fits in a constant multiple of $M$. The key and value blocks take $2B_cd \approx M/2$, the query block and its output take $2B_rd \le M/2$, and the tile takes $B_rB_c \le d\cdot M/4d = M/4$. (The cap $B_r \le d$ exists only to keep the tile small.)

Counting. The outer loop runs over the $n_c = \lceil T/B_c\rceil$ key blocks. Each key and value block is loaded once, $2Td$ in total. Inside, for each key block, the algorithm sweeps every query block, loading $\mathbf{Q}_i$, the partial output $\mathbf{O}_i$, and its statistics $\boldsymbol{\ell}_i$ and $\mathbf{m}_i$, and writing back $\mathbf{O}_i$, $\boldsymbol{\ell}_i$, $\mathbf{m}_i$. One sweep therefore moves $3Td + 4T$ values, and there are $n_c$ sweeps, so

Substituting, the second term is about $4Td/M\cdot 3Td = 12\,T^2d^2/M$. The first term, $2Td$, is at most of the same order, because $M \le Td$ means $T^2d^2/M \ge Td$. So the total is $\Theta(T^2d^2/M)$. $\blacksquare$

In plain words, the number of HBM accesses is the number of passes over the inner operand times the size of that operand. The number of passes is the length divided by the outer block size, and the outer block can hold about $M/d$ rows. So the traffic is $T/(M/d)$ passes of $Td$ values each, $T^2d^2/M$.

The FlashAttention-2 order swaps the roles. The query blocks are outside, with $B_r$ about $M/4d$ rows, and each makes one pass over $\mathbf{K}$ and $\mathbf{V}$, so the traffic is $Td + 2Td\cdot\lceil T/B_r\rceil + Td + T$, again $\Theta(T^2d^2/M)$, with the bonus that $\mathbf{O}$ is written once rather than read and written in every sweep. That is the formula our NumPy code measured. The backward pass obeys the same bound (Theorem 5 of the paper), by the same count over its tile loop.

How big is the saving? The ratio of standard to FlashAttention traffic is about $T^2\big/(T^2d^2/M) = M/d^2$, up to the constants. The plot shows it for $T = 4{,}096$ and $d = 64$. With an SM's worth of SRAM, around $M = 10^5$ numbers, FlashAttention moves an order of magnitude less. But notice the left of the plot. When $M$ is near $d^2$, the blocks are so thin that FlashAttention rereads $\mathbf{K}$ and $\mathbf{V}$ more times than standard attention reads its matrices, and it loses. FlashAttention needs SRAM much larger than $d^2$, so wide heads cost it more, and each GPU generation's larger shared memory helps.

def fa1_accesses(T, d, M):
    return 4 * T * T + 4 * T * d       # that is standard attention

def fa2_accesses(T, d, M):
    return 4 * T * T + 4 * T * d

No exact algorithm does better

Is $T^2d^2/M$ the best possible? Proposition 3 of the paper says there is no algorithm that computes exact attention with $o(T^2d^2/M)$ HBM accesses for every $M$ in $[d, Td]$. The proof takes the largest allowed SRAM, $M = Td$. There the bound is $T^2d^2/(Td) = Td$, and any algorithm must at least read its inputs, $\Omega(Td)$ numbers, so nothing can be asymptotically smaller at that end of the range. This is a weak form of optimality, since it pins down only one value of $M$. A lower bound that holds for each $M$ separately is left by the paper as future work. In practice the question is moot, because at realistic sizes FlashAttention is already compute-bound.

The theorem also explains why tiling attention looks like tiling a matrix multiply. Both sweep one operand once per block of the other, so both move about $n^3/\sqrt{M}$ or $T^2d^2/M$ values, a total that falls as fast memory grows (§11.5).

What it bought, measured

The FlashAttention paper reports these results on A100 GPUs, all for exact attention.

The last row is the one people forget. Memory linear in $T$ meant models could be trained on much longer sequences, and longer context improved quality. GPT-2 small trained with 4K context reached a perplexity of 17.5 on OpenWebText against 18.2 at 1K, and still trained 30% faster than Megatron-LM's 1K model.

Dropout, block-sparsity, and other real-kernel details

Real kernels handle a few more things, each of which fits the tile structure.

• Dropout on the weights. Storing a $T\times T$ dropout mask would defeat the purpose. The forward pass instead saves the state of the pseudo-random number generator, and the backward pass regenerates exactly the same mask tile by tile (Appendix B of the paper).
• Block-sparse attention. If a mask is known at block granularity, tiles marked empty are skipped, exactly as causal tiles are. Proposition 4 of the paper gives $\Theta(Td + T^2d^2s/M)$ HBM accesses, with $s$ the fraction of nonzero blocks. This is the hardware view of the sparse patterns in §12.6.
• Other score modifications. Biases such as ALiBi (§5.5), logit soft-capping, sliding windows, and document masks are all functions of $(r, l)$ and the score, computed inside the tile. PyTorch's FlexAttention lets you write such a function in Python and compiles it into a FlashAttention-style kernel (§12.4).

import numpy as np

def block_sparse_attention(Q, K, V, B, block_mask):
    raise NotImplementedError

The kernel, in Triton

On a GPU the NumPy loops become a kernel. Each program (a thread block) owns one query block of one head, which is exactly our outer loop, and the inner loop runs inside it. Triton, the subject of §11.6, lets us write it at the level of tiles, with tl.dot for the tile products on tensor cores. The forward kernel below is complete for contiguous inputs of shape (batch * heads, T, D). It cannot run in your browser, and production kernels add the details that make it fast (software pipelining, exp2 in place of exp, autotuned block sizes, and a separate path for the diagonal tiles), but every line maps to a step of the algorithm.

Notice the launch grid, one program per query block and per head. Parallelizing over query blocks, not only over batch and heads, is one of FlashAttention-2's central changes, and the reason it keeps the GPU busy when the batch is small (§12.4).

In PyTorch

To see the whole algorithm as a differentiable operation, wrap the tiled forward and backward passes in a torch.autograd.Function. The forward saves only $\mathbf{Q}$, $\mathbf{K}$, $\mathbf{V}$, $\mathbf{O}$, and the log-sum-exps, and the backward recomputes the tiles. gradcheck compares it with numerical derivatives, and the comparison with scaled_dot_product_attention checks both passes against PyTorch's own.

Read more

• Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. The algorithm (Section 3 and Appendix B), the IO theorem and lower bound (Section 3.2, Appendix C), and the measurements.
• Dao (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. Algorithms 1 and 2 are the clearest statement of the forward and backward passes with stored log-sum-exp.
• Tillet, Kung & Cox (2019). Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations. The tile-level language in which most FlashAttention variants are first prototyped.
• Aggarwal & Vitter (1988). The Input/Output Complexity of Sorting and Related Problems. The external memory model behind IO complexity, with matching lower bounds.
• Triton tutorials (2024). Fused Attention. A production-quality FlashAttention-2 forward and backward pass in Triton, worth reading line by line after this section.

Watch next

▶

▶

▶