Long Context LLMs: From 2K to 1M Tokens

The original Transformer handled 512 tokens. Gemini 1.5 Pro handles 1 million — a 2,000x increase achieved through a decade of fixing bottlenecks as they appeared. Transformer-XL introduced segment recurrence to break the 512-token barrier. RoPE enabled position encoding beyond training length. FlashAttention made dense attention fast enough for 32K contexts. Ring Attention distributed long sequences across GPUs for 128K-1M contexts. Each innovation unlocked new task categories: 4K isn’t enough for a research paper, 8K isn’t enough for a codebase, 32K isn’t enough for legal contracts with appendices. But long context has costs: quadratic memory scaling, attention dilution where important tokens drown in noise, and retrieval degradation as context grows. This post traces the full technical arc from 512 tokens to 1M+ and the tradeoffs at each step.

Why Long Context Matters

The Limitations of Short Context

The original GPT-2 had a context window of 1,024 tokens — roughly 750 words. GPT-3 extended this to 2,048 tokens. At these lengths, the model could handle a few paragraphs of conversation but nothing more. Any task requiring awareness of information spread across a longer document was out of reach.

Long Context vs. RAG

Retrieval-Augmented Generation (RAG) is often presented as an alternative to long context. Instead of feeding the entire document into the model, you retrieve the relevant chunks and include only those. This works well for factual lookup tasks, but it has fundamental limitations.

RAG requires knowing which chunks are relevant before generating the answer. For questions that depend on synthesizing information from multiple sections of a document — “How does the conclusion contradict the methodology?” or “Summarize the key themes across all chapters” — retrieval often fails because no single chunk contains the answer.

Long context eliminates this retrieval step entirely. The model sees everything and can attend to any part of the input. The tradeoff is computational cost: processing 128K tokens is far more expensive than processing 4K tokens of retrieved context.

Use Cases Enabled by Long Context

Document processing: Legal contracts, medical records, and financial reports often exceed 50K tokens. Long context allows end-to-end processing without chunking.

Code understanding: A medium-sized software project might have 100K-500K tokens of source code. With sufficient context, the model can understand cross-file dependencies, API contracts, and architectural patterns.

Multi-turn conversation: Long conversations accumulate context over time. A 128K context window can hold roughly 100 pages of conversation history, enabling the model to reference earlier parts of the discussion.

Few-shot learning with many examples: More context means more in-context examples, which directly improves task performance for pattern matching and classification.

The Quadratic Problem

The fundamental challenge with extending context length is the self-attention mechanism. Standard self-attention computes a score between every pair of tokens, giving it $O(n^2)$ time and memory complexity where $n$ is the sequence length.

This quadratic scaling means that naively doubling the context length quadruples the computation. Going from 2K to 1M tokens represents a 250,000x increase in attention FLOPs. Without architectural innovations to reduce this cost, long context would be completely impractical.

Historical Foundation: Transformer-XL Segment Recurrence

Transformer-XL (Dai et al., 2019) was among the first architectures to seriously tackle the context length limitation. Its key insight was that you do not need to process the entire sequence at once — you can process it in segments and carry information forward through a recurrence mechanism.

How Segment Recurrence Works

Instead of treating each segment independently (which loses all cross-segment information), Transformer-XL caches the hidden states from the previous segment and makes them available to the current segment during attention computation:

# Conceptual Transformer-XL forward pass
def transformer_xl_forward(current_segment, cached_states):
    """
    current_segment: [batch, seg_len, d_model]
    cached_states: [batch, mem_len, d_model] from previous segment
    """
    # Concatenate cached states with current segment for key/value
    extended_context = torch.cat([cached_states, current_segment], dim=1)

    # Query comes from current segment only
    Q = W_q(current_segment)          # [batch, seg_len, d_model]
    K = W_k(extended_context)          # [batch, seg_len + mem_len, d_model]
    V = W_v(extended_context)          # [batch, seg_len + mem_len, d_model]

    # Attention spans both current and cached segments
    attention = softmax(Q @ K.T / sqrt(d)) @ V

    # Cache current hidden states for next segment
    new_cache = current_segment.detach()  # Stop gradients through cache

    return attention, new_cache

