Learn: The Memory Wall & FlashAttention

What you're learning

Why the single most important kernel in modern LLMs — attention — is limited by moving bytes, not doing math, and how FlashAttention fixes it without changing the math at all. By the end you should be able to derive the online softmax recurrence from scratch, compute the HBM traffic of naive vs. tiled attention, and explain why every FlashAttention generation is a story of algorithm–hardware co-design rather than “more FLOPs.”

This lesson assumes Tier 1 and lessons 04-first-cuda-kernels (writing kernels, launch overhead) and 05-optimizing-gemm (SMEM tiling, register blocking). It is the payoff of 02-gpu-memory-hierarchy (the memory wall) and 03-performance-modeling-roofline (arithmetic intensity, the ridge point). If “ridge point” or “memory-bound” don’t feel automatic, reread those first.

---

1. Learning map

graph TD
    A["Memory wall + roofline<br/>(lessons 02–03)"] --> B["Kernel fusion<br/>fewer HBM round-trips"]
    B --> B1["Vertical fusion<br/>(op chains)"]
    B --> B2["Horizontal fusion<br/>(sibling ops)"]
    B --> B3["torch.compile / inductor<br/>(autofuser)"]
    A --> C["Attention is memory-bound<br/>materializing NxN S,P in HBM"]
    C --> D["FlashAttention<br/>(IO-aware, arXiv:2205.14135)"]
    D --> D1["Tile Q,K,V into SRAM<br/>never write NxN"]
    D --> D2["Online softmax<br/>running m, l + rescale"]
    D --> D3["Backward: recompute S<br/>trade FLOPs for memory"]
    D --> E["FlashAttention-2<br/>work partitioning (2307.08691)"]
    E --> F["FlashAttention-3<br/>Hopper async, FP8 (2407.08608)"]
    F --> G["FlashAttention-4<br/>Blackwell co-design (2603.05451)"]
    D2 --> H["Big idea:<br/>algorithm–hardware co-design"]
    D1 --> H
    G --> H

    style C fill:#a44,color:#fff
    style D fill:#44a,color:#fff
    style D2 fill:#4a4,color:#fff
    style H fill:#a83,color:#fff

---

2. Why this matters

Attention is where LLMs spend a large and growing fraction of their time as context lengths climb — and, done naively, it is one of the most bandwidth-wasteful operations in the whole model. The naive implementation materializes an $N\times N$ score matrix in HBM, so its memory traffic (and memory footprint) grows quadratically in sequence length. At $N=8192$ a single head’s score matrix is already 128 MiB in BF16; at $N=128\text{k}$ it is 32 GiB — larger than an H100’s entire 80 GB of HBM¹ if you keep a few of them around. Long context was, quite literally, blocked by this.

FlashAttention (Dao et al., 2022) removed that wall without approximation. It computes exactly the same attention, but by respecting the memory hierarchy — tiling the computation so the $N\times N$ matrix never touches HBM — it turned a memory-bound kernel into a near-compute-bound one, cut wall-clock 2–4×, and made million-token context tractable.² It is the canonical example of the thesis running through this whole curriculum: on modern accelerators you win by moving fewer bytes, not by doing fewer FLOPs.

The one-sentence version

FlashAttention does more arithmetic than naive attention (it recomputes things), yet runs far faster — because it pays FLOPs, which are cheap, to avoid HBM traffic, which is expensive. That trade is the entire game.

---

3. Kernel fusion: why many small kernels waste HBM

Before attention, understand the general disease it’s a cure for. Recall from lesson 03 the arithmetic intensity of an elementwise op: read the tensor, do a fixed few FLOPs per element, write it back. Intensity $I\approx O(1)$ FLOP/byte — far left of the H100 ridge point (~296 FLOP/byte for BF16).¹ Every such op runs at HBM bandwidth, and its runtime is essentially “time to read the tensor + time to write it.”

Now stack a typical activation chain: x → bias-add → GELU → dropout → residual-add. If each is its own kernel, each one:

1. reads the full activation tensor from HBM,
2. does a trivial amount of math,
3. writes the full tensor back to HBM.

For a tensor of $n$ BF16 elements, four separate kernels move $\approx 4\times (2n+2n)=16n$ bytes through HBM, to perform maybe a dozen FLOPs per element. The tensor is round-tripped four times. Each intermediate is born in HBM and immediately re-read by the next kernel — pure waste.

Kernel fusion folds the chain into one kernel: read $x$ once into registers/SMEM, apply bias → GELU → dropout → add in place while it’s on-chip, write once. Traffic drops from $16n$ to $4n$ bytes — a ~4× speedup on a purely bandwidth-bound op, for free. Nothing about the math changed; you just stopped spilling intermediates to HBM.

$$\text{unfused traffic}=\sum _{\text{ops}}(\text{read}+\text{write})\gg \text{fused traffic}={\text{read}}_{\text{first}}+{\text{write}}_{\text{last}}$$

Fusion raises arithmetic intensity

Fusion doesn’t add FLOPs and doesn’t remove them — it removes bytes. Same numerator, smaller denominator in $I=\text{FLOPs}/\text{bytes}$, so the fused kernel sits further right on the roofline. You are literally buying arithmetic intensity by keeping intermediates on-chip. There is also a second, smaller win: one kernel launch instead of four, so less launch/overhead latency (lesson 04) — this matters most for tiny tensors where launch cost rivals compute.

Vertical vs. horizontal fusion

