PAPERS-08-006 · Paper — FlashAttention

The paper

FlashAttention, introduced by Tri Dao, Daniel Y. Huang, et al. in 2022 at NeurIPS, is a landmark paper because it fundamentally reshaped how attention is computed, turning a memory‑bound bottleneck into an efficient kernel.

Before this paper

Before FlashAttention, standard attention required building the full N‑by‑N score matrix in GPU high‑bandwidth memory. This approach was memory‑bound and its memory usage grew quadratically with sequence length, making long contexts impractical. The huge matrix had to be written to slow memory, limiting speed and scalability. As a result, training and inference on long sequences were costly and often infeasible.

The key idea

The key insight is to make attention I/O‑aware by reordering the computation so the massive attention matrix never needs to be fully materialized in slow memory. By processing attention in small blocks that fit in fast on‑chip SRAM, the algorithm keeps only the necessary intermediate statistics. This avoids the quadratic memory blow‑up while still computing exact attention, not an approximation.

How it works

FlashAttention first tiles the queries, keys, and values into blocks that fit into the GPU's SRAM. For each block it computes the dot‑product scores, updates running softmax statistics, and multiplies by the corresponding value block, all while staying in fast memory. It repeats this block‑by‑block across the whole sequence, never writing the full score matrix to global memory. During the backward pass, it recomputes the needed blocks on the fly instead of storing them, further saving memory. The whole process yields the exact same results as the standard algorithm but with far lower memory traffic.

See it — never write the N x N matrix

flowchart LR
  A["Q, K, V"] --> B["Tile into blocks"]
  B --> C["Compute in fast<br/>on-chip SRAM"]
  C --> D["Never store full<br/>N x N in HBM"]
  D --> E["Exact, 2-4x faster"]

What they showed

The authors showed that FlashAttention ran attention two to four times faster than the traditional implementation and used up to an order of magnitude less memory, enabling much longer context windows.

What it changed

FlashAttention quickly became the default attention kernel in modern transformer training and inference pipelines. Its speed and memory efficiency allowed researchers and engineers to train larger models and use longer sequences without prohibitive hardware costs. Consequently, it set a new standard for efficient transformer implementations across the field. Next: QLoRA.

The one-line takeaway

FlashAttention proves that clever reordering of computation can eliminate the quadratic memory bottleneck without sacrificing exactness.

Read it / meet it

Reading the original FlashAttention paper reveals the elegant engineering behind the speedup and is rewarding for anyone interested in deep‑learning systems. Its ideas live on in the attention kernels of most large‑model stacks today.