The .detach() call is critical: gradients do not flow through the cached states, which keeps training memory manageable. The model learns to produce hidden states that are useful when attended to by future segments, even though it never directly optimizes for this objective through backpropagation across segments.

Relative Positional Encoding

Segment recurrence creates a problem for positional encoding. If you use absolute positions (position 0, 1, 2, …, 511 in each segment), then position 0 in the current segment and position 0 in the cached segment have the same encoding, even though they are semantically at very different positions in the overall sequence.

Transformer-XL solved this with relative positional encoding, which encodes the distance between tokens rather than their absolute position. This was a precursor to the rotary positional encoding (RoPE) used in most modern LLMs.

Limitations of Segment Recurrence

Despite its elegance, segment recurrence had practical drawbacks. The effective context grew linearly with the number of layers (each layer could “see” one more cached segment back), meaning deep networks were needed for long-range dependencies. The information from distant segments was also progressively degraded as it passed through multiple layers of processing. Modern approaches achieve much longer contexts more directly.

Sliding Window Attention

Sliding window attention, popularized by Mistral 7B (2023), takes a different approach to the quadratic problem. Instead of attending to all previous tokens, each layer only attends to a fixed window of the most recent $W$ tokens.

How It Works

def sliding_window_attention(Q, K, V, window_size=4096):
    """
    Each token attends only to the previous window_size tokens.
    """
    seq_len = Q.shape[1]

    # Create sliding window mask
    mask = torch.ones(seq_len, seq_len, dtype=torch.bool)
    for i in range(seq_len):
        # Token i can attend to tokens max(0, i-window_size+1) through i
        start = max(0, i - window_size + 1)
        mask[i, :start] = False
        mask[i, i+1:] = False  # Causal: can't attend to future

    scores = Q @ K.transpose(-1, -2) / math.sqrt(d_k)
    scores.masked_fill_(~mask, float('-inf'))

    return softmax(scores) @ V

With a window size of $W$, each token computes attention over at most $W$ key-value pairs instead of all $n$ pairs. This reduces the per-token cost from $O(n)$ to $O(W)$, making the total cost $O(n \cdot W)$ — linear in sequence length when $W$ is fixed.

Effective Context Through Layer Stacking

The key insight is that while each layer has a limited window, stacking $L$ layers creates an effective receptive field of $L \times W$ tokens. In Mistral 7B with $W = 4096$ and $L = 32$ layers, the theoretical receptive field is $32 \times 4096 = 131,072$ tokens, even though each layer only looks at 4,096 tokens.

Memory Savings

The most immediate benefit of sliding window attention is memory savings during inference. Instead of storing KV cache entries for all previous tokens (which grows linearly with sequence length), you only need to store the most recent $W$ entries. For Mistral 7B with a 4,096 window, the KV cache is capped at about 1 GB regardless of how long the sequence is, compared to 32 GB for a 128K full-attention model.

RoPE and Context Length Extension

Rotary Position Embedding (RoPE), introduced by Su et al. (2021), has become the standard positional encoding in modern LLMs (Llama, Mistral, Qwen, and many others). RoPE encodes position by rotating the query and key vectors in 2D subspaces at frequencies that depend on the position.

How RoPE Works

For a token at position $m$, RoPE applies a rotation to each pair of dimensions $(2i, 2i+1)$ in the query and key vectors:

$$\text{RoPE}(x, m) = \begin{pmatrix} x_{2i} \cos(m\theta_i) - x_{2i+1} \sin(m\theta_i) \\ x_{2i} \sin(m\theta_i) + x_{2i+1} \cos(m\theta_i) \end{pmatrix}$$

where $\theta_i = 10000^{-2i/d}$ defines the rotation frequency for dimension pair $i$. Low-frequency dimensions rotate slowly and encode coarse positional information; high-frequency dimensions rotate quickly and encode fine-grained position.

The inner product between two RoPE-encoded vectors depends only on the relative distance between them, making RoPE a form of relative positional encoding. This is what makes context extension possible.

