Learn: The Quadratic Bottleneck & IO-Aware Attention

What you're learning

The one idea that made long context tractable: naive attention is memory-bound, not compute-bound — it wastes its time moving the $s\times s$ score matrix through HBM, not computing it. The fix isn’t fewer FLOPs, it’s fewer bytes. By the end you’ll be able to derive the online (streaming) softmax recurrence — the mathematical heart of FlashAttention — and explain why it drops HBM traffic from $O(s^2)$ to $O(s)$ while computing exactly the same output.

First lesson of Tier 2. Assumes Tier 1: 01-scaled-dot-product-attention (SDPA, the $O(s^2)$ wall, FLOP counting), 02-multi-head-attention-and-the-block, and 03-positional-information. Vocabulary follows the attention-glossary.

This is the algorithm angle, not the kernel angle

This lesson is about why attention is memory-bound and the math that fixes it. The hardware/kernel treatment — warp specialization, TMA/wgmma overlap, FA-2/3/4 generations, roofline analysis — lives in 06-memory-wall-and-flashattention. We deliberately don’t repeat it here; cross-links point you there for the systems details.

---

1. Learning map

graph TD
    A["Recap: attention core is<br/>O(s²·d) FLOPs (lesson 01)"] --> B["But wall-clock ≠ FLOPs:<br/>it materializes s×s S,P in HBM"]
    B --> C["Arithmetic intensity ~ d,<br/>far below the ridge point<br/>⇒ MEMORY-BOUND"]
    C --> D["Realization: fix bytes moved,<br/>not FLOPs done (IO-awareness)"]
    D --> E["Obstacle: softmax needs a<br/>GLOBAL row normalizer"]
    E --> F["Online softmax:<br/>running m, ℓ, o + rescale α"]
    F --> G["Compute row output<br/>block-by-block, never form s×s"]
    G --> H["HBM traffic: O(s²) → O(s)<br/>same FLOPs, long context unlocked"]
    H --> I["Kernel/hardware details:<br/>see GPU lesson 06"]
    style B fill:#a44,color:#fff
    style C fill:#a44,color:#fff
    style F fill:#4a4,color:#fff
    style H fill:#44a,color:#fff

---

2. Intuition: attention is memory-bound, not compute-bound

Lesson 01 left you with the headline: the attention core costs $\approx 4b{s}^{2}d$ FLOPs — quadratic in sequence length $s$. Double the context, quadruple the compute. That’s the $O(s^2)$ wall, and it’s real. But it’s only half the story, and not even the binding half on modern hardware.

Here’s the realization. Count what the naive implementation actually does to memory when it runs as three library calls (S = Q@K.T/√d, P = softmax(S), O = P@V):

1. $S=Q{K}^{\top }/\sqrt{d}$ — a matmul that reads $Q,K$ ($O(sd)$) and writes the full $s\times s$ score matrix $S$ to HBM.
2. $P=\mathrm{softmax}(S)$ — an elementwise+reduction kernel that reads $S$ back ($s^2$), does a handful of FLOPs per element, and writes $P$ back ($s^2$).
3. $O=PV$ — a matmul that reads $P$ ($s^2$) and $V$, writes $O$ ($O(sd)$).

The $s\times s$ matrix is round-tripped through HBM roughly four times (write $S$, read $S$, write $P$, read $P$) — $\Theta (s^2)$ bytes moved — while carrying only $\Theta (s^{2}d)$ FLOPs split across separate kernels. The softmax in the middle is the smoking gun: it’s pure elementwise work with arithmetic intensity $\approx O(1)$ FLOP/byte, so it runs at HBM bandwidth and does nothing but shuttle the score matrix. Its entire runtime is “time to read $S$ + time to write $P$.”

So the attention block, as launched, is bottlenecked by moving the score matrix, not by the tensor cores computing it. Its effective arithmetic intensity collapses toward $\sim d$ (tens to ~128) — far below the hardware’s ridge point (~296 FLOP/byte on an H100).¹ Attention has plenty of FLOPs; it just can’t feed them because it’s starved on bandwidth. That’s what “memory-bound” means.²

The one-sentence version

Naive attention is slow not because it does too much math, but because it writes an $s\times s$ matrix to slow memory only to immediately read it back. The win is to never materialize that matrix in HBM — same FLOPs, a fraction of the bytes.

