FlashAttention Algorithm › Summary

FlashAttention achieves 2-4x speedup and dramatic memory reduction by restructuring the attention computation to exploit GPU memory hierarchy. The key ideas covered in this chapter:

• GPU memory hierarchy: HBM provides capacity while SRAM provides speed. Standard attention is slow because it repeatedly transfers the $n \times n$ attention matrix through HBM.

• Tiling for SRAM: By processing attention in blocks of size $B_r \times B_c$, FlashAttention keeps all intermediate values in fast SRAM, avoiding HBM transfers for the attention matrix.

• Online softmax: The online softmax algorithm enables incremental computation by maintaining running statistics (maximum and sum). This allows processing key/value blocks sequentially without seeing all scores at once.

• Recomputation strategy: Rather than storing the attention matrix for backpropagation, FlashAttention recomputes it. This trades $O(n^2)$ extra compute for $O(n^2)$ memory savings, a worthwhile trade given that memory is the bottleneck.

• Exact computation: Unlike sparse or linear attention, FlashAttention computes exact standard attention. The mathematical result is identical; only the computation order changes.

• Practical impact: FlashAttention has become the standard attention implementation, allowing longer sequences, faster training, and more efficient inference across production language models.

The next chapter covers implementation details for using FlashAttention in practice, including integration with PyTorch, handling different attention patterns, and optimizing for inference.