RoPE Scaling for Context Extension

A model trained with RoPE at context length $L_{\text{train}}$ has only seen rotation angles up to $L_{\text{train}} \times \theta_i$. When you try to use it at length $L > L_{\text{train}}$, the rotation angles exceed what was seen during training, causing performance to degrade.

Several techniques extend context by modifying the RoPE frequencies:

Position Interpolation (PI): Scale all positions by $L_{\text{train}} / L_{\text{target}}$ so that the maximum rotation angle stays within the training range. A model trained at 4K and extended to 32K would divide all positions by 8.

def position_interpolation(position, base_context=4096, target_context=32768):
    scale = base_context / target_context
    return position * scale  # Compress positions into original range

NTK-Aware Scaling: Instead of scaling all frequencies uniformly, modify the base frequency $\theta$ to spread the rotation angles more evenly. This preserves high-frequency (local) positional information while extending low-frequency (global) range.

def ntk_aware_rope(position, dim, base=10000, scale=4):
    """NTK-aware RoPE scaling."""
    # Increase the base frequency
    new_base = base * (scale ** (dim / (dim - 2)))
    theta = 1.0 / (new_base ** (torch.arange(0, dim, 2).float() / dim))
    return position * theta

YaRN (Yet another RoPE extensioN): Combines NTK scaling with an attention scaling factor and a temperature adjustment. YaRN further refines the approach by treating different frequency bands differently — high frequencies are interpolated, low frequencies are extrapolated, and middle frequencies are blended.

Ring Attention for Distributed Long Sequences

When context lengths reach hundreds of thousands or millions of tokens, a single GPU cannot hold the KV cache in memory. Ring attention (Liu et al., 2023) distributes the sequence across multiple devices and computes attention in a communication-efficient ring topology.

The Core Idea

Split the sequence into chunks and distribute them across $P$ devices arranged in a ring. Each device holds one chunk of queries and one chunk of key-value pairs. Attention is computed in $P$ rounds: in each round, every device computes the partial attention between its local queries and the current key-value chunk, then passes the key-value chunk to the next device in the ring.

def ring_attention(Q_local, K_local, V_local, num_devices, device_id):
    """
    Each device holds Q_local and starts with its own K_local, V_local.
    After P rounds, each device has computed full attention for its queries.
    """
    seq_chunk = Q_local.shape[1]  # tokens per device

    # Initialize accumulators for online softmax
    running_max = torch.full((Q_local.shape[0], seq_chunk, 1), float('-inf'))
    running_sum = torch.zeros_like(running_max)
    running_output = torch.zeros_like(Q_local)

    K_current, V_current = K_local, V_local

    for round_idx in range(num_devices):
        # Compute local attention scores
        scores = Q_local @ K_current.transpose(-1, -2) / math.sqrt(d_k)

        # Apply causal mask (only for appropriate chunks)
        if needs_causal_mask(device_id, round_idx, num_devices):
            scores = apply_causal_mask(scores)

        # Online softmax update (numerically stable accumulation)
        local_max = scores.max(dim=-1, keepdim=True).values
        new_max = torch.maximum(running_max, local_max)

        # Rescale previous accumulator
        correction = torch.exp(running_max - new_max)
        running_output = running_output * correction
        running_sum = running_sum * correction

        # Add current block's contribution
        local_exp = torch.exp(scores - new_max)
        running_output += local_exp @ V_current
        running_sum += local_exp.sum(dim=-1, keepdim=True)
        running_max = new_max

        # Ring communication: send K,V to next device, receive from previous
        K_current = ring_send_recv(K_current)
        V_current = ring_send_recv(V_current)

    return running_output / running_sum

Why Ring Topology Is Efficient

The key property of ring attention is that computation and communication overlap. While a device is computing attention with the current K,V block, it is simultaneously sending/receiving the next K,V block. If the computation time exceeds the communication time (which is typically true for large enough chunks), the communication cost is completely hidden.

Chunked Prefill for Long Prompts

When a user submits a long prompt (say, a 100K-token document), the model must process the entire prompt before generating the first output token. This prefill phase is compute-bound and can take many seconds. Chunked prefill splits the prompt into smaller chunks and processes them sequentially, which provides several benefits.

