FlashAttention for inference: tiled prefill and split-KV decoding

FlashAttention computes exact attention without ever writing the $N\times N$ matrix of attention scores to GPU main memory. It splits queries, keys and values into tiles that fit in on-chip SRAM, computes each tile's contribution, and merges the contributions with a running softmax that rescales earlier partial results whenever a larger score appears. The output is the same as standard attention up to floating-point rounding. For inference this matters in two different ways. In prefill, which is compute-bound, it removes a memory footprint that grows with the square of the prompt and keeps the tensor cores busy, which makes long prompts feasible and fast. In decode, which is memory-bound, the tiling idea is reused for a different purpose: splitting a long KV cache across many thread blocks so that a single sequence can use the whole GPU's bandwidth, plus native support for grouped-query heads and paged caches. Every major open-source serving engine runs a FlashAttention-family kernel or a close relative.

The bottleneck it attacks#

Standard attention for one head computes $S=Q{K}^{\top }/\sqrt{d}$, $P=\mathrm{softmax}(S)$ and $O=PV$ as separate kernels, writing $S$ and $P$, each $N\times N$, to high-bandwidth memory and reading them back. Dao, Fu, Ermon, Rudra and Ré (2022) argued that what made attention slow on GPUs was not FLOPs but these memory reads and writes, and that approximate attention methods had failed to deliver wall-clock speedups because they ignored IO. Their paper proves that standard attention needs $\Theta (Nd+N^2)$ HBM accesses, while FlashAttention needs $\Theta (N^2{d}^2{M}^{-1})$ for SRAM size $M$ with $d\le M\le Nd$, and that no exact algorithm can do asymptotically better across all SRAM sizes. For typical head dimensions of 64 to 128 and on-chip memory of order 100 KB per streaming multiprocessor, $d^2/M$ is well below 1, so far fewer bytes move.

Memory capacity is the more visible effect at inference time. For one Llama 3 8B layer (32 query heads), materialising FP16 scores for a single 32k-token prompt takes $32\times 32,768^2\times 2$ bytes = 68.7 GB, nearly a whole H100, for one layer of one request. FlashAttention keeps only two numbers per query row (the running maximum and the running sum), 8.4 MB in FP32 for the same layer and prompt.

The line of work: Milakov and Gimelshein (2018) showed the softmax normaliser can be computed in one pass with a running maximum. Rabe and Staats (2021) showed self-attention needs only $O(\sqrt{n})$ memory in practice by chunking with this trick. FlashAttention (2022) fused it into an IO-aware GPU kernel. FlashAttention-2 (Dao, 2023) reworked the loop order and parallelism to reach 50 to 73% of theoretical peak FLOP/s on A100, about 2× FlashAttention. FlashAttention-3 (Shah, Bikshandi, Zhang, Thakkar, Ramani and Dao, 2024) exploited Hopper's asynchronous tensor cores and memory engine to reach up to 740 TFLOP/s in FP16 (75% of H100 peak) and close to 1.2 PFLOP/s in FP8, with 2.6× lower FP8 error than a baseline FP8 attention. For decoding, Flash-Decoding (Dao, Haziza, Massa and Sizov, 2023) added the split-KV scheme described below, reporting up to 8× faster decoding for very long sequences.

The online softmax, with every symbol defined#

Take one query row $q\in {\mathbb{R}}^d$, keys and values split into blocks $K_1,\dots ,K_T$ and $V_1,\dots ,V_T$, and block scores $s_j=K_{j}q/\sqrt{d}$. Keep three running quantities: the maximum score seen so far $m$, the sum of exponentials relative to it $\ell$, and the unnormalised output $\tilde{o}$. Initialise $m^{(0)}=-\infty$, ${\ell }^{(0)}=0$, ${\tilde{o}}^{(0)}=0$, and for each block

After the last block, $o={\tilde{o}}^{(T)}/{\ell }^{(T)}$. The correction factor $c_j\le 1$ rescales everything accumulated so far whenever a new maximum appears; subtracting the running maximum keeps every exponential at most 1, so nothing overflows. By induction ${\ell }^{(j)}={\sum }_{t\le j}{\sum }_i{e}^{s_{t,i}-m^{(j)}}$ and ${\tilde{o}}^{(j)}$ is the matching weighted sum, so the result equals $\mathrm{softmax}(s{)}^{\top }V$ exactly in real arithmetic.

