Quadratic Attention Bottleneck › Memory Analysis: The Attention Matrix

Memory is often the first constraint that prevents you from extending sequence lengths, and the source of that constraint is the attention score matrix. To understand why, consider what must be stored during a single forward pass of a transformer with self-attention.

For every attention head, the model must compute and store an $n \times n$ matrix of attention scores before applying softmax. This matrix is not just used once: it drives the final weighted aggregation of value vectors, and during training it must be retained for the backward pass so that gradients can flow back through the softmax operation. The result is that the $n \times n$ attention matrix lives in GPU memory for the entire duration of the forward and backward pass, consuming resources proportional to the square of the sequence length.

Think of the attention matrix as a large lookup table. For a sequence of 1,024 tokens, the table has 1,024 rows and 1,024 columns, totaling roughly one million entries. For 8,192 tokens, the table grows to 67 million entries. For 32,768 tokens, it reaches one billion entries. At two bytes per entry (fp16), that single table consumes roughly 2 gigabytes for the 32K case, and that's just one head of one layer.

The total memory required to store attention matrices across all heads and all layers in a transformer can be expressed concisely. For a model with $h$ attention heads and $L$ layers, with each layer computing an $n \times n$ attention matrix per head at precision $b$ bytes per element:

$$M_{\text{attn}} = h \times n^2 \times b \times L$$

where:

• $M_{\text{attn}}$: total memory in bytes for all attention matrices
• $h$: number of attention heads (each head computes an independent $n \times n$ attention pattern)
• $n$: sequence length (number of tokens)
• $n^2$: the number of entries in each attention matrix, one score per query-key pair
• $b$: bytes per element (4 for fp32, 2 for fp16/bf16, often 2 in modern practice)
• $L$: number of transformer layers (each layer adds another full set of attention matrices)

Notice that this formula has $n^2$ as its only non-linear term. Every other factor ($h$, $b$, $L$) is a constant that scales linearly with model size. The sequence length $n$ is unique in appearing squared, which is why attention memory grows so much faster than model size as context lengths increase.

Let's make this concrete with a real calculation. Consider a LLaMA 7B-style model with 32 attention heads, 32 layers, and fp16 precision ($b = 2$ bytes). At a sequence length of 8,192 tokens:

$$M_{\text{attn}} = 32 \times 8192^2 \times 2 \times 32$$

Working through the calculation step by step:

$$\begin{aligned} 8192^2 &= 67{,}108{,}864 \text{ entries per head per layer} \\ \times 2 \text{ bytes} &= 134{,}217{,}728 \text{ bytes per head per layer} \\ \times 32 \text{ heads} &= 4{,}294{,}967{,}296 \text{ bytes per layer} \\ \times 32 \text{ layers} &= 137{,}438{,}953{,}472 \text{ bytes} \approx 128 \text{ GB} \end{aligned}$$

This is before counting model weights, activations, gradients, or optimizer states. A LLaMA 7B model at 8,192 tokens already needs 128 GB for attention matrices alone, far exceeding the 80 GB capacity of an H100 GPU. This is the memory wall in concrete form.