FlashAttention: Move Less, Not Compute Less

Fri, 25 Sep 2026 00:00:00 +0000

FlashAttention begins with a counterintuitive fact. It still computes exact attention. Its floating-point work still grows quadratically with sequence length, and the backward pass deliberately recomputes some values. Yet it is usually faster and reduces extra memory from quadratic to linear in sequence length.

Big-O arithmetic alone cannot explain that result. The answer lies in what a GPU actually spends time doing. Modern accelerators multiply matrices extremely quickly, but they dislike repeatedly writing a huge intermediate tensor to high-bandwidth memory (HBM), then reading it back for the next kernel. FlashAttention reorganizes attention from a sequence of mathematical operators into a dataflow designed around the memory hierarchy.

It does not change what the model learns. It changes how the same formula reaches the machine. That is why the work matters beyond one CUDA trick: model efficiency depends on both the computation graph and the journey each byte takes through hardware.

The Multiplication May Not Be the Slow Part

For one attention head:

$$S=QK^\top,\qquad P=\operatorname{softmax}(S),\qquad O=PV$$

With sequence length $N$ and head dimension $d$, $Q$, $K$, $V$, and $O$ have shape $N\times d$, while the score and probability matrices $S$ and $P$ are both $N\times N$. A conventional implementation lets separate kernels perform the first matrix multiplication, masking, softmax, dropout, and the second multiplication. It writes $S$ to HBM, reads it for softmax and writes $P$, then reads $P$ again to multiply by $V$.

The cost is not merely storing $N^2$ elements. At each kernel boundary, an intermediate may travel between on-chip compute and HBM. The original paper uses the A100 as an example: each streaming multiprocessor has about 192 KB of on-chip SRAM with an estimated 19 TB/s bandwidth, while the card has 40–80 GB of HBM at 1.5–2.0 TB/s. SRAM is roughly an order of magnitude faster but many orders of magnitude smaller. The useful question therefore becomes: can a small tile stay on chip long enough to finish the whole operation, without materializing the full $N\times N$ matrix in HBM?

Figure 1: standard attention materializes S and P in HBM; FlashAttention brings only small tiles into SRAM

Two meanings of memory are easy to conflate here. FlashAttention does not eliminate the input $Q$, $K$, $V$, or the final output. It reduces HBM traffic and storage for internal intermediates. The paper proves that, for a range of realistic SRAM sizes, standard attention performs $\Theta(Nd+N^2)$ HBM accesses while FlashAttention performs $\Theta(N^2d^2/M)$, where $M$ is the SRAM capacity. Arithmetic remains $O(N^2d)$, but movement is greatly reduced.

Tiling Is Easy; Softmax Couples the Entire Row

Matrix multiplication naturally supports tiling: split $Q$, $K$, and $V$ into pieces that fit in SRAM and multiply them block by block. Softmax appears to resist this treatment. A numerically stable softmax for a score vector $x$ first needs the maximum across the whole row, then the sum of all exponentials:

$$m=\max_i x_i,\qquad \ell=\sum_i e^{x_i-m},\qquad \operatorname{softmax}(x)_i=\frac{e^{x_i-m}}{\ell}$$

If the whole row must be saved before $m$ and $\ell$ are known, the quadratic intermediate has not gone away. FlashAttention uses online softmax instead. After reading each score block $x_b$, it updates only the running maximum $m$, normalization sum $\ell$, and an unnormalized output accumulator $o$:

