FlashAttention: Exact Attention at the Speed of Memory

Attention, written the way The Transformer’s Three Attention Patterns derived it, is three steps: scores $S = QK^\top / \sqrt{d}$, weights $P = \operatorname{softmax}(S)$ row by row, output $O = PV$. For a sequence of $n$ tokens, $S$ and $P$ are $n \times n$ matrices. At 32,768 tokens that’s a billion numbers per head, per layer: 2 GiB in 16-bit floats, for one head of one layer of one sequence.

Nothing about the output needs that matrix to exist all at once. Each output row is a weighted average of value rows, and FlashAttention (Dao et al., 2022) computes it while looking at only a block of keys at a time, never storing $S$ or $P$. It isn’t an approximation: the result is the same attention, to the last few bits of floating point. What it changes is how much the GPU has to move between its large slow memory and its small fast one, which How a GPU Runs Deep Learning showed is usually what decides the speed.

Table of contents

Where standard attention’s time goes

A GPU has two kinds of memory (How a GPU Runs Deep Learning). HBM, the GPU’s main memory, holds tens of gigabytes and delivers about 3.35 TB/s on an H100. SRAM, the shared memory inside each streaming multiprocessor, holds up to 228 KB per SM on an H100 and is roughly ten times faster: on the A100 the FlashAttention paper used, 192 KB per SM, it estimates about 19 TB/s across all SMs against 1.5–2.0 TB/s for HBM. Arithmetic happens on data in registers and SRAM; anything bigger lives in HBM and has to be brought in.

A straightforward implementation runs attention as separate operations, each reading its inputs from HBM and writing its result back: the matrix multiply writes $S$, the softmax reads $S$ and writes $P$, the second multiply reads $P$. Counted for one head, in 16-bit numbers, against what FlashAttention moves:

# Standard attention runs three separate kernels, and each one reads its
# inputs from GPU memory (HBM) and writes its result back:
#   S = Q K^T       reads Q, K (2nd), writes S (n^2)
#   P = softmax(S)  reads S, writes P (2 n^2)
#   O = P V         reads P, V (n^2 + nd), writes O (nd)
#
# FlashAttention, in the loop order FlashAttention-2 uses, keeps a block of
# B_r query rows in fast on-chip memory and streams K and V past it in
# blocks: Q and O cross HBM once, K and V once per block of query rows, and
# S and P never do. (This charges every re-read of K and V to HBM, a worst
# case: on a real GPU many of them hit the L2 cache.)
H100_SMEM = 228 * 1024                                    # bytes of shared memory per SM ("How a GPU Runs Deep Learning")

def standard(n, d):
    return 2 * (2 * n * d + n * n + 2 * n * n + n * n + n * d + n * d)

