Flash Attention: Definition and Examples

Flash Attention is an optimized implementation of the attention mechanism (self-attention) used in Transformer architectures. Developed by Tri Dao and his collaborators at Stanford in 2022, this algorithm solves a fundamental problem: classical attention computation has quadratic memory complexity with respect to sequence length, which severely limits the context sizes models can handle.

The central idea of Flash Attention relies on the concept of "tiling" (dividing into blocks) and intelligent management of the GPU memory hierarchy. Instead of materializing the full attention matrix in high-bandwidth memory (HBM), the algorithm performs calculations block by block directly in the GPU's fast memory (SRAM), then merges the results. This avoids costly round trips between different memory levels, which constitute the true bottleneck.

Flash Attention 2, released in 2023, further improved performance by optimizing parallelism and reducing non-essential operations. This version achieves approximately 50 to 70% of the theoretical maximum utilization of modern GPUs, compared to only 25 to 40% for standard attention. Flash Attention 3, released in 2024, pushes these optimizations further for Hopper GPUs (H100).

The impact of Flash Attention on the AI ecosystem is substantial: it has made possible the training of models with context windows of 100,000 tokens and beyond, where classical implementations were limited to a few thousand. Today, virtually all major language models (GPT-4, Claude, Llama, Mistral) use a variant of Flash Attention.