FlashAttention-3 and PagedAttention: IO-Aware Attention Kernels (2026)

Attention does not run slowly because the arithmetic is hard. It runs slowly because the arithmetic is cheap and the data is far away. A modern GPU can perform hundreds of floating point operations in the time it takes to fetch a single byte from its main memory, and naive attention fetches an enormous intermediate matrix several times over. FlashAttention fixed this in 2022 by restructuring the computation so that the matrix never leaves the chip, and nearly every serious training and inference stack now depends on it.

The story did not end there. FlashAttention-2 and FlashAttention-3 retuned the kernel for successive GPU generations, PagedAttention rebuilt how the key-value (KV) cache is stored, and Flash-Decoding fixed the single-query case that tiling alone could not. Together these attention kernels define what your GPU memory hierarchy can actually deliver.

This post derives the IO argument from numbers, walks the online softmax by hand, compares the kernel generations, explains block tables and split-KV decoding, and ends with runnable code and failure modes.

What this covers: the GPU memory hierarchy, IO complexity of attention, tiling with online softmax, FlashAttention-2, -3 and -4, PagedAttention, Flash-Decoding, GQA and MLA effects, a hardware matrix, a PyTorch comparison, and what goes wrong.

Context and Background

Self-attention computes, for each query, a weighted average of value vectors, where the weights come from a softmax over query-key dot products. For a sequence of length N and head dimension d, the score matrix is N by N. Compute grows with N squared times d, and the intermediate matrix grows with N squared. Early efficiency work attacked the FLOP count with sparse, low-rank and kernelized approximations, and most of it lost to the exact method in practice because it reduced arithmetic that was never the bottleneck.

The 2022 FlashAttention paper by Tri Dao, Daniel Fu, Stefano Ermon, Atri Rudra and Christopher Re made the opposite move. It kept attention exact and reduced memory traffic instead, and it framed the work as “IO-aware”: count reads and writes between the GPU’s large, slow high-bandwidth memory (HBM) and its small, fast on-chip SRAM, then minimize them. The abstract reports a 15% end-to-end speedup on BERT-large at sequence length 512 over the MLPerf 1.1 record, a 3x speedup on GPT-2 at length 1K, and 2.4x on Long Range Arena. It also enabled the first Transformers to beat chance on Path-X at length 16K.

Since then the kernel has become infrastructure. PyTorch exposes it through scaled_dot_product_attention, inference engines call it directly, and its successors target each new NVIDIA generation. If you operate an inference stack, such as the ones covered in our posts on continuous batching for LLM inference and the vLLM model runner v2 migration, the attention kernel is the component whose behavior you inherit rather than write.

There is a second half of the problem that FlashAttention does not touch: where the KV cache lives while a model generates text. A serving system holds the keys and values of every token of every in-flight request, and those bytes dominate memory at scale. PagedAttention, introduced with vLLM by Woosuk Kwon and colleagues at SOSP 2023, borrowed virtual memory from operating systems to manage that cache without fragmentation. The two ideas are complementary. One makes each attention computation cheap in memory traffic; the other makes the cache that feeds it dense and shareable.

A third concern is hardware drift. FlashAttention-2 reached only 35% utilization on H100, according to the FlashAttention-3 authors, because Hopper added asynchronous units that older kernels ignore. FlashAttention-4, posted in March 2026, argues Blackwell shifts the bottleneck again. Kernel design is a moving target, and understanding the principles lets you read each new paper as a variation on one theme: keep the tensor cores fed.

FlashAttention and the GPU Memory Hierarchy: Why Attention Is an IO Problem

Attention is memory-bound because its arithmetic intensity, the number of floating point operations per byte moved from HBM, sits far below the GPU’s ridge point. Standard attention writes and rereads an N by N score matrix in HBM; FlashAttention computes the same exact result in tiles held in on-chip SRAM, so HBM traffic falls from quadratic to linear in sequence length for the intermediates.

Figure 1: The GPU memory hierarchy from registers to host memory, and where naive attention versus FlashAttention keep their intermediates.

The figure shows capacity growing and speed falling as you move away from the tensor cores. Registers and shared memory sit inside each streaming multiprocessor (SM). The L2 cache is shared across SMs. HBM holds model weights, activations and the KV cache. Naive attention materializes its scores in the far tier; FlashAttention keeps its working tiles in the near tier.

The three tiers that matter