$$ \begin{aligned} m' &= \max(m,\max x_b)\\ \ell' &= e^{m-m'}\ell + \sum_j e^{x_{b,j}-m'}\\ o' &= e^{m-m'}o + \sum_j e^{x_{b,j}-m'}v_{b,j} \end{aligned} $$

The final result is $O=o/\ell$. When a new block introduces a larger maximum, the old accumulators are rescaled by $e^{m-m’}$. No matter how many blocks are used, the result is the same stable softmax attention, apart from ordinary floating-point rounding. This is neither approximate nor sparse attention.

Figure 2: online softmax maintains m, ℓ, and o block by block while preserving exact attention

In the fused kernel, a $K/V$ tile loaded into SRAM interacts with several $Q$ tiles. The kernel performs $QK^\top$, masking, the online softmax update, and multiplication by $V$ while the data is on chip. A local score tile is discarded after use; only the output tile and small normalization statistics survive. What had been a pipeline of separate operators becomes one I/O-aware kernel.

Why More Arithmetic Can Finish Sooner

Backpropagation normally needs $S$ and $P$ from the forward pass. A conventional implementation saves them in HBM. FlashAttention saves only the output and per-row softmax statistics. During the backward pass it reconstructs local $S$ and $P$ after the corresponding $Q$, $K$, and $V$ tiles have already reached SRAM.

The trade is explicit: spend extra matrix-multiply FLOPs to avoid large HBM writes, reads, and residency. On a GPU whose matrix-multiply throughput has grown much faster than its ability to move data, this can save both time and memory. The original paper reduces extra memory from $O(N^2)$ to $O(N)$. The FlashAttention-2 paper reports roughly 10–20× memory savings in the backward pass from not storing $S$ and $P$, along with 2–4× wall-clock speedups. Those ratios depend on sequence length, head dimension, masking, precision, and hardware; they are not universal promises.

This corrects a common intuition: recomputation is not necessarily slower. “Compute less” wins when arithmetic is the bottleneck. When memory access dominates, recomputing a value on chip can be cheaper than retrieving it from HBM. FlashAttention’s central idea is therefore not one softmax identity, but a hardware-appropriate choice of what to store, move, and recompute.

FA1 to FA3: Optimization Moves Closer to the Hardware

After the first generation removed the largest HBM round trips, attention still fell well short of the GPU’s matrix-multiply peak. Later versions kept the same backbone—tiling, online softmax, and recomputation—while attacking finer-grained idle time and synchronization.

FlashAttention-2 focuses on how GPU work is partitioned. On A100, the paper lists theoretical peaks of 312 TFLOPs/s for FP16/BF16 matrix multiplication and only 19.5 TFLOPs/s for non-matmul FP32 work. FA2 therefore reduces non-matmul rescaling, parallelizes along sequence length in addition to batch and heads to improve occupancy for long sequences and small batches, and repartitions work among warps to reduce shared-memory communication. The paper reports about 2× speed over FA1, up to 73% of theoretical peak in the A100 forward pass, and up to 225 TFLOPs/s per A100 in end-to-end GPT-style training.

FlashAttention-3 targets Hopper. On H100, Tensor Core multiplication, Tensor Memory Accelerator transfers, and ordinary CUDA-core work can proceed asynchronously. Software that still follows “load, then multiply, then softmax” makes those specialized units wait on one another. FA3 separates producer and consumer warps, uses a ping-pong pipeline to overlap transfer, the next block’s matrix multiplication, and the current block’s softmax, and combines FP8 with block quantization and incoherent processing to contain outlier-driven quantization error. The paper reports 1.5–2.0× FP16 forward speed over FA2 on H100, up to 740 TFLOPs/s; FP8 approaches 1.2 PFLOPs/s and shows 2.6× lower error than its baseline FP8 attention in the reported tests.

Figure 3: FA1 removes HBM traffic, FA2 improves work partitioning, and FA3 maps asynchronous pipelines and low precision to Hopper

The sequence is instructive. First redesign the algorithm’s data path, then divide work better across thread blocks and warps, and finally align the schedule with a particular GPU generation’s asynchronous units. Performance does not appear automatically because code “runs on a GPU.” The algorithm must renegotiate with each hardware generation.

What It Changes—and What It Does Not

FlashAttention is not a new attention pattern. Model weights, dense-attention semantics, and outputs do not change. A compatible model can swap kernels. Linformer, sparse attention, and linear attention alter or approximate which relationships are computed; FlashAttention changes only the execution order.

It does not remove quadratic arithmetic. Training and prefill over a long prompt still compute many pairwise token interactions. FlashAttention organizes that work in a hardware-friendly way and avoids quadratic intermediate activations, but it does not make arbitrary context free. Doubling context still roughly quadruples dense-attention FLOPs.

It pays a different bill from KV-cache compression. During single-token autoregressive decode, there is no large query-by-key matrix to materialize, and reading a growing KV cache often dominates. GQA reduces the number of KV heads, MLA compresses each token’s cached representation, and Mamba replaces the growing archive with fixed state. FlashAttention primarily optimizes dataflow inside the attention kernel. These techniques complement one another but do not substitute for one another. The official implementation’s causal, GQA/MQA, variable-length, and paged-attention paths also show that “attention” is not one fixed kernel shape.

Peak kernel speed is not application speed. A training step also includes MLPs, communication, the optimizer, and data loading. An online request also includes scheduling, sampling, and serving overhead. An attention microbenchmark, single-GPU TFLOPs/s, and end-to-end throughput answer different questions. Any performance claim should name the hardware, precision, sequence shape, causal mask, and baseline.

The Lasting Breakthrough Is Putting the Memory Hierarchy Into the Algorithm

FlashAttention did not invent the attention formula, tiling, or online softmax in isolation. Its breakthrough was assembling them into a complete, analyzable, implementable I/O-aware attention algorithm: never materialize the $N\times N$ intermediate in HBM, organize both passes around on-chip tiles, and ship open kernels that convert theoretical traffic reduction into wall-clock gains. The official repository later expanded across NVIDIA and AMD backends, while fused scaled-dot-product attention became a routine path in major frameworks. That kind of adoption says more about impact than a temporary benchmark lead.

The deeper lesson is that long-context capability is never just a model curve. Whether a window is usable depends on whether activations fit, bytes arrive quickly enough, kernels occupy the hardware, and inference systems manage the KV cache. FlashAttention did not solve every layer of that system, but it removed one previously expensive intermediate.

“Move less, not compute less” is therefore more than a performance slogan. It is a design method: ask not only how many FLOPs an algorithm performs, but where each byte lives, how many times it moves, when it is worth preserving, and when cheap recomputation can replace expensive traffic. As models grow, those questions only become more important.

References

Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022.
Dao, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, ICLR 2024.
Shah et al., FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision, 2024.
Dao AI Lab, FlashAttention official implementation.
Tri Dao, FlashAttention-3 technical blog, 2024.

GPU on Lee