Flash Attention: exact attention without the N N memory blow-up

If you have ever tried to push a transformer to a longer context and watched your GPU run out of memory, you have met the real bottleneck in attention. It is not the number of multiplications. It is the giant matrix that attention wants to write to memory. Flash Attention is the trick that makes that matrix disappear while computing the exact same answer. Where the memory goes Self-attention is one line: for queries Q, keys K and values V, the output is softmax(QKᵀ / √d) · V. Every query scores…