FlashAttention Explained: IO-Aware Attention, Tiling, Online Softmax and FP8 on Hopper

A 4,096-token attention head in fp16 produces a 33.5 MB score matrix that a naive implementation writes to GPU memory, reads back, writes again as probabilities, and reads once more. The arithmetic for that head takes about 27 microseconds on an A100; the memory traffic takes about 68. The tensor cores are idle most of the time, waiting on DRAM. FlashAttention is the family of kernels that fixes this by never letting the N-by-N matrix leave on-chip memory, and it is now the default attention path in PyTorch, vLLM, SGLang, TensorRT-LLM and nearly every training stack.

This matters in 2026 because context windows keep growing while memory bandwidth grows far more slowly than tensor-core throughput. The kernel you run decides whether a 128K-token prefill takes seconds or minutes, and whether FP8 attention is safe. The ideas are also reusable: the same tiling and rescaling trick appears in decoding kernels, FlexAttention and Blackwell’s FlashAttention-4.

You will leave with the memory-hierarchy arithmetic, a hand-worked online-softmax example, the backward-pass recomputation argument, an honest history of FlashAttention 1, 2, 3 and 4, and runnable PyTorch code.

What this covers: the GPU memory hierarchy, why standard attention is memory-bound, tiling and online softmax, backward recomputation, FA1 to FA4, FP8 and incoherent processing, FlexAttention, decoding variants, and when none of it helps.

Context and Background

Self-attention computes O = softmax(QKᵀ / √d) V. For a sequence of length N and head dimension d, Q, K and V are N×d matrices, and the intermediate score matrix S = QKᵀ is N×N. Compute is O(N²d) and the naive intermediate memory is O(N²). For years, the field read that quadratic cost as a FLOP problem and responded with approximations: sparse patterns, low-rank projections, kernel feature maps. Many reduced FLOPs and still did not run faster on real GPUs, because FLOPs were never the binding constraint.

The FlashAttention paper by Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra and Christopher Ré (arXiv:2205.14135) reframed the problem as one of IO-awareness: count the reads and writes between the large, slow memory (HBM, high-bandwidth memory) and the small, fast on-chip memory (SRAM, which on NVIDIA parts is the shared memory and register file inside each streaming multiprocessor). The abstract reports a 15% end-to-end speedup on BERT-large at sequence length 512 over the MLPerf 1.1 training record, 3x on GPT-2 at length 1K, and 2.4x on the Long Range Arena at lengths 1K to 4K. The result is exact attention, not an approximation, which is why it displaced most of the approximate methods in practice.

Since then the lineage has moved quickly. FlashAttention-2 (arXiv:2307.08691) improved work partitioning and reached about 2x over version 1. FlashAttention-3 (arXiv:2407.08608) rewrote the kernel for Hopper’s asynchronous hardware and added FP8. FlashAttention-4, described by the authors in a March 2026 write-up, targets Blackwell. Frameworks wrap these kernels: PyTorch exposes them through scaled_dot_product_attention, and serving engines such as those compared in vLLM vs SGLang vs TensorRT-LLM ship their own variants.

If you are new to the surrounding system, two companion posts help. Speculative decoding explains why decode latency, not prefill, often dominates user experience, and Mixture of Experts explains why attention becomes a larger share of the compute budget once the feed-forward layers are sparse. The original paper is the primary reference: FlashAttention on arXiv.

The numbers that frame everything

Three hardware facts drive the whole design. First, tensor-core throughput has grown much faster than DRAM bandwidth across generations. NVIDIA’s A100 80GB lists 312 TFLOPS of dense FP16 tensor throughput and 1,935 GB/s (PCIe) to 2,039 GB/s (SXM) of HBM bandwidth. The H100 SXM, per the FlashAttention-3 write-up on the PyTorch blog, peaks at 989 TFLOPS FP16 and 1,978 TFLOPS FP8 matmul. Second, non-matmul operations such as the softmax exponential run on special function units that are far slower: the same page cites 3.9 TFLOPS for special functions on H100, 256x lower than matmul. Third, on-chip SRAM is tiny, on the order of 100 to 230 KB per streaming multiprocessor, but roughly an order of magnitude faster than HBM. The original paper’s hardware figure quotes about 19 TB/s for SRAM versus 1.5 TB/s for HBM on an A100 40GB (taken from the paper’s memory-hierarchy figure; treat the SRAM aggregate as approximate).

The Core Idea: IO-Aware Attention and Why Standard Attention Is Memory-Bound