Three levels explain almost every attention kernel decision. Registers hold the accumulators that a warp updates every instruction. Shared memory, the programmable on-chip SRAM, holds the tiles that threads in a block cooperate on; on Hopper it is on the order of 228 KB per SM, and on Ampere A100 on the order of 192 KB per SM of combined L1 and shared memory, per NVIDIA’s architecture documentation. HBM is large but distant: an H100 SXM carries 80 GB at 3.35 TB/s according to NVIDIA’s product page.

The compute side is the other half of the ratio. H100’s dense BF16 tensor core throughput is about 989 teraFLOPS in NVIDIA’s datasheet; the 1,979 teraFLOPS shown on the product page is the figure with structured sparsity. As a consistency check, the FlashAttention-3 paper’s 740 TFLOPs/s at “75% utilization” implies a peak near 989, which matches the dense number. Dividing, the H100 ridge point is roughly 989e12 / 3.35e12, or about 295 FLOPs per byte. An A100 80 GB at 312 dense TFLOPs and roughly 2.0 TB/s has a ridge point near 153. Any kernel below those intensities is waiting on memory, not on math.

The arithmetic of waste: standard attention

Take one attention head with N = 4096 tokens, d = 128, in FP16. The two matmuls, QK-transpose and PV, cost about 2·N²·d each, so roughly 4·N²·d = 4 × 16,777,216 × 128, or 8.59 GFLOPs for the forward pass of this head. The mandatory traffic is reading Q, K and V and writing O: four tensors of 4096 × 128 × 2 bytes, 4.19 MB in total.

Naive implementation adds the score matrix. S has 4096² = 16.78 million entries, 33.5 MB in FP16. The sequence is: write S after the first matmul, read S for softmax, write P, read P for the second matmul. That is four passes, about 134 MB, which dwarfs the 4.19 MB of necessary traffic. Total traffic is about 138 MB, giving an arithmetic intensity of 8.59e9 / 138e6, about 62 FLOPs per byte.

Against the H100 ridge point of 295 that is a 4.7x shortfall. Memory time is 138 MB / 3.35 TB/s, about 41 microseconds; compute time at peak is 8.59 GFLOPs / 989 TFLOPs, about 8.7 microseconds. The head spends roughly 80% of its time moving a matrix that exists only to be consumed twice. These are illustrative first-principles numbers that ignore caches, launch overhead and the exact tensor-core efficiency, but the ratio is what matters, and it worsens as N grows because S scales quadratically while Q, K, V and O scale linearly.

The FlashAttention paper states the same result asymptotically. Standard attention needs Θ(N·d + N²) HBM accesses. FlashAttention needs Θ(N²·d² / M), where M is the SRAM size in elements, and the authors show this is optimal across a range of SRAM sizes. For typical d and M, d² / M is much smaller than one, so the quadratic term shrinks by a large factor. Equally important, the N by N matrix is never stored, so the memory footprint of the intermediates becomes linear in N.

Tiling and online softmax, worked by hand

The obstacle to tiling is softmax. Its denominator needs a sum over the entire row, so you cannot finish any output element until you have seen every key. The trick, borrowed from earlier work on online normalizer calculation, is to maintain a running maximum m and running sum l, and to rescale partial results whenever the maximum changes.

Figure 2: The FlashAttention forward loop. One query tile stays resident while K and V tiles stream through shared memory, with running statistics rescaling the accumulator.

Walk a row of four scores, [2, 1, 3, 0], split into two blocks of two. Real kernels subtract the running maximum before exponentiating to stay numerically stable, so the example below works with exponentials relative to a running max.

After block one, [2, 1], the running max is m₁ = 2 and the sum is l₁ = e⁰ + e⁻¹ = 1 + 0.3679 = 1.3679. The block’s unnormalized output accumulator o₁ holds the sum of e^(s − 2) times the value vectors.

Block two, [3, 0], has a local max of 3, which raises the global max to m₂ = 3. The old sum must be rescaled by e^(m₁ − m₂) = e⁻¹ = 0.3679, giving 0.5032. The new block contributes e⁰ + e⁻³ = 1 + 0.0498 = 1.0498. So l₂ = 0.5032 + 1.0498 = 1.5530. Direct computation, relative to max 3, gives e⁻¹ + e⁻² + e⁰ + e⁻³ = 0.3679 + 0.1353 + 1 + 0.0498 = 1.5530. They agree, and the same rescale factor multiplies the output accumulator so the final division by l is exact.

This is why FlashAttention is exact rather than approximate: it reorders the computation, but each row’s result equals standard softmax attention up to floating point rounding. The backward pass does not store the N by N probability matrix either. It stores only the output and a per-row log-sum-exp, then recomputes tiles of S and P on the fly. Recomputation costs extra FLOPs but saves far more memory traffic, a trade that favors the memory-bound regime.

