The FlashAttention paper contains a benchmark that looks like a misprint. On GPT-2 medium (sequence length 1024, batch size 64, A100), standard attention does 66.6 GFLOPs and runs in 41.7 ms. FlashAttention does 75.2 GFLOPs — 13% more arithmetic — and runs in 7.3 ms. The explanation is in the third column: standard attention moves 40.3 GB between GPU memory levels; FlashAttention moves 4.4 GB. Attention was never compute-bound. It was a memory-traffic problem wearing a FLOPs costume.
The hierarchy attention keeps ignoring
An A100 has 40–80 GB of HBM at 1.5–2.0 TB/s, and 192 KB of SRAM per streaming multiprocessor (108 SMs) at roughly 19 TB/s. SRAM is about an order of magnitude faster than HBM and four orders of magnitude smaller. Any kernel that streams large intermediates through HBM pays the slow tier's price on every byte.
Standard attention does exactly that, three times:
- Read Q, K from HBM; compute
S = QKᵀ; write S to HBM. - Read S back; compute
P = softmax(S); write P to HBM. - Read P and V; compute
O = PV; write O.
S and P are N×N matrices. At sequence length 4096 with fp16, each one is 32 MB per head per batch element — nowhere near fitting in 192 KB of SRAM, so they round-trip through HBM. Total HBM traffic is Θ(Nd + N²), where d is the head dimension. The N² term dominates, and it's pure intermediate state: nobody downstream wants S or P, only O.
The obvious fix is kernel fusion: compute a tile of S, softmax it, multiply by V, never materialize the big matrices. GEMM tiling is a solved problem. The thing that blocked it for years is the softmax in the middle.
Softmax is the anti-tiling operation
Softmax over a row needs the whole row before it can produce any output. The numerically safe version — the one every real library uses to avoid exp overflow — makes that worse, because it needs the row maximum first:
pass 1: m = max_j(x_j)
pass 2: l = Σ_j exp(x_j − m)
pass 3: y_j = exp(x_j − m) / lThree passes over data that, in attention, you were hoping never to store. If you tile QKᵀ into blocks, each block gives you a slice of a row — but you can't exponentiate-and-normalize that slice correctly, because the true row max and row sum depend on blocks you haven't seen yet.
The unblocking trick predates FlashAttention and wasn't invented for attention at all. Milakov and Gimelshein's 2018 online softmax paper showed you can maintain the max and the normalizer in a single pass, fixing up history as you go (they measured up to 1.3× on softmax itself). The update rule: when a new block arrives with block-max m_b and block-sum l_b,
m_new = max(m_old, m_b)
l_new = exp(m_old − m_new) · l_old + exp(m_b − m_new) · l_bThe exp(m_old − m_new) factor retroactively corrects every term you already accumulated, as if you'd known the true max all along. Exact, not approximate — the correction is just algebra on exp(a − m) = exp(a − m') · exp(m' − m).
FlashAttention's move: rescale the output, too
FlashAttention's contribution is pushing that correction one level further: not just the scalar normalizer, but the entire output accumulator. Q, K, V are split into blocks sized to fit SRAM. For each K/V block, the kernel computes the score tile, its block row-max and row-sum, and the tile's contribution P_b · V_b — then folds it into the running output:
O_new = exp(m_old − m_new) · O_old + exp(m_b − m_new) · (P_b · V_b)Per row, the kernel carries only the running statistics (m, l) and the d-dimensional output row. The N×N matrices exist one SRAM-sized tile at a time and are never written to HBM.
The paper proves the resulting I/O cost is Θ(N²d²/M) HBM accesses, where M is SRAM size, versus Θ(Nd + N²) for standard attention — up to 9× fewer accesses for typical d and M — and proves a matching lower bound: no exact attention algorithm can asymptotically beat this across all SRAM sizes. The GPT-2 measurement above (40.3 GB → 4.4 GB) is that theorem showing up in nvprof.
The backward pass: recompute instead of read
Training needs S and P again for gradients, and the classic implementation stores P — right back to O(N²) memory. FlashAttention's forward pass instead saves only O, the per-row statistics (m, l), and the dropout PRNG state. The backward pass regenerates S and P tile by tile in SRAM from Q, K, V using the saved statistics, and regenerates the dropout mask from the PRNG state.
That's deliberately spending FLOPs to avoid bytes, and it's the same bet as the forward pass: recomputing a tile from operands already in SRAM is faster than reading a 32 MB matrix back through HBM. Extra memory drops from O(N²) to O(N) — up to 20× more memory-efficient than exact baselines — which is what made 16K–64K contexts trainable at all. On Path-X (sequence length 16K), FlashAttention produced the first Transformer result better than chance, 61.4%; end to end it trained BERT-large 15% faster than the MLPerf 1.1 record and GPT-2 up to 3× faster than the HuggingFace implementation.
IO-optimal is not occupancy-optimal
The epilogue is instructive. For all its I/O optimality, the original FlashAttention kernel reached only 25–40% of the GPU's theoretical max FLOPs/s. Once HBM traffic stops being the bottleneck, the next ones surface: non-matmul FLOPs (those rescaling exp corrections run on the slow non-tensor-core units), too little parallelism for long sequences at small batch, and shared-memory shuffling between warps.
FlashAttention-2 attacked exactly those — fewer non-matmul operations, parallelizing a single head across thread blocks along the sequence dimension, and repartitioning work between warps — and roughly doubled throughput, reaching 50–73% of theoretical max on A100 and up to 225 TFLOPs/s per GPU (72% model FLOPs utilization) in end-to-end GPT training. Same math, same I/O pattern, 2× faster: once you fix the memory wall, you inherit an ordinary compute-kernel tuning problem.
The takeaway that outlives the kernel
The durable lesson isn't the specific tiling scheme — it's the accounting change. FLOP counts said attention was fine; counting HBM bytes said it was spending most of its time on a matrix nobody asked for. Approximate-attention papers spent years shaving FLOPs off a kernel that wasn't FLOPs-bound, which is why they rarely delivered wall-clock wins. When a kernel underperforms, profile bytes moved per memory tier before arithmetic. If the fix involves doing more math to move less data — recomputation, fusion, rescaling tricks — that's not a hack. On modern accelerators, that's the grain of the hardware.