One-sentence summary: FlashAttention is not approximate attention, and it does not remove the $N^2$ score interactions. Tiling and online Softmax avoid repeatedly writing a complete $N \times N$ intermediate matrix to device memory.

Naive attention writes full score and probability matrices to HBM, while FlashAttention merges on-chip tiles into the output

21.1 The Bottleneck Is Not Only “Too Much Arithmetic”

For one attention head, let $Q,K,V \in \mathbb{R}^{N \times d}$ :

O = \operatorname{softmax}\!\left(\frac{QK^T}{\sqrt d} + B_{\mathrm{mask}}\right)V

The leading arithmetic remains $\Theta(N^2d)$ . A naive implementation also materializes an $N \times N$ score matrix and may materialize another equally large probability matrix after Softmax.

At $N=4096$ , one FP16 score matrix for one head and one sample occupies:

4096^2 \times 2\ \text{bytes} = 32\ \text{MiB}

If a naive path retains that matrix for batch size 8 and 32 heads, the total is 8 GiB—before masks, probabilities, Dropout, and other backward state. This is an account of a naively materialized path, not a claim that every current PyTorch attention call really allocates 8 GiB.

The other half of the problem is movement:

read Q/K from HBM → compute S=QKᵀ → write S to HBM
read S from HBM   → mask/Softmax  → write P to HBM
read P/V from HBM → compute O=PV  → write O to HBM

Tensor Cores can multiply matrices quickly, but idle compute units still wait for data. FlashAttention attacks this I/O path.

21.2 Tiling: Keep the Large Sheet off the Floor

Keep the original desk analogy: HBM is a large bookshelf, while registers, shared memory, and L1 are the small work surface. Only one set of pages fits on the desk, so finish the current calculation—or reduce it to a compact running state—before fetching the next set.

FlashAttention brings one $Q_i$ tile and one $K_j,V_j$ pair on-chip:

Sᵢⱼ = QᵢKⱼᵀ / √d + mask
                 ↓
        online-Softmax state
                 ↓
         accumulate output Oᵢ

Q and K/V tiles form temporary scores on chip without writing the complete N by N matrix to HBM

Several hardware numbers are easily conflated. On A100, the whole GPU has 40 MB of L2; each SM has a combined L1/texture/shared-memory data cache of at most 192 KB, with a configurable shared-memory carveout. Those resources are not one freely shared “20 MB SRAM pool.” The 40 GB A100 has about 1550 GB/s HBM2 bandwidth and the 80 GB variant about 2039 GB/s. Different parts and generations do not support one timeless “SRAM is always 20 times faster” table.

The idealized tile sizes in Algorithm 1 of the original paper are:

B_c=\left\lceil\frac{M}{4d}\right\rceil,\qquad B_r=\min\!\left(\left\lceil\frac{M}{4d}\right\rceil,d\right)

Here $d$ is the per-head dimension, and $M$ is the number of scalar elements that fit in the paper's abstract on-chip memory—not a raw byte count to plug in without the dtype. Real kernels also balance dtype, register pressure, shared-memory carveout, alignment, warp partitioning, and GPU generation. The formula does not uniquely decree that every real tile must be 64.

21.3 Online Softmax: See One Tile, Preserve the Whole Row

The Softmax denominator spans a full row. Applying an independent Softmax to every tile and concatenating the results would be wrong. For one query row, maintain across processed key blocks:

$m$ : the running maximum score;
$\ell$ : the running exponential sum at that maximum's scale;
$a$ : the unnormalized, Value-weighted accumulator.

When a new score block $s$ and its $V_j$ arrive:

m'=\max(m,\max(s)),\qquad \alpha=e^{m-m'}