FlashAttention-2, FlashAttention-3 and FlashAttention-4: What Changed per Generation

Each generation keeps the same exact algorithm and changes how work maps to hardware. FlashAttention-2 fixed work partitioning on Ampere, FlashAttention-3 exploited Hopper’s asynchronous units and FP8, and FlashAttention-4 targets Blackwell’s asymmetric scaling. The recurring lesson is that after IO is minimized, the next bottleneck is keeping tensor cores busy while softmax and data movement run alongside.

FlashAttention-2: better parallelism and work partitioning

Tri Dao’s FlashAttention-2 paper reports that the original kernel reached only 25-40% of theoretical peak FLOPs/s, well short of optimized GEMM. The cause was not IO anymore but suboptimal partitioning across thread blocks and warps. The follow-up made three changes, per the abstract. It reduced non-matmul FLOPs, because tensor cores run matmuls many times faster than the scalar units that perform exponentials and rescales, so every extra non-matmul operation is disproportionately expensive. It parallelized across the sequence-length dimension in addition to batch and heads, which matters when batch times heads is too small to fill the GPU. And it redistributed work among warps within a block to cut shared-memory communication.

The reported result is about a 2x speedup over the first version, 50-73% of theoretical peak on A100, and up to 225 TFLOPs/s per A100 in end-to-end GPT-style training at 72% model FLOPs utilization. Those are the authors’ measurements on specific configurations, so treat them as upper-range indications. One structural detail worth knowing is the loop order: FlashAttention-2 puts the query tile in the outer loop and streams K and V in the inner loop, so each output tile is finalized by a single thread block without cross-block reductions.

FlashAttention-3: asynchrony and low precision on Hopper

The FlashAttention-3 paper by Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani and Tri Dao, posted in July 2024, starts from a hardware observation: FlashAttention-2 achieved only 35% utilization on H100. Hopper added the Tensor Memory Accelerator (TMA) for asynchronous bulk copies between HBM and shared memory, and warpgroup-level asynchronous matrix multiply instructions, WGMMA. A kernel written for Ampere issues loads and math in lockstep and leaves both units partly idle.

The paper names three techniques. The first is overlapping computation and data movement through warp specialization: some warps act as producers that issue TMA loads, while consumer warps run the matmuls, so memory latency hides behind math. The second is interleaving block-wise matmul and softmax. Softmax needs the special function unit, which has far lower throughput than the tensor cores; while one block’s softmax runs, the next block’s matmul can already be issued, either across two warpgroups in a “ping-pong” schedule or within a single warpgroup. The third is FP8 support with block quantization and incoherent processing, where a random orthogonal transform spreads outlier values across the vector before quantization so that per-block scales lose less accuracy.

The reported outcome on H100 is a 1.5-2.0x speedup over FlashAttention-2, FP16 reaching up to 740 TFLOPs/s at 75% utilization, and FP8 approaching 1.2 PFLOPs/s. The authors also report FP8 numerical error 2.6x lower than a baseline FP8 attention. Note the framing: “up to”, on H100, with their chosen shapes. Short sequences and small head counts see smaller gains, and the FlashAttention repository still describes the Hopper kernel as a beta with FP8 forward only, so check release notes before depending on a backward pass in FP8.

FlashAttention-4: Blackwell and the asymmetric-scaling problem

The FlashAttention-4 paper, submitted in March 2026 by Ted Zadouri, Markus Hoehnerbach, Jay Shah, Timmy Liu, Vijay Thakkar and Tri Dao, argues that on Blackwell tensor core throughput doubles while shared memory bandwidth and exponential units scale more slowly or not at all. The softmax exponential and the shared-memory traffic therefore become the limiting factors. The paper reports up to 1.3x over cuDNN 9.13 and up to 2.7x over Triton on B200 in BF16, reaching up to 1,613 TFLOPs/s at 71% utilization.

Its techniques follow the same theme as FlashAttention-3 but adapted: fully asynchronous MMA pipelines with larger tiles, software-emulated exponentials and conditional softmax rescaling to reduce non-matmul work, and use of tensor memory and a 2-CTA MMA mode to cut shared-memory traffic in the backward pass. The kernel is written in CuTe-DSL embedded in Python, with the authors reporting 20-30x faster compile times than C++ template approaches. I could not independently reproduce any of these numbers, so read them as the authors’ claims against the stated baselines.

Reading the generations as a table