Direct answer: FlashAttention computes exact softmax attention by splitting Q, K and V into tiles that fit in on-chip SRAM, fusing the matmul, softmax and weighted sum into one kernel, and keeping a running maximum and sum so the full N×N score matrix is never written to HBM. This cuts memory traffic from quadratic in N to linear in N for fixed d, which turns a memory-bound operation into a compute-bound one.

Figure 1: The fused FlashAttention loop. A Q tile is loaded once, K and V tiles stream through SRAM, an online softmax keeps running statistics, and only the output and a per-row logsumexp are written back to HBM.

Figure 1 shows the data flow of one thread block. It owns one tile of query rows. It loops over key and value tiles, computes the score tile on tensor cores, updates the running softmax statistics, folds the new contribution into an output accumulator held in registers, and finally writes one output tile and one small statistics vector back to HBM. Nothing of size N×N ever touches HBM.

Worked arithmetic: standard attention on one head

Take N = 4,096, d = 128, fp16 (2 bytes), one head, one sequence. The score matrix has N² = 16,777,216 entries, which is 33,554,432 bytes, or 33.5 MB.

A straightforward implementation runs three kernels: a GEMM for S = QKᵀ, a softmax, and a GEMM for O = PV. The HBM traffic is:

Write S: 33.5 MB.
Read S for softmax, write P: 67.1 MB.
Read P for the second GEMM: 33.5 MB.
Q, K, V reads and the O write: 4 × N × d × 2 bytes = 4.19 MB.

Total: about 138 MB. At the A100 SXM’s 2,039 GB/s that is about 67.6 microseconds, ignoring launch gaps. The matmul work is 4N²d = 4 × 16.78M × 128 ≈ 8.59 GFLOP, which at the 312 TFLOPS peak takes 27.5 microseconds. The memory floor is about 2.5x the compute floor, so the kernel can reach at most roughly 40% of tensor-core peak even with perfect overlap, and in practice much less.

Expressed as arithmetic intensity, the standard path does 8.59 GFLOP over 138 MB, about 62 FLOP per byte. The A100’s ridge point (peak FLOPS divided by peak bandwidth) is 312 / 2.039 ≈ 153 FLOP per byte. Anything below the ridge is memory-bound. Standard attention sits well below it.

A fused tiled kernel moves only Q, K, V and O, about 4.19 MB, plus a small logsumexp vector. Its HBM floor is about 2 microseconds, so it is firmly compute-bound: intensity jumps to roughly 2,000 FLOP per byte. Multiply by batch and heads (for example 8 sequences times 32 heads is 256 head-instances) and standard attention also needs 8.6 GB of scratch for S alone in a single layer’s forward pass, which is why memory, not just time, was a limit on context length before fused kernels.

Figure 2: Standard attention makes four N×N trips across the HBM boundary. The fused kernel keeps S and P in SRAM and writes only O and logsumexp.

The IO complexity result

The paper states the asymptotics. Standard attention needs Θ(Nd + N²) HBM accesses. FlashAttention needs Θ(N²d²M⁻¹) accesses, where M is the SRAM size in elements, under the condition d ≤ M ≤ Nd. For typical values (d = 64 to 128, M on the order of 100K elements) the factor d²/M is well below one, so FlashAttention reads and writes substantially fewer bytes than the N² term. The paper also gives a lower bound showing no exact attention algorithm can asymptotically beat this over the whole range of M, which is what “optimal for a range of SRAM sizes” in the abstract refers to.

Notice what FlashAttention does not change. The FLOP count is the same, in fact slightly higher in the backward pass because of recomputation, and the asymptotic compute is still O(N²d). The win is entirely in data movement. That is the thesis of IO-aware algorithm design: on modern accelerators, minimize bytes moved between memory levels first and FLOPs second.

Tile sizing under an SRAM budget

The original algorithm picks column block size Bc = ⌈M / 4d⌉ and row block size Br = min(Bc, d), so that the Q, K, V and O tiles fit in SRAM together. With d = 128 and, illustratively, 48K usable elements, Bc ≈ 94, which real kernels round to a hardware-friendly 64 or 128. A 128×128 fp16 tile is 32 KB. Holding a Q tile, a K tile, a V tile and an fp32 score tile of the same shape already costs 32 + 32 + 32 + 64 = 160 KB, close to the 164 KB maximum shared memory a single Ampere block can request. This is why head dimension limits tile sizes and why kernels for d = 256 behave differently from d = 64: bigger d forces smaller tiles or more register pressure, lowering reuse.

Tiling and Online Softmax: The Derivation

