The first FlashAttention was 2 to 4 times faster than standard attention, and still slow. Its forward pass reached only 30 to 50% of an A100's peak FLOP/s, and its backward pass 25 to 35%, while a well-tuned matrix multiplication reaches 80 to 90% (Dao 2023). The algorithm of §12.3 had fixed the memory traffic. What remained was a question of scheduling, of which processor does which tile, in what order, and what it waits for.

This section follows the two sequels. FlashAttention-2 (Dao 2023) rearranged the work to spend less time on operations that are not matrix multiplications and to keep every SM busy, which doubled the speed. FlashAttention-3 (Shah et al. 2024) rebuilt the kernel around the H100's asynchronous hardware and its FP8 tensor cores. Then we turn to practice, the kernels you will actually call from PyTorch, how to batch sequences of different lengths without padding, why decoding a single token needs yet another variant, and a preview of the paged caches that serving engines feed it.

Most of this is GPU engineering, and the details live in CUDA. But each change follows from a cost that you can compute on paper, and that is what an interviewer will ask you to do. Throughout, a thread block processes a query block of $B_r$ rows against key and value blocks of $B_c$ rows, as in §12.3, and $d$ is the head width, the $d_k$ of earlier chapters, written $d$ as in the FlashAttention papers.

Where the time went, three leaks

Profile the first FlashAttention on an A100 and three problems show up.

Non-matrix-multiply work is expensive. An A100 does 312 TFLOP/s of fp16 or bf16 matrix multiplication on its tensor cores, but only 19.5 TFLOP/s of other fp32 arithmetic. Each non-matmul FLOP costs 16 times as much as a matmul FLOP. The softmax's maximum, subtraction, exponential, sum, and the rescaling of the output are all non-matmul work, so even a small count of them takes a noticeable share of the time.

Not enough parallel work. The first FlashAttention gave each (sequence, head) pair to one thread block, which then looped over all the tiles. An A100 has 108 SMs. With a batch of 2 sequences and 16 heads there are only 32 thread blocks, and 70% of the GPU sits idle. Long sequences make it worse, because they force small batches to fit in memory.

Warps that must talk to each other. Inside a thread block, four or eight warps (groups of 32 threads, §11.1) share the tile. The first FlashAttention split the keys among the warps, so each warp computed part of every query's sum and the partial results had to be written to shared memory, synchronized, and added. Dao calls this the "split-K" scheme.

Each leak has a fix, and together they are FlashAttention-2.

FlashAttention-2, one fix per leak

Fewer non-matmul FLOPs

Two changes from §12.2 and §12.3 do it. First, keep the output unnormalized during the loop, rescaling $\tilde{\mathbf{O}}$ only by $e^{\mathbf{m} - \mathbf{m}^{\text{new}}}$ and dividing by $\boldsymbol{\ell}$ once at the end, instead of dividing the whole $B_r\times d$ block at every step. Second, store only the log-sum-exp $\operatorname{lse} = \mathbf{m} + \log\boldsymbol{\ell}$ for the backward pass, so each recomputed weight is one subtraction and one exponential, $e^{s - \operatorname{lse}}$. Neither changes the result. Together with skipping the masked causal tiles, they remove most of the non-matmul work that is not strictly needed.

How much does what remains cost? Per tile, the two products do $4B_rB_cd$ FLOPs, and the softmax does a handful of operations per score, say 5, so $5B_rB_c$. Weighting the non-matmul operations 16 times more,

which is 31% for $d = 64$ and 16% for $d = 128$. In plain words, the exponentials and their bookkeeping are a sizable fraction of the work even though they are a tiny fraction of the FLOPs, and the fraction shrinks as the head width grows, because the products grow with $d$ and the softmax does not. Every non-matmul operation removed is worth sixteen matmul operations.

The cell puts numbers on it for both GPUs, using each one's matmul rate and its rate for other arithmetic. (On the H100 the exponential alone is far slower than general fp32 arithmetic, a problem we return to with FlashAttention-3.)

Parallelism over the sequence