Generation	Primary target	Main idea	Reported headline (authors’ claims)
FlashAttention (2022)	A100 era	Tiling plus online softmax to cut HBM traffic	3x on GPT-2 at 1K, 15% on BERT-large vs MLPerf 1.1
FlashAttention-2 (2023)	Ampere, Ada, Hopper	Fewer non-matmul FLOPs, sequence parallelism, warp partitioning	About 2x over v1, up to 225 TFLOPs/s per A100 in training
FlashAttention-3 (2024)	Hopper	Warp specialization, softmax-matmul overlap, FP8	1.5-2.0x over v2, up to 740 TFLOPs/s FP16 on H100
FlashAttention-4 (2026)	Blackwell	Async MMA, emulated exp, tensor memory, 2-CTA	Up to 1.3x vs cuDNN 9.13, 1,613 TFLOPs/s on B200

The pattern is that algorithmic IO savings arrived once, in 2022, and everything since is a co-design with silicon. That is also why a kernel tuned for one architecture can lose to a generic one on another, and why the framework you use matters: PyTorch and inference engines choose among implementations at runtime.

A contrast from the other end of the scale

The same hierarchy thinking applies far below datacenter GPUs. A microcontroller running a quantized network has kilobytes of SRAM and megabytes of flash, and the whole discipline of scheduling layers to avoid spilling activations is the embedded cousin of IO-aware tiling. Our walk-through of TinyML on microcontrollers with TFLite Micro, INT8 quantization and CMSIS-NN makes the point at a scale of kilobytes instead of gigabytes: the principle is identical, keep data in the smallest memory that holds it and touch the larger one as rarely as possible.

PagedAttention: Virtual Memory for the KV Cache

PagedAttention stores each sequence’s keys and values in fixed-size blocks that need not be contiguous, and a per-sequence block table maps logical token positions to physical blocks. This removes the need to reserve a maximum-length slab per request, nearly eliminating internal and external fragmentation and allowing blocks to be shared between sequences.

Figure 3: Two sequences with separate block tables pointing into one pool of physical KV blocks. Physical block 7 is shared, for example as a common prompt prefix.

Why contiguous KV allocation wastes memory

During generation each new token appends one key and one value per layer per KV head. The final length is unknown when the request arrives, so a contiguous allocator must reserve for the maximum. The vLLM paper attributes the waste in prior systems to fragmentation and duplication, which limits batch size and therefore throughput.

Compute the per-token cost to see the scale. For a Llama-2-7B-style model with 32 layers, 32 KV heads, head dimension 128 and FP16, one token needs 2 (K and V) × 32 × 32 × 128 × 2 bytes = 524,288 bytes, which is 0.5 MiB. Reserving 4,096 tokens per request costs 2 GiB. If a chat request ends at 300 tokens, the reservation wasted about 93% of its slab. On an 80 GB GPU with, say, 14 GB of weights, a naive scheme could admit only a few dozen such requests even when the actual KV working set would fit hundreds.

Grouped-query attention changes the unit but not the logic. A Llama-3-8B-style layout with 32 layers, 8 KV heads and head dimension 128 needs 2 × 32 × 8 × 128 × 2 = 131,072 bytes, or 128 KiB, per token. That is 4x smaller, so more requests fit, yet reservation waste is still proportional to the over-reservation. These shapes are standard published configurations, but verify them against the model card for your specific checkpoint.

Block tables and what the kernel does

PagedAttention splits the cache into blocks of a fixed number of tokens; vLLM’s default is 16, which you can confirm in its documentation. A sequence of 300 tokens occupies ceil(300 / 16) = 19 blocks, covering 304 token slots, so the internal waste is 4 slots, about 1.3%. In the worst case the last block holds a single token and wastes 15 slots, so internal waste per sequence is bounded by one block minus one token no matter how long the sequence gets. External fragmentation vanishes because any free physical block can satisfy any request.

The cost moves into the kernel. A standard attention kernel assumes K and V for a sequence are contiguous. The paged kernel instead reads the block table, computes physical addresses block by block, and gathers keys and values. It adds an indirection per block and constrains tile sizes to align with the block size, which is why block size is a tuning parameter: smaller blocks waste less but increase table traffic and reduce coalescing; larger blocks do the opposite. The FlashAttention repository now supports a paged KV cache in its kernels, since version 2.5 per its README, so the two techniques compose rather than compete.

Because blocks are reference-counted, sequences can share them. Parallel sampling and beam search fork a sequence into many continuations that share the prompt blocks; when one continuation must write into a shared block, the system copies that block first, the same copy-on-write rule operating systems use. Prefix caching extends the idea across requests: if two prompts share a system prompt, the physical blocks holding its KV can be reused, saving both memory and prefill compute. The paper reports vLLM improving throughput by 2-4x over FasterTransformer and Orca at the same latency, with larger gains for longer sequences, larger models and more complex decoding.