Why Chunk the Prefill?

Memory management: Processing 128K tokens at once requires enormous temporary memory for the attention computation. Chunking reduces peak memory usage.

Scheduling flexibility: In a serving system handling multiple requests, a 128K-token prefill would monopolize the GPU for seconds. Chunking allows the scheduler to interleave prefill chunks from one request with decode steps from other requests, improving overall throughput and latency.

Pipeline efficiency: Chunks can be distributed across pipeline stages more evenly than a single monolithic prefill.

def chunked_prefill(prompt_tokens, model, chunk_size=8192):
    """
    Process a long prompt in chunks, building up the KV cache incrementally.
    """
    kv_cache = None
    num_chunks = (len(prompt_tokens) + chunk_size - 1) // chunk_size

    for i in range(num_chunks):
        start = i * chunk_size
        end = min(start + chunk_size, len(prompt_tokens))
        chunk = prompt_tokens[start:end]

        # Process chunk, attending to all previous KV cache entries + current chunk
        hidden, new_kv = model.forward(
            chunk,
            kv_cache=kv_cache,
            is_prefill=True
        )

        # Append new KV entries to cache
        if kv_cache is None:
            kv_cache = new_kv
        else:
            kv_cache = concatenate_kv(kv_cache, new_kv)

    return kv_cache  # Ready for autoregressive decoding

Memory Requirements at Scale

Understanding memory consumption is critical for capacity planning. The KV cache is the dominant memory consumer for long-context inference.

KV Cache Calculation

For a model with $L$ layers, $H$ attention heads, head dimension $d_h$, and sequence length $n$, the KV cache size is:

$$\text{KV Cache} = 2 \times L \times H \times d_h \times n \times \text{bytes\_per\_element}$$

The factor of 2 accounts for both keys and values.

Total Memory Budget

For a complete picture, the total GPU memory needed for inference includes model weights, KV cache, and activation memory:

How Gemini Reaches 1M Tokens

Google’s Gemini 1.5 Pro was the first production model to support 1 million tokens of context. While Google has not published full architectural details, the system likely combines several techniques:

Sparse/efficient attention: Some form of attention sparsity (possibly learned or structured) to reduce the quadratic cost. This might include a mix of local and global attention patterns.

Ring attention across TPU pods: Google’s TPU v5 pods provide massive interconnect bandwidth (up to 4.8 Tbps per chip via ICI). Ring attention across a TPU pod can distribute the 1M-token KV cache with minimal communication overhead.

Multi-query or grouped-query attention: Reducing KV heads is essential at 1M context. Even with 8 KV heads and 128-dim heads, a 100-layer model would need $2 \times 100 \times 8 \times 128 \times 10^6 \times 2 = 409$ GB for the KV cache alone.

Quantized KV cache: Quantizing the KV cache to INT8 or INT4 halves or quarters the memory requirement without significant quality loss.

Hierarchical attention: Some evidence suggests Gemini uses a form of hierarchical attention where early layers process local context and later layers have access to broader context, reducing the average attention span.

How Llama 3.1 Reaches 128K Tokens

Meta’s Llama 3.1 paper provides more detail about their long-context approach:

Training strategy: Llama 3.1 was trained in stages. The initial pretraining used 8K context. Then the model was continued-pretrained with progressively longer sequences: 16K, 32K, 64K, and finally 128K. This staged approach is more stable than training at 128K from the start.

RoPE frequency adjustment: Llama 3.1 uses a modified RoPE with an increased base frequency ($\theta = 500,000$ instead of the original $10,000$). This larger base allows the rotation angles to stay within a well-behaved range at longer positions.

GQA with 8 KV heads: All Llama 3.1 models use 8 KV heads regardless of model size, keeping the KV cache manageable at 128K context.

Long-context data: A significant fraction of the continued-pretraining data consisted of long documents (books, code repositories, long articles) to teach the model to actually use long context effectively.

When Long Context Hurts

Long context is not universally beneficial. There are several failure modes and costs that practitioners must understand.

