Flagship module · quadratic compute hid an IO problem · 7 beats

FlashAttention

Attention was called quadratic in compute for years while many standard GPU attention workloads were in fact limited by memory traffic despite that quadratic arithmetic cost. This module derives the rescaling identity that makes a tiled softmax exact, quantifies the bytes both implementations move, and draws the line — sharply — between what this kernel fixes and what it leaves entirely alone.

Core path ~32 min: 01 → 04 then 07 · Kernels deep dive +24 min: 05 → 06

Beat 01 / 07CONCEPT~5 min

What problem existed?

After this beat you can: You can explain why attention was slow for reasons its FLOP count does not reveal.

Standard attention is four operations: score, mask, softmax, weight. Written the obvious way, each one reads a matrix from HBM and writes another back.

The matrix is N × N. For a 16K sequence, 32 heads, and batch 8, that is 256 GiB of scores and probabilities per layer — moved several times over.

For years the field described attention as "quadratic in compute" and optimized accordingly: sparsity patterns, low-rank approximations, kernel tricks.

Many of those methods cut FLOPs substantially and delivered little wall-clock improvement.

That gap is the interesting part. If reducing the arithmetic does not reduce the time, the arithmetic was not what the hardware was waiting for.

An A100 does about 312 TFLOP/s in dense BF16 while moving about 2 TB/s from HBM. That is roughly 156 FLOPs per byte before compute becomes the limit.

Written the obvious way, attention moves so many bytes per FLOP that it sits far below that line. The workloads were bandwidth-bound despite their quadratic arithmetic cost.

The fix had to be an IO fix — not a FLOP fix. That is the whole module.

One precision on that claim: the later versions of this kernel spend their effort keeping tensor cores busy, which is the opposite lesson. FlashAttention-1 fixed the IO problem; FA2 and FA3 fixed the compute problem that was left after the IO problem was gone.

STANDARD ATTENTION, per head, per layer FLOPs 4 N² d (QKᵀ and PV, 2 FLOPs per MAC) HBM bytes Θ(N² ) per elementwise pass over the score matrix resident N² · bytes (kept for the backward pass) The FLOP term is quadratic and, for dense attention, unavoidable. The byte term is quadratic and — this is the whole module — avoidable.

Two quadratic terms were conflated for years. Only one of them was avoidable.

How this comes up in interviews

Try answering out loud before revealing.

Prove it before continuing

Sparse and low-rank attention variants cut FLOPs by large factors, yet wall-clock time barely moved. What was the hardware actually waiting on?

Answer the question above to continue — committing unlocks, right or wrong.