The Challenge of Attention at Long Contexts

A 128K Token Context Walks Into Production

Your company has just licensed access to GPT-4o at the 128K context tier. The use case is compelling: process complete customer contracts, extract key terms, identify risk clauses, and summarize obligations. No more chunking, no more retrieval, no more missed context from wrong chunk boundaries. Feed in the whole contract, get the answer.

The first test runs go beautifully. Short contracts - 20 pages, 15K tokens - come back in three seconds with accurate extractions. You're congratulating yourself on the architecture decision.

Then QA hands you the first large contract: a 340-page master services agreement. 96,000 tokens. You send it in. Fifteen seconds pass. Then thirty. The response arrives - but it costs $8.50 for a single API call. Your planned throughput of 5,000 contracts per day would cost$42,500 daily in LLM API fees alone. The business case evaporates.

You spend the next week learning why long context is expensive, what's happening inside the model, and whether there's a smarter approach. This lesson is that investigation.

---

The O(n²) Problem - Why Attention Doesn't Scale

The Standard Attention Mechanism

In the standard transformer attention mechanism, each token attends to every other token in the sequence. For a sequence of length $n$:

1. Compute query matrix $Q \in \mathbb{R}^{n \times d_k}$
2. Compute key matrix $K \in \mathbb{R}^{n \times d_k}$
3. Compute value matrix $V \in \mathbb{R}^{n \times d_v}$
4. Compute attention scores: $A = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) \in \mathbb{R}^{n \times n}$
5. Compute output: $O = AV \in \mathbb{R}^{n \times d_v}$

The attention matrix $A$ has $n \times n$ entries. At $n = 128,000$ tokens, this matrix has 16.4 billion entries. In float32 that's 65.5 GB - for a single attention head in a single layer.

Modern models use multiple heads and many layers. Llama-3-8B has 32 attention heads per layer and 32 layers. The attention matrices would require 32 × 32 × 65.5 GB = 67 terabytes of memory just to store all attention matrices simultaneously. This is physically impossible.

Compute Scaling

Beyond memory, compute is also quadratic. The matrix multiplication $QK^T$ requires $O(n^2 \cdot d_k)$ floating-point operations. For $n = 128K$ and $d_k = 128$:

$\text{FLOPs} = 2 \times n^2 \times d_k = 2 \times (128,000)^2 \times 128 \approx 4.2 \times 10^{12} \text{ FLOPs}$

Per attention head. Per layer. A 32-layer, 32-head model would need 32 × 32 × 4.2 trillion = 4.3 quadrillion FLOPs just for the attention computation on a single 128K token sequence. At A100 peak throughput (312 TFLOPS), this is nearly 14 seconds just for attention, excluding everything else.

These numbers explain why long-context inference is expensive and why approximate attention algorithms (sparse attention, sliding window attention) are needed for training at long contexts.

Concrete Numbers at Common Context Lengths

Context Length	Attention Matrix Size (per head, F32)	FLOPs (per head, d_k=128)
4K tokens	64 MB	4.3B
16K tokens	1 GB	68B
32K tokens	4 GB	274B
128K tokens	65.5 GB	4.4T

At 4K tokens, you can hold all attention matrices in GPU memory. At 128K, you physically cannot - the standard algorithm would require compute clusters to process a single sequence.

---

The KV Cache - Why Long Context Is Expensive at Inference Time

During training, the quadratic cost is paid once per batch. During inference (autoregressive decoding), a different problem emerges: the KV cache.

How the KV Cache Works

In autoregressive generation, each new token is generated by attending to all previous tokens. To avoid recomputing K and V for every previous token at each generation step, they are cached in memory. This is the KV cache.

For each token generated, you need:

• Key tensors: n_layers × n_heads × head_dim floats per token
• Value tensors: n_layers × n_heads × head_dim floats per token

KV Cache Size Formula

$\text{KV Cache Size} = 2 \times n_{layers} \times n_{heads} \times d_{head} \times n_{tokens} \times \text{bytes\_per\_element}$

Where factor of 2 accounts for both keys and values.

