FlashAttention, simplified

I computed the arithmetic intensity of standard attention and got 50 FLOP/byte, six times below an H100's ridge. The reason was the N×N attention matrix being written to and read from HBM. FlashAttention's entire trick is: don't do that.

It's one of the most influential kernel papers of the decade, and the core idea fits in a paragraph.

The problem

Standard attention materializes the N×N scores matrix in HBM (main GPU memory), then reads it back for softmax, then writes it again for the output. At N=4096 that's 86 MB of traffic. At N=32K it's gigabytes. The memory traffic, not the FLOPs, is what kills you.

Meanwhile, the GPU has a small, fast on-chip SRAM (shared memory) that's 100x faster than HBM. Standard attention barely uses it. FlashAttention's insight: use the fast memory aggressively, and never let the full N×N matrix touch HBM.

The trick

FlashAttention processes the attention computation in blocks (tiles), keeping only the running statistics it needs. The two hard parts:

1. Online softmax. You can't compute softmax until you've seen all the scores, but you don't want to store them all. Solution: maintain a running max and running sum, and rescale when a bigger max arrives. This is the "online" softmax trick.
2. Accumulate the output. Each block adds its contribution to the output, scaled by the running statistics, so the final output is exactly right even though the blocks were processed in order.

Here's the online-softmax core, simplified:

m = -inf   # running max
l = 0      # running sum of exp
acc = 0    # accumulated output
for block in tiles(Q, K, V):
    S = block_Q @ block_K.T      # scores for this tile
    m_new = max(m, rowmax(S))
    alpha = exp(m - m_new)
    l = alpha * l + rowsum(exp(S - m_new))
    acc = alpha * acc + exp(S - m_new) @ block_V
    m = m_new
O = acc / l   # rescale once at the end

The output is mathematically identical to standard attention, but the memory traffic drops from O(N²) to O(N) per head. That's the whole paper.

Why it matters for inference

• Longer context fits. Without the N² memory blowup, you can attend over much longer sequences on the same GPU.
• Prefill gets faster. It's a memory-bound optimization, so it helps the memory-bound parts of prefill too.
• It's the template. Every serious kernel now uses tiling + online softmax. If you read one kernel pattern, read this one.

The honest tradeoff

FlashAttention is not free. The online-softmax rescaling adds a little compute, and the tiling adds complexity. For very short sequences the overhead can exceed the savings, which is why engines only route to it above a threshold. And on the decode side, where the bottleneck is the weights not the attention, it helps less. It's a prefill and long-context win, and those are the expensive cases anyway.

The takeaway

FlashAttention didn't invent new math. It just stopped writing to memory. That's the whole game: the fastest memory access is the one you never make.

Tomorrow: profiling attention memory growth, and watching the N² curve bite in real time.