• Vertical (producer→consumer) fusion fuses ops along a dependency chain: the output of one is the input of the next (the bias→GELU→dropout example). This is the high-value case because it eliminates the intermediate tensors entirely. FlashAttention is an extreme vertical fusion — it fuses the entire matmul → softmax → matmul chain of attention into one kernel.
• Horizontal fusion fuses independent ops that share an input or can share a launch — e.g. computing the Q, K, and V projections as one batched GEMM instead of three, or applying the same elementwise op to several tensors in one grid. The win here is amortized launch overhead and better occupancy, not eliminated intermediates.

graph LR
    subgraph "Vertical fusion (chain)"
        direction TB
        v1["read x"] --> v2["bias"] --> v3["GELU"] --> v4["dropout"] --> v5["write y"]
    end
    subgraph "Horizontal fusion (siblings)"
        direction TB
        h0["read x once"] --> hq["→ Q proj"]
        h0 --> hk["→ K proj"]
        h0 --> hv["→ V proj"]
    end

torch.compile / Inductor as an autofuser

You rarely hand-write fusions for elementwise chains anymore. torch.compile traces your model into an FX graph and its backend TorchInductor performs automatic fusion: it groups compatible elementwise/reduction ops into fused Triton kernels (pointwise fusion, and reductions like softmax/layernorm folded with their neighbors), so the intermediates stay in registers/SMEM. This is why torch.compile gives large speedups on the “glue” between GEMMs with zero code changes. It generally will not discover a fundamentally new algorithm like FlashAttention on its own — that requires the algorithmic restructuring below — which is why attention is dispatched to a hand-written fused kernel (or F.scaled_dot_product_attention, which calls a FlashAttention backend). We go deep on Triton and Inductor in 07-triton-and-modern-kernels.

Quiz: Why does splitting a bias→GELU→dropout chain into three separate kernels waste HBM bandwidth, and roughly how much do you save by fusing? Would fusing help a large matmul the same way?

Answer

Why it wastes bandwidth. Each op is elementwise with arithmetic intensity $\sim O(1)$, so each is memory-bound: its runtime is dominated by reading the tensor from HBM and writing it back. Three separate kernels therefore round-trip the entire activation tensor three times — the intermediate after bias is written to HBM only to be immediately re-read by GELU, etc. For $n$ BF16 elements: unfused $\approx 3\times 4n=12n$ bytes; fused $\approx 4n$ bytes (read once, write once). That’s a ~3× traffic reduction → ~3× speedup, because the op is purely bandwidth-bound so time $\propto$ bytes.
Would it help a matmul? No, not in the same way. A large GEMM is compute-bound ($I\gg$ ridge point) — its runtime is set by tensor-core throughput, not HBM traffic, so removing a round-trip of its (small relative to FLOPs) I/O barely moves the needle. Fusion is a lever for the memory-bound ops, which is precisely why we fuse the elementwise/softmax “glue” around GEMMs rather than the GEMMs themselves. (You can fuse an elementwise epilogue into a GEMM — e.g. bias+GELU in the GEMM’s epilogue — and that helps, but the GEMM itself stays compute-bound.)

---

4. The attention memory problem

Standard scaled dot-product attention, for one head with $Q,K,V\in {\mathbb{R}}^{N\times d}$:

$$S=\frac{Q{K}^{\top }}{\sqrt{d}}\in {\mathbb{R}}^{N\times N},P=\text{softmax}(S)\in {\mathbb{R}}^{N\times N},O=PV\in {\mathbb{R}}^{N\times d}$$

The naive implementation (and what you get from three separate library calls) executes this as a sequence of kernels, each round-tripping through HBM:

1. GEMM $S=Q{K}^{\top }/\sqrt{d}$: read $Q,K$ ($\sim 2Nd$), write $S$ ($N^2$).
2. Softmax (a memory-bound elementwise+reduction kernel): read $S$ ($N^2$), write $P$ ($N^2$).
3. GEMM $O=PV$: read $P$ ($N^2$), read $V$ ($Nd$), write $O$ ($Nd$).

Count the HBM traffic in elements. The $N\times N$ matrix is written once, read once (softmax), written once (probs), and read once again (second GEMM): roughly $4{\mathbf{N}}^2$ element-accesses of $N^2$-sized traffic, versus only $O(Nd)$ for the actual inputs and output. For $N\gg d$ the traffic is $\Theta (N^2)$ — quadratic in sequence length — while the useful FLOPs are $\Theta (N^{2}d)$. The arithmetic intensity is therefore only $\sim d$ (tens to ~128), far below the ridge point.

Attention is memory-bound at typical shapes — the softmax is the smoking gun

The two matmuls ($Q{K}^{\top }$ and $PV$) are individually fine, but the softmax kernel in the middle is pure elementwise/reduction: it reads the entire $N\times N$ score matrix and writes it back, doing a handful of FLOPs per element ($I\approx O(1)$). It sits on the memory roof. And it forces the $N^2$ matrix to be materialized in HBM so it can be read back — the matmuls can’t hand it off on-chip because a separate kernel needs it. So the whole attention block, as launched, is dominated by $\Theta (N^2)$ HBM traffic and runs far below tensor-core peak. It is bottlenecked by moving the score matrix, not by the matmuls.

There’s a second, equally fatal problem: memory capacity. Materializing $S$ (and $P$) costs $O(N^2)$ HBM storage, which is what historically capped context length. This is the wall FlashAttention tears down.

Figure 1 — Naive attention’s HBM traffic is dominated by four $N\times N$ round-trips of the score/probability matrix (quadratic in $N$). FlashAttention never materializes those matrices in HBM, leaving only the linear-in-$N$ traffic of reading $Q,K,V$ and writing $O$. This $O(N^2)\to O(N)$ collapse in traffic is the whole win.