Cost Scaling

The most obvious cost is computational. Processing 128K tokens costs roughly $128K/4K = 32$x more than processing 4K tokens for the prefill phase. For API-based services, this translates directly to cost:

Attention Dilution

As context length increases, the attention distribution becomes more spread out. Each token has more keys to attend to, which can dilute the attention weights on the most relevant tokens. This manifests as:

• Decreased retrieval accuracy: The “needle in a haystack” task (finding a specific piece of information in a long context) becomes harder as the haystack grows.
• Lost in the middle: Several studies have shown that LLMs are better at using information at the beginning and end of the context than in the middle. This “U-shaped” attention pattern means that simply having long context does not guarantee the model will use all of it effectively.
• Decreased generation quality: For tasks like summarization, very long input can lead to less focused, more generic outputs as the model tries to attend to too much information simultaneously.

Latency Impact

Long context dramatically increases time-to-first-token (TTFT), which is often the most important latency metric for interactive applications:

When to Avoid Long Context

Flash Attention: The Enabler

No discussion of long context is complete without Flash Attention (Dao et al., 2022). While Flash Attention does not change the asymptotic complexity of attention ($O(n^2)$ remains), it reduces the constant factor dramatically by eliminating the need to materialize the full $n \times n$ attention matrix in GPU HBM.

How Flash Attention Helps Long Context

Standard attention computes the full $n \times n$ score matrix, stores it in HBM, applies softmax, then multiplies by values. For $n = 128K$, this matrix alone requires $128K \times 128K \times 2$ bytes $= 32$ GB in FP16 — more than the memory of most GPUs.

Flash Attention tiles the computation so that only small blocks of the attention matrix exist at any time, stored in fast SRAM (shared memory) rather than slow HBM. This reduces memory usage from $O(n^2)$ to $O(n)$ and also improves speed by 2-4x due to reduced HBM traffic.

Without Flash Attention, 128K context would be impossible on current hardware. With it, the bottleneck shifts from attention computation memory to KV cache storage.

Putting It All Together: A Modern Long-Context Stack

A production 128K-token serving system in 2025 typically combines:

1. RoPE with extended base frequency for positional encoding that generalizes beyond training length
2. GQA with 8 KV heads to keep the KV cache manageable
3. Flash Attention v2 for memory-efficient, fast attention computation
4. Paged KV cache (vLLM-style) for flexible memory management across requests
5. Chunked prefill for scheduling efficiency in multi-request serving
6. KV cache quantization (INT8 or FP8) to further reduce memory footprint
7. Tensor parallelism across multiple GPUs for larger models

For 1M-token systems, add:

8. Ring attention or sequence parallelism for distributing the sequence across devices
9. Sparse or hierarchical attention to reduce the effective quadratic cost
10. Specialized hardware (TPU v5, Grace Hopper) with high-bandwidth interconnects

Conclusion

The journey from 2K to 1M tokens has been driven by a stack of innovations, each addressing a different bottleneck:

• Positional encoding (Transformer-XL relative PE to RoPE to RoPE scaling) solved the problem of representing position at arbitrary lengths.
• Attention efficiency (Flash Attention, sliding window) solved the memory and compute cost of the attention operation itself.
• KV cache reduction (GQA, quantization, paged allocation) solved the memory cost of storing past context.
• Distribution (ring attention, sequence parallelism) solved the problem of fitting long sequences across multiple devices.
• Serving (chunked prefill, continuous batching) solved the problem of efficiently serving long-context requests alongside short ones.

The frontier continues to push forward. Context lengths of 10M tokens are being explored in research, and new architectures (state-space models, linear attention variants) promise to break the quadratic barrier entirely. But for now, the Transformer with these accumulated optimizations remains the dominant architecture, capable of processing context lengths that would have seemed impossible just a few years ago.

The practical lesson is clear: long context is a powerful capability, but it is not free. Every doubling of context length doubles the cost and increases latency. The best systems use long context surgically — for tasks that genuinely benefit from it — while falling back to cheaper approaches (RAG, summarization, truncation) when the task does not require global understanding over the full input.