Why architecture matters here
The architecture matters because it targets the actual bottleneck. On a modern GPU, arithmetic throughput vastly exceeds memory bandwidth, so for a memory-bound operation like attention, the number of HBM reads and writes — not the number of floating-point operations — determines runtime. Standard attention writes and re-reads an N×N matrix, incurring O(N²) HBM traffic. FlashAttention restructures the computation so that traffic is roughly O(N²/M) times smaller (where M relates to SRAM size), which is why it can be several times faster despite doing the same math.
It matters because the O(N²) memory of the naive approach is what caps sequence length. Materializing the full score matrix for a long context consumes memory quadratically, so doubling the context quadruples the memory just for attention scores, quickly exhausting the GPU. By never storing that matrix and using only O(N) memory, FlashAttention lets models train and run on far longer sequences with the same hardware — a direct enabler of the long-context models that have become standard.
It matters because it delivers these gains while remaining exact. Many earlier efficient-attention methods approximated the softmax (sparsity, low-rank, kernelization) and traded accuracy for speed, which made them risky to adopt. FlashAttention computes bit-for-bit the same result as standard attention (up to floating-point reordering), so it is a safe, transparent replacement: models get faster and fit longer contexts with no change in what they compute, which is why it was adopted so quickly and broadly.
The architecture: every piece explained
The memory hierarchy it exploits. A GPU has a small, extremely fast on-chip SRAM (shared memory, tens of KB per streaming multiprocessor) and a large, much slower off-chip HBM (tens of GB). Standard attention keeps intermediates in HBM; FlashAttention's whole design is organized around doing as much work as possible on data already resident in SRAM, because SRAM bandwidth is roughly an order of magnitude higher than HBM. Understanding attention as an IO problem across this hierarchy is the conceptual core.
Tiling into blocks. The Q, K, and V matrices are partitioned into blocks small enough that a block of Q plus a block of K and V fit together in SRAM. The algorithm loads a Q block once and streams K,V blocks through it, computing the partial attention scores for just that tile in fast memory. Because it works on tiles, it never needs the full N×N score matrix in existence at any moment — only one small block of scores lives in SRAM at a time.
The online (streaming) softmax. Softmax normally needs the whole row of scores to compute the max and the sum for numerical stability and normalization. FlashAttention uses an online softmax that maintains a running maximum and a running sum of exponentials as it processes each K,V block, rescaling the partial output accumulated so far whenever the running max changes. This lets it combine blocks incrementally and arrive at exactly the correct normalized result without ever holding the full row.
The output accumulator and recomputation. The output O (size N×d, linear in N) is accumulated in SRAM/registers block by block and written to HBM once at the end. For the backward pass, rather than having stored the N×N scores, FlashAttention recomputes the needed attention blocks on the fly from Q, K, V (which are cheap to keep), trading a little extra arithmetic for avoiding O(N²) storage — a favorable trade precisely because the operation is memory-bound.
End-to-end flow
The forward pass loops over blocks of the query matrix. For each Q block, the algorithm initializes an output accumulator and the running softmax statistics (running max and running sum) to their starting values, all held in fast on-chip memory. It then iterates over the blocks of K and V. For each K,V block, it loads the block into SRAM and computes the partial scores QKᵀ for just this tile — a small matrix that fits on chip.
With those partial scores in hand, it updates the online softmax: it finds the block's local maximum, adjusts the running maximum, rescales the output accumulated so far so that everything is normalized against the new maximum, exponentiates the scores relative to that maximum, and adds the weighted V contribution of this block to the accumulator. After every K,V block has streamed through, the accumulator holds the exact, fully normalized attention output for this Q block, which is written once to HBM.
Crucially, at no point was the full N×N score matrix written to or read from HBM — only tiles ever existed, and only in SRAM. The traffic to HBM is essentially reading Q, K, V once and writing O once, which is O(N) rather than O(N²). Because the operation was bottlenecked on exactly that traffic, collapsing it is what produces the multiple-times speedup, while the online softmax guarantees the numerical result matches standard attention.
The backward pass follows the same tiling discipline. Instead of having saved the N×N attention weights, it keeps only the small softmax statistics per row and recomputes the attention blocks from Q, K, V as gradients flow, again tile by tile in SRAM. This keeps training memory at O(N) rather than O(N²) for the attention activations, which — combined with the forward speedup — is what lets large models train on long sequences without running out of memory, all while remaining an exact computation.
A useful way to see why the recomputation is a bargain is to count the memory transactions. Standard attention writes the N×N scores to HBM, reads them back for the softmax, writes the normalized weights, and reads them again for the value multiply — several full O(N²) trips over the slowest link in the machine. FlashAttention replaces all of that with reading Q, K, and V once and writing O once, and in the backward pass it re-reads Q, K, V and recomputes the tiles in fast SRAM. The recomputation adds arithmetic, but arithmetic on data already in SRAM is enormously cheaper than the HBM round-trips it eliminates, so the net effect on a memory-bound operation is a large win in both time and memory. This is the general lesson FlashAttention popularized: on modern accelerators, restructuring an algorithm to move less data across the memory hierarchy often beats reducing the number of floating-point operations.