FlashAttention (2022) › The Problem

The traditional approach to computing attention in transformer models had relied on standard matrix multiplication operations that were memory-intensive and computationally expensive, especially for long sequences. The attention mechanism required storing intermediate results in memory, leading to quadratic memory complexity with respect to sequence length. When processing a sequence of length $n$, standard attention needed to store an $n \times n$ attention matrix, consuming memory that grew quadratically with sequence length. This limitation made it difficult to train or deploy models with long sequences, as the memory requirements grew rapidly with sequence length.

Consider a practical scenario: processing a document with 8,000 tokens. Standard attention would require storing a 8,000 × 8,000 matrix, consuming over 250 megabytes of memory just for the attention weights, even before accounting for the query, key, and value matrices. For sequences of 16,000 tokens or longer, which researchers wanted to enable for document understanding and long-form generation, the memory requirements became prohibitive, often exceeding available GPU memory.

Additionally, the standard attention computation was not optimized for modern hardware, leading to inefficient use of computational resources. Traditional implementations moved data between different levels of the memory hierarchy.from slow global memory to faster shared memory and registers.repeatedly and inefficiently. Each memory access had a cost in terms of latency and bandwidth, and the naive implementation of attention resulted in many unnecessary memory transfers. This meant that even when enough memory was available, the computation was slow because the GPU spent much of its time waiting for data to be transferred rather than performing actual computations.

The attention mechanism's memory bottleneck limited researchers' ability to train models that could understand and generate long documents, maintain context over extended conversations, or process entire books or research papers. As the field pushed toward longer contexts to enable more sophisticated applications, this limitation became one of the primary obstacles preventing further progress in long-context language models.