def flash(n, d, B_r):
    return 2 * (2 * n * d + (n // B_r) * 2 * n * d)       # Q and O once; K and V once per row block

d = 128
B_r = 128                                                 # 128 query rows of 128: 32 KB in 16 bits; fits with K, V blocks
                                                          # (the score tile and output accumulator live in registers)
print(f"one block of Q, K and V at {B_r} x {d} in 16 bits: {3 * B_r * d * 2 / 1024:.0f} KB of the {H100_SMEM // 1024} KB")
for n in (4096, 32768, 131072):
    s, f = standard(n, d), flash(n, d, B_r)
    print(f"n = {n:>7,}: standard {s / 2**20:9,.0f} MiB, flash {f / 2**20:7,.0f} MiB ({s / f:4.1f}x less); "
          f"the score matrix alone {n * n * 2 / 2**20:9,.0f} MiB per head")

Full script: traffic.py

The rest of traffic.py repeats the count at 32,768 tokens for row blocks of 64, 128 and 256, and computes what each method keeps for the backward pass.

one block of Q, K and V at 128 x 128 in 16 bits: 96 KB of the 228 KB
n =   4,096: standard       132 MiB, flash      66 MiB ( 2.0x less); the score matrix alone        32 MiB per head
n =  32,768: standard     8,224 MiB, flash   4,112 MiB ( 2.0x less); the score matrix alone     2,048 MiB per head
n = 131,072: standard   131,200 MiB, flash  65,600 MiB ( 2.0x less); the score matrix alone    32,768 MiB per head
n = 32,768: flash traffic with row blocks of 64: 8,208 MiB, 128: 4,112 MiB, 256: 2,064 MiB
n =  4,096: saved for backward per head: standard      32 MiB, flash   1.0 MiB
n = 32,768: saved for backward per head: standard   2,048 MiB, flash   8.1 MiB

Three things in that output:

• The score matrix dominates. At 4,096 tokens it’s 32 MiB per head; at 32,768, 2 GiB; at 131,072, 32 GiB, for one head; a 32-head layer computed at once would need about a terabyte. Everything else attention touches, $Q$, $K$, $V$ and $O$, is $n \times d$, a few MiB.
• The traffic falls, by a factor set by the block size. FlashAttention still reads $K$ and $V$ many times, once per block of query rows, so its traffic grows as $n^2 d / B_r$: with 128-row blocks it’s half the standard count here, with 256-row blocks a quarter. This count is a worst case: it charges every re-read of $K$ and $V$ to HBM, while thread blocks working on neighbouring query blocks stream the same tiles at about the same time, so many re-reads are served from the 50 MB L2 cache. Measured, the paper’s GPT-2 medium attention moved 4.4 GB against standard attention’s 40.3 GB. The bigger the on-chip memory, the bigger the blocks and the fewer the passes; the paper states it as $\Theta(n^2 d^2 / M)$ bytes for an SRAM of $M$ bytes, against $\Theta(nd + n^2)$ for standard attention.
• The storage falls much more. For training, standard attention keeps $P$ for the backward pass: $n^2$ numbers per head. FlashAttention keeps the output it produces anyway plus one number per row, about 64 KB at 32,768 tokens, instead of $P$‘s 2 GiB.

(This count is for unfused kernels in 16-bit numbers; real implementations fuse some steps and often keep scores in 32 bits. The shape of the comparison doesn’t change.)

The obstacle: softmax needs the whole row

The reason the standard version materializes $S$ is softmax. Row $i$ of the output is

and the denominator sums over every key. The first value’s weight can’t be finished until the last score has been seen. Worse, a stable softmax (Cross-Entropy and the Softmax Gradient) subtracts the row’s maximum before exponentiating, and the maximum is also known only at the end. So the obvious order is: compute the whole row of scores, find its max, exponentiate, sum, divide, multiply by $V$. Each step needs the whole row, and a row is $n$ numbers.

The online softmax

The way out is to keep a summary that can be corrected later. Process the keys in blocks, and for each query row keep three running quantities: $m$, the largest score seen so far; $\ell$, the sum of $e^{S_{ij} - m}$ over keys seen so far; and $o$, the sum of $e^{S_{ij} - m} V_j$ over the same keys. When a new block arrives with scores $s$:

The factor $e^{m - m'}$ is the correction: everything accumulated so far was computed relative to the old maximum, and multiplying by it re-expresses it relative to the new one. If the max didn’t change, the factor is 1. After the last block, $o / \ell$ is exactly the softmax-weighted average, because $o$ and $\ell$ are the numerator and denominator of the formula above, both scaled by the same $e^{-m}$, which cancels. (The running normalizer is Milakov and Gimelshein’s, 2018; carrying $o$ along with it is Rabe and Staats’, 2021, and FlashAttention’s.)

In code, the update is a short loop over blocks of keys. The output below shows a worked example first, one query and four keys in two blocks, then the same thing at full size, against standard attention:

def attention(Q, K, V):
    S = Q @ K.T / np.sqrt(Q.shape[1])                     # the full n x n score matrix
    P = np.exp(S - S.max(1, keepdims=True))
    return (P / P.sum(1, keepdims=True)) @ V

def tiled_attention(Q, K, V, block):
    n, d = Q.shape
    m = np.full(n, -np.inf)                               # running max per row
    l = np.zeros(n)                                       # running sum of exponentials
    o = np.zeros((n, V.shape[1]))                         # running (unnormalized) output
    for j in range(0, K.shape[0], block):
        S = Q @ K[j:j + block].T / np.sqrt(d)             # only n x block scores at a time
        m_new = np.maximum(m, S.max(1))
        scale = np.exp(m - m_new)                         # how much the old max overstated
        P = np.exp(S - m_new[:, None])
        l = l * scale + P.sum(1)
        o = o * scale[:, None] + P @ V[j:j + block]
        m = m_new
    return o / l[:, None]

Full script: online_softmax.py

attention is the standard version, with the whole score matrix; tiled_attention keeps only $m$, $\ell$ and $o$ per row. The rest of online_softmax.py runs the worked example block by block and compares the two functions on random inputs.

scores [2. 1. 4. 0.] -> softmax [0.1125 0.0414 0.831  0.0152]
after block 1: max 2, sum 1.3679, output so far 1.7358 -> 1.2689
after block 2: max 4, sum 1.2034, output so far 3.3082 -> 2.7489
standard attention gives 2.7489
n = 1024, d = 64, blocks of 128: largest difference from standard attention 7.8e-16
n = 4096, d = 128, blocks of 256: largest difference from standard attention 6.8e-16

In the worked example the first block sees scores 2 and 1, so $m = 2$ and the running answer is 1.27, a weighted average of the first two values. The second block brings a score of 4: the max rises to 4, the first block’s contributions are scaled down by $e^{2 - 4} = 0.135$, the new block is added, and the result, 2.7489, is the exact softmax average. At full size the tiled version and the standard one agree to $10^{-15}$: the same answer up to the order of floating-point additions. Step through it, with blocks of different sizes; the answer at the end never changes:

FlashAttention: tiles in fast memory

FlashAttention combines the online softmax with tiling. Here it is in the form FlashAttention-2 settled on, which is how it’s usually taught. Each block of GPU threads takes a block of $B_r$ query rows: the queries go into shared memory, their running $m$, $\ell$ and $o$ into registers. It then streams the keys and values through in blocks of $B_c$ rows: load a block of $K$ and $V$ into shared memory, compute that block’s $B_r \times B_c$ scores, update the running quantities, discard the scores, load the next block. When all the keys have passed, it divides $o$ by $\ell$ and writes the output rows to HBM. (The 2022 original put the loops the other way round: it streamed blocks of $K$ and $V$ in the outer loop and, for each one, revisited every block of queries, reading and rewriting their partial outputs in HBM. Both orders move $\Theta(n^2 d^2 / M)$ bytes; swapping them is part of what made version 2 faster.)

The full $n \times n$ matrix never exists anywhere. At any moment, a block of threads holds one $B_r \times B_c$ tile of it, in fast memory, for as long as it takes to use it. The figure also shows what a causal mask does to the work: a tile entirely above the diagonal, where every key comes after every query, is skipped outright; a tile on the diagonal is computed and masked element by element; a tile below it needs no mask. For a language model, that skips about half the tiles. (A row that is entirely masked needs care: with $m = -\infty$, $e^{m - m'}$ is undefined, so implementations special-case it.)

Measured

The core of the measurement times one head both ways and reads the process’s peak memory:

# Standard attention (score matrix, softmax, weighted sum: three operations)
# against PyTorch's fused scaled_dot_product_attention, which on the CPU
# runs a FlashAttention-style tiled kernel. One head of size 64, 32-bit
# floats. Each measurement runs in a fresh process so its peak memory can
# be read on its own. Run on an otherwise idle machine: under memory
# pressure the standard version's big matrices get swapped, and its times
# become meaningless.
def run(kind, n, d=64):
    torch.manual_seed(0)
    q, k, v = (torch.randn(1, 1, n, d) for _ in range(3))
    def standard():
        s = q @ k.transpose(-1, -2) / d ** 0.5
        return torch.softmax(s, -1) @ v
    def fused():
        return F.scaled_dot_product_attention(q, k, v)
    f = standard if kind == "standard" else fused
    f()
    t0 = time.perf_counter()
    for _ in range(3):
        out = f()
    t = (time.perf_counter() - t0) / 3
    peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * (1 if sys.platform == "darwin" else 1024) / 2**20   # bytes on macOS, KiB on Linux -> MiB
    print(f"{t * 1e3:.1f} {peak:.0f}")

Full script: sdpa_vs_standard.py

The script calls run once per method and length, each in a fresh process of its own, and prints one row per length, then compares the two outputs at 4,096 tokens.

n =  2,048: standard     3.9 ms, peak    228 MiB;  fused     1.1 ms, peak    197 MiB;  the n x n score matrix is    16 MiB
n =  8,192: standard    44.4 ms, peak    716 MiB;  fused    14.7 ms, peak    206 MiB;  the n x n score matrix is   256 MiB
n = 16,384: standard   175.6 ms, peak  2,263 MiB;  fused    57.7 ms, peak    215 MiB;  the n x n score matrix is 1,024 MiB
largest difference between the two at n = 4,096: 6.6e-07

PyTorch’s scaled_dot_product_attention picks a fused, tiled kernel when one is available; on the CPU it runs a FlashAttention-style implementation (_scaled_dot_product_flash_attention_for_cpu). Against the three-step version, on one head:

• The fused kernel is about 3 times faster at every length, 57 ms against 175 at 16,384 tokens. Both do the same $n^2$ work, and both times grow as $n^2$; the standard version also writes a 1 GiB score matrix to memory, reads it back for the softmax, writes the probabilities and reads them again.
• Its memory stays flat, about 200 MiB, most of it the Python process, while the standard version’s peak tracks the score matrix: 2.2 GiB at 16,384 tokens, the process plus two copies of a 1 GiB matrix.
• The two agree to $7 \times 10^{-7}$ in 32-bit floats, rounding error.

These timings come from the CPU of an Apple M4 laptop. ru_maxrss is in bytes on macOS and in KiB on Linux; the peak-memory line converts either to MiB.

(Three repeated runs agree to within a few percent. A first run, taken while a training job filled the machine’s memory, showed a far larger gap: the standard version’s large matrices were being swapped to disk. Timings of memory-heavy code mean nothing on a machine that’s under memory pressure. On GPUs the original paper reports GPT-2 training 3 times faster at 1,024 tokens than the Hugging Face implementation, and BERT-large training 15% faster than the MLPerf 1.1 record.)

The backward pass: recompute instead of store

Training needs gradients, and the backward pass of attention needs $P$: the softmax’s gradient (see Cross-Entropy and the Softmax Gradient) is $dS = P \circ (dP - D)$, where $dP = dO\, V^\top$ and $D_i$ is the row sum of $dP \circ P$. Storing $P$ is exactly what FlashAttention avoided, so it doesn’t. It keeps the output $O$ and one number per row, the log-sum-exp $m + \log \ell$. The row sum $D_i$ equals $dO_i \cdot O_i$, one dot product per row, which is why $O$ is kept. And each tile of $P$ is rebuilt as $e^{S - \text{LSE}}$ from $Q$ and $K$ when it’s needed, with no running max, since the log-sum-exp is already known.

That’s more arithmetic: one of the forward pass’s two matrix multiplies, $QK^\top$, done again. It’s still faster, because the extra multiply runs at full speed on data already in SRAM, while reading a stored $n \times n$ matrix from HBM would be limited by bandwidth (How a GPU Runs Deep Learning’s roofline). Trading compute for memory traffic is the same idea as activation checkpointing (see Backpropagation, by Hand), at the scale of one kernel.

FlashAttention 2 and 3, and Flash-Decoding

The math hasn’t changed since 2022; the implementations have, following the hardware.

• FlashAttention-2 (Dao, 2023) swapped the loops, so each thread block owns a block of query rows, which parallelizes a long single sequence across the GPU’s SMs as well as across batch and heads; kept $o$ unnormalized until the end and stored only the log-sum-exp; and cut the work that isn’t matrix multiplication, which can’t use the tensor cores (an A100 does 312 TFLOPS of 16-bit matrix multiply but 19.5 TFLOPS of other 32-bit arithmetic, so each such operation costs about 16 times as much). About twice as fast as the first version, 50–73% of an A100’s peak.
• FlashAttention-3 (Shah et al., 2024) is built for the H100. Warp specialization splits the work: producer warps issue the asynchronous memory copies (TMA) while consumer warpgroups run the tensor-core multiplies (WGMMA). Ping-pong scheduling lets one warpgroup’s softmax run while another’s multiplies do, and within a warpgroup the softmax of one tile overlaps the next multiply. For FP8 it adds block quantization and “incoherent processing”, a random rotation that spreads outliers before rounding. 740 TFLOPS in 16-bit, 75% of the GPU’s peak, and about 1.2 PFLOPS in FP8.
• Flash-Decoding (Dao et al., 2023) handles generation. When decoding, there is one query per sequence, so a kernel that parallelizes over query blocks has almost nothing to parallelize. Flash-Decoding splits the keys instead: blocks of the KV cache are processed in parallel, each producing its own $m$, $\ell$ and $o$, and the partial results are combined with the same rescaling as the online softmax. Up to 8 times faster decoding at long contexts. The KV cache’s layout in memory, which this kernel reads, is the subject of Paged Attention.

What FlashAttention doesn’t change

• The FLOPs, asymptotically. Attention is still $4 n^2 d$ per head in the forward pass, quadratic in the context. FlashAttention does somewhat more in training, since the backward pass recomputes $QK^\top$: about 25% more backward and 13–17% more overall. With a causal mask it does about half as much as a naive implementation, which computes every score and then masks half of them, because it skips the tiles above the diagonal. Either way it makes long contexts affordable in memory and in bandwidth, not in compute. Reducing the compute is what Attention Alternatives covers.
• The KV cache. Generation still keeps every earlier token’s keys and values (The KV Cache); FlashAttention only reads them more efficiently.
• The result. It’s exact attention, the same output up to the order of floating-point additions. Unlike sparse or linear attention, nothing about the model or its outputs changes, which is why it spread so quickly: it’s a drop-in replacement. (Its backward pass can be nondeterministic from run to run, where gradients are summed with atomic additions.)

Using it

In practice nobody writes this kernel. PyTorch’s F.scaled_dot_product_attention(q, k, v, is_causal=True) picks a fused backend by itself, FlashAttention, a memory-efficient kernel or cuDNN’s on NVIDIA GPUs, and falls back to the three-step version only when none applies; the flash_attn package exposes the kernels directly. Things worth knowing when using them:

• Head sizes are limited, to 256 in FlashAttention-2 and 3.
• Variable-length batches can be packed without padding: the kernels take the cumulative sequence lengths and keep sequences from attending to each other.
• Custom masks (sliding windows, document boundaries, prefixes) are what FlexAttention (Dong et al., 2024) is for: it compiles a mask or a score modification into a FlashAttention-style kernel.
• Contexts too long for one GPU are split across several with ring attention (Liu et al., 2023), which merges partial results with the same $(m, \ell, o)$ rescaling as the online softmax.
• The most readable implementation is Triton’s fused-attention tutorial, about a hundred lines.

Where this leads

• Anatomy of a Frontier Model, the capstone of Part V: Modern architecture, reads a frontier model’s design as the sum of this part’s ideas.
• Part XII: Inference engineering comes back to these kernels from the serving side: Paged Attention for the KV cache’s layout, and vLLM Architecture for the batching that decides how full each kernel’s work is.

Takeaways

• Standard attention writes an $n \times n$ score matrix to GPU memory and reads it back: 2 GiB per head at 32,768 tokens. At long contexts that traffic and storage, not the arithmetic, sets the speed.
• The online softmax keeps a running max, sum and output per row, and rescales them by $e^{m - m'}$ when the max rises, so attention can be computed one block of keys at a time with an exactly equal result.
• FlashAttention tiles the computation so a block of queries stays in fast on-chip memory while keys and values stream past; the score matrix never exists, and with a causal mask the tiles above the diagonal are skipped. Measured here: about 3 times faster on a CPU, with memory flat instead of growing as $n^2$.
• The backward pass recomputes the score tiles instead of storing them: more arithmetic, less traffic, a good trade on a GPU.
• Versions 2 and 3 follow the hardware (loop order and work partitioning, Hopper’s asynchrony, FP8); Flash-Decoding splits the KV cache for one-query decoding. None of them change attention’s quadratic FLOPs or the KV cache’s size.

References

Start with the FlashAttention paper’s section 3, then Tri Dao’s FlashAttention-2 paper for the loop order taught here. Triton’s fused-attention tutorial is the most readable implementation.

• Milakov and Gimelshein, Online normalizer calculation for softmax (2018): the one-pass softmax with a running max.
• Rabe and Staats, Self-attention Does Not Need O(n²) Memory (2021): attention in chunks with the same trick.
• Dao, Fu, Ermon, Rudra and Ré, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (NeurIPS, 2022): tiling, recomputation, and the IO analysis.
• Dao, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (ICLR, 2024): fewer non-matmul operations and better work partitioning.
• Shah, Bikshandi, Zhang, Thakkar, Ramani and Dao, FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (NeurIPS, 2024): Hopper’s asynchronous units and FP8.
• Dong, Feng, Guessous, Liang and He, Flex Attention: A Programming Model for Generating Optimized Attention Kernels (2024): FlexAttention, for custom masks compiled into a fused kernel.
• Liu, Zaharia and Abbeel, Ring Attention with Blockwise Transformers for Near-Infinite Context (2023): attention split across devices with the same merge.
• Dao, Haziza, Massa and Sizov, Flash-Decoding for long-context inference (Stanford CRFM, 2023): splitting the KV cache for decoding.