For Llama-3-8B (32 layers, 8 KV heads with GQA, head_dim=128) in BF16 (2 bytes):

$\text{KV Cache} = 2 \times 32 \times 8 \times 128 \times n_{tokens} \times 2 \text{ bytes}$ $= 131,072 \times n_{tokens} \text{ bytes}$ $\approx 128 \text{ KB per token}$

For 128K tokens: $128 \text{ KB} \times 128,000 = 16.4 \text{ GB}$ just for the KV cache.

This is in addition to the model weights (16 GB for 8B model in BF16) and activations. Total GPU memory requirement for 128K context with Llama-3-8B: approximately 35+ GB, exceeding the 80GB A100 when you account for multi-GPU overhead and the model itself.

def compute_kv_cache_size(
    n_layers: int,
    n_kv_heads: int,    # Note: GQA reduces this vs n_attention_heads
    head_dim: int,
    n_tokens: int,
    bytes_per_element: int = 2,  # 2 for BF16, 4 for FP32
) -> dict:
    """
    Compute KV cache memory requirements.

    For models using GQA (Grouped Query Attention), n_kv_heads < n_attention_heads.
    Llama-3-8B: 32 attention heads, 8 KV heads (GQA ratio 4:1)
    Llama-3-70B: 64 attention heads, 8 KV heads (GQA ratio 8:1)
    """
    # Size per token
    bytes_per_token = 2 * n_layers * n_kv_heads * head_dim * bytes_per_element

    # Total KV cache size
    total_bytes = bytes_per_token * n_tokens

    return {
        "bytes_per_token": bytes_per_token,
        "kb_per_token": bytes_per_token / 1024,
        "total_gb": total_bytes / (1024**3),
        "n_tokens": n_tokens,
    }

# Llama-3-8B at different context lengths
model_config = {"n_layers": 32, "n_kv_heads": 8, "head_dim": 128}

for ctx_len in [4_096, 16_384, 32_768, 128_000]:
    result = compute_kv_cache_size(
        **model_config,
        n_tokens=ctx_len,
        bytes_per_element=2  # BF16
    )
    print(f"{ctx_len:>8,} tokens: {result['kb_per_token']:.1f} KB/token, "
          f"total KV cache = {result['total_gb']:.2f} GB")

# Output:
#    4,096 tokens: 0.1 KB/token, total KV cache = 0.12 GB
#   16,384 tokens: 0.1 KB/token, total KV cache = 0.50 GB
#   32,768 tokens: 0.1 KB/token, total KV cache = 1.00 GB
#  128,000 tokens: 0.1 KB/token, total KV cache = 3.91 GB

# Wait - 3.91 GB for Llama-3-8B. That's manageable!
# The KV cache is much smaller with GQA (8 KV heads vs 32 attention heads)
# Without GQA, a 32-head model would be 4× larger = ~15.6 GB

# Original Llama-2-7B (no GQA, 32 KV heads):
result_no_gqa = compute_kv_cache_size(
    n_layers=32, n_kv_heads=32, head_dim=128,
    n_tokens=128_000, bytes_per_element=2
)
print(f"\nLlama-2-7B at 128K (no GQA): {result_no_gqa['total_gb']:.2f} GB KV cache")
# Output: Llama-2-7B at 128K (no GQA): 15.63 GB KV cache

:::note GQA Dramatically Reduces KV Cache Grouped Query Attention (GQA), used in Llama-3 and most modern models, reduces the number of KV heads while keeping the number of query heads the same. Llama-3-8B uses 8 KV heads for 32 query heads (4:1 ratio). This reduces KV cache size 4× compared to Multi-Head Attention (MHA), making long contexts much more feasible. :::

---

Perplexity Degradation - Why Long Contexts Fail Without Extension

Even if you could fit the attention matrices in memory, models trained on short contexts fail at long ones. This is the perplexity degradation problem.

The Distribution Shift Problem

During training, a model trained on 4K context length sees position indices from 0 to 4,095. At inference with 10K context, it encounters position index 8,000 - a value it has never seen during training. The position encoding module has no learned representation for positions outside the training range.