The scheduler is the other half. With dense memory, a continuous-batching scheduler can admit requests until blocks run out and then preempt or swap when they do, a policy explored in our continuous batching architecture post. Once GPU memory is the constraint, the next tier is offloading cold blocks to CPU memory or storage, which we cover in disaggregated serving with NVIDIA Dynamo and LMCache KV cache offload. Block granularity is what makes that offload tractable: you move pages, not whole sequences.

Decode-Time Kernels: Flash-Decoding, GQA and MLA

Prefill and decode stress the GPU in opposite ways. Prefill processes thousands of query tokens at once and can be compute-bound with tiled kernels. Decode processes one query per sequence against a long cache, so the work is a streaming read of the KV cache, and the kernel must parallelize across the cache length rather than across queries.

Figure 4: Split-KV decoding. Chunks of the KV cache are processed in parallel, each emitting a partial output and a log-sum-exp, then a small reduction combines them exactly.

Why decode starves the GPU

A FlashAttention forward kernel parallelizes over batch, heads and query blocks. In decode the query length is 1, so with batch size 1 and 32 heads you launch about 32 thread blocks on a GPU with 132 SMs (H100 SXM). Most of the chip idles while a single block per head walks a 32K-token cache serially. The Flash-Decoding write-up from Stanford CRFM, by Tri Dao and collaborators, describes exactly this under-utilization at small batch sizes and long contexts.

The fix, shown in Figure 4, adds a parallel axis over the keys and values. Split the cache into chunks, run the usual FlashAttention inner loop on each chunk in a separate block, and have each emit its partial output plus one extra scalar per row, the log-sum-exp. A second small kernel combines the partials using those scalars as weights: each partial output is scaled by exp(lse_i − lse_total), which reproduces full-softmax attention exactly. This is the same rescaling identity as the online softmax, applied across blocks instead of across loop iterations.

The post reports up to 8x faster end-to-end decoding on CodeLlama-34B at very long sequences, and up to about 50x faster attention alone than FlashAttention at those lengths. The micro-benchmark was an A100 with FP16, batch size 1, 16 query heads, 2 KV heads and head dimension 128. These are “up to” numbers at the extreme of the regime; with short prompts or large batches, where there is already enough parallelism, the benefit disappears. Modern FlashAttention releases and inference engines fold this idea into their kernels and pick a split count heuristically.

The decode roofline in numbers

Decode intensity is brutally low. For one KV head and one cached position, K and V in BF16 occupy 4d bytes, and each of the g query heads sharing that KV head performs 2d FLOPs for the dot product and 2d for the weighted sum, 4d in total. Intensity is therefore 4d·g / 4d = g FLOPs per byte. Multi-head attention has g = 1; a GQA model with 4 query heads per KV head has g = 4. Against a ridge point near 295, decode attention runs at about 0.3% to 1.4% of the compute roof, so the only meaningful metric is how close the kernel gets to memory bandwidth.

Use the Llama-3-8B-style figure of 128 KiB per token. A 32K-token context holds 131,072 × 32,768 = 4.29 GB of KV. Reading it once at 3.35 TB/s takes about 1.28 ms per generated token per sequence, before any weights. The 16 GB of FP16 weights add about 4.8 ms per step at the same bandwidth, shared across the whole batch. That asymmetry is why batching helps weights but not KV: weight reads amortize over the batch, KV reads scale with it. At batch 8 and 32K context, KV reads alone cost about 10 ms per step, more than twice the weight reads. These are first-principles lower bounds on an idealized H100, not benchmarks.

How GQA and MLA reshape the problem

Grouped-query attention shrinks the KV cache by sharing each KV head across g query heads. It raises decode intensity by g and cuts both memory capacity and bandwidth demand, so it is the main reason recent open models run long contexts affordably. The kernel consequence is that the query-group dimension can be packed into the tile to reuse each loaded K and V row g times, which good decode kernels do.

Multi-head latent attention (MLA), introduced in DeepSeek-V2, compresses keys and values into a low-rank latent that is cached instead of full K and V. The DeepSeek-V2 paper reports a 93.3% reduction in KV cache relative to its prior dense model. The published DeepSeek-V3 configuration caches a 512-dimensional latent plus a 64-dimensional decoupled rotary key per token per layer, which I recall as 576 values; confirm against the model config before quoting. MLA changes the kernel contract: during decode the attention is effectively computed in the latent space with a very large effective head dimension and a single shared KV head, so generic kernels tuned for head dimensions of 64 to 256 do not apply directly, and engines ship dedicated MLA decode kernels. The takeaway is that “attention kernel” is a family, and the cache layout decides which member you need.

