Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Context Length Challenges › Attention Memory Scaling

Michael BrenndoerferUpdated July 4, 20256 min read

Language AI HandbookAll chapters

Part 18 of 67Long Context

Chapter 126Context Length Challenges › Attention Memory Scaling

Reading section 3 of 11 · Attention Memory Scaling · 6 minSection 3 of 11 · 6 min

Section 3 of 11 · 6 min read

Context Length Challenges › Attention Memory ScalingOne 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

The attention mechanism computes pairwise interactions between all tokens. For each layer and each attention head, we must store an $n \times n$ matrix of attention weights, where every entry $(i, j)$ represents how much token $i$ attends to token $j$ . The key insight is that this pairwise structure is what makes attention so powerful: every token can directly influence every other token in a single forward pass. But that same all-pairs structure is precisely what makes scaling to long sequences so expensive. The total memory required to store these attention matrices is:

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

where:

$M_{\text{attn}}$ : total memory in bytes for all attention matrices
$h$ : number of attention heads (each head computes an independent attention pattern)
$L$ : number of transformer layers (each layer has its own set of attention matrices)
$n$ : sequence length in tokens
$n^2$ : the size of each attention matrix (one entry per query-key pair)
$b$ : bytes per element (2 for fp16/bf16, 4 for fp32)

The critical insight is that $n$ appears squared: doubling the sequence length quadruples the memory for attention matrices. This is the source of the quadratic memory scaling that limits context length.

But this formula only accounts for the attention weight matrices themselves. A complete picture of memory consumption includes:

Key-value cache (for inference): $2 \times L \times n \times d$
Activations (for training): Intermediate values needed for backpropagation
Gradients: Same size as activations

In[9]:

Code

def memory_breakdown_gb(n, d=4096, h=32, L=32, dtype_bytes=2, mode="training"):
    """
    Break down memory requirements for a transformer.

    Args:
        n: Sequence length
        d: Model dimension
        h: Number of attention heads
        L: Number of layers
        dtype_bytes: Bytes per element
        mode: "training" or "inference"

    Returns:
        Dictionary with memory breakdown in GB
    """
    bytes_to_gb = 1 / (1024**3)

    # Attention matrices: h * L * n^2 (but computed one layer at a time)
    # For training, we need to store all for backprop
    attn_matrices = h * L * n * n * dtype_bytes

    # KV cache: 2 (K and V) * L * n * d
    kv_cache = 2 * L * n * d * dtype_bytes

    # Activations per layer (rough estimate):
    # - Input embeddings: n * d
    # - Attention output: n * d
    # - FFN intermediate: n * 4d (assuming 4x expansion)
    # - Total per layer: ~6 * n * d
    activations_per_layer = 6 * n * d * dtype_bytes

    if mode == "training":
        # Store activations for all layers for backprop
        activations = L * activations_per_layer
        gradients = activations  # Same size as activations
        return {
            "attention_matrices": attn_matrices * bytes_to_gb,
            "activations": activations * bytes_to_gb,
            "gradients": gradients * bytes_to_gb,
            "total": (attn_matrices + activations + gradients) * bytes_to_gb,
        }
    else:  # inference
        return {
            "attention_matrices": 0,  # Computed on-the-fly with FlashAttention
            "kv_cache": kv_cache * bytes_to_gb,
            "total": kv_cache * bytes_to_gb,
        }


# Analyze for LLaMA 7B style model
sequence_lengths = [2048, 8192, 32768, 131072]

Out[10]:

Console

Memory Breakdown for LLaMA 7B Style Model (d=4096, h=32, L=32):

Training Mode:
   Seq Len   Attn Matrices     Activations    Gradients      Total
-----------------------------------------------------------------
        2K             8.0 GB          3.0 GB        3.0 GB     14.0 GB
        8K           128.0 GB         12.0 GB       12.0 GB    152.0 GB
       32K          2048.0 GB         48.0 GB       48.0 GB   2144.0 GB
      131K         32768.0 GB        192.0 GB      192.0 GB  33152.0 GB

Inference Mode:
   Seq Len        KV Cache
------------------------------
        2K             1.0 GB
        8K             4.0 GB
       32K            16.0 GB
      131K            64.0 GB

For training, the numbers quickly become untenable. A 32K sequence requires over 500 GB for attention matrices alone, far exceeding any single GPU. Even with gradient checkpointing (recomputing activations during backprop instead of storing them), the attention matrices remain the bottleneck.

For inference, the KV cache dominates. At 128K tokens, storing keys and values for all layers consumes over 60 GB. This explains why long-context inference requires careful memory management and often multiple GPUs.

Out[11]:

Visualization

Stacked area chart showing training memory components with attention matrices exploding quadratically. — Training memory breakdown showing attention matrices dominating at long sequences. The quadratic growth makes training beyond 32K impractical without specialized techniques.

Line plot showing linear growth of KV cache memory with sequence length. — Inference memory is dominated by the KV cache, which scales linearly with sequence length. This is more manageable but still limits context windows.

These visualizations reveal why long-context training and inference require fundamentally different approaches. Training hits the quadratic wall with attention matrices. Inference hits a linear wall with KV cache storage. Solutions like FlashAttention address the first problem by never materializing the full attention matrix. Solutions like KV cache compression and sliding window attention address the second.