FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Attention on a GPU is limited by memory traffic. FlashAttention computes the same result while moving far less data.

Standard attention writes an N×N matrix to the GPU's main memory and reads it back. FlashAttention cuts the work into tiles that fit in on-chip memory and fuses every step into one kernel. On GPT-2 the attention kernel runs 7.6× faster, and its extra memory grows linearly with sequence length instead of quadratically.

Attention's cost grows with the square of the sequence length $N$, and the usual explanation points at the arithmetic: every query multiplied against every key. FlashAttention starts from a measurement that says otherwise. On an A100 GPU, the standard attention kernel for GPT-2 spends most of its time moving an $N \times N$ matrix between two kinds of memory, with the arithmetic units idle while the data moves. Reorganize the computation so that matrix never reaches the slow memory, and the same attention runs several times faster in a fraction of the memory.

The paper combines two known techniques. Tiling splits the inputs into blocks small enough for the GPU's on-chip memory and computes attention block by block, so the $N \times N$ matrix is never assembled anywhere. Recomputation discards the score blocks after the forward pass and rebuilds them during the backward pass instead of saving them. Written as a single hand-tuned GPU kernel, the combination does about 13% more arithmetic than standard attention and still finishes 5.7× sooner on the paper's GPT-2 medium benchmark.

The page goes in dependency order: how GPU memory is laid out, where standard attention spends its memory traffic, the streaming softmax that makes tiling possible, the tiled algorithm and its block size, the backward pass, the count of memory accesses (which stays quadratic), and the results, including the long-context benchmarks the memory savings opened up.

Attention is limited by memory

Attention for one head takes three matrices, each of shape $N \times d$, where $N$ is the number of tokens and $d$ is the head dimension (usually 64 or 128): queries $Q$, keys $K$ and values $V$. The paper writes the computation as three steps:

Row $i$ of $S$ holds the dot products of query $i$ with every key. The softmax runs along each row and turns the scores into weights that are positive and sum to one. Row $i$ of the output is the weighted average of the value vectors. (Implementations scale $S$ by $1/\sqrt{d}$ before the softmax; the paper leaves the scale out of (1) and adds it in its appendix as a constant $\tau$, typically $1/\sqrt{d}$.) If you have not met attention before, the Transformer explainer builds it from scratch.

Put numbers on it with the paper's example, GPT-2: $N = 1024$ and $d = 64$. For one head, $Q$ holds 65,536 values, which is 128 KB in 16-bit floating point. $S$ holds 1,048,576 values ($1024^2$), 2 MB, sixteen times the size of $Q$ (the ratio is $N/d$). $P$ is the same size again. A training batch of 64 sequences with 16 heads has 1,024 such heads, so the score matrices of a single attention layer come to 2 GB.

A GPU operation is either compute-bound or memory-bound. A compute-bound operation spends its time on arithmetic; a memory-bound one spends it waiting for data to arrive from memory. The deciding quantity is arithmetic intensity, the number of arithmetic operations per byte read or written. A large matrix multiply reuses each loaded number across a whole row or column of outputs, so its intensity is high and it is compute-bound. A softmax does a handful of operations per number (subtract, exponentiate, add, divide), so it is memory-bound, and so are masking and dropout. Over the past decade GPU arithmetic throughput has grown much faster than memory bandwidth, which moves more operations into the memory-bound column every hardware generation.

The paper therefore measures attention by the bytes that cross between memory levels rather than by its floating-point operations. The authors call an algorithm that accounts for these reads and writes IO-aware, a practice that database systems and numerical libraries have followed for decades.

Two kinds of GPU memory