This causes a dramatic increase in perplexity (the model becomes "confused") for tokens that appear beyond the training context length. The model may produce repetitive, incoherent, or hallucinated output when the relevant context is located at positions beyond its training horizon.

import torch
import numpy as np
import matplotlib.pyplot as plt

def simulate_perplexity_degradation(
    trained_max_len: int = 4096,
    eval_lengths: list[int] = None,
) -> list[float]:
    """
    Simulate the perplexity degradation pattern seen when evaluating
    a model at context lengths beyond its training maximum.

    This is a simplified model of the real phenomenon - actual degradation
    depends on the specific model and position encoding.
    """
    if eval_lengths is None:
        eval_lengths = [512, 1024, 2048, 4096, 5000, 6000, 8000, 12000, 16000]

    # Approximate model: perplexity is stable within training range,
    # then increases rapidly beyond it
    baseline_ppl = 7.5  # typical for a well-trained 7B model
    perplexities = []

    for length in eval_lengths:
        if length <= trained_max_len:
            # Within training distribution: near-baseline perplexity
            ppl = baseline_ppl + np.random.uniform(-0.3, 0.3)
        else:
            # Beyond training range: exponential degradation
            overshoot_ratio = length / trained_max_len
            # The exact shape depends on the position encoding type
            # Absolute encodings degrade more sharply than relative ones
            ppl = baseline_ppl * (1 + 0.5 * (overshoot_ratio - 1)**2)
        perplexities.append(ppl)

    return perplexities

# Real data from Peng et al. 2023 (Position Interpolation paper):
# Llama-7B (trained on 2K) evaluated at various lengths:
# 2K tokens: 7.4 PPL (baseline)
# 4K tokens: 8.5 PPL
# 8K tokens: 25.0 PPL
# 16K tokens: >1000 PPL (model effectively random)

The sharp degradation beyond the training length is the core problem that context window extension techniques (Lesson 03) must solve.

---

FlashAttention - IO-Aware Exact Attention

FlashAttention (Dao et al. 2022, "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness") is the key engineering contribution that makes long-context training and inference tractable.

What FlashAttention Is NOT

FlashAttention does not change the mathematical result of attention. It computes exactly the same output as standard attention. It does not approximate or skip any computation. It is purely an engineering optimization.

What FlashAttention IS

FlashAttention is a tiling algorithm that restructures the computation to be compatible with GPU memory hierarchy. Modern GPUs have fast on-chip SRAM (shared memory) and slower off-chip HBM (high bandwidth memory). Standard attention moves data between HBM and SRAM many times (once per attention computation). FlashAttention reorganizes the computation to minimize this data movement.

The IO-Awareness Insight

Standard attention requires materializing the full $n \times n$ attention matrix in HBM. This requires:

• Write Q, K, V matrices to HBM: $3nd$ elements
• Read Q, K to compute $QK^T$, write $S = QK^T$ to HBM: $n^2$ elements
• Read $S$ to compute softmax, write $P$: $n^2$ elements
• Read $P, V$ to compute $PV$, write output: $n^2 + nd$ elements

Total HBM accesses: $O(n^2)$ - dominated by reading/writing the $n \times n$ matrix.

FlashAttention processes the attention matrix in blocks that fit in SRAM. It computes the softmax incrementally (using the online softmax trick) without materializing the full matrix. The attention matrix is never written to HBM - only the final output is.

Total HBM accesses with FlashAttention: $O(n \cdot d)$ - linear in sequence length.

Standard Attention IO:
──────────────────────────────────────────────────────
SRAM (fast)     |  HBM (slow, ~10× slower than SRAM)
─────────────   |  ─────────────────────────────────────
Compute block   |  ← Read Q, K, V
Compute block   |  ← Read S (full n×n matrix) - BOTTLENECK
Compute block   |  → Write S
Compute block   |  ← Read P (full n×n matrix) - BOTTLENECK
Compute block   |  → Write P
Compute block   |  ← Read P, V
                |  → Write Output O