FlashAttention-2 uses the loop order of §12.3, query blocks outside, and makes each query block its own thread block. The launch grid becomes (number of query blocks) × (batch × heads), so a batch of 2 sequences of 8,192 tokens with 16 heads and $B_r = 128$ launches $64\times32 = 2{,}048$ thread blocks, enough to fill 108 SMs nineteen times over. The query blocks are independent (each has its own running statistics), so no communication is needed between them. Dao credits this loop order and the sequence-length parallelism to Phil Tillet's Triton implementation.

The cell compares the two launch grids for a few shapes on an A100, assuming one thread block per SM at a time and blocks of equal cost.

For a large batch of short sequences the two are equivalent, and FlashAttention-1 was designed for that regime. For long sequences and small batches, which is where attention costs the most, FlashAttention-2's grid fills the GPU and FlashAttention-1's does not.

The backward pass parallelizes the same way over key blocks, since $\mathrm{d}\mathbf{K}_j$ and $\mathrm{d}\mathbf{V}_j$ belong to one key block. The price is that many thread blocks add into the same $\mathrm{d}\mathbf{Q}_i$, which is done with atomic additions in HBM, as in §12.3.

Split the queries, not the keys

Within a thread block, FlashAttention-2 gives each warp a slice of the queries and lets every warp read all of the current key and value block. Now each warp owns complete rows. Its running maximum, sum, and output never need anyone else's, so the warps never exchange partial results through shared memory. In the language of §12.2, split-K forces a merge $\oplus$ of partial softmax summaries across warps at every tile, and split-Q needs none, because a row's summary lives entirely in one warp.

What it bought

Dao (2023) reports about a 2 times speedup over FlashAttention, 1.7 to 3.0 times depending on the setting, reaching 50 to 73% of the A100's theoretical peak, with up to 73% in the forward pass and 63% in the backward pass. Training GPT-style models end to end, it reached up to 225 TFLOP/s per A100, 72% model FLOPs utilization. Block sizes are typically $\{64, 128\}\times\{64, 128\}$, with 4 or 8 warps per thread block, tuned per head width, because larger blocks reduce shared-memory traffic until they run out of registers.

def grid_utilization(batch, heads, T, Br, n_sm, version):
    raise NotImplementedError

Decoding needs a different split

Generation breaks FlashAttention-2's parallelism. During decoding each sequence has a single new query (§6.5), so there is one query block per sequence and head, and the grid is batch × heads again. Dao, Haziza, Massa, and Sizov (2023) note that with a batch of one, FlashAttention "will use less than 1% of the GPU". Meanwhile each of those few thread blocks streams the entire KV cache, possibly hundreds of thousands of keys, by itself.

Flash-Decoding splits the keys instead. The cache is cut into chunks, each chunk is processed by its own thread block, which produces a partial output and its log-sum-exp, and a final small kernel merges the partials. The merge is exactly the log-sum-exp form from §12.2,

In plain words, each chunk reports what it found and how much softmax mass it holds, and the final output is the mass-weighted average of the chunks' outputs. This is split-K again, the scheme FlashAttention-2 removed from inside a thread block, now used across thread blocks because at decode time there is nothing else to parallelize over.

Here is the merge at work on one decoding step, with the keys split into eight chunks that could each run on a different SM.

On CodeLlama-34B, Dao and colleagues measured up to 8 times faster generation for very long sequences, with the attention operation itself up to 50 times faster than FlashAttention. The same reduction is what the serving engines that decode many sequences at once build on, with the cache stored in fixed-size pages, as the next subsection describes.

import numpy as np

def partial_attention(q, K, V):
    s = K @ q / np.sqrt(len(q)); w = np.exp(s); return w @ V / w.sum(), np.log(w.sum())

def combine(outs, lses):
    return np.mean(outs, axis=0), np.max(lses)        # wrong: equal weights

Paged KV caches, a preview of serving

How a serving system schedules hundreds of requests and manages their memory belongs to continuous batching and paged memory management, a later subject in its own right. One idea from it changes the attention kernel, though, so we preview it here. A serving engine decodes hundreds of sequences at once, and it does not know in advance how long each one will grow. Reserving a contiguous buffer of the maximum length for every sequence wastes most of the memory. Kwon et al. (2023) measured that earlier serving systems used only about 20 to 40% of their KV cache memory for actual tokens, with the rest lost to over-reservation and fragmentation.