The same algebra gives the rule for combining independent partial results, which is what split-KV decoding needs. If chunk $k$ produced a normalised output $o_k$ and its log-sum-exp ${\lambda }_k=m_k+\log {\ell }_k$, then

which is a softmax over chunks weighted by how much probability mass each chunk holds.

A worked example by hand

One query, four keys with scores $s=(2,1,4,3)$, scalar values $v=(1,0,2,-1)$, two blocks of two.

Block 1, scores $(2,1)$: $m=2$; $\ell =e^0+e^{-1}=1.3679$; $\tilde{o}=1\cdot e^0+0\cdot e^{-1}=1.0000$. The interim output would be $1/1.3679=0.7311$.

Block 2, scores $(4,3)$: new maximum $m=4$, so $c=e^{2-4}=0.1353$. Then $\ell =0.1353\times 1.3679+e^0+e^{-1}=0.1851+1.3679=1.5530$ and $\tilde{o}=0.1353\times 1+2e^0-1\cdot e^{-1}=0.1353+2-0.3679=1.7675$.

Result: $o=1.7675/1.5530=1.1381$. Directly, $e^s=(7.389,2.718,54.598,20.086)$ with sum 84.791, and $(7.389+0+109.196-20.086)/84.791=1.1381$. Our numpy run reproduces both to machine precision.

Prefill: tiles, loop order and the tensor cores#

for each (sequence, head, query tile Q_i of B_r rows) in parallel:       # one thread block each
    load Q_i into SRAM; m = -inf; l = 0; acc = 0
    for each key/value tile j up to the causal limit:                     # B_c columns per tile
        load K_j, V_j into SRAM
        S = Q_i K_j^T / sqrt(d)            (tensor cores; apply causal mask on the diagonal tile)
        m_new = max(m, rowmax(S)); P = exp(S - m_new); c = exp(m - m_new)
        l = c * l + rowsum(P); acc = c * acc + P V_j   (tensor cores)
        m = m_new
    write O_i = acc / l and the row log-sum-exp m + log(l) to HBM

This is FlashAttention-2's loop order: the outer parallel loop is over query tiles, so each thread block owns its output tile and no two blocks write the same rows. Causal masking lets tiles strictly above the diagonal be skipped entirely, halving the work. Keys and values are re-read once per query tile, often from L2 cache rather than HBM, and the score matrix never leaves the chip. Because prefill is compute-bound (see the roofline), the goal is to keep the matrix units saturated, which is what FlashAttention-2's work partitioning and FlashAttention-3's overlap of softmax with matrix multiplies achieve.

How much prefill time attention takes depends on prompt length. Its FLOPs are $2L{n}_h{d}_h{N}^2$ against $2PN$ for the linear layers, so the attention share is $L{n}_h{d}_{h}N/(P+L{n}_h{d}_{h}N)$. For Llama 3 8B ($L{n}_h{d}_h=131,072$) that is 3.2% at 2,000 tokens, 11.8% at 8k, 34.8% at 32k and 68.1% at 128k. A faster attention kernel barely moves TTFT for chat-length prompts and dominates it for long documents.

Decode: a different bottleneck#

At decode time each sequence contributes one query row per head. The score "matrix" is a vector of length $S$, so there is nothing large to avoid materialising, and the cache must be read exactly once whatever the algorithm. The arithmetic intensity is fixed by the model, as derived in the roofline article: $2n_h/(n_{kv}{b}_{kv})$ FLOP per byte, 4 for Llama 3 8B in BF16. Decode attention is memory-bound on every GPU, and the only goal is to stream the cache at full bandwidth. Three things stand in the way.

Too little parallelism. A kernel that assigns one thread block per (sequence, head) launches $B\times n_h$ blocks, or $B\times n_{kv}$ if query heads sharing a key-value head are processed together. At batch 1 with 8 key-value heads that is 8 thread blocks on a GPU with 132 streaming multiprocessors, so most of the chip, and most of its memory bandwidth, sits idle. Flash-Decoding splits the key-value sequence into chunks, gives each chunk its own thread block, writes per-chunk outputs with their log-sum-exp, and merges them in a small reduction kernel using the combination rule above. With 1,024-token chunks a 32k-token sequence yields 32 chunks per key-value head, 256 blocks at batch 1. The Flash-Decoding authors report attention up to 50× faster than FlashAttention for long sequences at small batch and end-to-end decoding up to 8× faster, with runtime roughly constant up to 32k tokens on CodeLlama 34B at batch 1. At large batch there is already enough parallelism and splitting adds only reduction overhead, so kernels choose the number of splits from batch size, heads and length.

