FlashAttention Implementation › Summary

FlashAttention maps attention computation from a memory-bound operation into one that achieves near-peak hardware utilization. The key insights and takeaways from this chapter:

• GPU memory hierarchy: Understanding the dramatic speed difference between HBM and SRAM (10x) is essential. Standard attention repeatedly transfers data to slow HBM, while FlashAttention keeps data in fast SRAM.

• Tiled computation: By processing attention in small blocks that fit in SRAM, FlashAttention never materializes the full $n \times n$ attention matrix. This reduces memory complexity from $O(n^2)$ to $O(n)$.

• Online softmax: The online softmax algorithm maintains running statistics (maximum and sum of exponentials) that allow correct normalization without seeing all scores at once. Previous accumulations are corrected as new tiles reveal larger values.

• FlashAttention-2 improvements: Reduced non-matrix-multiply operations, better work partitioning across threads and warps, and improved memory access patterns combine to give approximately 2x speedup over FlashAttention-1.

• PyTorch integration: The easiest path to using FlashAttention is through torch.nn.functional.scaled_dot_product_attention, which automatically selects FlashAttention when conditions permit (CUDA, fp16/bf16, supported head dimensions).

• Limitations: FlashAttention requires recent NVIDIA GPUs (Ampere or newer for full support), has constraints on attention mask patterns, and may produce results that differ from standard attention at floating-point precision. While it eliminates the attention matrix memory cost, activations and KV cache still scale linearly with sequence length.

FlashAttention has become the de facto standard for attention computation in modern LLMs. Models like LLaMA 2, Mistral, and most recent transformers use it by default. Understanding its implementation details helps you appreciate why it works, when it applies, and what to expect when deploying models that rely on it.