PagedAttention, the kernel behind vLLM, borrows the operating system's answer. The cache is a pool of fixed-size pages (blocks of 16 tokens' keys and values by default) that can sit anywhere in HBM, and each sequence has a block table that maps its logical blocks (tokens 0 to 15, 16 to 31, and so on) to physical pages. A new page is taken from the free list only when the sequence's last page fills, so at most one partly filled page per sequence is wasted, under 4% of the cache in their measurements. Sequences that share a prompt can share its pages and copy one only when they would write to it. With the memory freed for larger batches, vLLM reported 2 to 4 times the throughput of earlier systems at the same latency.

The attention kernel barely changes. Instead of reading keys at consecutive addresses, it looks up each block's page in the table and reads the tile there. That is a gather at page granularity, so each page is still one contiguous, coalesced read, and everything else (the online softmax, and the split over keys of Flash-Decoding) is untouched. The flash-attn package's flash_attn_with_kvcache accepts a block table for exactly this purpose. Admission, preemption, eviction, prefix caching, and how pages are scheduled across a batch are the serving side of the story, and they are left to the treatment of serving.

import numpy as np

def append_token(k_pages, v_pages, block_table, seq_len, k, v, free_pages):
    raise NotImplementedError

def gather(pages, block_table, seq_len):
    return pages[block_table].reshape(-1, pages.shape[-1])      # keeps the garbage slots

FlashAttention-3, built for asynchrony

FlashAttention-2 reaches only about 35% of an H100's peak (Shah et al. 2024). The H100 changed the rules in three ways that a kernel written for the A100 does not exploit.

• Asynchronous tensor cores. The WGMMA instruction runs a matrix multiply for a warpgroup (four warps) and returns immediately. The warps can do other work while the tensor cores compute.
• An asynchronous copy engine. The Tensor Memory Accelerator (TMA) moves tiles between HBM and shared memory on its own, without threads computing addresses.
• FP8 tensor cores at twice the bf16 rate, about 1,979 dense TFLOP/s.

And it made one old problem much worse. Shah et al. compute it. The H100 SXM5 has 989 TFLOP/s of fp16 matrix multiply but only 3.9 TFLOP/s of special functions such as the exponential. In the forward pass with $d = 128$, a tile does $4B_rB_cd = 512\,B_rB_c$ matmul FLOPs and $B_rB_c$ exponentials, so there are 512 times more matmul FLOPs, but the exponential is 256 times slower. If they ran one after the other, the exponentials would take half as long as the products. In FP8 the products get twice as fast and the exponentials do not.

The answer is to run them at the same time, and FlashAttention-3 has three techniques for it.

1. Warp specialization. Some warps become producers that only issue TMA loads of $\mathbf{K}_j$ and $\mathbf{V}_j$ into a circular buffer of shared memory. The rest are consumers that only compute. Producers stay a few tiles ahead, so consumers rarely wait for memory. Registers are reallocated toward the consumers, since a producer needs almost none.
2. Ping-pong between warpgroups. With two consumer warpgroups, the kernel schedules them so that while warpgroup 1 does the softmax of its tile on the special-function units, warpgroup 2 runs its matrix products on the tensor cores, and then they swap.
3. Pipelining inside a warpgroup. Within one warpgroup, the softmax of tile $j$ overlaps with the product $\mathbf{Q}_i\mathbf{K}_{j+1}^\top$ for the next tile, which is issued asynchronously before the softmax starts. This takes extra registers to hold two tiles of scores at once.

In the paper's ablation (one fixed problem size with $d = 128$), the full kernel ran at 661 TFLOP/s, 582 without the softmax-matmul pipelining, and 570 without warp specialization.

FP8, and the outlier problem

FP8 (the e4m3 format has 4 exponent bits and 3 mantissa bits, with a largest value of 448) doubles throughput but has very little precision and range, and attention's inputs have outliers, a few coordinates of $\mathbf{Q}$ and $\mathbf{K}$ much larger than the rest. FlashAttention-3 uses two tricks.