Quiz: Standard attention does $\Theta (N^{2}d)$ FLOPs — genuinely a lot of compute. Why is it nonetheless memory-bound in practice, and which sub-step is the culprit?

Answer

Because the FLOPs and the bytes don’t live in the same kernel. The $\Theta (N^{2}d)$ FLOPs are in the two matmuls, but the algorithm as launched forces the $N\times N$ score matrix through HBM four times (write $S$, read $S$, write $P$, read $P$), which is $\Theta (N^2)$ traffic with a tiny amount of math attached — the arithmetic intensity of the whole block collapses toward $\sim d$, far below the ~296 FLOP/byte ridge point. The softmax is the culprit: it’s an elementwise/reduction kernel ($I\approx O(1)$) that sits on the memory roof and forces the $N^2$ matrix to be materialized in HBM so it can be read back. So even though attention has plenty of FLOPs, wall-clock is set by shuttling the score matrix, not by the tensor cores. (This is also why lower head dim $d$ makes it worse: intensity $\sim d$ shrinks.) FlashAttention fixes exactly this by never letting $S,P$ reach HBM.

---

5. FlashAttention: the IO-aware algorithm

FlashAttention (Dao, Fu, Ermon, Rudra, Ré — arXiv:2205.14135) is “IO-aware”: it is designed around minimizing HBM reads/writes, the same way tiled GEMM (lesson 05) is designed around reusing SMEM tiles.² Two ideas do all the work:

1. Tile $Q,K,V$ into blocks small enough to fit in SRAM/SMEM. Compute attention block-by-block. The block-level scores $S_{ij}$ live only in SRAM — they are produced, consumed, and discarded on-chip, and never written to HBM.
2. Online (streaming) softmax.³ Softmax needs a global normalizer over the whole row (${\sum }_j{e}^{s_j}$), which seems to require seeing all of $S$ at once. FlashAttention computes it incrementally with a running max and running denominator, applying a small correction each time a new block shifts the max. This is the mathematical heart of the method.

5.1 Deriving the online softmax

Start from safe softmax for one query row with scores $s_1,\dots ,s_N$ (here $s_j=q^{\top }{k}_j/\sqrt{d}$). To avoid overflow of $e^{s_j}$ we subtract the row max $m={\max }_j{s}_j$:

$$o=\sum _{j=1}^N\frac{e^{s_j-m}}{{\sum }_{l=1}^N{e}^{s_l-m}}{v}_j,m=\underset{j}{\max }{s}_j.$$

The obstacle: both $m$ and the denominator $\ell ={\sum }_l{e}^{s_l-m}$ depend on the entire row, but we want to stream over blocks $B_1,B_2,\dots ,B_T$ of the key/value sequence, seeing only one block at a time. Can we maintain the running state and correct it when a later block reveals a bigger max?

Maintain three running quantities after processing blocks $\mathrm{1..}t$:

• $m^{(t)}$ — running max of all scores seen so far,
• ${\ell }^{(t)}$ — running denominator, normalized to the current max $m^{(t)}$,
• ${\mathbf{o}}^{(t)}\in {\mathbb{R}}^d$ — running unnormalized output accumulator, also referred to $m^{(t)}$.

Initialize $m^{(0)}=-\infty$, ${\ell }^{(0)}=0$, ${\mathbf{o}}^{(0)}=0$. For block $B_t$, let ${\tilde{m}}^{(t)}={\max }_{j\in B_t}{s}_j$ be the block-local max. The recurrence is:

$$m^{(t)}=\max (m^{(t-1)},{\tilde{m}}^{(t)})$$ $$\boxed{{\ell }^{(t)}=\underset{{\alpha }^{(t)}\text{(rescale)}}{\underbrace{e^{m^{(t-1)}-m^{(t)}}}}{\ell }^{(t-1)}+\sum _{j\in B_t}{e}^{s_j-m^{(t)}}}$$ $$\boxed{{\mathbf{o}}^{(t)}=\underset{{\alpha }^{(t)}}{\underbrace{e^{m^{(t-1)}-m^{(t)}}}}{\mathbf{o}}^{(t-1)}+\sum _{j\in B_t}{e}^{s_j-m^{(t)}}{v}_j}$$

After the last block, normalize once: $o={\mathbf{o}}^{(T)}/{\ell }^{(T)}$.

Why the correction factor ${\alpha }^{(t)}=e^{m^{(t-1)}-m^{(t)}}$? The old accumulators ${\ell }^{(t-1)},{\mathbf{o}}^{(t-1)}$ were computed with exponentials referred to the old max $m^{(t-1)}$. When the max jumps to $m^{(t)}$, every previously accumulated term $e^{s_j-m^{(t-1)}}$ must be re-referenced to the new max:

$$e^{s_j-m^{(t-1)}}\cdot e^{m^{(t-1)}-m^{(t)}}=e^{s_j-m^{(t)}}.$$

Multiplying the whole running sum by ${\alpha }^{(t)}=e^{m^{(t-1)}-m^{(t)}}$ does exactly this to every stored term at once — that’s why a single scalar multiply corrects the entire accumulator. If the new block doesn’t raise the max ($m^{(t)}=m^{(t-1)}$), then ${\alpha }^{(t)}=e^0=1$ and no rescale happens. Note ${\alpha }^{(t)}\le 1$ always (we subtract a $\ge$ value), so it never overflows.

The invariant that makes it exact