There’s a second, equally fatal consequence of materializing $S$: capacity. Holding an $s\times s$ matrix costs $O(s^2)$ HBM storage. At $s=8\text{k}$ that’s ~128 MiB per head in BF16; at $s=128\text{k}$ it’s ~32 GiB² — larger than an entire H100’s 80 GB if you keep a few around. That storage wall is, quite literally, what capped context length before 2022.

The fix is IO-awareness: design the algorithm around minimizing HBM reads/writes, the way tiled GEMM is designed around SRAM reuse. Compute attention block-by-block, keeping each block’s scores on-chip, producing, consuming, and discarding them before they ever reach HBM. The obstacle is softmax: its normalizer ${\sum }_j{e}^{s_j}$ is a sum over the entire row, which seems to demand seeing the whole row at once. Section 3 shows why it doesn’t.

Quiz (retrieval): Attention genuinely does $\Theta (s^{2}d)$ FLOPs — 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 (s^{2}d)$ FLOPs are in the two matmuls, but running the op as three separate kernels forces the $s\times s$ score/probability matrix through HBM four times (write $S$, read $S$, write $P$, read $P$) — that’s $\Theta (s^2)$ traffic with only a handful of FLOPs attached per byte. The effective arithmetic intensity of the block collapses toward $\sim d$, far below the ridge point, so wall-clock is set by HBM bandwidth, not tensor-core throughput. The softmax is the culprit: it’s an elementwise/reduction kernel ($I\approx O(1)$) that both sits on the memory roof and forces the $s^2$ matrix to be materialized so a later kernel can read it back. (Note lower head dim $d$ makes it worse — intensity $\sim d$ shrinks.)

---

3. The math: online (streaming) softmax

We want to compute, for a single query row with scores $s_1,\dots ,s_s$ (here $s_j=q^{\top }{k}_j/\sqrt{d}$) and value vectors $v_j\in {\mathbb{R}}^d$, the softmax-weighted output

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

where subtracting the row max $m$ is the standard safe-softmax trick to avoid overflow of $e^{s_j}$ (every exponent becomes $\le 0$). The obstacle for streaming: both $m$ and the denominator $\ell ={\sum }_l{e}^{s_l-m}$ depend on the whole row, but we want to sweep over blocks $B_1,B_2,\dots ,B_T$ of the key/value sequence, seeing only one block at a time and keeping $S$ off HBM.

The trick (Milakov & Gimelshein, 2018)³ is to maintain a small running state and correct it whenever a later block reveals a larger max. Track three quantities after processing blocks $\mathrm{1..}t$:

• $m^{(t)}$ — running max of all scores seen so far;
• ${\ell }^{(t)}$ — running denominator, referred 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$. 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 were built with exponentials referred to the old max $m^{(t-1)}$. When the max jumps to $m^{(t)}$, every previously stored term $e^{s_j-m^{(t-1)}}$ must be re-referenced to the new max. 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 running sum by the scalar ${\alpha }^{(t)}=e^{m^{(t-1)}-m^{(t)}}$ applies this correction to every stored term at once — that’s why a single multiply fixes the entire accumulator. If the new block doesn’t raise the max ($m^{(t)}=m^{(t-1)}$), then ${\alpha }^{(t)}=e^0=1$ and nothing changes. And because $m^{(t)}\ge m^{(t-1)}$, the exponent is $\le 0$, so ${\alpha }^{(t)}\in (0,1]$ — a shrink factor, never an overflow.

The invariant that makes it exact

At every step $t$ the recurrence maintains 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; the inductive step is exactly the rescale algebra above), after the final block ${\mathbf{o}}^{(T)}/{\ell }^{(T)}$ equals full softmax attention bit-for-bit up to floating-point associativity. FlashAttention is exact, not an approximation — you can drop it into any model with zero accuracy change. That is the crucial selling point that separates it from the sparse/linear-attention approximations.

Because $\ell$ and $\mathbf{o}$ are just a running scalar and a running $d$-vector, each block’s scores $s_j$ can be thrown away the instant they update the state — they never need to be stored. That is the mechanical reason the $s\times s$ matrix never reaches HBM. The state is $O(d)$ per query row, independent of $s$.

Quiz (derive it): Justify the rescale factor ${\alpha }^{(t)}=e^{m^{(t-1)}-m^{(t)}}$ from first principles. What breaks if you omit it, and why can ${\alpha }^{(t)}$ never overflow?

Answer