A GPU has a large main memory called HBM (high-bandwidth memory), where tensors live between operations, and a small on-chip memory called SRAM, located next to the arithmetic units. On the NVIDIA A100 the paper uses, HBM is 40 GB (80 GB on the larger model) with a bandwidth of 1.5–2.0 TB/s. The chip has 108 streaming multiprocessors (SMs, the GPU's independent compute units), each with 192 KB of SRAM, about 20 MB in total, at an estimated 19 TB/s. The fast memory is about two thousand times smaller than the slow one and about ten times faster.

An analogy for the two levels: a cook's countertop and a walk-in fridge. The countertop holds little, but anything on it is within reach; each trip to the fridge costs a walk. A cook who works quickly plans the ingredients so each one comes out of the fridge as few times as possible. SRAM is the countertop and HBM is the fridge. The analogy stops at timing: GPU loads overlap with arithmetic, so a kernel only waits when the data cannot arrive as fast as it is consumed.

The SRAM in this paper is the per-SM scratchpad, the combined L1 cache and shared memory that a kernel can address directly. The A100's 40 MB L2 cache is a different memory: shared across the chip, larger, and slower. The paper's block sizes use a per-SM budget of around 100 KB.

Every GPU operation, or kernel, loads its inputs from HBM, computes in registers and SRAM, and writes its result to HBM. For memory-bound work the standard remedy is kernel fusion: if several operations touch the same data, load it once, apply all of them on chip, and write the result once. Compilers already fuse simple elementwise chains. Attention resists this in two ways. The softmax needs a complete row of scores before it can normalize, so the steps cannot be cut into independent tiles without extra machinery. And in training the intermediate matrices have to be saved for the backward pass, which sends them to HBM anyway and removes most of the benefit of fusing. The fastest attention kernel the authors knew of when they started, NVIDIA's Apex FMHA, fused the whole forward computation and still wrote the attention matrix to HBM for the backward pass.

Standard attention, pass by pass

The paper's Algorithm 0 is attention as a framework runs it: three kernels, one for each step of (1).

# standard attention (Algorithm 0): three kernels, one head
S = (Q @ K.T) / sqrt(d)   # read Q, K;        write S (N x N) to HBM
P = softmax(S)            # read S (N x N);   write P (N x N) to HBM
O = P @ V                 # read P (N x N), V; write O (N x d) to HBM

Follow the $N \times N$ data. The first kernel writes $S$ to HBM. The second reads $S$ back and writes $P$. The third reads $P$. That is two writes and two reads of an $N \times N$ matrix, four passes of 2 MB each per head in the GPT-2 example, against 0.5 MB for reading $Q$, $K$, $V$ and writing $O$. A real Transformer adds more passes: a mask applied to $S$ and dropout applied to $P$ each read and write the full matrix again when they run as separate kernels, which the paper notes makes the problem worse.

The arithmetic between those passes is cheap. The softmax performs a few operations per value it reads, so the second kernel runs at the speed of HBM, not the speed of the arithmetic units. The goal of FlashAttention is to make none of the four $N \times N$ passes: read $Q$, $K$, $V$ in blocks, do all three steps on chip, and write only $O$. Toggle the figure between the two schemes and watch where the $N \times N$ matrix lives.

Two requirements follow. The softmax has to be computed without ever holding a full row of scores, and the backward pass has to produce gradients without a saved copy of $S$ or $P$; a streaming softmax handles the first and recomputation the second.

Softmax one block at a time

The softmax couples a whole row. To avoid overflow in $e^x$, implementations subtract the row maximum before exponentiating, then divide by the sum of the exponentials. For a vector $x$ of length $B$, the paper writes:

Both $m$ and $\ell$ depend on every entry of the row. With only the first block of a row in hand, you cannot finish: a later block may contain a larger score, which changes the maximum, and every exponential computed so far was taken relative to the old maximum.

Instead of the row itself, the algorithm carries two numbers per row: the running maximum $m$ and the running normalizer $\ell$. Split the row into two blocks, $x = [x^{(1)}, x^{(2)}]$. Each block has its own $m$ and $\ell$, computed with (2). The paper combines them as:

The factor $e^{\,m(x^{(1)}) - m(x)}$ converts the first block's sum from its old reference point to the new one, because exponents add:

Multiplying by the factor turns each term the first block computed into the term it would have had with the true maximum. The factor is at most 1, and it equals 1 when the maximum did not change. Nothing is approximated; after every block, $(m, \ell)$ are exactly the maximum and normalizer of the scores seen so far.

A worked example with numbers: one query row has four scores in two blocks, $[1, 3 \mid 2, 5]$, and each key has a one-number value, $[10, 20 \mid 30, 40]$. The first block gives its maximum, its sum, and a partial output (the weighted average of 10 and 20):

The second block has its own maximum 5 and sum $e^{-3} + e^{0}$, which is 1.0498. The maximum grows from 3 to 5, so the first block's sum and output are rescaled by $e^{3-5}$, which is 0.1353:

The full softmax of $[1, 3, 2, 5]$ gives weights 0.0152, 0.1125, 0.0414, 0.8310, whose sum of exponentials (shifted by the maximum 5) is 1.2034 and whose weighted average of the values is 36.88. The streamed result matches to the digits shown. The output update multiplies the old partial output by $\ell_{\text{old}}$ to undo its normalization, rescales it, adds the new block's unnormalized contribution, and divides by the new $\ell$.

Figure 3 runs the same procedure on sixteen scores in four blocks. Step through with the slider and look for the blocks that raise the running maximum: those are the ones where the accumulated sum is rescaled.

This streaming softmax is older than FlashAttention. Milakov and Gimelshein published it at NVIDIA in 2018 to compute an ordinary softmax with fewer memory accesses, and measured up to 1.3× faster softmax kernels. It works like Welford's running-variance update: keep a small summary and correct it as each new chunk arrives. Rabe and Staats (2021) used the same blockwise softmax to show attention does not need quadratic memory (their practical implementation needs O(√n) extra memory). They optimized the peak memory, not the number of reads and writes, and they kept a separate partial output per block and combined them at the end. The paper reports their method runs at about the speed of standard attention or slightly slower. FlashAttention updates a single output in place after every block and orders the loads to minimize HBM traffic.

The tiled kernel

With the streaming softmax available, attention can be computed tile by tile. Cut $Q$ into $T_r$ row blocks of $B_r$ rows, and $K$ and $V$ into $T_c$ column blocks of $B_c$ rows. Given an SRAM of size $M$, Algorithm 1 sets:

The factor 4 splits the budget among the blocks that share SRAM at the same time: $Q_i$, $K_j$, $V_j$, and the output and score tiles. For the GPT-2 head, suppose $M$ is a 100 KB SRAM counted in 16-bit values, 51,200 of them. Then (4) gives:

That is 6 column blocks and 16 row blocks. The largest tile of scores ever held is $64 \times 200$, 12,800 values instead of a million.

The algorithm is two nested loops. The outer loop walks the column blocks of $K$ and $V$; the inner loop walks the row blocks of $Q$. For each pair it loads $Q_i$ and the running statistics for those rows, computes the $B_r \times B_c$ score tile on chip, and folds it into the running output with the streaming softmax:

# FlashAttention forward (Algorithm 1), v1 loop order, one head
# HBM holds Q, K, V, O (N x d) and l, m (N); nothing of size N x N
O[:] = 0;  l[:] = 0;  m[:] = -inf
for j in range(T_c):                  # OUTER: column blocks of K, V
    load K_j, V_j                     # B_c x d each, into SRAM
    for i in range(T_r):              # INNER: row blocks of Q
        load Q_i, O_i, l_i, m_i       # B_r rows each
        S_ij  = Q_i @ K_j.T / sqrt(d) # B_r x B_c tile, on chip only
        m_ij  = rowmax(S_ij)
        P_ij  = exp(S_ij - m_ij)      # shifted by the tile's own max
        l_ij  = rowsum(P_ij)
        m_new = max(m_i, m_ij)
        l_new = exp(m_i - m_new) * l_i + exp(m_ij - m_new) * l_ij
        O_i   = (l_i * exp(m_i - m_new) * O_i          # per-row scalars
                 + exp(m_ij - m_new) * (P_ij @ V_j)) / l_new
        store O_i, l_i = l_new, m_i = m_new           # back to HBM
# kept for the backward pass: O, l, m (and the dropout RNG state)

Written in the paper's matrix notation, where $\tilde m_{ij}$ and $\tilde P_{ij}$ are the tile's own row maxima and shifted exponentials, the output update (line 12 of Algorithm 1) is:

It is the scalar update from the worked example, one row per entry of the diagonal matrices: $\mathrm{diag}(\ell_i)$ undoes the old normalization of each row, the first exponential moves the old sum to the new maximum, the second does the same for the tile, and $\mathrm{diag}(\ell_i^{\text{new}})^{-1}$ normalizes again. Theorem 1 of the paper states that this procedure returns exactly $\text{softmax}(QK^\top)V$, uses $O(N^2 d)$ floating-point operations, and needs $O(N)$ memory beyond the inputs and output, namely the vectors $\ell$ and $m$.

Figure 4 steps through the loops on a 4×4 grid of tiles and counts the block transfers to and from HBM. Watch the output column on the right: in the paper's loop order every output block stays partial until the last column, and each one is read and written back once per column. Then switch to the order FlashAttention-2 used.

Because every step now happens on a tile in SRAM, the whole of attention fits in one CUDA kernel: the score matrix multiply, the mask, the softmax, dropout, and the multiply by $V$. The appendix version of the algorithm adds the scale $\tau$, the mask (which sets masked scores to $-\infty$) and dropout inside the same loop body, so none of them adds a pass over HBM.

With $Q$ on the inner loop, the v1 kernel reads and rewrites each output block, with its $\ell$ and $m$, once per column block. FlashAttention-2 (Dao, 2023) swapped the loops, putting $Q$ on the outside, so each output block stays on chip until it is finished and is written once. It also parallelized the work across the sequence length and reported about 2× the speed of FlashAttention. This page follows the paper's v1 order throughout; diagrams of the modern kernel show the swapped order.

Choosing the block size

Each pass of the outer loop reads all of $Q$ and reads and writes all of $O$, and there are $T_c = N/B_c$ passes. A wider column block therefore means fewer passes and less HBM traffic, in proportion to $1/B_c$. The limit is SRAM: the tiles for one step have to fit on chip together, and their size grows with $B_c$. The formula $B_c = \lceil M/4d \rceil$ picks the widest block that fits.

The paper measured this trade-off directly. Varying $B_c$ in the forward pass, it found that HBM accesses fall as the block grows and runtime falls with them. Beyond a block size of about 256, runtime is limited by other costs, such as the arithmetic, and larger blocks also stop fitting in SRAM.

The backward pass recomputes

Training needs gradients, and the gradient of attention uses $P$, the $N \times N$ matrix the forward pass avoided storing. Storing it now would bring back the quadratic memory. FlashAttention instead keeps only the output $O$, the per-row statistics $m$ and $\ell$, and the state of the random-number generator used for dropout. In the backward pass it loads blocks of $Q$, $K$, $V$ again and recomputes each score tile in SRAM. Because $m$ and $\ell$ are the final row maxima and normalizers, each recomputed tile gives exact probabilities directly as $\exp(S_{ij} - m_i)$ divided row by row by $\ell_i$, with no second streaming pass; the saved generator state regenerates the same dropout mask.

# FlashAttention backward, one tile (Algorithm 4, masking and dropout left out)
S_ij  = Q_i @ K_j.T / sqrt(d)         # recomputed on chip from Q, K
P_ij  = exp(S_ij - m_i) / l_i         # exact probabilities from saved m, l
dV_j += P_ij.T @ dO_i
dP_ij = dO_i @ V_j.T
D_i   = rowsum(dO_i * O_i)            # eq. (6): length-d dot product per row
dS_ij = P_ij * (dP_ij - D_i)
dQ_i += dS_ij @ K_j / sqrt(d)         # written back to HBM
dK_j += dS_ij.T @ Q_i / sqrt(d)       # kept on chip until column j is done

The softmax Jacobian gives the gradient with respect to a score as $dS_{ij}$ equal to $P_{ij}(dP_{ij} - D_i)$, where $D_i$ is the sum of $P_{ik}\, dP_{ik}$ over a whole row of length $N$, the kind of full-row reduction tiling cannot do in one tile. Each $dP_{ik}$ is $do_i^\top v_k$, the dot product of the output gradient for row $i$ with value $k$. Pull $do_i$ out of the sum, and what remains is the weighted sum of values $\sum_k P_{ik} v_k$, which is the output row $o_i$ itself:

The $N$-length reduction becomes a dot product of two length-$d$ vectors: $o_i$, saved by the forward pass, and $do_i$, the incoming gradient. For GPT-2 that is 64 multiply-adds per row instead of 1,024, and the backward pass stays tile by tile like the forward pass.

Discarding intermediates and recomputing them is the idea behind gradient checkpointing (Chen et al., 2016). There it is a trade: less memory, more compute, slower training. Here the recomputation makes the backward pass faster, because that pass was also limited by HBM traffic. Recomputing a score tile on chip costs a small matrix multiply; reading the stored $P$ back from HBM costs a full $N \times N$ pass. The paper states the result as "even with more FLOPs, our recomputation speeds up the backward pass." The activation recomputation explainer covers the same trade at the scale of whole Transformer layers.

Counting HBM accesses

The paper checks the argument on GPT-2 medium (sequence length 1024, head dimension 64, 16 heads, batch 64), forward and backward pass together, on one A100. FlashAttention does 75.2 GFLOPs against 66.6 for standard attention, about 13% more, because of the recomputation. It reads and writes 4.4 GB of HBM against 40.3 GB, about 9× less. It finishes in 7.3 ms against 41.7 ms, 5.7× faster.

The attention kernel alone is up to 7.6× faster than the PyTorch implementation on GPT-2. Whole-model training speedups are smaller, 15% to 3× in the results below, because attention is one part of each layer.

Theorem 2 counts the traffic in general, in the external-memory model of Aggarwal and Vitter: a fast memory of size $M$, a slow memory of unlimited size, and cost equal to the number of values moved between them. For $d \le M \le Nd$:

The FlashAttention count comes straight from the loop structure. Each value of $K$ and $V$ is loaded once. The outer loop makes $T_c = N/B_c$ passes, each reading $Q$ and reading and writing $O$, which is $\Theta(Nd)$ per pass. With $B_c = \Theta(M/d)$, the total is $\Theta(Nd \cdot Nd/M)$, which is $\Theta(N^2 d^2 / M)$. Compared term by term with the standard $N^2$, FlashAttention moves $d^2/M$ as much. The paper's typical values are $d$ from 64 to 128 and $M$ around 100 KB, where $d^2$ (4,096 to 16,384) is many times smaller than $M$. The backward pass has the same two complexities (Theorem 5).

Both expressions in (7) still contain $N^2$. FlashAttention reduces the constant in front of the quadratic term; it does not change the growth rate. Double the sequence length and both kernels move about four times as much data, which matches the paper's benchmark: runtime "grows quadratically with sequence length" for FlashAttention as for standard attention. What becomes linear is the extra memory, the $O(N)$ of Theorem 1.

Figure 7 counts every value that Algorithm 0 and Algorithm 1 move in the forward pass for one head with $d = 64$, constants included. At the GPT-2 setting ($N = 1024$, a 100 KB SRAM) standard attention moves $4N^2 + 4Nd$ values, 4,456,448 of them. FlashAttention moves 1,402,880: 67,584 to initialize $O$, $\ell$, $m$, 131,072 to load $K$ and $V$, and 6 passes of 200,704. That is 3.2× less for the forward pass alone of the bare algorithms; the measured 9× also covers the backward pass. Move the length slider to see how both counts grow when $N$ doubles, then take the SRAM slider to both ends.

Proposition 3 adds a lower bound: no exact attention algorithm achieves $o(N^2 d^2 / M)$ HBM accesses for every $M$ in the range $[d, Nd]$. The proof uses the top of the range, $M = \Theta(Nd)$, where the bound reduces to $\Omega(Nd)$, the cost of reading the inputs once. It shows no algorithm beats FlashAttention asymptotically across all SRAM sizes; it does not show FlashAttention is optimal at the 100 KB an A100 actually has. The authors leave a bound in terms of $M$ for future work.

Exact, not approximate

Most earlier work on attention's $N^2$ cost changed the function being computed. Sparse attention scores only some query-key pairs; low-rank methods such as Linformer project the keys to a shorter sequence; kernel methods such as Performer replace the softmax with a feature map that can be computed in linear time. These are approximate attention: fewer operations, a different output, and usually some loss in model quality. The paper observes that many of them show no wall-clock speedup over standard attention, because they cut floating-point operations while the time was going to memory access. In its own benchmarks, approximate methods only overtake FlashAttention at sequence lengths between 512 and 1024.

FlashAttention computes $\text{softmax}(QK^\top)V$, the same function, which is what the paper means by exact. In 16-bit floating point the output can differ from a standard kernel in the last bits, because the streaming softmax and the recomputation add the same terms in a different order. That difference is rounding. In training, models with FlashAttention reach the same perplexity as the same models with standard attention and follow the same training and validation curves; the paper calls FlashAttention "as numerically stable as the baselines."

The paper also presents one approximate variant, block-sparse FlashAttention, covered with the results below. The exactness statement covers Algorithm 1 and Theorem 1 only.

Results and limits

The first group of results is speed with no change to the model. On BERT-large at sequence length 512, FlashAttention reaches the MLPerf 1.1 target accuracy in 17.4 ± 1.4 minutes against 20.0 ± 1.5 for the MLPerf 1.1 record, a 15% speedup (the record time divided by the new one). MLPerf is the industry benchmark for training throughput. On GPT-2, training on 8 A100s is up to 3× faster than the HuggingFace implementation and 1.7× faster than Megatron-LM, NVIDIA's optimized training library: GPT-2 small takes 2.7 days against 9.5 and 4.7, GPT-2 medium 6.9 days against 21.0 and 11.5. Perplexity matches: 18.2 for all three on GPT-2 small, 14.3 on medium against 14.3 for Megatron-LM and 14.2 for HuggingFace. On the Long-Range Arena, a benchmark suite of tasks with sequences of 1K to 4K tokens, FlashAttention is 2.4× faster than a standard Transformer with matching average accuracy (59.8 against 59.3).

The second group comes from memory. The kernel's memory footprint grows linearly with sequence length, and the paper measures it at up to 20× less than exact-attention baselines and 2× less than Linformer. On a 40 GB A100, every other method except Linformer runs out of memory before 64K tokens. Drag the slider to longer sequences and watch the two footprints separate.

The memory went into longer contexts. GPT-2 small with a 4K context, trained with FlashAttention, is still 30% faster than Megatron-LM's GPT-2 small at 1K and reaches 0.7 better perplexity (17.5 against 18.2). On two long-document classification tasks, longer inputs lift results by 6.4 points on average: MIMIC-III clinical notes gain 4.3 points at 16K tokens over 512, and ECtHR legal cases gain 8.5 points at 8K over 512 (8K is ECtHR's best length, not 16K). These gains come from the longer context; at the same length the perplexities match, as in the GPT-2 comparison above.

Path-X is a Long-Range Arena task: decide whether two points in a 128×128 image are connected by a path, with the image fed in as a sequence of 16K pixels. Before this paper every Transformer either ran out of memory or scored at chance. FlashAttention reaches 61.4% accuracy, the first Transformer to do better than chance. Path-256 is the same task at 256×256, a 64K-token sequence, where block-sparse FlashAttention reaches 63.1%.

Block-sparse FlashAttention takes a mask over blocks and skips the tiles the mask zeroes out. It is the one approximate method in the paper. With $s$ the fraction of nonzero blocks, its HBM accesses are:

It runs 2 to 4× faster than dense FlashAttention and, in the authors' benchmarks, faster than every exact, sparse and approximate attention implementation they tested, at every sequence length.

The authors close with three limitations and future directions. Every new attention variant needs a new CUDA kernel, written far below the level of PyTorch and not always portable to the next GPU generation; they call for a compiler that turns a high-level description of attention into an IO-aware kernel, as Halide does for image processing. They expect the IO-aware approach to extend beyond attention, since every layer of a network reads and writes HBM, and hope for IO-aware versions of other modules. The IO analysis covers a single GPU, and splitting attention across several GPUs adds another level of data movement to account for.