The recurrence maintains, at every step $t$, the exact partial softmax over the prefix seen so far:
${\ell }^{(t)}={\sum }_{j\in B_1\cup \cdots \cup B_t}{e}^{s_j-m^{(t)}},{\mathbf{o}}^{(t)}={\sum }_{j\in B_1\cup \cdots \cup B_t}{e}^{s_j-m^{(t)}}{v}_j.$
By induction (base case trivial; inductive step is the algebra above), after the final block ${\mathbf{o}}^{(T)}/{\ell }^{(T)}$ equals the full softmax attention output bit-for-bit up to floating-point associativity — FlashAttention is exact, not an approximation. This is the crucial selling point: you can drop it into any model with no accuracy change.

Because $\ell$ and $\mathbf{o}$ are just running scalars/vectors, the block scores $S_{ij}$ can be thrown away immediately after they update the state — they never need to be stored. That is why the $N\times N$ matrix never reaches HBM.

Figure 2 — For each query block $Q_i$, FlashAttention streams over the K/V blocks. Each score tile $S_{it}$ is formed, consumed to update the running $(m,\ell ,\mathbf{o})$ state, and discarded — all in SRAM. Only $Q,K,V$ (in) and $O$ (out) cross HBM; the $O(d)$-sized running state stays in registers. The whole matmul→softmax→matmul chain is fused into one kernel.

5.2 The HBM traffic reduction: $O(N^2)\to O(N)$

Naive attention’s HBM traffic is $\Theta (N^2+Nd)$ — dominated by the $N\times N$ score/probability round-trips. FlashAttention never materializes those matrices, so the only mandatory HBM traffic is reading $Q,K,V$ and writing $O$: that’s $\Theta (Nd)$ — linear in $N$. The score matrix’s contribution to HBM traffic goes to zero.

There is a re-read subtlety: because you tile, whichever operand is in the outer loop gets re-read once per inner-loop pass. The paper’s exact bound is that FlashAttention performs $\Theta \left(N^2{d}^2/M\right)$ HBM accesses (with $M$ = SRAM size in elements), versus $\Theta (N^2+Nd)$ for standard attention — a reduction by a factor $\sim M/d^2$. Since $M$ (tens of thousands of elements) is much larger than $d^2$ (e.g. $64^2=4096$), this is a large multiplicative win; empirically the original paper measured ~9× fewer HBM accesses for GPT-2 and 2–4× wall-clock speedup.² The headline mental model, though, is the clean one: you replaced quadratic score-matrix traffic with linear activation traffic.

5.3 The backward pass: recomputation trades FLOPs for memory

The backward pass needs the attention probabilities $P=\text{softmax}(S)$ to compute gradients (e.g. $dV=P^{\top }dO$, and the softmax Jacobian for $dS$). Naive backprop would store the $N\times N$ matrix $P$ from the forward pass to reuse it — reintroducing exactly the $O(N^2)$ HBM footprint we just eliminated.

FlashAttention’s fix is recomputation (a form of activation checkpointing specialized to attention): during backward, it re-derives the block scores $S_{ij}$ and probabilities on the fly in SRAM from $Q,K,V$, using the stored per-row softmax statistics. Cleverly, it only needs to save the logsumexp $L_i=m_i+\log {\ell }_i$ per row (an $O(N)$ vector, not $O(N^2)$) from the forward pass — with $L_i$ in hand, a block’s probabilities are $P_{ij}=\exp (S_{ij}-L_i)$, recomputable exactly without re-running the softmax reduction.² This trades extra FLOPs (recomputing $S$) for a massive memory saving (no $N^2$ matrix stored), and because attention is memory-bound, the recomputed backward is still faster than the store-and-reload version. It is the same FLOPs-for-bytes bargain as the forward pass.

Recomputation is not free FLOPs — it's cheaper than the memory it saves

Recomputing $S$ in the backward pass genuinely does more arithmetic than a store-based backward. That sounds backwards until you internalize the memory wall: the recomputed FLOPs run on idle tensor cores that were waiting on HBM anyway, whereas storing/reloading the $N^2$ matrix would saturate the scarce bandwidth. On a memory-bound op, spending FLOPs to save bytes is a win. This is the same logic as gradient checkpointing (lesson 03’s HFU/MFU gap) — but here it also speeds you up, not just saves memory.

Quiz: In the online-softmax recurrence, justify the correction factor ${\alpha }^{(t)}=e^{m^{(t-1)}-m^{(t)}}$ from first principles. What goes wrong if you omit it, and why is ${\alpha }^{(t)}$ never a source of overflow?

Answer

Justification. The running accumulators after block $t-1$ store exponentials referred to the then-current max $m^{(t-1)}$: ${\ell }^{(t-1)}={\sum }_{j<B_t}{e}^{s_j-m^{(t-1)}}$ and likewise ${\mathbf{o}}^{(t-1)}$. When block $t$ reveals a larger max, the running max updates to $m^{(t)}=\max (m^{(t-1)},{\tilde{m}}^{(t)})$, and every stored term must be re-referenced to $m^{(t)}$. For a single term: $e^{s_j-m^{(t-1)}}\cdot e^{m^{(t-1)}-m^{(t)}}=e^{s_j-m^{(t)}}$. Multiplying the whole accumulator by the scalar ${\alpha }^{(t)}=e^{m^{(t-1)}-m^{(t)}}$ applies this correction to all stored terms simultaneously, restoring the invariant ${\ell }^{(t)}={\sum }_{j\le B_t}{e}^{s_j-m^{(t)}}$. Then the new block’s terms ${\sum }_{j\in B_t}{e}^{s_j-m^{(t)}}$ (already referred to $m^{(t)}$) are added.
If you omit it: the old terms stay scaled by $e^{s_j-m^{(t-1)}}$ while new terms use $e^{s_j-m^{(t)}}$ — you’d be summing exponentials with inconsistent reference points, so the denominator and numerator no longer correspond and the result is wrong (unless the max never changed). The final normalization $\mathbf{o}/\ell$ would not equal the true softmax.
No overflow: $m^{(t)}=\max (m^{(t-1)},{\tilde{m}}^{(t)})\ge m^{(t-1)}$, so the exponent $m^{(t-1)}-m^{(t)}\le 0$, hence ${\alpha }^{(t)}=e^{(\le 0)}\in (0,1]$. It’s a shrink factor, never a blow-up. Likewise every $e^{s_j-m^{(t)}}\le 1$ because $s_j\le m^{(t)}$. Numerical stability is preserved throughout — which is the whole point of the safe-softmax max subtraction, now maintained incrementally.