Justification. After block $t-1$ the accumulators store exponentials referred to the then-current max: ${\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: $e^{s_j-m^{(t-1)}}\cdot e^{m^{(t-1)}-m^{(t)}}=e^{s_j-m^{(t)}}$. Multiplying the whole accumulator by ${\alpha }^{(t)}=e^{m^{(t-1)}-m^{(t)}}$ does this 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 at the new reference) are added.
If omitted: old terms stay scaled by $e^{s_j-m^{(t-1)}}$ while new terms use $e^{s_j-m^{(t)}}$ — you’d sum exponentials with inconsistent reference points, so numerator and denominator no longer correspond and $\mathbf{o}/\ell$ is wrong (unless the max happened never to change).
No overflow: $m^{(t)}\ge m^{(t-1)}\Rightarrow m^{(t-1)}-m^{(t)}\le 0\Rightarrow {\alpha }^{(t)}\in (0,1]$. Every $e^{s_j-m^{(t)}}\le 1$ too, since $s_j\le m^{(t)}$. Safe-softmax stability is preserved, now incrementally.

Quiz (retrieval): Why does maintaining $(m,\ell ,\mathbf{o})$ let you avoid ever writing the $s\times s$ score matrix to HBM? How large is the running state per query row?

Answer

Softmax’s only global dependency on the row is through the max $m$ and the denominator $\ell$. The online recurrence carries both as running scalars, plus a running unnormalized output vector $\mathbf{o}$. Once a block’s scores $s_j$ have been folded into $(m,\ell ,\mathbf{o})$, they carry no further information and can be discarded immediately — so a block of scores is produced, consumed, and thrown away on-chip, never written out. The state is one scalar $m$, one scalar $\ell$, and one $d$-vector $\mathbf{o}$: $O(d)$ per query row, independent of sequence length $s$. That’s the whole reason the $s\times s$ matrix never needs to exist in HBM.

---

4. Illustration: tiled block-by-block attention

For each query block $Q_i$, stream over the key/value blocks $B_1,\dots ,B_T$, keeping only the $O(d)$ running state on-chip. The $s\times s$ scores are never assembled.

flowchart LR
    subgraph HBM["HBM (slow, large) — only linear-in-s traffic crosses here"]
        Qi["Q_i block"]
        KV["K_1..K_T, V_1..V_T"]
        Oi["O_i (written once)"]
    end
    subgraph SRAM["SRAM / on-chip — loop t = 1..T"]
        S1["1. score tile<br/>S_it = Q_i K_tᵀ/√d<br/>(stays on-chip, never to HBM)"]
        S2["2. update state<br/>m ← max(m, m̃_t)<br/>α = e^(m_old−m)<br/>ℓ ← αℓ + Σ e^(S−m)<br/>o ← αo + Σ e^(S−m) V_t"]
        S3["3. after last block<br/>O_i = o / ℓ"]
    end
    Qi --> S1
    KV --> S1
    S1 --> S2
    S2 -->|"next K/V block"| S1
    S2 --> S3
    S3 --> Oi

Figure 1 — For each query block, 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 on-chip. Only $Q,K,V$ (in) and $O$ (out) cross HBM. The $s\times s$ matrix is never materialized.

Before vs. after, in HBM traffic:

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

The tiling reorders the softmax sum — associativity, not identity

“Bit-for-bit exact” in the invariant means up to floating-point associativity. Streaming accumulates the exponential sums in a different order than a single-pass softmax, so results can differ in the last bit or two. This is the ordinary non-associativity of float addition (the same caveat as any tiled reduction), not an approximation of attention — accuracy is unaffected in practice.

---

5. FLOPs & memory: same math, a fraction of the bytes

The crucial accounting, side by side. Let $b$ = batch, $s$ = sequence length, $d$ = head dim.

	Naive (3 kernels)	IO-aware / tiled
Attention-core FLOPs	$\approx 4b{s}^{2}d$	$\approx 4b{s}^{2}d$ (same — forward)
HBM traffic	$\Theta (s^2)$ (four $s\times s$ round-trips)	$\Theta (sd)$ (Q,K,V in, O out)
Peak HBM footprint of $S,P$	$\Theta (s^2)$	$\Theta (d)$ running state / row
Bound by	HBM bandwidth (memory-bound)	near tensor-core peak (compute-bound)

The headline: the FLOP count does not change. Online softmax is not a cheaper algorithm in arithmetic — it’s the same $\Theta (s^{2}d)$ multiply-adds. What changes is the bytes moved: HBM traffic drops from quadratic to linear in $s$, and the $s\times s$ storage footprint vanishes. Because the op was memory-bound, cutting the bytes is what cuts the wall-clock — the original paper measured ~9× fewer HBM accesses and 2–4× wall-clock speedup on GPT-2, computing identical outputs.²

