FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Inside the paper

The problem

Standard attention on long sequences creates large intermediate matrices and moves them between levels of GPU memory. Reducing arithmetic alone does not necessarily reduce elapsed time when memory traffic is the bottleneck.

The idea

FlashAttention reorganizes exact attention around the limited fast memory available on a GPU. It processes blocks of the inputs and combines partial softmax results without writing the full attention matrix to slower memory.

How it works

• Divide the query, key and value matrices into blocks sized for fast on-chip memory.
• Load blocks and calculate portions of attention scores and output on chip, maintaining running quantities needed to normalize the softmax correctly.
• Write the resulting output and compact normalization information to high-bandwidth memory without materializing the full score matrix there.
• During the backward pass, recompute attention blocks from the saved information instead of storing the entire intermediate matrix.

What it showed

The paper reports faster attention and end-to-end training on several benchmarks, and uses longer sequences to improve results on tasks in its experiments. Its analysis also bounds memory accesses for the proposed algorithm under the stated memory model.

Read it for

Read Section 2 for the hardware and standard attention baseline, then Algorithm 1 and Section 3 for the blockwise calculation. Section 4 gives the measured results; the detailed proofs can wait until you need the bounds.

Written in our own words from the paper itself. A draft the editors have not yet reviewed.