Understanding FlashAttention: Memory Complexity

The Computational Steps of Standard Self-Attention

At its core, the self-attention mechanism allows a model to weigh the importance of different tokens in a sequence when producing a representation for each token. The canonical formula for scaled dot-product attention is:

Where $Q$ (Query), $K$ (Key), and $V$ (Value) are matrices derived from the input sequence, and $d_k$ is the dimension of the key vectors.

For a sequence of length $n$, the matrices $Q$, $K$, and $V$ each have a shape of $n\times d_k$. A naive implementation of this formula on a GPU follows a straightforward, multi-kernel or multi-step process:

1. Compute Scores: Calculate the dot product between every query and every key to get the attention scores. This is a matrix multiplication: $S=Q{K}^T$. The resulting score matrix $S$ has dimensions $n\times n$.
2. Apply Softmax: Apply a softmax function row-wise to the score matrix $S$ to get the attention probabilities, $P=\text{softmax}(S)$. This matrix $P$ also has dimensions $n\times n$.
3. Compute Output: Multiply the probability matrix $P$ by the value matrix $V$ to get the final output. This is another matrix multiplication: $O=PV$.

Each of these steps typically involves a separate kernel launch, leading to multiple trips to and from the GPU's main memory, which, as you know, is a high-latency operation.

To see this in action, watch the following short animation, which visualizes this sequence of operations. Pay close attention to how intermediate matrices ($S$ and $A$, which is our $P$) are written to and read from High Bandwidth Memory (HBM).

The diagram below offers a static view of this same process, highlighting the data flow between the compute units and the slow, off-chip HBM.

The Quadratic Memory Bottleneck

The fundamental problem with the standard attention mechanism lies in the materialization of the $n\times n$ attention matrix. Let's analyze why this is so detrimental.

For a sequence of length $n$ and a model with $h$ attention heads, the attention probability matrix $P$ has a shape of (h, n, n). The memory required to store this single tensor is proportional to $h\times n^2$. This is the source of the quadratic memory complexity.

To appreciate the scale of this problem, consider the numbers. In early models, a sequence length of $n=512$ was common. For a 12-head model using 16-bit floats (2 bytes), the memory for the attention matrix is:

$12\text{ heads}\times 512\times 512\times 2\text{ bytes/float}\approx 6.3\text{ MB}$

This is manageable. However, modern Large Language Models (LLMs) operate on much longer contexts. If we increase the sequence length to $n=8192$:

$12\text{ heads}\times 8192\times 8192\times 2\text{ bytes/float}\approx 1.6\text{ GB}$

And at $n=32768$:

$12\text{ heads}\times 32768\times 32768\times 2\text{ bytes/float}\approx 25.7\text{ GB}$

As you can see, the memory consumption explodes. For a sequence length of 32k, storing this single intermediate matrix can consume the entire memory of a consumer GPU and a significant portion of a high-end datacenter GPU like an A100. This is before accounting for model weights, gradients, optimizer states, and other activations.

The following resource provides a detailed analysis of this issue, connecting it to a typical PyTorch implementation.

The graph below visually contrasts the linear, $O(n)$, complexity of a Recurrent Neural Network (RNN) with the quadratic, $O(n^2)$, complexity of self-attention. While attention offers superior modeling capabilities by allowing direct interaction between any two tokens, it comes at a steep computational and memory cost.

The I/O Problem: More Than Just Memory Size

For a systems programmer like yourself, it's clear that the issue isn't just the size of the memory allocation, but the cost of accessing it. The $n\times n$ matrix is too large to fit in the fast on-chip SRAM. Consequently, it must be written to and read from the much slower global memory (HBM).

This creates a severe I/O bottleneck. The standard algorithm is not I/O-aware; it is memory-bound. The GPU's powerful compute units spend most of their time idle, waiting for data to be shuttled back and forth from HBM.

Tri Dao, the creator of FlashAttention, explains this exact issue. He highlights that the quadratic scaling of reads and writes to HBM is the primary obstacle to processing long sequences efficiently.

The arithmetic intensity of standard attention is very low. The number of memory operations scales as $O(n^2)$, while the number of floating-point operations scales as $O(n^2{d}_k)$. For long sequences where $n\gg d_k$, the kernel spends far more time waiting on memory than doing useful computation, making it a classic memory-bound problem.

Conclusion

In this lesson, we have pinpointed the fundamental performance limitations of the standard self-attention algorithm. By breaking down its computational steps, we've identified the materialization of the $n\times n$ attention matrix as the primary culprit.

Key Takeaways:

• Quadratic Memory Complexity: Standard attention requires storing an intermediate $n\times n$ matrix, leading to memory requirements that scale with $O(n^2)$. This makes processing long sequences prohibitively expensive in terms of memory capacity.
• I/O Bottleneck: The large attention matrix does not fit in fast on-chip SRAM and must be repeatedly written to and read from slower global memory (HBM). This makes the algorithm memory-bound, leaving the GPU's compute units underutilized.
• The Attention Bottleneck: The combination of quadratic memory and I/O scaling creates the "attention bottleneck," which is the primary factor limiting the context length of Transformer models.

You now have a firm grasp of the problem that FlashAttention was designed to solve. It's not about approximating attention; it's about reformulating the exact same computation to be I/O-aware, avoiding the costly round-trips to HBM.

In our next lesson, we will explore the elegant solution. You will learn about the two core techniques FlashAttention employs—tiling and online softmax—to compute exact attention without ever materializing the full $n\times n$ matrix in global memory, thereby breaking the quadratic memory barrier.