FlashAttention

A transformer's attention mechanism computes a weighted sum of value vectors, where the weights come from comparing each query to every key. For a sequence of length N, that means N² comparisons, and the standard implementation writes all N² scores into memory as a dense matrix, applies softmax to each row, writes the result back, then uses it to weight the values. The matrices are large — for a 4096-token sequence, 16 million floats per head — and they live in high-bandwidth memory (HBM), the main DRAM pool on a GPU.

Modern GPUs have a three-tier memory hierarchy. At the top, each streaming multiprocessor has a small amount of SRAM, often called shared memory, with latencies under ten cycles and bandwidth measured in tens of terabytes per second. Below that sits HBM, with latencies in the hundreds of cycles and bandwidth around one or two terabytes per second. At the bottom, if the GPU runs out of its own memory, data must cross the PCIe bus to system RAM, which is slower still. The gap between SRAM and HBM is the critical one: moving a float from HBM into a register where the GPU can actually use it costs roughly a hundred times more energy and takes fifty times longer than performing a multiply-add on it once it arrives.

The standard attention implementation is written as if memory were free. It computes Q times K-transpose and writes the result — the full attention matrix — to HBM. It reads that matrix back, applies softmax row by row, and writes the softmaxed matrix back to HBM. It reads the softmaxed matrix again, multiplies it by V, and writes the final output. For a sequence of 4096 tokens with 128-dimensional heads, each of those reads and writes moves 64 megabytes. The arithmetic itself — a few tens of millions of FLOPs — finishes quickly, but the GPU spends most of its time stalled, waiting for data to arrive from HBM.

FlashAttention reorganises the computation so that intermediate matrices never leave SRAM. It divides Q, K and V into blocks sized to fit comfortably in on-chip memory, then computes attention one block of queries at a time. For each query block, it loops over the key and value blocks, computing partial attention scores and accumulating weighted sums in SRAM. The full N-by-N attention matrix is never materialised. Only the final output is written back to HBM.

The obstacle is softmax. To compute softmax over a row, you need the maximum value in that row, which is used to shift all the exponentials for numerical stability. If you are processing the row in chunks, you do not know the global maximum until you have seen every chunk. FlashAttention handles this with online softmax: it tracks a running maximum and a running sum of exponentials, and whenever a new block reveals a larger maximum, it rescales the contributions from previous blocks. The rescaling factor is exact, so the final result is bitwise identical to what you would get from a single global softmax over the full row. This technique was known in the numerical computing literature, but applying it to attention required working out how to thread it through the backward pass as well, which the FlashAttention paper describes in detail.

The result is an algorithm that performs the same arithmetic as standard attention but moves far less data. Memory traffic drops from quadratic in sequence length to linear, because the intermediate attention matrices are never written out. On an A100 GPU with a sequence length of 2048, the original FlashAttention paper reports wall-clock speedups of 2.4× for the forward pass and similar gains for the backward pass. The improvement grows with sequence length, because the memory bottleneck becomes more severe as N increases.

FlashAttention-2, published a year later, refines the algorithm further. It changes how work is partitioned across the GPU's thread blocks to reduce synchronisation overhead and improve parallelism, particularly for shorter sequences where the original version did not fully saturate the hardware. It also reduces the number of non-matmul FLOPs — the comparisons, exponentials and divisions involved in softmax — by tuning the order of operations. The paper reports speedups of around 2× over the original FlashAttention on typical workloads, which compounds with the earlier gains over the standard implementation.