---

6. FlashAttention-2: better work partitioning

FlashAttention-1 already cut HBM traffic, but the forward pass still only reached ~30–50% of A100 peak FLOP/s (the backward worse, ~25–35%) — not because of HBM traffic anymore, but because of suboptimal work partitioning across thread blocks and warps, leaving GPU resources idle. FlashAttention-2 (Dao — arXiv:2307.08691, ICLR 2024) is a scheduling rewrite that keeps the same math and IO story but feeds the hardware better.⁴ Three changes:

1. Fewer non-matmul FLOPs. Non-matmul ops (the softmax rescales, exponentials) run on the general-purpose units, whose throughput can be ~16× lower than the tensor cores. FA-2 restructures the algorithm to minimize them — notably, it defers the per-block rescaling of the output accumulator, keeping $\mathbf{o}$ unnormalized through the loop and dividing by $\ell$ once at the very end instead of rescaling on every block. Same result, far fewer non-matmul operations, so more of the time is spent in tensor-core matmuls.
2. Parallelize over sequence length. FA-1 parallelized over batch × heads; when sequences are long (and batch is therefore small), that’s too few thread blocks to fill the GPU. FA-2 also parallelizes across the sequence-length (query-block) dimension, raising occupancy for the long-context regime that matters most.
3. Better warp scheduling within a block. FA-1 split the K/V dimension across warps (“split-K”), forcing warps to communicate partial results through SMEM. FA-2 instead splits the Q dimension across warps so each warp owns a slice of the output independently — eliminating inter-warp SMEM communication and syncs.

Result: ~2× over FlashAttention-1, reaching 50–73% of A100 peak FLOP/s and ~72% MFU in end-to-end GPT training — i.e. attention finally approaches GEMM-level efficiency.⁴ Note that none of this changed the IO complexity; FA-2 is entirely about mapping the same tiled algorithm onto warps/blocks so the hardware isn’t idle. (This is the lesson-05 GEMM lesson applied to attention: the algorithm was right, the occupancy and warp layout were the bottleneck.)

---

7. FlashAttention-3: Hopper asynchrony and FP8

FlashAttention-2 was designed for a synchronous execution model and left ~65% of an H100 on the table (only ~35% utilization on Hopper). Hopper (H100) added hardware that FA-2’s model doesn’t exploit: the Tensor Memory Accelerator (TMA) for asynchronous bulk HBM↔SMEM copies, asynchronous warpgroup matmul (wgmma) instructions, and FP8 tensor cores. FlashAttention-3 (Shah, Bikshandi, Zhang, Thakkar, Ramani, Dao — arXiv:2407.08608, NeurIPS 2024) is a Hopper-native redesign around asynchrony and low precision.⁵ Three techniques:

1. Warp-specialization / producer–consumer async. Split warps into producers (issue TMA loads of the next K/V tile into SMEM) and consumers (run wgmma + softmax on the current tile), coordinated by async barriers. This overlaps data movement with computation — the TMA fetches the next block while tensor cores chew on the current one, so neither unit stalls.
2. Interleaved (ping-pong) matmul and softmax. The $Q{K}^{\top }$ matmul (tensor cores) and the softmax (exp on the multifunction/special-function units) use different hardware units, so FA-3 schedules them to overlap: while one warpgroup does softmax on block $t$, another runs the matmul for block $t+1$. This hides the non-matmul softmax latency behind matmul throughput — attacking exactly the “non-matmul is slow” problem FA-2 identified, but by overlap rather than reduction.
3. FP8 with block quantization + incoherent processing. To use Hopper’s 2× FP8 throughput, FA-3 quantizes in blocks and applies an incoherent-processing (random orthogonal/Hadamard rotation) trick to spread outliers and cut quantization error, addressing FP8’s tiny mantissa.

Reported performance: 1.5–2.0× over FA-2 on H100, with FP16 reaching up to ~740 TFLOP/s (~75% utilization) and FP8 approaching ~1.2 PFLOP/s (the camera-ready version reports up to ~840 TFLOP/s BF16, ~85% utilization). FA-3 also improves accuracy vs a naive FP8 attention baseline thanks to the incoherent processing.⁵

Each generation targets a different bottleneck — read them as a sequence

FA-1 killed HBM traffic (IO). FA-2 killed idle warps/blocks (occupancy & non-matmul FLOPs). FA-3 killed synchronous stalls (async overlap) and unlocked FP8. The math — tiling + online softmax — is unchanged since 2022. What changes is which hardware resource is the binding constraint, generation to generation. That’s the co-design loop in action.

---

8. FlashAttention-4: Blackwell co-design (current SOTA, Sept 2026)

