FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning

FlashAttention-2 computes the same attention as FlashAttention about twice as fast, by changing how the work is split across the GPU.

FlashAttention reached 25 to 40% of an A100's peak arithmetic rate. Three changes raise that to 50 to 73%: fewer operations outside matrix multiplies; a separate group of GPU threads (a thread block) for each chunk of query rows instead of one per attention head; and a split of each chunk of work among groups of 32 threads (warps) so they no longer pass partial results to each other through on-chip memory.

FlashAttention (2022) made attention fast by moving less data. It computes softmax($QK^\top$)$V$ in tiles small enough for a GPU's on-chip memory, so the $N \times N$ attention matrix never goes to the GPU's main memory, and it ran 2 to 4× faster than the attention implementations of the time. It still used the A100's arithmetic poorly: its forward pass ran at 30 to 50% of the GPU's peak FLOPs/s and its backward pass at 25 to 35%, (the paper's abstract summarizes these as 25 to 40%), while a well-tuned matrix multiply (GEMM) runs at 80 to 90%. The paper traces the gap to how the work was divided among the GPU's thread blocks and warps.

FlashAttention-2 is a single-author technical report by Tri Dao, published in July 2023 and at ICLR 2024. It keeps FlashAttention's math and its memory traffic analysis and changes how the computation is laid out on the hardware. The changes are:

1. rewrite the streaming softmax so the kernel does fewer operations that are not matrix multiplies, because on an A100 each of those costs about 16 times as much as a matrix-multiply FLOP;
2. put the loop over query blocks outside the loop over key blocks (FlashAttention had them the other way round), and give each block of query rows its own thread block (a group of threads that runs on one of the GPU's 108 processors), so a long sequence with a small batch still occupies all of them;
3. inside a thread block, split each tile between warps (bundles of 32 threads) by query rows instead of by keys, so the warps never have to add up each other's partial results.

The output is still exactly softmax($QK^\top$)$V$. In the paper's benchmarks FlashAttention-2 runs around 2× faster than FlashAttention, reaches 73% of the A100's peak in the forward pass, and trains a 2.7B-parameter GPT model end to end at 225 TFLOPs/s per GPU. The page starts with a short recap of FlashAttention, then explains the GPU cost model behind the first change, walks through the forward and backward algorithms and the two levels of work partitioning, and ends with the results.

What FlashAttention already did

One attention head takes queries $Q$, keys $K$ and values $V$, each an $N \times d$ matrix, where $N$ is the sequence length and $d$ the head dimension (64 or 128 in practice). The paper writes the computation as:

The softmax runs along each row of $S$, and row $i$ of $O$ is a weighted average of the value vectors with weights from row $i$ of $P$. Like the paper, this page leaves out the $1/\sqrt{d}$ scale on $S$, the causal mask, and dropout until they matter. The Transformer explainer covers why attention is computed this way.

A GPU has two kinds of memory that matter here. HBM is the large main memory: 40 or 80 GB on an A100, at 1.5 to 2.0 TB/s. SRAM is on-chip memory inside each of the A100's 108 streaming multiprocessors (SMs), 192 KB per SM, roughly ten times faster. A standard implementation writes $S$ and $P$ to HBM and reads them back, which for $N = 8192$ is two 64M-entry matrices per head. FlashAttention instead loads a block of queries and a block of keys and values into SRAM, computes that tile of $S$, folds it into the output, and discards it.

Tiling the softmax is hard: normalizing a row needs the maximum and the sum over the whole row, and a tile sees only part of it. The online softmax of Milakov and Gimelshein (2018) gets around this. Split one block of query rows into two key blocks, $S^{(1)}$ and $S^{(2)}$, each $B_r \times B_c$. After the first block, keep the row maximum $m^{(1)}$, the row sum $\ell^{(1)}$ of $e^{S^{(1)} - m^{(1)}}$, and a normalized partial output $O^{(1)}$. When the second block arrives, the maximum may grow, and every term computed under the old maximum is too large by a factor $e^{m^{(2)} - m^{(1)}}$. FlashAttention corrects that and renormalizes:

Here $\mathrm{diag}(\ell)$ scales row $i$ by $\ell_i$, and $e^{m^{(1)}-m^{(2)}}$ is a per-row factor at most 1. Multiplying $O^{(1)}$ by $\ell^{(1)}$ undoes its old normalization, the exponential moves it to the new maximum, and the division by $\ell^{(2)}$ normalizes the sum again. Repeated over any number of blocks, the update ends at exactly (1).

For the backward pass, FlashAttention saves only $O$, $m$ and $\ell$ and recomputes each tile of $S$ and $P$ from $Q$ and $K$. That does about 13% more arithmetic and still finishes sooner, because the recomputation replaces reading an $N \times N$ matrix from HBM. Extra memory drops from $O(N^2)$ to $O(N)$, a 10 to 20× saving. FlashAttention-2 keeps all of this. Its changes are about what happens once the data is on the chip.

Why non-matmul FLOPs cost 16× on a GPU

An A100 has two kinds of arithmetic units. Tensor Cores do matrix multiplies in FP16 or BF16 at up to 312 TFLOPs/s. Everything else (exponentials, maxima, sums, elementwise multiplies and divides) is limited to the 19.5 TFLOPs/s of non-matmul FP32 arithmetic. One non-matmul FLOP therefore takes as long as 16 matmul FLOPs. Kernel speed is quoted as a fraction of 312 TFLOPs/s; a kernel that spends a fifth of its time on non-matmul work, with nothing overlapping it, cannot exceed 80% of peak, however few FLOPs that work contains.

Count the work in one tile of the forward pass, per query row, for a key block of $B_c$ keys. The two matrix multiplies, $Q_iK_j^\top$ and $\tilde P V_j$, each cost $2B_c d$ FLOPs per row (one multiply and one add per term), $4B_c d$ in total. The softmax on the tile costs about $4B_c$: a max, a subtraction, an exponential and an addition per score. Then the output row of length $d$ is rescaled by (2). In FlashAttention that is four operations per element: multiply the old output by $\ell^{(1)}e^{m^{(1)}-m^{(2)}}$, multiply the new product by its exponential factor (Algorithm 1 of FlashAttention shifts each tile by its own max), add, and divide by $\ell^{(2)}$.

With $B_c = 128$ and $d = 128$, that is 65,536 matmul FLOPs, 512 softmax operations and 512 rescaling operations per row per key block. The non-matmul part is 1.5% of the FLOPs. Weighted by 16, it is 20% of the time if the two kinds of work run one after the other. Figure 1 shows how that split moves with the head dimension: the softmax cost is fixed per score while the matmul cost grows with $d$, so small heads spend a larger share of their time outside the Tensor Cores.

At $d = 64$ and $B_c = 128$, non-matmul work is 2.3% of FlashAttention's FLOPs and 27% of its time. FlashAttention-2 removes half of the rescaling work, which brings the time share to 24%. The exponentials themselves cannot be removed. By this page's count, halving the rescaling saves 4 to 10% of a tile's time across head dimensions 32 to 256, so most of the 2× has to come from the other two changes and from the new implementation.

How FlashAttention-2 works: the forward pass

The paper makes two changes to (2). The first: do not divide by $\ell$ after every block. Keep an unscaled accumulator $\tilde O$, the numerator of the softmax average, and divide once at the end. The two-block update becomes:

Here $\tilde O^{(1)}$ is $e^{S^{(1)}-m^{(1)}}V^{(1)}$, and $m$ and $\ell$ are updated as in (2). Each block now costs two operations per output element (one multiply, one add) instead of four, and the divide happens once per row for the whole sequence. The factor $e^{m^{(1)}-m^{(2)}}$ multiplies; the paper prints it with an inverse in Algorithm 1 and in its two-block example, which the official kernel does not do (see the Provenance panel). Since the new maximum is never smaller than the old one, the factor is at most 1 and shrinks the old terms.

The second change: save one statistic per row for the backward pass, the logsumexp, instead of the two numbers $m$ and $\ell$:

This works because dividing by $\ell_i$ is subtracting its logarithm in the exponent:

One subtraction and one exponential per score recover the exact probabilities, with no division, from one stored number per row.

Figure 2 runs both versions on one query row with eight scores in four blocks of two keys and a single number as each value ($d = 1$). Step through and compare the two accumulators. FlashAttention's $O$ is a valid weighted average after every block. FlashAttention-2's $\tilde O$ becomes an average only at the last step; after block 4 it is 5.105, and dividing by $\ell = 2.463$ gives 2.073, the exact answer.

The row ends with $m$ = 4.0 and $\ell$ = 2.463, so the logsumexp saved for the backward pass and the probability it gives for the last key (score 4.0), which matches a full softmax over the eight scores, are:

The full forward pass also swaps FlashAttention's loops. FlashAttention's Algorithm 1 had the outer loop over key and value blocks and the inner loop over query blocks, so every output block was read from HBM, updated and written back once per key block. FlashAttention-2 puts the query blocks outside. A thread block loads $Q_i$ once, keeps $\tilde O_i$, $m_i$ and $\ell_i$ on chip while it walks across all of $K$ and $V$, and writes $O_i$ and $L_i$ once. No iteration of the outer loop needs anything from another (row block 5 never reads row block 4), which the parallelism change in a later section relies on. In pseudocode:

# FlashAttention-2 forward (Algorithm 1), one head; scale, mask, dropout omitted
for i in range(T_r):                      # OUTER: one thread block per row block
    load Q_i                              # B_r x d, into SRAM
    O = zeros(B_r, d)                     # all three live on chip for the whole row
    l = zeros(B_r)
    m = full(B_r, -inf)
    for j in range(T_c):                  # INNER loop: walk across K, V
        load K_j, V_j                     # B_c x d each
        S = Q_i @ K_j.T                   # B_r x B_c, never leaves the chip
        m_new = maximum(m, rowmax(S))
        P = exp(S - m_new[:, None])       # unnormalized probabilities
        a = exp(m - m_new)                # a <= 1, shrinks the old terms
        l = a * l + rowsum(P)
        O = a[:, None] * O + P @ V_j      # multiply and add; no divide
        m = m_new
    O = O / l[:, None]                    # the only divide, once per row
    write O to HBM as O_i                 # once per row block
    write L_i = m + log(l) to HBM         # one number per query row

With $N = 16384$, $d = 128$ and the A100 block sizes from the code ($B_r = 128$, $B_c = 64$), one head has $T_r = 128$ row blocks and $T_c = 256$ key blocks. Each tile is 4,194,304 matmul FLOPs and 8,192 exponentials. Each row block streams 256 tiles and writes back 128 × 128 output values and 128 logsumexps. The forward pass uses $O(N^2 d)$ FLOPs and $O(N)$ memory beyond the inputs and output, the same as FlashAttention.

Skipping masked blocks

In autoregressive language modeling, token $i$ may attend only to tokens $j \le i$. A causal mask sets $S_{ij} = -\infty$ for $j > i$, so those probabilities are zero. Because the kernel already works in blocks, it can skip a whole tile when every key in it comes after every query in it. For a long sequence that is about half the tiles, and the paper measures a 1.7 to 1.8× speedup over the same computation without the mask.

The tiles strictly below the diagonal need no mask at all, since every key precedes every query. With square blocks, only the diagonal tile of each row needs the elementwise mask. Of the 36 tiles that survive in an 8 × 8 grid, 8 are masked and 28 are computed as if unmasked. The code walks each row from its last key block backward, so the masked diagonal tile is handled first and the remaining steps run without masking logic. The next section returns to the masked grid, where it makes rows uneven in length.

The backward pass

Given the gradient of the loss with respect to the output, $dO \in \mathbb{R}^{N \times d}$, the chain rule through (1) gives:

The softmax gradient for one row $p$ of $P$ is $(\mathrm{diag}(p) - pp^\top)\,dp$. Written out per element:

The sum in $D_i$ runs over all $N$ keys, which a tile cannot see. Since $dP_{ij}$ is the dot product $dO_i \cdot V_j$, it can be rewritten:

So $D$ is a row-wise dot product of two $N \times d$ matrices that are already in HBM, one number per query row, computed before the main loop. This trick comes from FlashAttention; the backward pass of FlashAttention-2 is almost the same algorithm, with the probabilities rebuilt from $L$ by (4):

# FlashAttention-2 backward (Algorithm 2), one head; scale, mask, dropout omitted
D = rowsum(dO * O)                        # length N: one number per query row
for j in range(T_c):                      # OUTER: one thread block per column
    load K_j, V_j
    dK_j = zeros(B_c, d); dV_j = zeros(B_c, d)   # accumulated on chip
    for i in range(T_r):
        load Q_i, O_i, dO_i, L_i, D_i
        S = Q_i @ K_j.T                   # matmul 1: recomputed scores
        P = exp(S - L_i[:, None])         # exact probabilities from L alone
        dV_j += P.T @ dO_i                # matmul 2
        dP = dO_i @ V_j.T                 # matmul 3
        dS = P * (dP - D_i[:, None])
        atomic_add(dQ_i, dS @ K_j)        # matmul 4: other columns add here too
        dK_j += dS.T @ Q_i                # matmul 5
    write dK_j, dV_j to HBM               # once per column block

The loops run the other way around from the forward pass. The outer loop is over key and value blocks, so each thread block owns one $dK_j$ and one $dV_j$, accumulates them on chip down its column, and writes them once. The cost of that choice falls on $dQ$: row block $i$ of $dQ$ receives a contribution from every column block. Different thread blocks add into the same $dQ_i$ in HBM. A plain read, add and write from two thread blocks at once could lose one of the two updates, so they use atomic adds, which the hardware performs as one indivisible step (in fp32, in the official code).

The backward pass does five matrix multiplies per tile against two in the forward pass, which is why the benchmark counts its FLOPs as 2.5 times the forward count. It also keeps more blocks in SRAM at once, and at best it reaches 63% of peak against the forward pass's 73%.

Parallelism over the sequence length

A CUDA kernel runs as a grid of thread blocks. The GPU schedules each thread block onto one SM, where it runs to completion; a thread block is made of warps of 32 threads. Threads in a warp exchange values through fast register shuffles, and warps in the same thread block exchange values through shared memory, which is the SM's SRAM. Occupancy is the fraction of the GPU's resources a kernel keeps busy; a kernel that launches fewer thread blocks than the GPU has SMs leaves the rest of the SMs idle. A kernel that launches more runs them in waves: the first 108 start, and each SM takes the next waiting thread block when its current one finishes.

FlashAttention, as the paper describes it, uses one thread block per attention head per sequence: batch size × number of heads thread blocks in total. That fills an A100 when the product is large (the paper says about 80 or more). Long sequences break it, because with a fixed budget of tokens per batch a longer sequence means a smaller batch. In the paper's benchmark setting (16k tokens per batch, hidden size 2048, $d = 128$, so 16 heads), a 16k-token sequence comes as a batch of one and FlashAttention launches 16 thread blocks for 108 SMs.

FlashAttention-2 makes each iteration of the outer loop its own thread block. The forward pass launches $\lceil N/B_r \rceil$ × batch × heads thread blocks, which the code builds as a three-dimensional grid with the row-block index first. The 16k-token example becomes 128 × 1 × 16 = 2,048 thread blocks, and since the row blocks of a head never communicate, nothing else in the algorithm changes. In this setting the count is 2,048 at every sequence length, because 16k tokens always make 128 row blocks per head.

At $N = 512$ the batch is 32, and FlashAttention's 512 thread blocks keep 95% of the SMs busy on average. At $N = 16384$, FlashAttention's 16 thread blocks occupy 15% of the GPU, and in this model FlashAttention-2 finishes 6.7× sooner. The measured gap is much smaller, for a reason the model leaves out: the released FlashAttention code had already added a heuristic that splits each head across several thread blocks to fill the SMs (the num_splits parameter in v1.0.9), and the measured FlashAttention baseline reaches 110 TFLOPs/s at 16k tokens with $d = 64$ (32 heads), more than the ≈92 TFLOPs/s (32 × 312/108) that one thread block per head on 32 of the 108 SMs could deliver. The paper's description of FlashAttention as one block per head is the original algorithm, not the baseline in its plots. In FlashAttention-2 the query loop is outside, so the row blocks are independent by construction and each writes its own rows of the output.

The loop swap and the sequence-length parallelism were first implemented by Phil Tillet in the Triton version of FlashAttention, which the paper credits. The backward pass is parallelized the same way, with one thread block per column block and the atomic adds into $dQ$ described above.

Figure 4 shows both partitions on an 8 × 8 grid of tiles. In the forward pass each color is a thread block that owns a row, so the output column on the right fills one row at a time as each block finishes. In the backward pass each color owns a column, and every tile it finishes adds into the $dQ$ entry of that row. Turn on the causal mask and watch the row lengths: row block 1 has one tile and row block 8 has eight, so with causal attention the forward thread blocks finish at different times.

With the mask off, the forward pass computes 64 tiles and writes 8 output blocks once each; the backward pass computes the same 64 tiles and performs 64 atomic adds into $dQ$, 8 per row. With the mask on, 28 tiles are skipped in both passes and the backward pass does 36 atomic adds.

Splitting the work between warps

Inside a thread block, 4 or 8 warps share each tile. FlashAttention split the tile by keys: each of 4 warps took a quarter of $K_j$ and $V_j$, while all of them read $Q_i$. The paper calls this the "split-K" scheme, after the GEMM technique of splitting the dimension being summed over, which in the product $\tilde P V_j$ is the key dimension. Walk it through for $B_r = B_c = 128$ and $d = 64$: warp 1 computes $S$ for all 128 query rows against keys 1 to 32, a $128 \times 32$ slice. It multiplies its slice of $\tilde P$ by rows 1 to 32 of $V_j$ and gets a $128 \times 64$ result, but that result covers only 32 of the 128 keys. The full output is the sum of the four warps' results, so each warp writes its 8,192 values to shared memory, all four synchronize, and the sum is formed: 32,768 values written and read back per key block. The row maximum and row sum of the softmax are split across the four warps the same way.

FlashAttention-2 splits by query rows instead. Each warp takes 32 of the 128 rows of $Q_i$, and all warps read all of $K_j$ and $V_j$. A warp then owns complete rows of $S$, complete rows of the softmax statistics, and complete rows of the output, so no warp needs another warp's result. In the code this is a single layout choice: the matrix-multiply tile stacks its warps along the query dimension.

The backward pass also avoids split-K, but its dependencies among $Q$, $K$, $V$, $O$, $dO$, $dQ$, $dK$ and $dV$ still require some synchronization between warps. The paper reports a speedup from the reduced shared-memory traffic in both passes but does not measure it separately.

Block sizes

Larger blocks mean fewer loads and stores through shared memory per FLOP, but each thread holds more of the tile in registers and the thread block needs more shared memory. Past some size the registers spill to slower memory, or the shared memory needed exceeds what the SM has and the kernel cannot launch. The paper uses blocks of $\{64, 128\} \times \{64, 128\}$, tuned by hand per head dimension and GPU, and leaves auto-tuning to future work. The current code uses $128 \times 128$ for $d = 64$ and $128 \times 64$ for $d = 128$ on an A100, with 4 warps; a comment beside the $d = 64$ choice records that $256 \times 64$ runs 85% slower at 2k tokens because of register spilling.

Results

The benchmarks use one A100 80GB SXM4, sequence lengths from 512 to 16k with the batch set so each run holds 16k tokens, hidden size 2048, and head dimension 64 or 128. Speed is reported in TFLOPs/s, with the forward FLOPs counted as:

per sequence: two matrix multiplies of $2N^2d$ FLOPs each. With the causal mask the count is halved, and the backward pass is counted as 2.5 times the forward. Since every implementation in a panel is assigned the same count, a ratio of TFLOPs/s is a ratio of runtimes. The baselines are a standard PyTorch implementation, FlashAttention, the FlashAttention implementation in xformers ("cutlass"), and FlashAttention in Triton.

The highest forward number is 227 TFLOPs/s (no mask, $d = 128$, 2k tokens), 73% of the 312 peak; the paper rounds it to 230. The highest backward number is 196 TFLOPs/s, 63%. Over FlashAttention, the forward-plus-backward ratio ranges from 1.45× (no mask, $d = 64$, 512 tokens) to 2.46× (no mask, $d = 128$, 2k tokens), and the forward pass alone reaches 3.39× (224 against 66 TFLOPs/s at 1k). The paper summarizes this as "around 2×" and as 1.7 to 3.0×. The forward ratio is largest at $d = 128$: at 16k tokens, going from $d = 64$ to $d = 128$ slows FlashAttention's forward pass from 104 TFLOPs/s to 73 and speeds FlashAttention-2's from 192 to 223.

Against PyTorch, the forward-plus-backward speedup is 2.3 to 9.2×, close to the paper's "3 to 10×". In the causal forward pass PyTorch runs at 10 TFLOPs/s and FlashAttention-2 at 181, an 18× gap at 8k tokens. PyTorch computes the full matrix and then masks it, so it does twice the work it is credited with.

End to end, the paper trains GPT-style models of 1.3B and 2.7B parameters on 8 A100s (Table 1, TFLOPs/s per GPU):

At 2k context FlashAttention-2 is only 1.04× (1.3B) and 1.08× (2.7B) faster than FlashAttention, because attention is a small part of each layer and the matrix multiplies of the MLP and projections dominate. At 8k it is 1.29× faster in both models. Against training without FlashAttention, the speedup at 8k is 3.1× for the 1.3B model and 2.8× for the 2.7B model; the paper quotes the 2.8×.

The model FLOPs come from the formula used by Megatron-LM, per sequence of length $s$:

The first term is the forward and backward matrix multiplies against the weights, the second is attention. The paper notes that the second term is not halved for the causal mask, as the attention benchmark's count is, and keeps it unhalved for consistency with the literature. That affects how to read 225 TFLOPs/s as 72% of peak. Take the GPT-3 paper's shape for the 2.7B model, 32 layers of width 2560, at $s = 8192$: per token the weight term is about $1.62 \times 10^{10}$ FLOPs and the attention term $8.05 \times 10^9$. Halving the attention term lowers the total by 17%, so the 225 TFLOPs/s of counted work corresponds to about 188 TFLOPs/s of work actually performed, 60% of peak.

Run unchanged on an H100, without its Tensor Memory Accelerator (TMA, a hardware unit for asynchronous copies between HBM and shared memory) or fourth-generation Tensor Cores, the same kernel reaches up to 335 TFLOPs/s. The paper expects another 1.5 to 2× from using those features.

The kernel ships as the flash-attn Python package for NVIDIA Ampere, Ada and Hopper GPUs, in fp16 or bf16, with head dimensions up to 256. Inputs are laid out as (batch, sequence, heads, head dimension). Passing fewer key and value heads than query heads runs grouped-query attention without copying $K$ and $V$: the kernel maps each query head to its shared key-value head by index, and the backward pass sums the gradients of the query heads that share one, the variant covered in the Multi-Query Attention explainer.

# pip install flash-attn --no-build-isolation   (Ampere/Ada/Hopper; fp16 or bf16)
import torch
from flash_attn import flash_attn_func

q = torch.randn(1, 16384, 16, 128, device="cuda", dtype=torch.bfloat16)
k = torch.randn(1, 16384, 4, 128, device="cuda", dtype=torch.bfloat16)  # 4 KV heads
v = torch.randn_like(k)
out = flash_attn_func(q, k, v, causal=True)   # shape (1, 16384, 16, 128)
# softmax_scale defaults to 1/sqrt(d); q's head count must be a multiple of k's

Limits and what came next

The paper is a systems report. It has no ablation: the 2× over FlashAttention is measured for the whole package, a new implementation built on NVIDIA's CUTLASS 3.x library, and the three changes are never measured one at a time. By the count in Figure 1 the non-matmul change is worth a few percent, and the occupancy argument applies most clearly to the original FlashAttention algorithm rather than the released baseline.

The numbers are for the A100. Block sizes are tuned by hand for each head dimension, and the kernel uses none of the H100's new features. The backward pass, at 63% of peak at best, remains the slower half. The authors list their next steps as H100 features (TMA, fourth-generation Tensor Cores, FP8), other hardware such as AMD GPUs, combining the kernel with sparse or local attention patterns, and compilers that can produce such kernels from a high-level description.

FlashAttention-3 (Shah, Bikshandi, Zhang, Thakkar, Ramani and Dao, 2024) did the H100 part. It overlaps the softmax of one tile with the matrix multiplies of another, using warp specialization and the H100's asynchronous Tensor Cores and TMA, and reaches up to 740 TFLOPs/s in FP16 (75% of the H100's peak) and close to 1.2 PFLOPs/s in FP8.