p=e^{s-m'},\qquad \ell'=\alpha\ell+\sum p,\qquad a'=\alpha a+pV_j

Return $o=a/\ell$ after the final block. If a later tile contains a larger score, $\alpha$ rescales everything already accumulated to the new numerical scale.

Online Softmax merges score blocks through a running maximum, denominator, and unnormalized output accumulator

This unmasked PyTorch version is for checking the recurrence, not for speed:

import math
import torch


def online_attention(q, k, v, block_size):
    # q: [Nq, d], k: [Nk, d], v: [Nk, dv]
    scale = 1.0 / math.sqrt(q.size(-1))
    rows = q.size(0)
    m = torch.full((rows,), -torch.inf, device=q.device, dtype=q.dtype)
    denominator = torch.zeros_like(m)
    accumulator = torch.zeros(
        rows, v.size(-1), device=q.device, dtype=q.dtype
    )

    for start in range(0, k.size(0), block_size):
        k_block = k[start:start + block_size]
        v_block = v[start:start + block_size]
        scores = q @ k_block.transpose(-2, -1) * scale

        block_max = scores.max(dim=-1).values
        new_m = torch.maximum(m, block_max)
        correction = torch.exp(m - new_m)
        probabilities = torch.exp(scores - new_m[:, None])

        denominator = (
            correction * denominator + probabilities.sum(dim=-1)
        )
        accumulator = (
            correction[:, None] * accumulator + probabilities @ v_block
        )
        m = new_m

    return accumulator / denominator[:, None]

For the same $Q,K,V$ , it should match softmax(Q @ K.T / sqrt(d)) @ V within floating-point tolerance. “Exact attention” means the algorithm does not introduce sparse or low-rank approximation; reordered finite-precision operations are not promised to be bit-for-bit identical.

21.4 Keep Arithmetic, Extra Memory, and HBM I/O Separate

FlashAttention preserves quadratic score arithmetic while reducing extra memory from quadratic to linear and lowering HBM traffic

Question	Naively materialized attention	FlashAttention
Leading arithmetic	$\Theta(N^2d)$	$\Theta(N^2d)$
Extra intermediate memory, fixed d	$\Theta(N^2)$	$O(N)$
HBM accesses in the original abstract machine	$\Theta(Nd+N^2)$	$\Theta(N^2d^2/M)$

The third row has conditions: $M$ counts on-chip elements, and the paper analyzes $d \le M \le Nd$ . It cannot be casually simplified to $O(N^2d/M)$ . Nor does the asymptotic drop from quadratic to linear extra memory mean real peak memory is exactly N times smaller; Q/K/V, outputs, weights, optimizer state, allocator behavior, and the rest of the model remain.

In backward, FlashAttention recomputes local score tiles from saved statistics instead of retaining the full probability matrix. This adds some FLOPs, but it does not imply that backward must be slower overall. Saving HBM traffic can still make the kernel and end-to-end training faster.

21.5 FA1 to FA4: One Core Idea, Hardware-Specific Kernels

Algorithm and hardware emphasis from FlashAttention 1 through FlashAttention 4 as of July 2026

FA1 (2022) established I/O-aware tiling plus online Softmax without materializing the full matrix.
FA2 (2023) reduced non-matmul FLOPs and improved thread-block and warp work partitioning. On the paper's A100 workloads, it reported roughly 2× FA1, 50–73% of theoretical peak, and up to 225 TFLOPs/s/GPU at 72% model FLOPs utilization for GPT-style training. Those are bounded paper measurements, not a guaranteed speedup for any model.
FA3 (2024; still labelled a Hopper beta in the official repository) uses H100/H800 paths such as TMA, WGMMA, and FP8.
FA4 (2026) targets the asymmetric scaling of Blackwell B200/GB200 and uses CuTe DSL with pipeline co-design. As of 2026-07 the official repository provides flash-attn-4 pre-releases; head-dimension, dtype, and backward coverage must still be checked against the current release rather than inferred from FA2.

A larger version number is not universally faster on every GPU. FA3 targets Hopper and FA4 targets Blackwell; Ampere and Ada may still use FA2.

21.6 In Practice: Let the Framework Dispatch First

PyTorch's scaled_dot_product_attention selects among available backends according to device, dtype, shape, masks, and other constraints. With the Chapter 18 layout, arrange Q/K/V as [batch, heads, time, head_dim]:

import torch.nn.functional as F


def attention(q, k, v, dropout_p, training):
    return F.scaled_dot_product_attention(
        q,
        k,
        v,
        attn_mask=None,
        dropout_p=dropout_p if training else 0.0,
        is_causal=True,
    )

scaled_dot_product_attention always applies Dropout according to the passed dropout_p; it does not inspect an enclosing module's training flag. Pass 0.0 explicitly in evaluation.

To prove that one input is eligible for the Flash backend, restrict the backend in diagnostic code:

from torch.nn.attention import SDPBackend, sdpa_kernel

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    output = F.scaled_dot_product_attention(
        q, k, v, dropout_p=0.0, is_causal=True
    )

PyTorch reports why a fused kernel cannot run. Production code should not necessarily disable every fallback merely to claim the word “Flash.”

The official Dao-AILab/flash-attention package is another direct route. Its hardware support is no longer accurately described as “new NVIDIA only”: as of 2026-07 the FA2 CUDA path covers Ampere/Ada/Hopper, while the official repository also provides CK and Triton AMD backends for ROCm. Check the current PyTorch, CUDA/ROCm, GPU, dtype, and head-dimension matrix before installing.

21.7 How to Benchmark Honestly

At minimum, freeze and report:

GPU model and power mode; PyTorch, CUDA/ROCm, and kernel versions;
dtype, batch, heads, query length, key length, and head dimension;
causal/noncausal mode, mask, Dropout, and forward versus forward+backward;
warmup, timing method, median/percentiles, and whether this is a kernel microbenchmark or whole-model wall time;
peak memory plus output and gradient error.

For short sequences or small batches, launch and layout overhead can erase the gain. An attention kernel running 2× faster does not imply that the whole model—or per-token latency—runs 2× faster.

FlashAttention and KV caching also solve different problems: FlashAttention reduces I/O within one attention computation; a KV cache avoids recomputing old Keys and Values across autoregressive decoding steps.

21.8 Chapter Summary

FlashAttention preserves the mathematical attention result; its core contribution is an I/O-aware implementation.
It still performs leading $\Theta(N^2d)$ arithmetic, but does not materialize the full $N \times N$ intermediate matrix.
Online Softmax rescales both the denominator and Value accumulator, which makes tile results composable.
$d$ is head dimension and $M$ is abstract on-chip capacity; model width and raw byte counts cannot be substituted blindly.
FA1–FA4 kernels and hardware support keep changing, so every speed number needs a workload and version boundary.
PyTorch SDPA can dispatch automatically, but callers must still pass dropout_p=0.0 for evaluation.

Chapter Checklist

Separate compute complexity, extra-memory complexity, and HBM I/O
Derive the online-Softmax tile merge with $m,\ell,a$
Explain why “exact” is not a bit-for-bit promise
Call PyTorch SDPA correctly and verify the backend actually selected
Design a benchmark that does not present kernel speed as whole-model speed

Primary Sources

See You in the Next Chapter

FlashAttention makes the current attention computation cheaper, but Chapter 20's autoregressive loop still recomputes old Keys and Values at every step. Chapter 22 removes that cross-step redundancy with a KV Cache.