Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Attention Complexity › Memory Requirements

Michael BrenndoerferUpdated May 29, 20256 min read

Language AI HandbookAll chapters

Part 13 of 67Self-Attention

Chapter 95Attention Complexity › Memory Requirements

Reading section 3 of 11 · Memory Requirements · 6 minSection 3 of 11 · 6 min

Section 3 of 11 · 6 min read

Attention Complexity › Memory RequirementsOne section · Sections in this chapter3 / 11 ↓

Sections in this chapter11

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Computation is only half the story. Memory consumption often becomes the limiting factor before compute does. Self-attention requires storing several large tensors, and the memory profile of these tensors follows the same quadratic pattern as the computation, with some additional complications that make memory management especially challenging in practice.

Think of memory pressure as a different kind of constraint than compute. With compute, the question is how long something takes. With memory, the question is whether it fits at all. A computation that is slow can still complete given enough time; a computation that exceeds available memory simply fails, often abruptly and without a clear error message. This binary nature of memory constraints makes them especially important to understand before designing or deploying attention-based systems.

During a forward pass, the key memory consumer is the attention weight matrix itself. Every element of this $n \times n$ matrix must be materialized in memory before the output can be computed in the standard implementation. For training, the situation is worse: the entire matrix must be kept in memory throughout the forward pass so that gradients can be computed during the backward pass. For each layer, this quadratic memory cost persists throughout the backward computation rather than appearing as a momentary spike.

The memory bottleneck is what motivated the development of FlashAttention, which we will discuss later in this chapter. The key insight behind FlashAttention is that the forward and backward passes can be restructured to process the attention matrix in small blocks that fit in the GPU's fast on-chip memory (called SRAM), rather than writing the full $n \times n$ matrix to the slower off-chip memory (HBM). This dramatically reduces memory transfers, even though it performs the same total number of floating-point operations.

Attention Matrix StorageLink Copied

The attention weight matrix has shape $n \times n$ , where $n$ is the sequence length. This is the dominant memory consumer in standard attention implementations. For a single attention head, this matrix contains $n^2$ elements. With $h$ heads, we store $h \cdot n^2$ values. Using 32-bit floats (4 bytes each), the memory requirement is:

M_{\text{attn}} = 4 \cdot h \cdot n^2 \text{ bytes}

where:

$M_{\text{attn}}$ : memory required for attention matrices
$h$ : number of attention heads
$n$ : sequence length
$4$ : bytes per element for 32-bit floats (fp32)

For 16-bit floats (common in modern training with fp16 or bf16), this halves to $2 \cdot h \cdot n^2$ bytes. The quadratic dependence on $n$ means that doubling the sequence length quadruples the memory requirement. This is the same quadratic relationship we saw for compute, but memory is a harder constraint: you can always wait longer for computation, but you cannot wait your way around a memory limit.

Total Memory BreakdownLink Copied

A complete forward pass through attention requires storing several tensors simultaneously. Understanding this breakdown helps you predict memory usage and design systems that fit within hardware constraints. A complete forward pass through attention requires storing:

Query, Key, Value matrices: Each is $n \times d$ , totaling $3nd$ elements. These scale linearly with sequence length.
Attention scores/weights: $n^2$ per head, $hn^2$ total. These scale quadratically and dominate for long sequences.
Output before projection: $n \times d$ elements. Linear in sequence length.
Intermediate activations for backpropagation: Often requires keeping the attention matrix and pre-softmax scores in memory for the duration of the backward pass.

During training, gradient computation doubles many of these requirements. The optimizer state, which stores momentum and variance estimates for every parameter in Adam-style optimization, adds further overhead that scales with model size rather than sequence length. The total memory footprint during training is commonly estimated as 12 to 20 bytes per parameter for the model and optimizer states, plus the activation memory (dominated by the attention matrices) which scales with batch size and sequence length.

In[9]:

Code

def attention_memory_gb(n, d, n_heads, n_layers=1, dtype_bytes=2):
    """
    Estimate memory for attention computation.

    Args:
        n: Sequence length
        d: Model dimension
        n_heads: Number of attention heads
        n_layers: Number of transformer layers
        dtype_bytes: 2 for fp16/bf16, 4 for fp32

    Returns:
        Memory in gigabytes
    """
    # Per layer memory
    # Q, K, V: 3 * n * d
    qkv_mem = 3 * n * d * dtype_bytes

    # Attention matrices (per head): n * n, total: n_heads * n * n
    attn_mem = n_heads * n * n * dtype_bytes

    # Output: n * d
    output_mem = n * d * dtype_bytes

    # For training, need to store attention weights for backward pass
    # This is often the dominant term
    layer_mem = qkv_mem + attn_mem + output_mem

    total_bytes = layer_mem * n_layers
    return total_bytes / (1024**3)


# Typical model configurations
configs = [
    {"name": "GPT-2 Small", "d": 768, "heads": 12, "layers": 12},
    {"name": "GPT-2 Medium", "d": 1024, "heads": 16, "layers": 24},
    {"name": "GPT-2 Large", "d": 1280, "heads": 20, "layers": 36},
    {"name": "13B Model", "d": 5120, "heads": 40, "layers": 40},
]

seq_lengths_mem = [512, 2048, 8192, 32768]

Out[10]:

Console

Memory Requirements for Attention (fp16, per sample):

Model              Seq Len  Memory (GB)
----------------------------------------
GPT-2 Small            512         0.11
GPT-2 Small          2,048         1.27
GPT-2 Small          8,192        18.56

GPT-2 Medium           512         0.28
GPT-2 Medium         2,048         3.38
GPT-2 Medium         8,192        49.50

The memory requirements reveal the quadratic scaling in action. For GPT-2 Small, moving from 512 to 2048 tokens increases memory by roughly 16x (from ~0.2GB to ~3GB). At 8192 tokens, even the smaller model requires substantial memory. For GPT-2 Medium with its additional layers and heads, the numbers are even larger. These estimates cover only attention matrices and exclude model weights, optimizer states, and gradient storage needed during training.

Out[11]:

Visualization

Heatmap showing memory in GB with sequence length on x-axis and model size on y-axis, darker colors indicating higher memory. — Memory requirements for attention matrices across different sequence lengths and model sizes. The quadratic scaling in sequence length creates a steep cliff where memory becomes prohibitive.

The heatmap shows that memory grows quadratically with sequence length but only linearly with model size. Doubling from 2048 to 4096 tokens has a larger impact than moving from GPT-2 Small to GPT-2 Large. For truly long sequences, even small models hit memory limits.