As of September 2026 the newest generation is FlashAttention-4 (Zadouri et al. — arXiv:2603.05451, published March 2026; first shown at Hot Chips, August 2025), purpose-built for NVIDIA Blackwell (B200/GB200).⁶ Its motivating observation is asymmetric hardware scaling: from Hopper to Blackwell, BF16 tensor-core throughput jumps ~2.25× (to ~2.25 PFLOP/s), but the special-function units (SFUs) that compute exp and the shared-memory bandwidth did not scale.⁶ So on Blackwell the tensor cores are no longer the bottleneck — the new limits are (a) the SFU exp throughput in the forward pass and (b) SMEM bandwidth in the backward pass. FA-4’s techniques follow directly:

• Software-emulated exponentials. Instead of routing exp through the scarce SFUs (MUFU.EX2), FA-4 computes it via a polynomial approximation on the abundant FMA units — moving the softmax bottleneck onto general-purpose compute Blackwell has in surplus.
• Conditional online-softmax rescaling. The rescale-on-every-max-change from §5.1 is skipped unless the max shift is large enough to threaten numerical stability, cutting rescaling operations ~10×.⁶ (A direct optimization of the recurrence you derived above.)
• Tensor Memory (TMEM) for the backward pass. Blackwell’s per-SM 256 KB TMEM (wired into the tensor cores) holds backward intermediates, relieving SMEM-bandwidth pressure; combined with the new 2-CTA MMA mode to further cut SMEM traffic and halve atomic reductions.
• CuTe-DSL in Python. The kernel is written in a Python-embedded CuTe DSL, giving ~20–30× faster compile times than C++ templates while keeping full control.

Reported: up to ~1.6 PFLOP/s BF16 (~71% utilization) on B200, 1.3× over cuDNN 9.13 and 2.7× over the Triton backend for the forward pass.⁶ FA-4 also backs PyTorch’s FlexAttention as a JIT target on Hopper/Blackwell, so custom variants (ALiBi, sliding-window, doc-masking, soft-capping) compile into FA-4 kernels. Caveat: FA-4 requires Hopper or Blackwell; older GPUs stay on FA-2/FA-3.

SOTA moves fast — verify before quoting

The FlashAttention line has shipped a new generation roughly yearly (2022 → 2023 → 2024 → 2026), each retargeting the then-current GPU’s binding constraint. The exact TFLOP/s numbers above are the papers’ reported figures and depend on head dim, causal masking, and precision; treat them as ballparks and re-check the paper/repo for your shape. The ideas (IO-awareness, online softmax, async overlap, precision) are the durable takeaways.

---

9. The broader lesson: algorithm–hardware co-design

Step back. Across four generations, FlashAttention never reduced the FLOP count of attention — the forward pass still does $\Theta (N^{2}d)$ multiply-adds, and the backward adds FLOPs via recomputation. Every speedup came from respecting the memory hierarchy:

• FA-1: don’t materialize $N^2$ in HBM (IO).
• FA-2: map the work onto warps/blocks so the tensor cores stay busy (occupancy).
• FA-3: overlap async copy/compute and exploit FP8 (Hopper).
• FA-4: move exp off the SFUs and intermediates into TMEM (Blackwell).

This is algorithm–hardware co-design: the algorithm is (re)shaped to the specific costs and capabilities of the memory/compute hierarchy it runs on. It is the same principle as tiled GEMM (lesson 05) and elementwise fusion (§3) — do more work per byte pulled from HBM, and match the computation’s structure to the units that execute it. The deepest lesson of this entire GPU tier: on the far side of the memory wall, performance engineering is data-movement engineering. FLOPs are cheap; bytes and the units that produce them are what you optimize around.

The transferable heuristic

When a kernel is slow, don’t ask “how do I do fewer FLOPs?” Ask “what’s the binding hardware resource — HBM bandwidth, SMEM bandwidth, SFU throughput, occupancy, or launch overhead? — and how do I restructure the algorithm to stop starving it?” FlashAttention is the textbook worked example of answering that question four times in a row.

---

10. Practice problems

Practice problem 1: HBM traffic — naive vs. FlashAttention

Problem: One attention head, sequence length $N=8192$, head dim $d=128$, BF16 ($b=2$ bytes). Ignore causal masking. (a) Estimate the HBM traffic of naive attention, counting the score/probability round-trips. (b) Estimate FlashAttention’s ideal HBM traffic (Q,K,V read once, O written once). (c) Give the ratio, and state how it scales with $N$.

Worked solution:
(a) Naive. The $N\times N$ matrix is written after $Q{K}^{\top }$, read by softmax, written as $P$, and read by the $PV$ matmul — 4 round-trips of an $N^2$ matrix:
${\text{traffic}}_{N^2}\approx 4N^{2}b=4\cdot 8192^2\cdot 2=4\cdot 6.71\times 10^7\cdot 2\approx 5.4\times 10^8\text{bytes}\approx 512\text{MiB}.$
The input/output traffic ($Q,K,V,O$) is $4Ndb=4\cdot 8192\cdot 128\cdot 2\approx 8.4\times 10^6$ bytes — negligible (~1.6%) beside the $N^2$ term. So naive $\approx 5.4\times 10^8$ bytes, dominated by the score matrix.
(b) FlashAttention (ideal). No $N^2$ traffic. Read $Q,K,V$ once and write $O$ once:
${\text{traffic}}_{\text{flash}}\approx 4Ndb=4\cdot 8192\cdot 128\cdot 2\approx 8.4\times 10^6\text{bytes}\approx 8\text{MiB}.$
(In reality there are bounded K/V re-reads from tiling, giving the paper’s $\Theta (N^2{d}^2/M)$; still far below the naive $N^2$. The ideal number captures the asymptotics.)
(c) Ratio. $\frac{5.4\times 10^8}{8.4\times 10^6}\approx 64\times$ less HBM traffic here. More importantly the scaling differs: naive is $\Theta (N^2)$, flash is $\Theta (Nd)$ — linear in $N$. So the advantage grows with sequence length: double $N$ and naive’s traffic 4בs while flash’s only 2בs. This is why FlashAttention is what made long context feasible — and note the naive version’s 512 MiB is also 512 MiB of HBM you must hold, the capacity wall that capped context length before FA.