The obstacle to tiling is softmax. For a row of scores x₁…x_N, softmax needs the global maximum for numerical stability and the global sum of exponentials for the denominator. Both depend on all N elements, so a naive tiled kernel would need two passes over the keys before it could produce any output.

Safe softmax, then online softmax

The numerically safe form is softmax(x)ᵢ = exp(xᵢ − m) / Σⱼ exp(xⱼ − m) with m = max(x). Computing it directly takes three passes: find m, compute the sum, then normalize. Milakov and Gimelshein (arXiv:1805.02867) showed that the maximum and the denominator can be computed in a single pass by maintaining a running pair (m, ℓ). When a new element x arrives:

m_new = max(m, x)
ℓ_new = ℓ · exp(m − m_new) + exp(x − m_new)

The old sum is rescaled by exp(m − m_new) so that it is expressed relative to the new maximum. That paper reports softmax up to about 1.3x faster and, fused with top-K, about 5x, because it removes memory passes. FlashAttention generalizes the same recurrence from scalars to tiles and from a sum to a weighted sum of value vectors.

The tile recurrence

For one query row, process key tiles j = 1, 2, …. Maintain a running max m, running denominator ℓ, and an unnormalized output accumulator o (a d-vector). For tile j with scores s⁽ʲ⁾:

m_new = max(m, rowmax(s_j))
p_j   = exp(s_j − m_new)                 # tile-local probabilities, unnormalized
ℓ_new = exp(m − m_new) · ℓ + rowsum(p_j)
o_new = exp(m − m_new) · o + p_j · V_j

After the last tile, the output is o / ℓ. The correction factor exp(m − m_new) is the only coupling between tiles, and it is a scalar per row. Because this is algebraically identical to the three-pass computation, the result is exact up to floating-point reordering, not an approximation.

Worked numeric example

Use four scores [2, 1, 3, 0] with scalar values v = [1, 2, 3, 4], split into tile A = [2, 1] and tile B = [3, 0].

Tile A: m = 2. The exponentials relative to m are e⁰ = 1 and e⁻¹ = 0.3679, so ℓ = 1.3679. The unnormalized output is 1·1 + 0.3679·2 = 1.7358.

Tile B arrives with local max 3, so m_new = 3 and the correction factor is e^(2−3) = 0.3679. The tile’s exponentials relative to 3 are e⁰ = 1 and e⁻³ = 0.0498. Then:

ℓ = 1.3679 × 0.3679 + (1 + 0.0498) = 0.5032 + 1.0498 = 1.5530.
o = 1.7358 × 0.3679 + (1·3 + 0.0498·4) = 0.6386 + 3.1991 = 3.8377.
Final output = 3.8377 / 1.5530 = 2.4712.

The direct computation gives weights e⁻¹, e⁻², e⁰, e⁻³ = 0.3679, 0.1353, 1, 0.0498 (relative to the global max 3), a sum of 1.5530, a weighted value sum of 0.3679 + 0.2707 + 3 + 0.1991 = 3.8377, and the same answer 2.4712. The two paths agree, which is the whole proof of concept in four numbers.

Why the statistics are cheap to store

The forward pass must save something for the backward pass. Instead of saving the N×N probability matrix, FlashAttention saves the output O and one scalar per row, the logsumexp L = m + log ℓ. That is N extra floats per head, negligible next to N². Given L, any probability can be recomputed as P = exp(S − L) from the scores alone, with no second reduction.

Backward Pass: Recomputation Beats Storage

Training needs gradients of Q, K and V. Standard attention stores P (N×N) during the forward pass and reads it back. FlashAttention instead recomputes S and P tile by tile in the backward pass from Q, K, the saved L, and the output. This is gradient checkpointing applied inside a single kernel, and it is a strict trade: extra FLOPs for much less memory traffic.

The backward algebra for one head is, with dO the incoming gradient:

dV = Pᵀ dO
dP = dO Vᵀ
dS = P ∘ (dP − D)      where D_i = rowsum(dO ∘ O)
dQ = dS K / √d
dK = dSᵀ Q / √d

The key simplification is that D_i = Σ dO_i ∘ O_i is a per-row scalar that can be computed from the output and its gradient without ever forming P, which is what lets the softmax Jacobian be applied tile by tile. The forward pass does two matmuls (QKᵀ and PV); the backward pass does five (recomputing QKᵀ, plus dV, dP, dQ, dK). That is why the FlashAttention-4 write-up lists backward-pass tensor-core work at 2,560 cycles per tile versus 1,024 for the forward pass, a 2.5x ratio.