Block quantization. Instead of one scale per tensor, keep one scale per block of $\mathbf{Q}$, $\mathbf{K}$, and $\mathbf{V}$, the blocks the kernel loads anyway, so an outlier only inflates the scale of its own block.

Incoherent processing. Multiply $\mathbf{Q}$ and $\mathbf{K}$ by the same random orthogonal matrix $\mathbf{R}$ before quantizing. Since $\mathbf{R}\mathbf{R}^\top = \mathbf{I}$,

so the scores are unchanged in exact arithmetic, while each outlier's magnitude is spread over all $d$ coordinates. FlashAttention-3 uses $\mathbf{R}$ = a Hadamard matrix times a random diagonal of $\pm1$, which costs $O(d\log d)$ per vector with the fast Walsh-Hadamard transform and can be fused with the rotary embedding. The idea comes from quantizing weights in QuIP (Chee et al. 2023).

On a test with deliberately injected outliers (each entry $\mathcal{N}(0,1)$, plus an independent $\mathcal{N}(0, 100)$ term for 0.1% of entries), FP8 FlashAttention-3 had a root mean squared error of $9.1\times10^{-3}$, against $2.4\times10^{-2}$ for a baseline FP8 attention with per-tensor scaling, 2.6 times lower. In that ablation incoherent processing did most of the work. Without it the error was back to $2.4\times10^{-2}$, while without block quantization it was $9.3\times10^{-3}$. There is one more practical wrinkle. FP8 matrix multiplies on Hopper accept operands only in one memory layout, so the kernel transposes tiles of $\mathbf{V}$ in shared memory and reshuffles the registers that hold $\mathbf{P}$ before the second product.

What it bought

FlashAttention-3 is 1.5 to 2.0 times faster than FlashAttention-2 on H100 in fp16, reaching up to 740 TFLOP/s (75% of peak), and close to 1.2 PFLOP/s in FP8. The backward pass is 1.5 to 1.75 times faster. These are kernel-level numbers on one GPU. The lesson that generalizes is that on modern hardware a kernel is fast when its memory traffic, its tensor-core work, and its other arithmetic all overlap, and the designer's job is to arrange the dependencies so they can.

And on Blackwell

The pattern continues with each GPU generation. FlashAttention-4 (2025) targets NVIDIA's Blackwell GPUs, and public accounts of it come mainly from its open-source code and its authors' talks, so treat the details as provisional. Two of its reported ideas follow directly from this section and §12.2. It computes part of the exponentials in software on the ordinary floating-point units, writing $2^x$ as a power of two with an integer exponent times $2^f$ for the fractional part $f$, approximated by a cubic polynomial, to take load off the special-function units. And it rescales the accumulator only when the running maximum has grown by enough to threaten overflow, which reportedly cuts the number of rescalings about tenfold. That shortcut is exact. As §12.2 pointed out, any reference point works as long as $\boldsymbol{\ell}$ and $\tilde{\mathbf{O}}$ share it, so a stale maximum only makes some exponentials somewhat larger than 1, which fp32 accumulators absorb, and the final division by $\boldsymbol{\ell}$ cancels the choice.

def exp_to_matmul_time(d, matmul_flops, exp_flops):
    return 1 / (4 * d)            # counts FLOPs, ignores the rates

def overlap_speedup(d, matmul_flops, exp_flops):
    return 1.0

Using it, the real way

You will rarely write an attention kernel. You will call one, and the details of how you call it decide whether you get FlashAttention or something much slower.

In PyTorch the entry point is torch.nn.functional.scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False), with tensors shaped (batch, heads, T, d_k). It chooses among several backends, which you can list and restrict with the torch.nn.attention.sdpa_kernel context manager and the SDPBackend enum.

Three details trip people up.

