FlashAttention, Playable

EP 01

The memory wall — your GPU is a chef waiting on deliveries

A modern GPU can do ~1,000 trillion multiply-adds per second, but can only fetch ~3 TB/s from its main memory (HBM). So the question for any operation isn't "how much math?" — it's how much math per byte fetched (arithmetic intensity). Drag the slider and watch the utilization gauge: heavy-math ops keep the chef cooking; naive attention keeps him standing at the loading dock.

✓ good case: big matmul, ~300 ops/byte — silicon at 90%+, the chef never waits✗ bad case: naive attention's softmax phase, ~4 ops/byte — compute idles at ~10%, you paid for a kitchen and ran a warehouse

arithmetic intensity 4 ops/byte —

The gauge implements the roofline model: achievable FLOPs = min(peak compute, intensity × bandwidth), with H100-class numbers (989 TFLOPs, 3.35 TB/s). The break-even intensity ≈ 295 ops/byte falls out of the division.

On modern accelerators, most "slow" code isn't compute-bound — it's memory-bound. The whole kernel-optimization discipline is moving work above the roofline's ridge. FlashAttention is this idea applied to one specific, spectacularly important operation.

EP 02

The N² tourist — a giant matrix that only exists in transit

Naive attention computes the N×N score matrix, writes all of it to slow memory, then immediately reads it back to apply softmax and multiply by V. It's a tourist: it visits HBM once and contributes nothing. Drag the sequence length and read the bill — at 32k tokens, that round-trip is gigabytes per attention head, and it grows with the square.

✓ good case: 512 tokens — the matrix detour costs 1 MB, who cares✗ bad case: 32k tokens — 2 GB written and re-read per head-batch, memory traffic dwarfs the actual math

sequence length 512 —

Traffic = 2 × N² × 2 bytes (write + read of fp16 scores) per head; the animation shows the round trip. FLOPs for the same work = 2·N²·d — the ratio between the bars is the whole problem.

The scores matrix is scaffolding, not product — nothing downstream needs it once the output is computed. Every byte of it that touches HBM is pure waste. The question FlashAttention asked: what if it never leaves the chip at all?

EP 03

Tiling — cook from the pantry, not the warehouse

Next to its huge slow memory, a GPU has a tiny fast one: SRAM — ~200 KB per compute block, but ~50× faster. The trick: chop Q, K, V into tiles that fit in the pantry, compute a block of attention entirely on-chip, fold it into the running output, and move on. The N² matrix still gets computed — it just never exists in one place. Pick a tile size: too big won't fit, too small drowns in loading overhead.

✓ good case: tiles sized to fill SRAM — HBM traffic drops ~20× vs naive✗ bad case: tiles bigger than SRAM — spills back to slow memory, congratulations you rebuilt naive attention

tile size 64×64 —

Traffic counters accumulate live during the animation: naive counts every score byte twice; tiled counts only Q/K/V/O block loads. The ~N²·d/M ratio (M = SRAM size) is measured off your run, not asserted.

Tiling is the oldest trick in high-performance computing — it's how matmul libraries have worked for decades. FlashAttention's contribution wasn't inventing tiles; it was noticing that attention could be tiled at all, despite one operation that seems to forbid it. That operation is the next episode.

EP 04

The softmax puzzle — how to average what you haven't seen

Here's why nobody tiled attention for five years: softmax needs the whole row — its max (for numerical stability) and its total sum — before any single weight is final. But tiles arrive one at a time. The fix is the online softmax: keep a running max and running sum, and every time a new tile brings a bigger max, rescale everything you've accumulated so far. Step through a row below and watch the books stay balanced — then run the naive streamer that skips the rescale and watch it silently produce garbage.

✓ good case: online softmax — final output matches the all-at-once computation to 15 decimal places✗ bad case: streaming without rescale — a later tile with a bigger max makes earlier work wrong by 40%+, with no error message

—

Real numbers: the row's scores are random each reset; both the tile-by-tile ledger (running max m, running sum ℓ, rescale factor) and the ground-truth full softmax are computed live and compared at the end.

This little accounting identity is the entire theorem. Once softmax can be computed in one streaming pass, attention becomes tileable, the N² tourist dies, and the memory wall falls. The best systems ideas often look like this: one algebraic observation that unlocks a decade-old hardware trick.

EP 05

Race them yourself — a real benchmark in your tab

No simulation this time. The button below runs two real implementations of attention in JavaScript, right here: a naive one that genuinely allocates the full N×N score matrix, and a streaming-tiled one with the online softmax from EP 04 that never holds more than one tile. Wall-clock times from your own CPU, memory allocated counted exactly, outputs cross-checked for equality.

✓ good case: tiled at N=2048 — same answer, a fraction of the memory, faster past the cache cliff✗ bad case: naive at N=2048 — 4.2M floats (34 MB) allocated for a matrix used exactly once

—

Honesty label: your browser has no HBM/SRAM split, so speedups here come from cache behavior and allocation cost — smaller than on a GPU (where FlashAttention gains 2-10×), but the memory asymptotics (O(N²) vs O(N)) are identical and measured exactly. Outputs are verified equal to 1e-10 every run.

What this trick actually bought the world: context windows. Quadratic memory was the wall between 2k-token models and today's 100k-1M-token ones — FlashAttention (and its descendants v2, v3, plus paged attention for serving) turned that wall into a line item. When you paste a whole codebase into a chat, this kernel is why it doesn't crash.