The practical consequence is that the memory for attention in training drops from O(N²) to O(N) per head, which is what made 32K and 128K training contexts feasible without enormous activation memory. The backward pass is also harder to make fast: the kernel must accumulate dQ across key tiles with atomic adds into HBM, which introduces nondeterminism. Deterministic modes serialize those reductions at a throughput cost; the FA4 write-up reports its deterministic backward reaching roughly 85 to 90% of the nondeterministic throughput.

From FlashAttention-1 to FlashAttention-4: What Changed on Each GPU Generation

Figure 3: Each FlashAttention generation answered the bottleneck of its hardware. Utilization figures are the authors’ reported peaks on the stated GPU and are not directly comparable across rows.

The central lesson of the lineage is that the algorithm stayed essentially fixed (tiling plus online softmax) while the kernel was repeatedly re-engineered around whatever resource had become scarce. Read the four generations as a sequence of bottlenecks.

FlashAttention-1: prove IO-awareness works

Version 1 established the fused tiled kernel on Ampere-era GPUs. It parallelized over batch and heads, looped over key and value blocks in the outer loop and query blocks in the inner loop, and so had to write partial outputs back to HBM repeatedly. The FlashAttention-2 paper reports that this left the kernel at roughly 25 to 40% of theoretical peak FLOPs on A100 (as I read that paper; confirm against the PDF before quoting).

FlashAttention-2: better parallelism and work partitioning

FlashAttention-2 (arXiv:2307.08691) made three changes. It reduced non-matmul FLOPs, because on an A100 each non-matmul FLOP costs about 16 times more than a matmul FLOP (the specialized tensor cores versus general units), so shaving rescaling work mattered. It parallelized across the sequence dimension in addition to batch and heads, which fixed low occupancy for long sequences with small batches. And it swapped the loop order so that each thread block owns a query tile and iterates over keys and values, with the output accumulator staying in registers, and it split work among warps to avoid shared-memory communication.

The paper’s abstract reports roughly a 2x speedup over version 1, 50 to 73% of theoretical maximum FLOPs per second on A100, and up to 225 TFLOPs/s per A100 in end-to-end GPT-style training, stated as 72% model FLOPs utilization. For most practitioners “FlashAttention” in PyTorch on an A100 or consumer Ampere/Ada GPU means this generation.

FlashAttention-3: asynchrony and low precision on Hopper

Hopper changed the shape of the problem. Its tensor cores became so fast that FlashAttention-2 reached only about 35% of the H100’s theoretical maximum, per the PyTorch blog’s account. Three hardware features had to be used to go further, and the FlashAttention-3 paper (arXiv:2407.08608, Shah, Bikshandi, Zhang, Thakkar, Ramani, Dao) exploits all of them:

WGMMA (warpgroup matrix multiply-accumulate) is Hopper’s tensor-core instruction operating on a warpgroup of four warps. The blog notes that the older mma.sync path reaches only about two thirds of peak Hopper tensor throughput.
TMA (Tensor Memory Accelerator) is a hardware unit that copies tiles between global and shared memory and handles index math and bounds predication, freeing registers for bigger tiles.
Warp specialization assigns producer warps to issue TMA loads and consumer warps to issue WGMMA, so data movement and math overlap.

On top of that, FA3 hides softmax latency by overlapping it with matmul. In pingpong scheduling, two warpgroups alternate: while one runs its GEMMs, the other runs softmax, enforced with barriers. In intra-warpgroup overlap, parts of one warpgroup’s softmax run while its own GEMMs are in flight, at the cost of more registers. The blog’s FP16 forward progression for head dimension 128 shows how much each step contributed: about 350 TFLOPS for FA2, roughly 540 to 570 after rewriting for WGMMA and TMA, about 620 with pingpong, and about 640 to 660 with intra-warpgroup overlap.

The headline results reported by the authors: up to 740 TFLOPS in FP16, about 75% utilization of the H100, 1.5 to 2.0x faster than FlashAttention-2, and close to 1.2 PFLOPS in FP8. The reason softmax matters so much: with 989 TFLOPS of matmul against 3.9 TFLOPS of special-function throughput, the exponentials in a tile cost a non-trivial fraction of the matmul time unless hidden.

FlashAttention-4: co-design for Blackwell

The FlashAttention-4 write-up on Tri Dao’s blog (March 2026) describes the next bottleneck shift. On B200, BF16 tensor throughput rises from about 1 PFLOP on H100 to 2.25 PFLOPs, but the special function unit count and shared-memory bandwidth are unchanged. The blog’s roofline puts forward-pass cost per tile at 1,024 cycles for tensor cores, 1,024 for exponentials, and 768 for shared memory, so the exponential unit becomes as limiting as the matmul. In the backward pass shared-memory traffic (3,328 cycles) exceeds even tensor-core time (2,560).