Reading shared heads more than once. With GQA, $g=n_h/n_{kv}$ query heads share each key-value head. A naive kernel that expands keys and values to $n_h$ heads, or runs one block per query head, reads each cache entry $g$ times. GQA-aware kernels load each key-value tile once and multiply it against all $g$ query rows, which both cuts bytes by $g$ and turns a matrix-vector product into a small matrix-matrix product.

Scattered memory. Paged caches store keys and values in blocks listed in a block table. FlashAttention added paged-cache support in version 2.5, and FlashInfer (Ye et al., 2025) builds its kernels around block-sparse layouts that cover paged caches, shared prefixes and tree-structured speculative decoding; FlashInfer reports 29 to 69% lower inter-token latency than compiler backends on a serving benchmark and 28 to 30% lower latency for long-context inference.

A quick bound shows the stakes. One 32k-token Llama 3 8B sequence holds 4.29 GB of cache, read every step. At the H100's 3.35 TB/s that is 1.28 ms per step at best; a kernel that reaches only a sixth of the chip's bandwidth because it launched 8 thread blocks spends several times that. Decode attention kernels are judged by achieved bandwidth, not FLOP/s.

Toy run: tiling and splitting are exact#

We implemented three versions in numpy (seed 0, FP64): standard attention that materialises the full score matrix, a tiled causal prefill with 64×64 tiles and the online softmax, and split-KV decoding with 1,024-token chunks merged by log-sum-exp. For a 2,048-token causal prefill with $d=128$, the largest absolute difference between tiled and standard outputs was $1.1\times 10^{-15}$. For four query heads decoding against a 32,768-token cache, split-KV and standard differed by at most $3.0\times 10^{-16}$. Both are rounding noise: neither technique approximates anything. The hand example above came out of the same script, matching to the fourth decimal.

Interaction with post-trained models#

Exact, but not bit-identical. FlashAttention changes the order of floating-point additions, so logits differ from a naive implementation in the last bits, and different kernels (prefill versus decode, split versus unsplit, different split counts at different batch sizes) can disagree slightly with each other. This is harmless for serving, but it matters in reinforcement learning: if a rollout engine and a trainer use different attention kernels, the log-probabilities used in importance ratios for GRPO or PPO disagree before any update has happened. Measure the mismatch between rollout and trainer log-probabilities on the same tokens before blaming the algorithm for instability.

FP8 attention needs an accuracy check. FlashAttention-3's FP8 path reduces error relative to naive FP8 attention through block quantisation and incoherent processing, but it is still lower precision than BF16. A fine-tune whose behaviour depends on precise long-range retrieval should be evaluated with the exact kernel and dtype you will serve. See FP8 and FP4 inference.

Architecture features must be supported by the kernel. Head dimensions up to 256, grouped-query attention, sliding windows, ALiBi, logit soft-capping and paged caches are supported in current FlashAttention releases according to its documentation, but each new variant arrives in kernels later than in model code. A post-training change that alters attention (a different head size after pruning, a custom mask, an added bias) can push a model off the fast path and onto a slow fallback without any error.

Packed fine-tuning data and variable-length batches. Fine-tuning pipelines that pack several examples into one sequence rely on variable-length attention so examples do not attend to each other; inference engines use the same variable-length interface to batch prompts of different lengths in prefill without padding. If a fine-tune was trained with packing but without per-example masking, it has seen cross-example attention during training, which no inference kernel will reproduce. That is a training bug that only shows up as a quality gap.

Adapters do not change the attention kernel. LoRA modifies the projections that produce $Q$, $K$ and $V$, not the attention computation, so FlashAttention is unaffected by adapters, merged or not.

Tuning knobs#

Failure modes and how to detect them#

Silent fallback. Symptom: prefill of long prompts is several times slower than expected, or memory use spikes during prefill. The engine fell back to a math implementation that materialises scores, usually because of an unsupported head dimension, dtype, mask or GPU. Check the engine's start-up log for the selected attention backend and profile one long prefill.

Low decode bandwidth at small batch. Symptom: single-user TPOT on long contexts far above (weights + cache bytes) / bandwidth. Splitting is off or capped. Check that the decode kernel supports split-KV for your configuration.