FlashAttention IO:
──────────────────────────────────────────────────────
SRAM (fast)     |  HBM (slow)
─────────────   |  ─────────────────────────────────────
Load Q block    |  ← Read Q block
Load K block    |  ← Read K block
Compute tiled   |  (n×n matrix stays in SRAM blocks)
attention       |
Write output    |  → Write Output O (only this!)

FlashAttention Performance Impact

Metric	Standard Attention	FlashAttention	FlashAttention-2
Memory (n=128K)	~65 GB per head	~2 GB per head	~1.5 GB per head
Speed vs standard	1×	2-4× faster	3-9× faster
Max context length	Limited by HBM	2-4× longer	3-5× longer
Result accuracy	Exact	Exact	Exact

FlashAttention doesn't eliminate the $O(n^2)$ compute complexity - it still must compute all $n^2$ attention scores. But it eliminates the memory bottleneck by keeping intermediate results in fast SRAM rather than writing them to slow HBM.

# Using FlashAttention via PyTorch 2.0's SDPA
import torch
import torch.nn.functional as F

def standard_attention(Q, K, V, scale=None):
    """Standard attention - materializes full attention matrix in memory."""
    if scale is None:
        scale = Q.shape[-1] ** -0.5
    # This materializes n×n matrix - memory bottleneck at long contexts
    scores = torch.matmul(Q, K.transpose(-2, -1)) * scale
    weights = F.softmax(scores, dim=-1)
    return torch.matmul(weights, V)

def flash_attention(Q, K, V, scale=None):
    """
    FlashAttention via PyTorch 2.0 scaled_dot_product_attention.

    PyTorch automatically selects FlashAttention kernel when:
    1. GPU supports it (Ampere A100, Hopper H100, Ada RTX 30xx/40xx)
    2. Inputs are in the right dtype (float16 or bfloat16)
    3. Sequence length is within supported range
    """
    # PyTorch 2.0+ uses FlashAttention kernel automatically
    with torch.backends.cuda.sdp_kernel(
        enable_flash=True,
        enable_math=False,
        enable_mem_efficient=False,
    ):
        output = F.scaled_dot_product_attention(Q, K, V, scale=scale)
    return output

# Memory comparison for n=32K tokens
batch_size = 1
n_heads = 32
seq_len = 32768
head_dim = 128
dtype = torch.bfloat16

Q = torch.randn(batch_size, n_heads, seq_len, head_dim, dtype=dtype, device="cuda")
K = torch.randn(batch_size, n_heads, seq_len, head_dim, dtype=dtype, device="cuda")
V = torch.randn(batch_size, n_heads, seq_len, head_dim, dtype=dtype, device="cuda")

# Standard attention would require storing seq_len × seq_len × n_heads matrix:
attention_matrix_gb = (seq_len * seq_len * n_heads * 2) / (1024**3)  # BF16 = 2 bytes
print(f"Standard attention matrix size: {attention_matrix_gb:.1f} GB")  # ~64 GB

# FlashAttention never materializes this - uses O(seq_len * head_dim) memory
print("FlashAttention: attention matrix never materialized in HBM")

# In practice, use with huggingface transformers:
# model = AutoModelForCausalLM.from_pretrained(
#     "meta-llama/Meta-Llama-3-8B",
#     attn_implementation="flash_attention_2",  # enables Flash Attention 2
#     torch_dtype=torch.bfloat16,
# )

FlashAttention-2 and FlashAttention-3

FlashAttention-2 (Dao 2023) improved upon the original by:

• Better parallelization across thread blocks (2× speedup on A100)
• Smarter work partitioning between Q and KV dimensions
• Reduced number of non-matrix-multiply operations

FlashAttention-3 (Shah et al. 2024) targeted NVIDIA H100 specifically:

• Exploits H100's asynchronous memory copy operations
• Overlaps computation and memory access
• Achieves ~75% of H100's theoretical peak on long sequences

---

Why FlashAttention Doesn't Solve Everything

FlashAttention solves the memory IO problem but doesn't solve:

1. $O(n^2)$ compute: You still must compute $n^2$ dot products. FlashAttention makes this faster per operation, but the count is unchanged.

2. KV cache growth: The KV cache grows linearly with sequence length regardless of the attention algorithm.