The reported techniques follow directly from that roofline:

Software exponential emulation. Part of the base-2 exponential is computed on the FMA units with a cubic polynomial (Horner evaluation after Cody-Waite range reduction), alongside the hardware unit, to use idle FMA throughput.
Conditional rescaling. The output accumulator is only rescaled when the row maximum rises by more than a threshold, skipping most of the correction work. Final normalization uses the true statistics, so the result stays exact in the sense the authors describe.
Tensor memory (TMEM) and 2-CTA MMA. Blackwell adds 256 KB of tensor memory per SM; the forward pass stages probabilities through it with a ping-pong over two query tiles and a dedicated correction warpgroup. The backward pass pairs two CTAs (thread blocks) to split the accumulator and exchange half of dS over distributed shared memory, which the authors say roughly halves shared-memory traffic for one operand and halves global dQ atomic adds.
CuTe-DSL implementation. The kernel is written in CUTLASS’s Python DSL, with reported compile-time reductions of about 20 to 30x versus C++ templates.

The reported result: up to 1,605 TFLOPs/s on B200 in BF16 at 71% utilization, 1.1 to 1.3x faster than cuDNN 9.13 and 2.1 to 2.7x faster than Triton in the forward pass. The same page candidly notes that the latest cuDNN (9.19.1.2) reportedly offers similar performance, and that cuDNN 9.13 and 9.14 already absorbed some FA4 techniques. Treat that as the honest takeaway: on Blackwell the vendor library and the research kernel are converging.

Reading the numbers responsibly

These figures are the authors’ own measurements on specific GPUs, sequence lengths and head dimensions, usually at long sequences where the kernel is saturated. Your workload (short prompts, grouped-query attention, small batches, head dimension 64, a different driver) can land far from them. Always benchmark your own shapes. A 2x claim over a previous kernel says nothing about your end-to-end step time if attention is 25% of it.

FP8 Attention and Incoherent Processing

Hopper and Blackwell tensor cores run FP8 at twice the rate of FP16 on paper (1,978 versus 989 TFLOPS on H100 SXM). Naively casting attention inputs to FP8 is dangerous because Q and K activations in large models contain outlier channels. Scores are products of those entries, so a few large values blow up quantization error for everything else in the row.

FlashAttention-3 addresses this with two techniques the paper describes. Incoherent processing multiplies Q and K by the same random orthogonal matrix before quantization. Since (QM)(KM)ᵀ = QKᵀ for orthogonal M, the scores are mathematically unchanged, but the rotation spreads outlier energy across all channels. The authors build M from a Hadamard transform with random sign flips, which costs O(d log d) per vector rather than O(d²) and can be fused with rotary embeddings. In their test with 0.1% large outliers, the technique reduced quantization error by 2.6x. The blog also reports FP8 attention with 2.6x smaller numerical error than a baseline FP8 attention. Block quantization (per-block scaling factors rather than one per tensor) is described in the paper rather than the blog; verify details in the PDF if you depend on them.

The practical rule: FP8 attention is a throughput option for inference and some training, not a drop-in default. Validate perplexity or downstream eval deltas on your model, because outlier structure is model-specific. Whether FA3’s FP8 path is exposed in your framework version, and in which dtype (E4M3 for inputs is typical), is a deployment detail to check in the release notes rather than assume.

Walk-Through: Using FlashAttention Through PyTorch

You rarely call the CUDA kernels directly. PyTorch’s scaled_dot_product_attention (SDPA) dispatches to a FlashAttention-family kernel, a memory-efficient (xFormers-style) kernel, a cuDNN kernel, or a plain math fallback. The sdpa_kernel context manager restricts which backends are allowed, which is the standard way to confirm what you are actually running.

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

assert torch.cuda.is_available()
B, H, N, D = 4, 16, 4096, 128
q, k, v = (torch.randn(B, H, N, D, device="cuda", dtype=torch.bfloat16)
           for _ in range(3))

def bench(backends, iters=20):
    with sdpa_kernel(backends):
        for _ in range(3):                       # warmup
            F.scaled_dot_product_attention(q, k, v, is_causal=True)
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        for _ in range(iters):
            F.scaled_dot_product_attention(q, k, v, is_causal=True)
        torch.cuda.synchronize()
        dt = (time.perf_counter() - t0) / iters
    flops = 4 * B * H * N * N * D / 2            # causal ~ half the work
    return dt * 1e3, flops / dt / 1e12