Why this unlocked long context — the scaling, not the constant

Naive traffic is $\Theta (s^2)$; tiled traffic is $\Theta (sd)$. So the advantage grows with context: double $s$ and naive’s traffic 4בs while the tiled version’s only 2בs. Equally important, the $\Theta (s^2)$ storage of the score matrix — the capacity wall that once OOM’d long sequences — is gone. Same FLOPs, linear bytes: that combination is what turned 128k+ token context from impossible to routine.

"Fewer FLOPs" is the wrong instinct here

Faced with the $O(s^2)$ wall, the tempting move is to approximate attention to cut FLOPs (sparse, low-rank, linear attention). Those change the math and can cost accuracy. FlashAttention shows a different, often better lever: keep the exact math, cut the bytes. On a memory-bound op, bytes — not FLOPs — are what you’re actually paying for. (The backward pass even adds FLOPs via recomputation and still wins — see 06-memory-wall-and-flashattention §5.3.)

For the hardware mechanics that turn this IO complexity into realized TFLOP/s — SRAM tiling specifics, the exact $\Theta (s^2{d}^2/M)$ IO bound, warp specialization, and the FA-2/FA-3/FA-4 generations — see the systems lesson 06-memory-wall-and-flashattention. FlashAttention-2⁴ in particular re-derives the recurrence to defer per-block rescaling to a single final divide, cutting non-matmul work; that’s an optimization of the exact recurrence in §3.

Try it: watch the two curves diverge

Open the interactive Attention FLOPs & KV-cache calculator and set $d=128$, one head, $b=1$. Sweep $s$ from 2k → 8k → 32k → 128k. Watch two things: (1) prefill FLOPs grow ~quadratically — that cost is the same for naive and IO-aware attention; (2) the $s\times s$ score-matrix footprint (the thing naive attention writes to HBM) explode as $s^2$, while the IO-aware version holds only $O(d)$ state per row. The gap between “FLOPs you must do” and “bytes you must move” is exactly the memory-bound story of §2.

Quiz (retrieval): Online softmax computes the identical attention output. So where does the speedup come from, and by what factor does HBM traffic grow when you 4× the context?

Answer

The speedup is not from doing less arithmetic — the forward FLOP count is unchanged at $\approx 4b{s}^{2}d$. It comes entirely from moving fewer bytes: the $s\times s$ score/probability matrix is never written to or read from HBM, so mandatory traffic drops from $\Theta (s^2)$ to $\Theta (sd)$. Since the op was memory-bound, less traffic ⇒ less wall-clock (~9× fewer HBM accesses, 2–4× faster on GPT-2²). On 4×-ing the context: naive HBM traffic scales as $s^2$, so 16×; tiled traffic scales as $sd$, so 4×. The gap between them widens with $s$ — which is exactly why long context became feasible.

---

6. Primary source & where to go deeper

• FlashAttention (primary): Dao, Fu, Ermon, Rudra, Ré, “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness” — arXiv:2205.14135 (NeurIPS 2022).² Read §3 (the tiling + online-softmax algorithm) and the IO-complexity analysis. This is the paper that framed attention as an IO problem.
• Online softmax (the recurrence itself): Milakov & Gimelshein, “Online normalizer calculation for softmax” — arXiv:1805.02867 (2018).³ The single-pass running-max + running-denominator trick, four years before FlashAttention adapted it.
• Kernel/hardware angle (companion): 06-memory-wall-and-flashattention — roofline, SRAM tiling, recomputation backward, and FA-2/3/4.

Ask your teacher

Stuck on the induction that makes online softmax exact, the role of ${\alpha }^{(t)}$, or the FLOPs-vs-bytes accounting? Ask — I’m your teacher for this course and can go deeper on any piece, hand-run the recurrence on a new example, or build you a harder problem.

---

7. Practice

Problem 1: HBM traffic, naive vs. IO-aware

One head, $s=8192$, $d=128$, BF16 (2 bytes). Ignore causal masking. (a) Estimate naive attention’s HBM traffic from the score-matrix round-trips. (b) Estimate the IO-aware ideal (Q,K,V read once, O written once). (c) Give the ratio and state how it scales with $s$.

