Quantized kernels

01 · The memory wall

Attention's intermediates grow as N²

Standard attention computes softmax(QKᵀ)V in three steps, and the middle product — the score matrix S, one attention weight per query-key pair — has N² entries per head. At N = 4096 context, 32 heads, FP16: 4096² × 32 × 2 B ≈ 1.1 GB of intermediates per layer, allocated, written, read, and freed for a matrix that is never the final answer. Push to 128k context and N² grows 1,024× — the wall that, circa 2022, made long context effectively impossible regardless of compute.

Day 23's fusion arithmetic applies at terrifying scale: those scores get written to DRAM and read back — multiple times — in the softmax. This is a kernel-shaped problem: the math is fine; the data flow materializes a tensor that never needed to exist.

02 · The insight

Softmax is row-wise, so it can stream

Day 21's softmax needed the whole row at once for one reason: the max. But watch what happens when a row arrives in blocks. Keep a running max m and running sum l. A new block arrives with a bigger max m'? Rescale what you have by exp(m - m') — a factor ≤ 1, numerically safe — and fold the new block in. The same trick extends to the output accumulation, since attention output is a weighted sum of V rows: keep a running weighted sum, rescale it in lockstep with l, divide once at the end. This is the online softmax: correct for any block order, never needing the full row, never materializing S.

The payoff is categorical: memory drops from O(N²) to O(N) — FlashAttention holds only the Q block, the current K/V block, and the accumulators. The 1.1 GB of intermediates becomes kilobytes of SRAM. Long context stopped being a memory problem the day this shipped.

03 · The structure

Day 15's tiling and day 21's softmax, composed

The kernel structure is the two ideas you already own, nested. Outer loop: one block of Q rows per program, loaded to SRAM and kept there. Inner loop: sweep blocks of K and V through SRAM; per block, compute scores against the resident Q, run the online-softmax update, accumulate. The update, in Triton-flavored pseudocode:

# state: acc (output accum), m (running max), l (running sum)
s = tl.dot(q, tl.trans(k)) * scale   # scores for this K-block
m_new = tl.maximum(m, tl.max(s, axis=1))
alpha = tl.exp(m - m_new)             # rescale what we hold
p = tl.exp(s - m_new)                 # this block's weights
l = l * alpha + tl.sum(p, axis=1)
acc = acc * alpha + tl.dot(p, v)      # …and the output sum
m = m_new

Every mechanism has a birthday in this journey: the SRAM staging is day 15's tiling with its barriers; the max-subtraction is day 21's stability trick promoted to a streaming recurrence; the masked-block accumulation logic descends from day 16's reductions; the tile-shape-versus-SRAM negotiation is day 08's budget; the intensity it achieves is why it sits right of naive attention on day 11's chart. Nothing here is exotic — it's composition under memory constraints, which is what kernel engineering is.

04 · Why nobody wrote it in 2017

Exactness, not approximation — and a mental shift

Two things made this late. First, the naive formulation had everyone convinced softmax needed the full row — the online rescaling looks fiddly until you notice the factors are ≤ 1 and everything is exact. FlashAttention's output is bit-for-bit the same computation, not an approximation; there's no accuracy footnote. Second, writing it demanded the mental shift this journey has been building: treat DRAM traffic, not FLOPs, as the thing being optimized. FlashAttention-2 and -3 refine the parallelism and hardware mapping; the founding idea is the one sentence you now own — if an intermediate is only read once downstream, don't let it reach DRAM.

Next: day 25 — quantized kernels: what actually changes when numbers shrink from FP32 to INT8 and FP8, beyond "smaller numbers."