3. Positional encoding out-of-distribution: FlashAttention doesn't help models generalize to positions beyond their training maximum.

4. Lost-in-the-middle attention patterns: FlashAttention computes exactly the same attention scores as standard attention - if those scores concentrate on tokens at the boundaries, FlashAttention doesn't help.

These remaining problems require the techniques covered in the next lessons: RoPE scaling for position encoding, YaRN/LongRoPE for context extension, and structural improvements to attention patterns.

---

The Position Encoding Bottleneck

The reason context extension is hard boils down to position encoding. Standard absolute position encoding (APE) - as used in the original transformer - learns a separate embedding for each position index. Training at 2048 positions means positions 0-2047 have learned embeddings. Position 2049 has no learned embedding at all.

Different position encoding schemes have different extrapolation behaviors:

RoPE's degradation is "fixable" because it has a principled mathematical structure (frequency-based rotation) that can be systematically modified to extend the context range. This is why virtually all context extension research focuses on modifying RoPE. ALiBi has less systematic structure for extension. APE cannot be extended at all without retraining.

---

Production Cost Model

Understanding attention scaling translates directly to cost modeling:

def estimate_inference_cost(
    model_name: str,
    n_input_tokens: int,
    n_output_tokens: int,
    cost_per_million_input: float,  # in USD
    cost_per_million_output: float,
) -> dict:
    """
    Estimate inference cost for a single API call.

    Prices as of 2024 (approximate, check current pricing):
    GPT-4o:           $5/M input,   $15/M output
    GPT-4o-mini:      $0.15/M input, $0.6/M output
    Claude-3.5-Sonnet: $3/M input,  $15/M output
    Gemini-1.5-Pro:   $3.5/M input, $10.5/M output
    """
    input_cost = (n_input_tokens / 1_000_000) * cost_per_million_input
    output_cost = (n_output_tokens / 1_000_000) * cost_per_million_output
    total_cost = input_cost + output_cost

    return {
        "model": model_name,
        "input_tokens": n_input_tokens,
        "output_tokens": n_output_tokens,
        "input_cost_usd": round(input_cost, 4),
        "output_cost_usd": round(output_cost, 4),
        "total_cost_usd": round(total_cost, 4),
    }

# Compare cost for a 96K token document analysis (the contract use case)
configs = [
    ("GPT-4o (128K)", 96_000, 2_000, 5.0, 15.0),
    ("GPT-4o-mini (128K)", 96_000, 2_000, 0.15, 0.60),
    ("Claude-3.5-Sonnet", 96_000, 2_000, 3.0, 15.0),
    ("Gemini-1.5-Pro", 96_000, 2_000, 3.5, 10.5),
]

print("Cost per API call for 96K input tokens, 2K output:")
for name, inp, out, cin, cout in configs:
    result = estimate_inference_cost(name, inp, out, cin, cout)
    print(f"  {name:30s}: ${result['total_cost_usd']:.4f}")

# At 5000 calls/day:
print("\nDaily cost at 5000 calls/day:")
for name, inp, out, cin, cout in configs:
    result = estimate_inference_cost(name, inp, out, cin, cout)
    daily = result["total_cost_usd"] * 5000
    print(f"  {name:30s}: ${daily:,.2f}/day")

The cost model directly motivates the alternatives covered in Lessons 05 and 06: RAG (retrieve relevant chunks instead of stuffing everything in context) and context compression (reduce the number of tokens before passing to the expensive model).

---

Common Mistakes

:::danger Assuming that fitting in the context window means good performance A model with a 128K context window can process 128K tokens without error. But that doesn't mean it attends to all of them with equal reliability. Information placed in the middle of a very long context is systematically under-attended compared to information at the beginning or end (the "lost in the middle" problem). Fitting your content in the context window is necessary but not sufficient for good performance. :::

:::warning Confusing KV cache with attention computation The KV cache and the $O(n^2)$ attention computation are different bottlenecks. The KV cache is a memory cost at inference time - it grows linearly with context length but persists across generation steps. The $O(n^2)$ attention computation is a compute cost - it's paid once per forward pass and grows quadratically. FlashAttention reduces the attention compute cost; GQA reduces the KV cache memory cost. You need both to make long context practical. :::