Solution.
(a) The $s\times s$ matrix is written ($S$), read (softmax), written ($P$), read ($PV$) — 4 round-trips: $4s^{2}b=4\cdot 8192^2\cdot 2\approx 5.4\times 10^8$ bytes $\approx 512$ MiB. Input/output traffic $4sdb=4\cdot 8192\cdot 128\cdot 2\approx 8.4\times 10^6$ bytes is ~1.6% — negligible beside the $s^2$ term.
(b) No $s^2$ traffic; read $Q,K,V$ once, write $O$ once: $\approx 4sdb\approx 8.4\times 10^6$ bytes $\approx 8$ MiB. (Real tiling adds bounded K/V re-reads, giving the paper’s $\Theta (s^2{d}^2/M)$; still far below $s^2$.)
(c) $\frac{5.4\times 10^8}{8.4\times 10^6}\approx 64\times$ less traffic here. Scaling: naive is $\Theta (s^2)$, IO-aware is $\Theta (sd)$ — so the ratio grows $\propto s/d$; the advantage widens with context. (And the naive 512 MiB is 512 MiB you must hold — the old capacity wall.)

Problem 2: hand-run the online-softmax recurrence

A query attends to 4 keys in two blocks of two. Post-scaled scores: $B_1=\{s_1=1,s_2=3\}$, $B_2=\{s_3=4,s_4=2\}$. Scalar values $v_1=1,v_2=2,v_3=3,v_4=4$. Run the recurrence tracking $(m,\ell ,\mathbf{o})$ and verify it equals the direct softmax-weighted sum.

Solution. Init: $m^{(0)}=-\infty ,{\ell }^{(0)}=0,{\mathbf{o}}^{(0)}=0$.
Block 1 ($\tilde{m}=3$): $m^{(1)}=3$, ${\alpha }^{(1)}=e^{-\infty }=0$. ${\ell }^{(1)}=e^{1-3}+e^0=0.1353+1=1.1353$. ${\mathbf{o}}^{(1)}=e^{-2}\cdot 1+e^0\cdot 2=2.1353$.
Block 2 ($\tilde{m}=4$): $m^{(2)}=4$, ${\alpha }^{(2)}=e^{3-4}=0.3679$. Rescale then add $e^{4-4}=1$ (for $s_3$) and $e^{2-4}=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.5413=4.3269$.
Normalize: $O=4.3269/1.5530=2.786$.
Direct check (global max 4): weights $\propto (e^{-3},e^{-1},e^0,e^{-2})=(0.0498,0.3679,1,0.1353)$, sum $=1.5530$ ✓. Weighted value sum $=0.0498+0.7358+3+0.5413=4.3269$ ✓. $O=2.786$ ✓. Block 2’s larger max triggered the ${\alpha }^{(2)}=e^{-1}$ rescale that corrected block 1’s contribution.

---

8. What’s next

1. 05-mqa-gqa-mla (next) — the other attention memory problem. FlashAttention fixed training/prefill HBM traffic; at autoregressive inference the binding constraint is the KV cache (linear in $s$ but read every step). MQA/GQA/MLA shrink it by sharing or compressing keys/values. Different bottleneck, same “move fewer bytes” philosophy.
2. 03-positional-information (back) — revisit how positions enter attention. Position schemes (especially RoPE) must be applied inside the fused kernel, block-by-block, since the scores never materialize — a concrete constraint the IO-aware algorithm imposes on everything upstream.

Further threads: linear attention & SSMs (change the algorithm to escape $O(s^2)$ entirely, trading exactness for asymptotics — the opposite bet to FlashAttention), and the kernel-level continuation in 06-memory-wall-and-flashattention.

---

Reference: attention-glossary | Topic: attention-mechanisms | Filed: 2026-09-10

---

References

Footnotes

1. NVIDIA H100 Tensor Core GPU — BF16 arithmetic-intensity ridge point ≈ 296 FLOP/byte (≈ 989 TFLOP/s tensor-core ÷ ~3.35 TB/s HBM3). See NVIDIA H100 architecture whitepaper; the ridge-point figure and its use in the attention roofline are detailed in 06-memory-wall-and-flashattention. ↩

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 IO-aware framing, tiling + online-softmax algorithm, the exactness result, the ~9× fewer HBM accesses and 2–4× wall-clock speedup on GPT-2, and the $s^2$ storage/capacity wall. ↩ ↩² ↩³ ↩⁴ ↩⁵

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. Restructures the recurrence to defer per-block output rescaling to a single final normalization, cutting non-matmul work; work partitioning / warp scheduling covered in the GPU-systems lesson. ↩