FlashAttention (2022) › The Solution

FlashAttention addressed these limitations by redesigning the attention computation to be IO-aware and memory-efficient. The algorithm computed attention by processing the input in blocks and using techniques including tiling and recomputation to minimize memory usage while maintaining exact attention computation. This approach reduced memory complexity from quadratic to linear with respect to sequence length, enabling the training and deployment of models with much longer sequences.

The key innovation of FlashAttention was its use of tiling to process attention computation in blocks, reducing memory requirements while maintaining numerical accuracy. The algorithm divided the input into smaller blocks and computed attention for each block separately, using techniques including online softmax and online attention to maintain exact computation. Rather than computing the full attention matrix at once, FlashAttention processed queries and keys in tiles, computing attention scores incrementally and only storing what was necessary for the final output.

This tiling approach worked by dividing the input sequence into blocks and processing attention computation block by block. For each query block, the algorithm would iterate through key blocks, computing attention scores incrementally. The critical insight was that by using online algorithms for softmax and attention computation, the algorithm could maintain numerically exact results without ever storing the full attention matrix. The online softmax algorithm allowed computing the softmax normalization factor incrementally as new attention scores were computed, while the online attention computation enabled accumulating the attention-weighted values without storing intermediate scores.

The algorithm's architecture was designed to be hardware-aware, optimizing for modern GPU architectures and memory hierarchies. FlashAttention used techniques including shared memory, register blocking, and memory coalescing to maximize computational efficiency and minimize memory access overhead. The implementation carefully managed data movement between different levels of GPU memory.global memory, shared memory, and registers.to minimize the number of slow memory accesses. By keeping frequently accessed data in fast shared memory and registers, the algorithm could compute attention much more efficiently than naive implementations.

The algorithm was also designed to be numerically stable and exact. This ensured that the results were identical to standard attention computation. This was central because any numerical differences could affect model training and deployment. The online softmax and attention algorithms used by FlashAttention were carefully designed to maintain the same numerical properties as standard attention computation, just computed in a more memory-efficient order.