for name, be in [("flash", [SDPBackend.FLASH_ATTENTION]),
                 ("efficient", [SDPBackend.EFFICIENT_ATTENTION]),
                 ("math", [SDPBackend.MATH])]:
    try:
        ms, tf = bench(be)
        print(f"{name:10s} {ms:8.2f} ms  {tf:7.1f} TFLOPs")
    except RuntimeError as e:
        print(f"{name:10s} unavailable: {str(e)[:60]}")

Three details matter. First, the flash backend requires fp16 or bf16 on a supported NVIDIA GPU; fp32 inputs will fall back or error when you restrict to it. Second, the math backend materializes the full N×N matrix, so at N = 4,096 with 64 head-instances it allocates gigabytes, which is the cleanest demonstration of the problem in this post. Third, if you restrict to a single backend and the call raises, the error message usually states which constraint (dtype, head dimension, mask type) failed, which is more informative than a silent fallback.

Which kernel generation PyTorch uses under the hood depends on the PyTorch version and GPU. Check the release notes for your version instead of assuming FA3 or FA4 is reachable through SDPA; the standalone flash-attn package from the Dao-AILab repository tracks the newest kernels faster than framework releases do.

Correctness check against the math path

Always diff a fused kernel against the reference once:

with sdpa_kernel([SDPBackend.MATH]):
    ref = F.scaled_dot_product_attention(q.float(), k.float(), v.float(), is_causal=True)
with sdpa_kernel([SDPBackend.FLASH_ATTENTION]):
    out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print((out.float() - ref).abs().max())   # expect bf16-level error, about 1e-2 or lower

Differences at bf16 rounding level are normal because tile order changes summation order. A large error indicates a mask, dtype or layout bug, not an algorithmic one.

FlexAttention: Programmable Attention Without Writing a Kernel

Figure 4: A practical selection flow. Standard masks go to the FlashAttention path, custom masks or score changes to FlexAttention, and long-cache decoding to split-KV or paged kernels.

A hand-tuned fused kernel is fast but rigid. Every new attention variant (ALiBi, sliding window, document masks, soft-capping) historically needed either a new kernel or a fall back to the slow math path. The PyTorch team called this the “software lottery.” FlexAttention (PyTorch blog) lets you express the variant as a short Python function that torch.compile fuses into a FlashAttention-style kernel, with the backward pass generated by autograd.

There are two hooks. score_mod(score, b, h, q_idx, kv_idx) edits each score before softmax, enough for relative position bias, ALiBi or soft-capping (tanh(score/cap) * cap, used by Gemma 2 and Grok-1 according to the blog). mask_mod(b, h, q_idx, kv_idx) returns a boolean and is converted by create_block_mask into a BlockMask that records which key blocks each query block must visit. Fully masked blocks are skipped entirely, which gives real FLOP savings for sliding windows and packed documents, something a dense boolean mask passed to a standard kernel cannot do.

from torch.nn.attention.flex_attention import flex_attention, create_block_mask

def sliding_causal(b, h, q_idx, kv_idx):
    return (q_idx >= kv_idx) & (q_idx - kv_idx <= 1024)

block_mask = create_block_mask(sliding_causal, B=None, H=None, Q_LEN=8192, KV_LEN=8192)
flex = torch.compile(flex_attention)
out = flex(q, k, v, block_mask=block_mask)

The blog reports about 90% of FlashAttention-2 forward and 85% of backward performance on A100 for causal attention, 2x faster for a mask_mod causal mask than the equivalent score_mod, and a 71% sample-packing throughput improvement in torchtune. It also lists early limitations: sequence lengths multiples of 128, and grouped-query attention initially requiring replicated key/value heads. Those were launch-time statements; check current docs before relying on them. The BlockMask metadata at 1M tokens and block size 128 is about 60 MB per the same page, a number worth knowing when you push context length.

The decision rule is simple. If your mask is dense causal or full, use SDPA. If it is anything else (windowed, prefix-LM, document-packed, biased), try FlexAttention before writing Triton or CUDA.

Decoding, Paged and Split-KV Variants

Everything above optimizes prefill and training, where many query rows share a K/V pass. Autoregressive decoding is different: each step has a single query token (or a handful with speculative decoding) attending to the entire KV cache. The matmul is a matrix-vector product with arithmetic intensity near one FLOP per byte of cache read, so decoding is memory-bound no matter how clever the tiling. The goal flips from “avoid writing S” to “saturate HBM bandwidth while reading the cache”.