Practice problem 2: Work an online-softmax update by hand

Problem: A single query attends to 4 keys, processed in two blocks of two. The raw scores (post-$1/\sqrt{d}$) are $B_1=\{s_1=1,s_2=3\}$ and $B_2=\{s_3=4,s_4=2\}$. The value vectors are scalars for simplicity: $v_1=1,v_2=2,v_3=3,v_4=4$. Run the online-softmax recurrence block by block, tracking $(m,\ell ,\mathbf{o})$, and verify the final output equals the direct softmax-weighted sum.

Worked solution:
Init: $m^{(0)}=-\infty ,{\ell }^{(0)}=0,{\mathbf{o}}^{(0)}=0$.

Block 1 ($s_1=1,s_2=3$): block max ${\tilde{m}}^{(1)}=3$; $m^{(1)}=\max (-\infty ,3)=3$; ${\alpha }^{(1)}=e^{-\infty -3}=0$.
${\ell }^{(1)}=0\cdot 0+(e^{1-3}+e^{3-3})=e^{-2}+e^0=0.1353+1=1.1353$.
${\mathbf{o}}^{(1)}=0+(e^{-2}\cdot 1+e^0\cdot 2)=0.1353+2=2.1353$.

Block 2 ($s_3=4,s_4=2$): block max ${\tilde{m}}^{(2)}=4$; $m^{(2)}=\max (3,4)=4$; ${\alpha }^{(2)}=e^{3-4}=e^{-1}=0.3679$.
Rescale old state, then add new terms $e^{4-4}=1$ (for $s_3$) and $e^{2-4}=e^{-2}=0.1353$ (for $s_4$):
${\ell }^{(2)}=0.3679\cdot 1.1353+(1+0.1353)=0.4177+1.1353=1.5530$.
${\mathbf{o}}^{(2)}=0.3679\cdot 2.1353+(1\cdot 3+0.1353\cdot 4)=0.7856+(3+0.5413)=0.7856+3.5413=4.3269$.

Normalize: $O={\mathbf{o}}^{(2)}/{\ell }^{(2)}=4.3269/1.5530=2.786$.

Direct check (global softmax, max $=4$): weights $\propto (e^{1-4},e^{3-4},e^{4-4},e^{2-4})=(0.0498,0.3679,1,0.1353)$, sum $=1.5530$ ✓ (matches ${\ell }^{(2)}$). Weighted sum of values $=(0.0498\cdot 1+0.3679\cdot 2+1\cdot 3+0.1353\cdot 4)=0.0498+0.7358+3+0.5413=4.3269$ ✓ (matches ${\mathbf{o}}^{(2)}$). Output $=4.3269/1.5530=2.786$ ✓. The streaming computation reproduces the exact softmax — note how block 2’s larger max ($4>3$) triggered the ${\alpha }^{(2)}=e^{-1}$ rescale that corrected block 1’s contribution.

---

11. Practice & resources

Hands-on

1. Implement online softmax in NumPy. Write a function that computes softmax-weighted attention for one row by streaming over key/value blocks, maintaining $(m,\ell ,\mathbf{o})$ and applying the rescale. Assert it matches scipy.special.softmax @ V to floating-point tolerance for random inputs and for adversarial cases (one huge score, all-equal scores). This cements the recurrence far better than reading it.
2. Work the Triton FlashAttention tutorial. The official Triton tutorial implements a real fused FlashAttention forward (and backward) kernel. Read it after you can derive the online softmax by hand — the tutorial’s m_i, l_i, acc are exactly your $(m,\ell ,\mathbf{o})$. Modify the block sizes and profile. (Deep-dived in 07-triton-and-modern-kernels.)
3. Profile attention memory, naive vs. Flash. In PyTorch, run attention two ways at $N=4096,8192,16384$: (a) explicit S = Q@K.T; P = S.softmax(-1); O = P@V, and (b) F.scaled_dot_product_attention (FlashAttention backend). Compare peak memory (torch.cuda.max_memory_allocated) and wall-clock. You should see naive’s memory grow $\propto N^2$ and OOM first, while SDPA stays roughly linear. Then nsys/ncu the two and confirm naive is dominated by the softmax/elementwise kernels round-tripping HBM.
4. Measure the HBM traffic gap. In Nsight Compute, read dram__bytes.sum for the naive attention’s softmax kernel and for the fused SDPA kernel at fixed $N$. Confirm the ratio tracks your Problem-1 estimate.

Real resources