:::tip Use GQA models for long context workloads Models that use Grouped Query Attention (GQA) - including Llama-3, Mistral, Phi-3 - have significantly smaller KV caches than models using standard Multi-Head Attention. For a 32-head model with 8 KV heads, the KV cache is 4× smaller. When choosing models for long context applications, prefer GQA models. :::

---

Interview Q&A

Q: Why is attention O(n²) in memory and compute, and why does this matter for long contexts?

A: In standard transformer attention, each of the $n$ tokens attends to every other token, producing an $n \times n$ attention matrix. This matrix requires $O(n^2)$ memory to store and $O(n^2 \cdot d)$ operations to compute. At $n = 128K$ tokens, the attention matrix for a single head is approximately 65 GB in float32. With 32 heads and 32 layers, storing all attention matrices simultaneously would require tens of terabytes. This is physically impossible, so naive attention algorithms cannot operate at long context lengths. The quadratic scaling also means that doubling the context length quadruples the compute cost - a fundamental economic barrier to long context.

Q: What is the KV cache and how does it grow with context length?

A: The KV cache stores the key and value tensors from all previous tokens during autoregressive generation, enabling each new token to attend to all previous tokens without recomputing their keys and values. The cache size is $2 \times n_{layers} \times n_{kv\_heads} \times d_{head} \times n_{tokens} \times \text{bytes}$. For Llama-3-8B (32 layers, 8 KV heads with GQA, head_dim=128) in BF16, this is approximately 128 KB per token. At 128K tokens, the KV cache requires about 16 GB - a significant portion of available GPU memory. GQA reduces the number of KV heads (and thus KV cache size) by a factor of 4-8× compared to standard multi-head attention.

Q: How does FlashAttention work and what does it not solve?

A: FlashAttention restructures the attention computation to minimize data movement between GPU HBM (slow, large) and SRAM (fast, small). Standard attention materializes the full $n \times n$ attention matrix in HBM - a massive memory allocation that requires $O(n^2)$ reads and writes. FlashAttention uses a tiling algorithm to compute attention in blocks that fit in SRAM, using the online softmax trick to produce the correct output without ever writing the full attention matrix to HBM. HBM accesses drop from $O(n^2)$ to $O(n \cdot d)$. However, FlashAttention does not reduce the $O(n^2)$ compute complexity - all attention scores are still computed. It also doesn't help with the KV cache, positional encoding out-of-distribution, or the lost-in-the-middle problem.

Q: Why do models fail catastrophically when given inputs longer than their training context?

A: Position encodings are learned or computed for specific position indices. A model trained on sequences up to 4K tokens has only seen position indices 0-4095. At inference with an 8K token input, tokens at positions 4096-8191 receive positional encodings that are either undefined (for absolute encodings) or have very different statistical properties from training (for relative encodings like RoPE). The model's attention patterns, which depend on positional information, become unreliable. Perplexity degrades sharply - often by 10-100× - beyond the training context length. For RoPE-based models, this manifests as attention weights becoming noisy and incoherent, since the model has no learned representation for the rotational angles corresponding to those large position indices.

Q: What is the per-token KV cache cost for a 70B model, and how does GQA affect it?

A: For Llama-3-70B (80 layers, 8 KV heads via GQA, head_dim=128) in BF16: per-token KV cache = $2 \times 80 \times 8 \times 128 \times 2 = 327,680$ bytes ≈ 320 KB per token. Without GQA (if it used 64 KV heads like 64 query heads), the cost would be $2 \times 80 \times 64 \times 128 \times 2 = 2,621,440$ bytes ≈ 2.5 MB per token - 8× more. At 128K context, GQA reduces KV cache from 320 GB to 40 GB for a 70B model. This is the difference between requiring a 16× A100 cluster and fitting on a more reasonable 4× A100 setup.

:::tip 🎮 Interactive Playground

Visualize this concept: Try the Attention Complexity & Long Context demo on the EngineersOfAI Playground - no code required.

:::