Hardware Support Matrix and Backend Selection

Which kernel you actually get depends on GPU generation, dtype and framework. The table summarizes what the FlashAttention README states, plus how PyTorch typically dispatches.

Hardware	FlashAttention-2	FlashAttention-3	FlashAttention-4	Notes
Ampere A100, Ada RTX 40	Yes, FP16 and BF16	No	No	BF16 needs Ampere or newer
Hopper H100, H800	Yes	Yes, beta; FP16 and BF16 forward and backward, FP8 forward only	Yes per README	FA3 needs CUDA 12.3 or newer, 12.8 recommended
Blackwell B200, GB200	Limited; check release	Not targeted	Yes, CuTe-DSL	Verify your installed version
Turing T4, RTX 20	Not in main repo	No	No	Community fork with a subset of features
AMD Instinct	Composable Kernel or Triton backend	Via Triton FA3 interface for FP8	No	Sliding window in progress on Triton

Treat the Blackwell FlashAttention-2 cell as unverified; I did not confirm it from primary documentation. Head dimension limits also apply: the FlashAttention-2 README states head dimensions up to 256, while the FA3 and FA4 README sections do not list limits, so check your model’s head size, including any non-standard sizes such as 192 or 576, before assuming a fast path exists.

In PyTorch, torch.nn.functional.scaled_dot_product_attention chooses among the FlashAttention-derived kernel, a memory-efficient kernel from xFormers lineage, a cuDNN backend on supported builds, and a plain math fallback. The torch.nn.attention.sdpa_kernel context manager lets you force a backend, which is the right way to test whether your inputs qualify. A silent fall back to the math path is among the most common reasons a model is slower than expected.

Runnable Example: Measuring the IO Gap

Two short programs make the earlier claims testable. The first is a NumPy reference that implements the tiled online-softmax forward pass and checks it against standard attention. I ran it on CPU; the maximum absolute difference over a 256 by 64 random problem was about 6e-16, confirming the exactness argument. It also includes a toy block-table gather mirroring Figure 3.

import numpy as np

def attention_ref(Q, K, V):
    S = Q @ K.T / np.sqrt(Q.shape[-1])
    S -= S.max(axis=-1, keepdims=True)
    P = np.exp(S)
    P /= P.sum(axis=-1, keepdims=True)
    return P @ V

def tiled_attention(Q, K, V, block=64):
    n, d = Q.shape
    O = np.zeros_like(Q)
    scale = 1.0 / np.sqrt(d)
    for qs in range(0, n, block):
        q = Q[qs:qs + block]
        m = np.full(q.shape[0], -np.inf)      # running max
        l = np.zeros(q.shape[0])              # running sum
        acc = np.zeros_like(q)                # unnormalized output
        for ks in range(0, n, block):
            s = (q @ K[ks:ks + block].T) * scale
            m_new = np.maximum(m, s.max(axis=1))
            p = np.exp(s - m_new[:, None])
            alpha = np.exp(m - m_new)         # rescale old state
            l = l * alpha + p.sum(axis=1)
            acc = acc * alpha[:, None] + p @ V[ks:ks + block]
            m = m_new
        O[qs:qs + block] = acc / l[:, None]
    return O

def paged_gather(pool, table, length):
    out = np.concatenate([pool[b] for b in table], axis=0)
    return out[:length]

rng = np.random.default_rng(0)
Q, K, V = (rng.standard_normal((256, 64)) for _ in range(3))
print(np.abs(attention_ref(Q, K, V) - tiled_attention(Q, K, V)).max())

pool = rng.standard_normal((8, 16, 64))   # 8 physical blocks of 16 tokens
print(paged_gather(pool, [5, 1, 6], length=40).shape)

The second program times naive attention against each PyTorch SDPA backend on a CUDA GPU and records peak memory. I have not run it on a GPU in this session, so no timings are quoted; run it on your hardware and expect naive peak memory to grow quadratically while the fused backends stay roughly linear.

import time
import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

def naive(q, k, v):
    s = (q @ k.transpose(-2, -1)) * (q.shape[-1] ** -0.5)
    n = s.shape[-1]
    mask = torch.ones(n, n, device=q.device, dtype=torch.bool).triu(1)
    s = s.masked_fill(mask, float("-inf"))
    return torch.softmax(s, dim=-1) @ v