• FlashAttention (primary): Dao, Fu, Ermon, Rudra, Ré, “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness” — arXiv:2205.14135 (NeurIPS 2022). The IO-complexity theorem and the tiling+online-softmax algorithm. Read §3 (algorithm) and the IO-complexity analysis. arxiv.org/abs/2205.14135
• FlashAttention-2: Dao, “FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning” — arXiv:2307.08691 (ICLR 2024). Work partitioning, seq-length parallelism, warp scheduling. arxiv.org/abs/2307.08691
• FlashAttention-3: Shah, Bikshandi, Zhang, Thakkar, Ramani, Dao, “FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision” — arXiv:2407.08608 (NeurIPS 2024). Hopper warp-specialization, TMA/wgmma overlap, FP8. arxiv.org/abs/2407.08608
• FlashAttention-4: Zadouri et al. (Tri Dao co-author), “FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling” — arXiv:2603.05451 (2026). Blackwell: software exp, conditional rescaling, TMEM, 2-CTA MMA. Companion blog: tridao.me/blog/2026/flash4; reverse-engineering write-up: modal.com/blog/reverse-engineer-flash-attention-4.
• Code / repo: github.com/Dao-AILab/flash-attention — reference CUDA/CuTe implementations for all generations.
• Triton FlashAttention tutorial: triton-lang.org → Tutorials → “Fused Attention” — a readable, hackable fused kernel; the fastest way to see the recurrence in real code.
• GPU MODE lecture on FlashAttention: the GPU MODE (formerly CUDA MODE) lecture series has a dedicated FlashAttention session walking the algorithm and the Triton kernel. github.com/gpu-mode/lectures
• The mental model (must-read): Horace He, “Making Deep Learning Go Brrrr From First Principles” — compute vs. bandwidth vs. overhead, and why fusion wins. The prose companion to this whole lesson. horace.io/brrr_intro.html
• Online softmax (background): Milakov & Gimelshein, “Online normalizer calculation for softmax” — arXiv:1805.02867. The streaming-softmax trick that FlashAttention builds on, predating it by four years.

---

12. What’s next

Two threads, in order:

1. 07-triton-and-modern-kernels (next) — write the fused kernels yourself. Triton is the language FlashAttention’s tutorial and much of the modern kernel ecosystem is written in; you’ll implement a fused attention forward and see torch.compile/Inductor generate fusions automatically. The natural continuation of §3 and the Triton hands-on above.
2. 05-optimizing-gemm (back) — reread the SMEM tiling and warp-partitioning material now that you’ve seen FA-2 apply exactly that thinking to attention. FlashAttention is “tiled GEMM with an online-softmax epilogue fused in the middle”; the two lessons are the same discipline.

Further threads to pull on: paged attention / KV-cache management (the inference-time memory problem, a cousin of this one), linear-attention & SSMs (changing the algorithm to escape $O(N^2)$ entirely rather than just its IO), and FP8/FP4 attention numerics (the precision frontier FA-3/FA-4 opened).

---

Reference topic: gpu-systems-for-llms | Concepts: memory-wall, arithmetic-intensity, gpu-memory-hierarchy | Filed: 2026-09-02

---

References

Footnotes

1. NVIDIA H100 Tensor Core GPU — 80 GB HBM3, 989 TFLOP/s FP16/BF16 tensor-core (dense), ~3.9 TFLOP/s special-function (exp) throughput; BF16 arithmetic-intensity ridge ≈ 296 FLOP/byte (≈ 989 TFLOP/s ÷ ~3.35 TB/s HBM3). NVIDIA, “NVIDIA H100 Tensor Core GPU Architecture” whitepaper; SFU vs tensor-core throughput figures also stated in the FA-3 paper (arXiv:2407.08608). ↩ ↩²

2. [established] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré, “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness” — arXiv:2205.14135 (NeurIPS 2022). arxiv.org/abs/2205.14135. Source of the $\Theta (N^2{d}^2/M)$ IO-complexity bound, the ~9× fewer HBM accesses (GPT-2) and 2–4× wall-clock speedup, and the recomputation/logsumexp backward pass. ↩ ↩² ↩³ ↩⁴

3. [established] Maxim Milakov, Natalia Gimelshein, “Online normalizer calculation for softmax” — arXiv:1805.02867 (2018). arxiv.org/abs/1805.02867. The single-pass streaming-softmax recurrence (running max + running denominator with rescale) that FlashAttention adapts. ↩

4. [established] Tri Dao, “FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning” — arXiv:2307.08691 (ICLR 2024). arxiv.org/abs/2307.08691. Source of the ~2× over FA-1, 50–73% of A100 peak, ~72% MFU figures, and the FA-1 ~30–50% (fwd) / 25–35% (bwd) A100 baselines. ↩ ↩²

5. [recent] Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao, “FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision” — arXiv:2407.08608 (NeurIPS 2024). arxiv.org/abs/2407.08608. Source of the 1.5–2.0× over FA-2, ~740 TFLOP/s FP16 (~75% util) and ~1.2 PFLOP/s FP8 (arXiv v1/v2); the camera-ready ~840 TFLOP/s BF16 (~85%) figures appear in the revised abstract (OpenReview / AI-at-Meta). FA-2’s ~35% H100 utilization is stated here. ↩ ↩²

6. [recent] Ted Zadouri, Markus Hoehnerbach, Jay Shah, Timmy Liu, Vijay Thakkar, Tri Dao, “FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling” — arXiv:2603.05451 (March 2026). arxiv.org/abs/2603.05451. Companion blog: tridao.me/blog/2026/flash4; reverse-engineering write-up: modal.com/blog/reverse-engineer-flash-attention-4. Source of the asymmetric-scaling motivation (B200 ~2.25 PFLOP/s BF16 tensor cores; SFU/SMEM did not scale), software-emulated exp, conditional rescaling (~10× fewer rescales), TMEM / 2-CTA MMA, CuTe-DSL compile-time claim, and the ~1.6 PFLOP/s BF16 (~71% util) / 1.3× cuDNN 9.13 / 2.7× Triton results. Note: treat the exact TFLOP/s figures as shape/precision-dependent — verify against the primary source when quoting. ↩ ↩² ↩³ ↩⁴