• Passing a mask can cost you the kernel. FlashAttention kernels support causal masks, sliding windows, and variable lengths natively, but not an arbitrary dense mask. Passing attn_mask for something that is_causal=True could express may silently move you to a slower backend. Check with sdpa_kernel([SDPBackend.FLASH_ATTENTION]), which raises an error instead of falling back.
• is_causal is aligned to the top left. With 4 new queries and 128 keys, is_causal=True lets the first query see only key 0, which is wrong for chunked prefill or decoding with a cache. Use torch.nn.attention.bias.causal_lower_right there (§4.4 derives the bottom-right rule). The flash-attn package's causal=True is aligned bottom-right, the opposite convention.
• Grouped-query attention. enable_gqa=True (PyTorch 2.5 and later) accepts fewer key-value heads than query heads and shares them in groups, without copying $\mathbf{K}$ and $\mathbf{V}$ (§12.5).

For masks and score modifications that SDPA cannot express, PyTorch 2.5 added FlexAttention. You write the mask as a Python predicate on (b, h, q_idx, kv_idx), or a score modification as a function of the score and those indices, and torch.compile generates a FlashAttention-style kernel. create_block_mask evaluates the predicate ahead of time to find which 128 by 128 tiles are entirely masked, so the kernel skips them, which is block-sparse FlashAttention (Exercise 12.3d) with the sparsity pattern written as code.

The printed block mask shows which 128 by 128 tiles the kernel would visit. With a 64-token window, only the diagonal tiles and the ones just below them survive.

The standalone flash-attn package exposes the kernels directly, with tensors in (batch, T, heads, d_k) layout, and adds options SDPA lacks. The block below shows the three calls you will meet in model code. It needs a CUDA GPU.

Variable-length batches without padding

Real batches hold sequences of different lengths. Padding them to the longest wastes compute on pad tokens, and a padding mask (§4.4) only stops the pads from mattering, not from costing. With lengths $L_1, \dots, L_B$, padded attention computes $B\cdot(\max_b L_b)^2$ scores per head, while the useful work is $\sum_b L_b^2$. For lengths 1,000, 3,000, and 96 that is 27 million scores for 10 million useful ones.

The variable-length ("varlen") kernels instead take the sequences packed end to end into one tensor of shape (total_tokens, heads, d_k), plus a vector of cumulative offsets cu_seqlens $= [0, L_1, L_1 + L_2, \dots]$ (int32, of length $B + 1$). Sequence $b$ occupies rows cu_seqlens[b] to cu_seqlens[b + 1]. Mathematically this is attention with the block-diagonal document mask of §4.4. The kernel never visits the tiles between different sequences, because each thread block works within one sequence's offsets, so the cost is $\sum_b L_b^2$ and there is no mask to build at all. This is also how packed pretraining with document masking is run at scale, and how encoders such as ModernBERT drop padding entirely (§7.1).

PyTorch's equivalent is a nested tensor with the jagged layout, which SDPA accepts directly. Here is the same idea on a CPU.

import numpy as np

def make_cu_seqlens(lengths):
    return np.cumsum(lengths).astype(np.int32)

def varlen_attention(q, k, v, cu_seqlens, causal=False):
    S = np.einsum("thd,shd->hts", q, k) / np.sqrt(q.shape[-1])     # attends across sequences
    P = np.exp(S - S.max(-1, keepdims=True)); P /= P.sum(-1, keepdims=True)

Read more

• Dao (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. Section 3 covers the three changes, with the non-matmul FLOP argument and the warp partitioning figure.
• Shah et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. Warp specialization, ping-pong scheduling, the exponential bottleneck, and FP8 with incoherent processing.
• Dao, Haziza, Massa & Sizov (2023). Flash-Decoding for long-context inference. The split-KV kernel for decoding, with benchmarks on CodeLlama-34B.
• Kwon et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. Block tables, copy on write, and the measurements of wasted cache memory behind vLLM.
• Dong et al. (2024). Flex Attention: A Programming Model for Generating Optimized Attention Kernels. How FlexAttention turns mask and score functions into block-sparse fused kernels.
• PyTorch documentation. torch.nn.functional.scaled_dot_product_attention. Backends, masks, and the exact semantics of is_causal and enable_gqa.
• Chee et al. (2023). QuIP: 2-Bit Quantization of Large Language Models With Guarantees. The incoherence processing idea that FlashAttention-3 borrows for FP8.

Watch next

▶

▶