def bench(fn, iters=20):
    for _ in range(3):
        fn()
    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    t0 = time.perf_counter()
    for _ in range(iters):
        fn()
    torch.cuda.synchronize()
    ms = (time.perf_counter() - t0) / iters * 1e3
    return ms, torch.cuda.max_memory_allocated() / 2**20

B, H, D = 4, 16, 128
for N in (1024, 2048, 4096, 8192):
    q, k, v = (torch.randn(B, H, N, D, device="cuda", dtype=torch.bfloat16)
               for _ in range(3))
    runs = {"naive": lambda: naive(q, k, v)}
    for name in ("FLASH_ATTENTION", "EFFICIENT_ATTENTION", "CUDNN_ATTENTION"):
        be = getattr(SDPBackend, name, None)
        if be is None:
            continue
        def run(be=be):
            with sdpa_kernel(be):
                return F.scaled_dot_product_attention(q, k, v, is_causal=True)
        runs[name] = run
    for name, fn in runs.items():
        try:
            ms, mib = bench(fn)
            print(f"N={N:5d} {name:20s} {ms:8.2f} ms  peak {mib:9.0f} MiB")
        except (RuntimeError, torch.cuda.OutOfMemoryError) as e:
            print(f"N={N:5d} {name:20s} unavailable: {str(e)[:60]}")
        torch.cuda.empty_cache()

Interpret the output with the roofline in mind. At N = 4096 with this shape, the naive score tensor alone is 4 × 16 × 4096² × 2 bytes, about 2.1 GB, and the causal mask and softmax output each add comparable temporaries. The fused backends should show peak memory near the size of Q, K, V and O. Time ratios will depend on GPU, driver and PyTorch version, and the backend that wins on A100 may differ from the one that wins on H100, so measure rather than assume.

Trade-offs, Gotchas, and What Goes Wrong

Fused attention is a strong default, but it is not free of sharp edges. Most production problems come from the boundary between the kernel’s assumptions and the model’s actual requirements.

Numerics are exact in theory, not bitwise. The online softmax reorders summation, so results differ from a reference at the level of floating point rounding, and different tile sizes or backends give slightly different outputs. Kernels accumulate in FP32 and keep the running max and sum in FP32, but if you port the algorithm to a lower precision accumulator you will see drift that grows with sequence length. Backward passes that use atomic additions for gradient accumulation can also be nondeterministic from run to run, which matters when you debug a loss spike or need reproducible training. Check whether your build exposes a deterministic option and what it costs.

Arbitrary masks and biases are awkward. Causal masks, sliding windows and ALiBi-style biases have fast paths because they can be computed from indices. A dense, arbitrary attention bias breaks the “never materialize N by N” premise, so frameworks either fall back to slower kernels or force you to pay for the full tensor. PyTorch’s FlexAttention exists to express such patterns as compiled score modifications, though I did not benchmark it here. A fully masked row, such as a left-padded token that attends to nothing, can yield NaN in a naive softmax; kernels handle it in different ways, so test padding explicitly.

Variable-length batches need packing. Padding to the longest sequence wastes both FLOPs and bandwidth. The varlen interfaces take concatenated sequences plus cumulative-length offsets and skip padding entirely, but you must build those offsets correctly, and a bug there produces plausible-looking wrong attention across sequence boundaries rather than a crash. Unit-test against a per-sequence reference.

Paging has a price. The block-table indirection and block-aligned tiling cost some kernel efficiency versus a contiguous cache, and a poorly chosen block size can erase the memory savings you wanted. Prefix sharing assumes identical token prefixes at identical positions; a tiny difference in a templated prompt, such as a timestamp near the top, defeats reuse. Swapping or preempting blocks under memory pressure adds latency tails that look like random slowness unless you monitor preemption counters.

Short sequences and small shapes may not benefit. The IO savings scale with N squared. At short lengths, kernel launch overhead and the cost of non-matmul work dominate, and a simple fused kernel or even the math path can tie. Head dimensions outside supported sets, unusual dtypes, or older GPUs silently drop you to a slower backend.

Version churn is real. The FlashAttention-3 Hopper kernel is described as beta in the repository, FP8 is forward only, and FlashAttention-4 is new enough that its feature coverage may lag. Pin library versions, record which backend ran, and rerun accuracy checks when you upgrade, particularly for FP8, where quantization error interacts with long contexts.

Finally, a conceptual trap: faster attention does not make long context free. Compute still grows quadratically in prefill, and the KV cache still grows linearly with context and batch. The kernels shift the practical ceiling; they do not remove it.

Practical Recommendations

