FlashAttention.
An IO-aware attention algorithm that tiles computation in fast SRAM memory to reduce data movement, enabling efficient training and inference on long sequences.
FlashAttention is a memory-efficient implementation of the attention mechanism that addresses the quadratic memory complexity bottleneck in transformer models. Traditional attention implementations require storing large intermediate matrices in GPU memory, which becomes prohibitive for long sequences. FlashAttention solves this by recognizing that the bottleneck is not computation but memory bandwidth - the cost of moving data between different levels of the GPU memory hierarchy.
The algorithm works by tiling the attention computation and using the GPU's fast SRAM (shared memory) as a cache to minimize data movement to and from the slower high-bandwidth memory (HBM). It computes attention in blocks, loading only the necessary tiles of queries, keys, and values into SRAM, performing the computation there, and then moving results back. This approach maintains mathematical equivalence to standard attention while dramatically reducing memory traffic and enabling much longer context lengths.
FlashAttention enables training and inference on sequences that would otherwise be impossible due to memory constraints, making it crucial for long-context applications like document analysis and extended conversations. The algorithm typically provides 2-4x speedups and allows models to handle context windows an order of magnitude larger than before. However, the implementation complexity is significantly higher than standard attention, and the benefits are most pronounced on newer GPU architectures with sufficient SRAM capacity.