A standard FlashAttention grid parallelizes over batch, heads and query blocks. With one query and a small batch, that is a handful of thread blocks on a GPU with over a hundred SMs, leaving most of the chip idle while a single block streams a 100K-token cache. Flash-Decoding (Dao and collaborators, October 2023 blog) fixes this by also splitting the key/value sequence across thread blocks. Each block computes a partial output and its logsumexp over its slice of the cache, and a small final kernel combines the partials with the same rescaling identity used in the tile loop. The associativity of the online-softmax merge is what makes this legal: (m, ℓ, o) triples from independent chunks combine exactly.

PagedAttention, introduced with vLLM, attacks a different problem: KV-cache fragmentation. It stores the cache in fixed-size non-contiguous blocks addressed through a block table, like virtual memory pages, so sequences can grow without preallocating maximum length and prefixes can be shared. Paged kernels read keys and values through the block table, adding an indirection but removing most memory waste. In production serving engines the two ideas combine: a paged cache plus split-KV decoding kernels. How engines schedule and offload that cache is a topic in vLLM vs SGLang vs TensorRT-LLM, and the tokens-per-step economics change again with speculative decoding, which turns decode steps into small multi-query verifications that look more like short prefills.

Two modeling choices shape decode cost more than kernel tweaks. Grouped-query and multi-query attention shrink the KV cache by sharing key/value heads across query heads, and in Mixture of Experts models where the feed-forward cost per token is low, attention and its cache become the dominant per-token memory traffic. Kernel work and architecture work compound.

Worked decode example

Take a 70B-class model with 80 layers, 8 KV heads, head dimension 128, fp16 cache, serving a 32K-token context. Cache bytes per token per layer are 2 (K and V) × 8 × 128 × 2 bytes = 4,096. Across 80 layers that is 327,680 bytes per token, so a 32,768-token context holds about 10.7 GB. Each decode step must read all of it. At 3.35 TB/s of HBM bandwidth (the H100 SXM figure from NVIDIA’s published spec, quoted from memory) the floor is about 3.2 ms per token for the cache alone, before reading any weights. This is an illustrative calculation with assumed model dimensions, not a benchmark, but it shows why cache size and bandwidth, not FLOPs, set decode speed.

When FlashAttention Does Not Help

FlashAttention is not magic, and the honest list of limits is long.

Short sequences. At a few hundred tokens the N×N matrix is small, launch overhead and the feed-forward and projection layers dominate, and the end-to-end gain shrinks toward zero. The paper’s own BERT-large result at length 512 is 15%, not 3x.

Decoding with tiny batches. As above, the problem is bandwidth and parallelism over the cache, not S materialization. Plain FlashAttention helps little; split-KV kernels help a lot.

Attention is not your bottleneck. If profiling shows 20% of step time in attention, a 2x faster kernel gives 11% overall. Amdahl’s law applies. Profile first.

Unusual masks, biases and head dimensions. Supported combinations of dtype, head dimension, mask type and GPU architecture are finite. Outside them you silently get the memory-efficient or math path. Verify with sdpa_kernel.

Numerics. The kernel is exact in exact arithmetic, but reordered summation and low-precision accumulators change bits. Training runs that depend on bit-reproducibility must use deterministic backward modes and accept some throughput loss. FP8 adds real approximation error.

Hardware coupling. The Hopper and Blackwell kernels depend on architecture-specific instructions (WGMMA, TMA, TMEM). They do not run on A100, consumer cards or other vendors; those use the Ampere-era FA2 path or vendor ports. The best number in a paper is the best number on one chip.

Quadratic compute remains. FlashAttention leaves O(N²) FLOPs untouched. At 1M tokens, exact attention is still enormous, which is why sparse, windowed and linear-attention hybrids remain active research even though IO-aware kernels beat most older approximations.

Memory-efficiency claims are about attention only. Activation memory elsewhere in the transformer, optimizer state and the KV cache itself are unaffected.

Trade-offs, Gotchas, and What Goes Wrong

The recurring operational failures are mundane. A model silently falls back from the flash kernel to the math kernel after someone adds an attention bias or changes dtype to fp32, memory usage jumps by gigabytes, and a job that fit now runs out of memory. Pin the backend in tests with sdpa_kernel so a regression raises an error instead of slowing down.

Padding-heavy batches waste work: FlashAttention variants for variable-length (“varlen”) inputs avoid computing on pad tokens, but only if you pack sequences with cumulative-length metadata. Dense padded batches with a boolean mask pay for the full rectangle.