Start from measurement, not from the newest kernel. Profile a representative request mix and decide whether you are prefill-bound, decode-bound or capacity-bound, because each points to a different lever. Prefill-bound workloads benefit most from FlashAttention-3 class kernels on Hopper or FlashAttention-4 on Blackwell. Decode-bound workloads benefit from split-KV kernels, GQA or MLA models and quantized KV. Capacity-bound workloads benefit from paging, prefix caching and offload.

Treat the framework default as a hypothesis. Force each backend with sdpa_kernel on your actual shapes, confirm that the fused path is taken, and keep the benchmark script from this post in CI so a dependency upgrade that silently changes dispatch gets caught. Log the attention backend, block size, KV dtype and maximum batched tokens alongside latency percentiles.

Use the roofline arithmetic as a sanity check before buying hardware or changing models. If decode is bandwidth-bound, a GPU with higher HBM bandwidth or a model with fewer KV heads helps more than extra FLOPs. If your contexts are long and batches are small, check that split-KV is active.

A short checklist:

Confirm GPU architecture, CUDA version and library versions, and match them to the support matrix above.
Verify the fused kernel runs for your head dimension, dtype and mask type; fail loudly rather than falling back.
Choose KV dtype and block size deliberately, and measure accuracy when quantizing KV or running FP8.
Test padding, fully masked rows and variable-length packing against a slow reference.
Enable prefix caching only after checking that your prompts share exact prefixes.
Re-benchmark after every major framework upgrade, including the model runner changes discussed in our vLLM model runner v2 migration guide.

Frequently Asked Questions

What is FlashAttention and why is it faster than standard attention?

FlashAttention is an exact attention algorithm that computes softmax attention in tiles held in on-chip SRAM instead of writing the full N by N score matrix to GPU main memory. Because attention is limited by memory bandwidth, cutting reads and writes to HBM is what speeds it up. It uses an online softmax to combine tiles correctly, and it recomputes intermediates in the backward pass instead of storing them, reducing memory use from quadratic to linear in sequence length.

What is the difference between FlashAttention and PagedAttention?

They solve different problems. FlashAttention is a compute kernel that makes each attention calculation efficient in memory traffic. PagedAttention is a memory management scheme for the KV cache that stores keys and values in fixed-size, non-contiguous blocks addressed through block tables, which removes fragmentation and allows sharing. In practice they combine: a paged kernel reads the cache through a block table while using FlashAttention-style tiling, and FlashAttention kernels now support a paged KV cache.

What is new in FlashAttention-3 compared with FlashAttention-2?

FlashAttention-3 targets NVIDIA Hopper GPUs. It uses warp specialization so producer warps issue asynchronous copies while consumer warps run matrix multiplies, interleaves softmax with matmul so the slow exponential unit overlaps with tensor core work, and adds FP8 support using block quantization and incoherent processing. The authors report 1.5 to 2.0 times the speed of FlashAttention-2 on H100, with FP16 reaching up to 740 TFLOPs/s. Those are the paper’s measurements, so verify them on your shapes.

Does FlashAttention work on any GPU?

No. The mainline FlashAttention-2 CUDA kernels support Ampere, Ada and Hopper GPUs in FP16 and BF16, and older Turing cards are only covered by a community fork with fewer features. FlashAttention-3 is Hopper only, and FlashAttention-4 targets Hopper and Blackwell according to the repository. AMD GPUs use Composable Kernel or Triton backends. PyTorch will fall back to other kernels where the fast path is unavailable, so confirm which backend actually ran.

Why is LLM decoding memory-bound even with FlashAttention?

During decoding each new token needs a single query to attend to the whole cached history, so every generated token must stream the KV cache from HBM while performing very little arithmetic per byte. For grouped-query attention with g query heads per KV head, intensity is roughly g FLOPs per byte, far below an H100’s ridge point near 295. Throughput therefore tracks memory bandwidth, which is why GQA, MLA, KV quantization and split-KV kernels matter.

How do GQA and MLA change attention kernel design?

Grouped-query attention shares each KV head across several query heads, shrinking the cache and raising arithmetic intensity, so kernels pack the query group into a tile to reuse each loaded key and value. Multi-head latent attention caches a compressed latent rather than full keys and values, which cuts cache size much further but needs dedicated kernels because the effective head dimension and sharing pattern differ from standard attention. Always match the kernel to the model’s cache layout.

FlashAttention-3 and PagedAttention: IO-Aware Attention Kernels and the GPU Memory Hierarchy (2026)

FlashAttention-3 and PagedAttention: IO-Aware Attention Kernels (2026)

Context and Background