Mismatch between prefill and decode kernels. Symptom: tiny logit differences between the first generated token and later ones that make deterministic tests flaky. Expected behaviour; compare with tolerances, not equality.

Numerical issues in FP16. Models trained in BF16 can produce attention logits or values that overflow FP16. The online softmax protects the exponentials, but the matrix products themselves can still overflow. Serve BF16-trained fine-tunes in BF16.

When it matters#

Minimal implementation#

In PyTorch, scaled_dot_product_attention dispatches to a FlashAttention-style fused kernel on supported GPUs, which covers prefill in simple settings. For a decode step against a paged cache, the flash-attn package's flash_attn_with_kvcache takes the cache, a block table and per-sequence lengths, appends the new keys and values in place, and handles GQA and split-KV internally.

import torch
import torch.nn.functional as F
from flash_attn import flash_attn_with_kvcache

# prefill: q, k, v of shape (batch, heads, seq, head_dim); fused kernel, no N x N tensor
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)

# decode: one new token per sequence against a paged cache
# q_new: (batch, 1, q_heads, head_dim); k_new, v_new: (batch, 1, kv_heads, head_dim)
# k_cache, v_cache: (num_blocks, block_size, kv_heads, head_dim)
# block_table: (batch, max_blocks) int32; cache_seqlens: (batch,) int32 tokens already cached
out = flash_attn_with_kvcache(
    q_new, k_cache, v_cache, k=k_new, v=v_new,
    cache_seqlens=cache_seqlens, block_table=block_table, causal=True,
)
cache_seqlens += 1

# the merge rule behind split-KV, for reference
def merge(outs, lses):            # outs: (splits, ..., d); lses: (splits, ...)
    w = torch.softmax(lses, dim=0)
    return (w.unsqueeze(-1) * outs).sum(0)

Serving engines wrap these calls: vLLM, SGLang and TensorRT-LLM select FlashAttention, FlashInfer or their own kernels per GPU generation. When POFE evaluates a fine-tune on long-context tasks, it runs the same attention backend that the deployment target will use, so the numbers on the evaluation report are the numbers you will serve.

Further reading#

Dao et al. (2022) is the original paper, with the IO-complexity analysis and lower bound. Dao (2023) gives the FlashAttention-2 loop order and partitioning, and Shah et al. (2024) the Hopper-specific FlashAttention-3 with FP8. The Flash-Decoding post by Dao, Haziza, Massa and Sizov (2023) explains split-KV decoding. Milakov and Gimelshein (2018) and Rabe and Staats (2021) are the algorithmic ancestors. Ye et al. (2025) describes FlashInfer's serving-oriented kernels for paged and shared-prefix caches. The FlashAttention repository documents current support for paged caches, GQA, sliding windows and soft-capping.

Frequently asked questions#

References#

1. Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arxiv.org/abs/2205.14135
2. Tri Dao (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arxiv.org/abs/2307.08691
3. Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arxiv.org/abs/2407.08608
4. Tri Dao, Daniel Haziza, Francisco Massa, Grigory Sizov (2023). Flash-Decoding for long-context inference. crfm.stanford.edu/2023/10/12/flashdecoding.html
5. Maxim Milakov, Natalia Gimelshein (2018). Online normalizer calculation for softmax. arxiv.org/abs/1805.02867
6. Markus N. Rabe, Charles Staats (2021). Self-attention Does Not Need O(n^2) Memory. arxiv.org/abs/2112.05682
7. Zihao Ye, Lequn Chen, Ruihang Lai, Wuwei Lin, Yineng Zhang, Stephanie Wang, Tianqi Chen, Baris Kasikci, Vinod Grover, Arvind Krishnamurthy, Luis Ceze (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arxiv.org/abs/2501.01005
8. Jordan Juravsky, Bradley Brown, Ryan Ehrlich, Daniel Y. Fu, Christopher Ré, Azalia Mirhoseini (2024). Hydragen: High-Throughput LLM Inference with Shared Prefixes. arxiv.org/abs/2402.05099
9. Tri Dao and contributors (2025). flash-attention repository README. github.com/Dao-AILab/flash-attention
10. NVIDIA (2022). NVIDIA Hopper Architecture In-Depth. developer.nvidia.com/blog/nvidia-hopper-architecture-in-depth/

Related reading