Flash Attention in transformers

Module 16 · Advanced AI · 23 concepts

In depth / 15 lines

Flash Attention, in more detail

Flash Attention is a faster and more memory-efficient way to compute exactly the same attention.

Standard attention computes a score for every pair of positions.

For a sequence of 8,000 tokens that matrix holds 64 million entries per head.

It is written to GPU memory, read back for the softmax, and written again.

That traffic dominates the cost, since the arithmetic itself is comparatively cheap.

Flash Attention never builds the full matrix in memory.

It processes the sequence in tiles that fit in fast on-chip SRAM.

Each tile is computed, used and discarded without a round trip to main memory.

The softmax is computed incrementally as tiles are processed, which is the clever part.

The output is mathematically identical, so this is not an approximation.

Memory use drops from growing with the square of the length to growing linearly.

That alone made much longer context windows practical.

Speedups of two to four times are typical, and larger on long sequences.

It is now built into the major frameworks and used by default.

It is a good illustration that on modern hardware, memory movement is the thing to optimise.

Deep dive

What is underneath Flash Attention

The cost that is not the arithmetic

S = \frac{QK^{\top}}{\sqrt{d_k}} \in \mathbb{R}^{N \times N}

A standard implementation writes that matrix to memory, reads it back for the softmax, writes the result, and reads it again for the multiply by $V$ .

O(N^{2}) \text{ memory traffic}

With a sequence of 8,192 and 32 heads, the attention matrix alone is gigabytes, and every byte crosses the memory bus several times.

I = \frac{\mathrm{FLOPs}}{\text{bytes}} \text{ is low}

Attention is memory-bound, so the speedup does not come from doing less arithmetic. FlashAttention does slightly more and moves far less.

Tiling, and the problem it creates

Q = [Q_1, \dots, Q_{T_r}], \qquad K, V = [K_1, \dots, K_{T_c}], \; [V_1, \dots, V_{T_c}]

Blocks of the three matrices are loaded into shared memory, where the arithmetic is an order of magnitude faster than in main memory.

\mathrm{softmax}(\mathbf{x})_i = \frac{e^{x_i - m}}{\sum_j e^{x_j - m}}, \qquad m = \max_j x_j

The obstacle is that softmax needs the maximum and the sum over the whole row, and a tile only ever sees part of it.

Online softmax

m^{(k)} = \max\big(m^{(k-1)}, \; \tilde{m}^{(k)}\big)

\ell^{(k)} = e^{m^{(k-1)} - m^{(k)}} \ell^{(k-1)} + e^{\tilde{m}^{(k)} - m^{(k)}} \tilde{\ell}^{(k)}

Keep a running maximum and a running sum, and rescale them whenever a new block raises the maximum.

O^{(k)} = e^{m^{(k-1)} - m^{(k)}} O^{(k-1)} + \tilde{P}^{(k)} V_k

The output accumulator is rescaled by the same factor, so the result after the last block is exactly the softmax over the whole row.

It is not an approximation. The output is numerically identical to the standard implementation, which is why it could be adopted without retraining anything.

One block row, in order

Load a block of $Q$ into shared memory and initialise the accumulators.
Loop over blocks of $K$ and $V$ , loading each into shared memory.
Compute the scores for that tile and its local maximum and sum.
Rescale the running statistics and the output accumulator, then add the tile's contribution.
After the last tile, divide by the running sum and write the output once.

O(N) \text{ memory instead of } O(N^{2})

The attention matrix is never materialised, which is where the memory saving comes from and why context lengths that previously did not fit now do.

The backward pass

The trick is to recompute $S$ from $Q$ and $K$ instead of storing it.

Backpropagation needs the attention weights, and storing them would undo everything.

They are recomputed inside the same tiling, which costs extra arithmetic and still wins, because arithmetic was never the constraint.

2 \text{ to } 4\times \text{ faster, } 10 \text{ to } 20\times \text{ less memory}

The later versions improve the work partitioning across warps and reduce the non-matmul operations, which matter because the matrix units are far faster than the general ones.

Nothing about the model changes. It is a different implementation of the same function, which is the cleanest kind of optimisation there is.

Test yourself

One question on this card. Every option explains itself once you pick, right or wrong.

Question 1

How does Flash Attention speed up attention?

Advanced AI

23 concepts in this module, then a graded test with every answer explained.

Study this module →