Deterministic training is a real trade. The backward pass accumulates dQ with atomics; non-deterministic order gives run-to-run bit differences. If you need reproducibility, enable the deterministic option and measure the throughput loss on your shapes (the FA4 authors report 85 to 90% of non-deterministic speed on their kernel).

Version drift bites too. Kernel packages are compiled per CUDA version and GPU architecture; mismatched wheels either fail to import or quietly select a slower path. Treat the attention package like a driver: pin versions and test on the target GPU.

Finally, beware chasing peak utilization numbers. A kernel at 75% of tensor-core peak for head dimension 128 and 8K sequences says little about your 512-token, head-dimension-64 workload on a different GPU. Benchmark with your own shapes, include the backward pass for training, and compare against end-to-end step time.

Practical Recommendations

For most teams the right plan is to use the framework path and verify it. Call scaled_dot_product_attention with bf16 or fp16 inputs, pin the intended backend in a unit test, and compare outputs once against the math reference. That captures nearly all the benefit of FlashAttention-2 class kernels on Ampere and Ada parts with no custom code.

On Hopper, check whether your framework version exposes an FA3-class kernel, and if not, evaluate the standalone flash-attn package. Gains are largest for long-sequence prefill and training with head dimension 128. Evaluate FP8 only with an accuracy gate on your own model. On Blackwell, benchmark FA4-style kernels against the current cuDNN attention path, because the FA4 write-up itself notes cuDNN has converged.

For serving, treat decode as a bandwidth problem: use paged KV caches, split-KV decoding, grouped-query attention where the model allows, and quantized caches if accuracy permits. For nonstandard masks, reach for FlexAttention first.

A short checklist:

Confirm the active backend with sdpa_kernel; fail tests on fallback.
Use bf16 or fp16, head dimensions your kernel supports, and packed varlen batches instead of padding.
Profile attention’s share of step time before optimizing it.
Benchmark your shapes, forward and backward, against end-to-end step time.
Gate FP8 attention behind an accuracy evaluation.
Pin kernel package, CUDA and driver versions per GPU architecture.
Decide on deterministic backward explicitly and measure its cost.

Frequently Asked Questions

What is FlashAttention and how does it work?

FlashAttention is an exact attention algorithm that avoids materializing the N×N score matrix in GPU memory. It loads tiles of Q, K and V into on-chip SRAM, computes scores, applies an online softmax that tracks a running maximum and sum, accumulates the output, and writes only the result and a per-row logsumexp back to HBM. The same math as standard attention, far fewer memory transfers.

Is FlashAttention an approximation of attention?

No. It computes the same softmax attention as the standard implementation. The online softmax recurrence is algebraically identical to the three-pass version, so differences come only from floating-point reordering, which typically stays at bf16 or fp16 rounding level. Approximation enters only if you opt into FP8 attention or pair it with sparse or windowed masks that change what attention computes.

What is the difference between FlashAttention-2, FlashAttention-3 and FlashAttention-4?

FA2 improved work partitioning and parallelism on Ampere GPUs, reaching up to about 73% of peak per its paper. FA3 targets Hopper, using WGMMA, TMA, warp specialization, softmax-matmul overlap and FP8, reporting up to 740 TFLOPS in FP16 and near 1.2 PFLOPS in FP8. FA4 targets Blackwell with tensor memory, 2-CTA MMA and softmax exponential emulation, reporting up to 1,605 TFLOPs in BF16.

Does FlashAttention help during inference decoding?

Only partly. Decoding processes one query against a long KV cache, which is memory-bound regardless of tiling, and plain FlashAttention under-uses the GPU for small batches. Flash-Decoding splits the key/value sequence across thread blocks and merges partial results with the online-softmax rescaling identity, and PagedAttention reduces cache fragmentation. Together these, plus grouped-query attention, matter more for decode than the original kernel.

How do I check which attention kernel PyTorch is using?

Wrap the call in torch.nn.attention.sdpa_kernel with a single backend such as SDPBackend.FLASH_ATTENTION. If that backend cannot handle your dtype, head dimension, mask or GPU, PyTorch raises an error that names the failing constraint instead of silently falling back. You can also profile with the PyTorch profiler and look for flash or efficient attention kernel names in the trace.

Should I use FP8 attention?

Treat it as an optimization to validate, not a default. FP8 doubles nominal tensor throughput on Hopper (1,978 versus 989 TFLOPS), but outlier activations can cause large quantization error. FlashAttention-3 mitigates this with incoherent processing, a random orthogonal Hadamard-based rotation of Q and K, and reports 2.6x smaller error than baseline FP8 attention. Always confirm downstream accuracy on your model before shipping.