Towards Infinite Context: How LLMs Are Breaking the Context Limit

ChatGPT was launched in 2023 with 4k context length, fast forward to 2025 Gemini 3 has more than 1M context window. So how did we bridge this 250x gap? This post goes through the pieces that made it possible: position embeddings that generalize, attention mechanisms that scale, memory systems that reach beyond the context window, and the architectural changes that make near-infinite context workable. Long context is a complex problem that can be tackled from multiple perspectives. I've broken it down into five key areas, each addressing a different aspect of the challenge. Fair warning this one is gonna be long!!!

The Context Length Problem

Why is extending context hard? Three reasons, and they're independent of each other.

1. Quadratic Attention Complexity

Standard attention is $O(n^2)$ in sequence length, and here's why that matters. The attention mechanism computes a similarity score between every query token and every key token. With $n$ tokens you have $n$ queries, each compared against $n$ keys, which is $n \times n = n^2$ comparisons. Each comparison is a dot product over $d$ dimensions, so $n^2 \times d$ floating-point operations.

At 1M tokens, the numbers get ridiculous:

For a typical model with $d = 128$ per head and 32 heads, that's roughly $4 \times 10^{14}$ operations just for attention. Worse, you need to store the full $n \times n$ attention matrix in memory. At 1M tokens in FP16 that is about 2TB for a single attention matrix, and you have dozens of layers.

FlashAttention helps by never materializing the full matrix, but the computation is still $O(n^2)$. The quadratic cost is still there; the optimizations just make it hurt less.

2. Position Encoding Generalization

Models trained on 4K tokens don't automatically work on 100K tokens. Position encodings must generalize to unseen lengths.

During training, the model learns position embeddings only for the sequence lengths it has seen. A model trained on sequences up to 4,096 tokens has never encountered position 10,000; it's out-of-distribution. Feed it a 100K token sequence and things break down. Absolute position embeddings assign garbage values to positions they've never seen. Learned relative distances that worked perfectly for "distance 4,000" mean nothing at "distance 50,000". And attention scores become unstable because the model can't compute relative positions between tokens far outside its training range.

The fix is position encoding schemes that extrapolate by construction: mathematical functions (rotations in RoPE, linear biases in ALiBi) that are well-defined for any position, not just the ones seen during training. Without that, performance degrades on longer sequences.

3. Memory and Information Retrieval

Even with the computational cost solved, there's a retrieval problem. With 1M tokens, most of the context is irrelevant to any given query, and standard attention treats every token the same. Each query attends to all 1M keys, and the signal gets diluted by noise.

The problem is identifying which tokens matter. That takes attention patterns that prioritize recent or semantically important tokens over the rest, retrieval mechanisms that can locate the relevant chunks quickly, memory systems that compress or summarize the less-relevant history, and some form of selective attention that learns to ignore most of the context.

Without those, a 1M token context can be less useful than a shorter one. The model struggles to focus, and performance drops despite having more information available.

Part 1: Position Encoding for Length Generalization

The first breakthrough: making position embeddings that extrapolate beyond training length.

Absolute Position Embeddings (The Old Way)

Original transformers add learned position embeddings:

If you train with positions 0-2047, position 2048 is out-of-distribution garbage.

Rotary Position Embeddings (RoPE)

The idea that changed everything downstream. Instead of adding positions, RoPE rotates query and key vectors:

where $R_{\Theta,m}$ is a rotation matrix based on position $m$:

The attention score between positions $m$ and $n$ depends only on their relative distance:

Fig 1: RoPE Encoding Diagram

import torch
 
def precompute_rope_frequencies(dim: int, max_seq_len: int, base: float = 10000.0):
    """Precompute RoPE rotation frequencies"""
    # Compute theta values for each dimension pair
    freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
    
    # Compute position * frequency for all positions
    positions = torch.arange(max_seq_len)
    angles = torch.outer(positions, freqs)  # [seq_len, dim/2]
    
    # Return cos and sin
    return torch.cos(angles), torch.sin(angles)
 
def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
    """Apply rotary embeddings to queries or keys"""
    # x: [batch, seq, heads, dim]
    # Split into pairs for rotation
    x1, x2 = x[..., ::2], x[..., 1::2]
    
    # Apply rotation
    rotated = torch.stack([
        x1 * cos - x2 * sin,
        x1 * sin + x2 * cos
    ], dim=-1).flatten(-2)
    
    return rotated

Why this works: rotations are well-defined for any position. The model learns to interpret relative rotations, and relative rotations generalize beyond the training length.

RoPE Scaling: Extending Trained Models

Even RoPE has limits. Models trained on 4K struggle at 32K. There are a few ways to push it.

Position Interpolation (Linear Scaling)

Simply rescale positions to fit within training range:

A model trained on 4K tokens, when scaled to 32K, sees position 32000 as position 4000.

def scaled_rope(position: int, scale_factor: float):
    """Linear position interpolation"""
    return position / scale_factor

The cost is resolution. Compressed positions push nearby tokens together until they become hard to tell apart.

NTK-Aware Scaling

Instead of scaling positions, scale the frequency base:

where $\alpha$ is the scaling factor. This preserves high-frequency (local) information while extending low-frequency (global) range.

def ntk_scaled_rope(dim: int, max_seq_len: int, base: float = 10000.0, scale: float = 1.0):
    """NTK-aware RoPE scaling"""
    # Scale the base frequency
    base = base * (scale ** (dim / (dim - 2)))
    
    freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
    positions = torch.arange(max_seq_len)
    angles = torch.outer(positions, freqs)
    
    return torch.cos(angles), torch.sin(angles)

YaRN (Yet another RoPE extensioN)

The current best. Combines NTK scaling with attention temperature adjustment:

where temperature $t$ is tuned per-layer to restore attention entropy.

ALiBi: Attention with Linear Biases

Alternative to RoPE: add a linear bias based on distance directly to attention scores:

where $m$ is a head-specific slope.

def create_alibi_bias(num_heads: int, seq_len: int):
    """Create ALiBi attention bias matrix"""
    # Slopes: geometric sequence from 2^(-8/n) to 2^(-8)
    slopes = 2 ** (-8 * torch.arange(1, num_heads + 1) / num_heads)
    
    # Distance matrix
    positions = torch.arange(seq_len)
    distances = positions.unsqueeze(0) - positions.unsqueeze(1)  # [seq, seq]
    distances = distances.abs()
    
    # Bias: [num_heads, seq, seq]
    bias = -slopes.view(-1, 1, 1) * distances.unsqueeze(0)
    
    return bias

The upside is that it extrapolates naturally: a linear decay extends to any length. The downside is that linear decay is an assumption, and some heads would prefer non-monotonic attention patterns that ALiBi can't express.

Part 2: Efficient Attention for Long Context

Position encodings solve how to represent long contexts mathematically; RoPE and ALiBi let us extrapolate far beyond training length. But being able to represent a million tokens doesn't mean we can afford to compute attention over them. The quadratic cost that was fine at short lengths is what breaks at scale. I've discussed this in details in my other article "on solving quadratic complexity of attention".

Sliding Window Attention

In Nutshell: each token only attends to a local window of nearby tokens instead of the entire sequence. Think of it like reading a book through a small window. You can only see a few pages at a time, but you can slide the window to read the whole book.

The bet is that most of the information needed for the next token is local. A word usually depends on nearby words, not on tokens thousands of positions away. Restricting attention to a fixed window of size $w$ brings the cost from $O(n^2)$ down to $O(n \cdot w)$: linear in sequence length, constant in window size.

Each token at position $i$ attends only to tokens within its window:

The window slides along the sequence, giving each position a "local receptive field". With a 4K window, token 10,000 can directly attend to tokens 8,000-12,000, but not to token 1,000 or token 50,000. Information can still travel long distances, but it has to hop through layers, and each layer can move it by at most the window size.

Comparison of regular causal attention and sliding window attention patterns

def sliding_window_attention(Q, K, V, window_size: int):
    """Sliding window attention with O(n*w) complexity"""
    batch, seq_len, num_heads, head_dim = Q.shape
    
    outputs = []
    for i in range(seq_len):
        start = max(0, i - window_size // 2)
        end = min(seq_len, i + window_size // 2 + 1)
        
        q_i = Q[:, i:i+1, :, :]
        k_window = K[:, start:end, :, :]
        v_window = V[:, start:end, :, :]
        
        scores = torch.matmul(q_i, k_window.transpose(-2, -1)) / (head_dim ** 0.5)
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, v_window)
        outputs.append(out)
    
    return torch.cat(outputs, dim=1)

Mistral uses a 4K window; Gemma 2 alternates local and global layers.

The trade-off is that local attention is fast and memory-cheap, but long-range dependencies need depth. For a 1M token context with a 4K window, information needs roughly 250 layer hops to get from the beginning to the end. That's why Mistral uses pure sliding windows while Gemma 2 interleaves global attention layers to recover some long-range modeling.

Landmark Attention

Instead of attending to all tokens, insert special "landmark" tokens that summarize chunks of the sequence. Landmarks work like a table of contents: you check the summary to find which chapter is relevant, then read that section.

The mechanism is hierarchical attention. First route through the landmarks to identify relevant chunks, then retrieve detailed information from those chunks. Two stages, each much cheaper than full attention.

The sequence is structured with landmarks inserted between chunks:

Each landmark token $\texttt{<Li>}$ is created by pooling (typically mean pooling) the tokens in its chunk, giving a compressed representation. A query first attends to the landmarks, a small set of summary tokens, to work out which chunks are relevant. Then it attends within those chunks for the details.

The saving is large. Instead of comparing against all $n$ tokens, the query compares against $k$ landmarks (where $k \ll n$), then attends into only the selected chunks. For a 1M token sequence in 1000 chunks, the cost drops from $O(n^2)$ to roughly $O(k^2 + c \cdot w^2)$, where $c$ is the number of selected chunks and $w$ is the chunk size.

Comparison of standard attention and attention with landmarks

class LandmarkAttention(nn.Module):
    def __init__(self, d_model: int, chunk_size: int = 512):
        super().__init__()
        self.chunk_size = chunk_size
        self.landmark_proj = nn.Linear(d_model, d_model)
    
    def forward(self, x):
        batch, seq_len, d = x.shape
        
        # Create landmark tokens by pooling chunks
        num_chunks = seq_len // self.chunk_size
        chunks = x.view(batch, num_chunks, self.chunk_size, d)
        landmarks = chunks.mean(dim=2)  # [batch, num_chunks, d]
        landmarks = self.landmark_proj(landmarks)
        
        # Two-stage attention:
        # 1. Query attends to landmarks to find relevant chunks
        landmark_scores = torch.matmul(Q, landmarks.transpose(-2, -1))
        top_chunks = landmark_scores.topk(k=3, dim=-1).indices  # Select top-k chunks
        
        # 2. Query attends within selected chunks
        selected_chunks = chunks.gather(1, top_chunks.unsqueeze(-1).expand(-1, -1, -1, d))
        chunk_attention = F.scaled_dot_product_attention(Q, selected_chunks, selected_chunks)
        
        return chunk_attention

What you get is random access to any part of the context. You can jump straight to the relevant chunk without processing everything in between, which is exactly what question answering over a long document needs.

What you give up depends on how well the landmarks summarize their chunks. If a landmark misses something important, its chunk gets skipped even when it holds the answer. Chunk size matters too: too small and there are too many landmarks, too large and each landmark says too little.

Ring Attention (Distributed Long Context)

For contexts that genuinely don't fit on a single GPU, Ring Attention distributes the computation across devices. Attention needs queries (Q) to see keys and values (K, V), but nothing says they have to live on the same GPU. You can split the sequence across GPUs and pass K/V blocks around in a ring.

Here's how it works. Split a long sequence across $N$ GPUs. Each GPU holds its local chunk of queries, keys and values, but each query needs to attend to keys from all GPUs, not just its own. So each GPU computes attention with its local K/V, then passes that K/V block to the next GPU in the ring. After $N$ steps, every GPU has seen every K/V block.

The process:

1. Partition: Split the sequence across $N$ GPUs, so each GPU gets $\frac{n}{N}$ tokens
2. Local computation: Each GPU computes attention with its local Q and local K/V
3. Ring pass: Each GPU sends its K/V block to the next GPU and receives K/V from the previous GPU
4. Accumulate: Each GPU accumulates attention outputs across all $N$ steps

With 8 GPUs each handling 128K tokens you get a 1M token effective context. Memory per GPU stays constant no matter how long the total context is; you just do more communication rounds.

Ring Attention: Circular data flow of K/V blocks across multiple GPUs

def ring_attention_step(Q_local, K_local, V_local, rank: int, world_size: int):
    """One step of ring attention (simplified)"""
    # Each GPU has Q for its partition, but needs K,V from all partitions
    
    O_local = torch.zeros_like(Q_local)
    L_local = torch.zeros(Q_local.shape[:-1])  # log-sum-exp
    
    K_recv, V_recv = K_local, V_local
    
    for step in range(world_size):
        # Compute attention with current K, V block
        scores = Q_local @ K_recv.T / sqrt(d)
        
        # Online softmax update
        m_new = torch.maximum(L_local, scores.max(-1))
        O_local = O_local * exp(L_local - m_new) + softmax(scores) @ V_recv
        L_local = m_new
        
        # Ring: send K,V to next GPU, receive from previous
        K_recv = ring_send_recv(K_recv, rank, world_size)
        V_recv = ring_send_recv(V_recv, rank, world_size)
    
    return O_local / exp(L_local)

Attention Sinks (StreamingLLM)

A surprising discovery: LLMs put a lot of attention on the first few tokens, regardless of what those tokens are. These "attention sinks" turn out to matter for model stability, and understanding why says something about how transformers work.

The observation is that the first few tokens (often just the BOS token and maybe 2-3 more) consistently receive disproportionately high attention scores across all layers and heads. It isn't that these tokens are semantically important. It's that during training the model learned it needs somewhere to dump excess attention, and without that somewhere, attention distributions become unstable.

Why would that be? Softmax has to put its mass somewhere. When a query doesn't strongly match any key, it still has to distribute its attention, and the model learned to use the initial tokens as the default destination. That gives every attention distribution a stable baseline.

The Problem with Sliding Window

When you naively slide a window over long text, you eventually evict the initial tokens:

Window at t=0:    [BOS, tok1, tok2, tok3, tok4, ...]
Window at t=1000: [tok997, tok998, tok999, tok1000, ...] ← BOS is gone!

Without the sink tokens, attention scores become unstable and perplexity explodes. The model tries to redistribute the attention that used to go to the sinks and there's nowhere stable to put it. Attention becomes erratic, with tokens randomly getting very high or very low weight, and generation quality falls apart.

The Solution: Keep the Sinks

StreamingLLM's fix is simple: keep a small number of initial "sink" tokens permanently in the KV cache, then add a sliding window for recent tokens. That gives the model its stability anchors and access to recent context at the same time:

The sink tokens absorb excess attention and provide the baseline; the sliding window provides recent context. Together they allow unbounded streaming with stable perplexity, and you only need about 4 sink tokens no matter how long the sequence gets.

class StreamingLLMCache:
    def __init__(
        self, 
        num_sink_tokens: int = 4, 
        window_size: int = 1024,
        num_layers: int = 32
    ):
        self.num_sink = num_sink_tokens
        self.window_size = window_size
        
        # KV cache structure: sink tokens + sliding window
        self.sink_cache = None  # [layers, 2, batch, num_sink, head_dim]
        self.window_cache = None  # [layers, 2, batch, window_size, head_dim]
        self.window_pos = 0
    
    def update(self, new_k: torch.Tensor, new_v: torch.Tensor, layer: int):
        """Add new KV, maintaining sink + window structure"""
        if self.sink_cache is None:
            # First tokens become sink tokens
            self.sink_cache = ...
            return
        
        # Add to sliding window (circular buffer)
        pos = self.window_pos % self.window_size
        self.window_cache[layer, 0, :, pos, :] = new_k
        self.window_cache[layer, 1, :, pos, :] = new_v
        self.window_pos += 1
    
    def get_kv(self, layer: int):
        """Return sink + recent window for attention"""
        sink_k = self.sink_cache[layer, 0]
        sink_v = self.sink_cache[layer, 1]
        
        # Get window in correct order
        if self.window_pos < self.window_size:
            window_k = self.window_cache[layer, 0, :, :self.window_pos, :]
            window_v = self.window_cache[layer, 1, :, :self.window_pos, :]
        else:
            # Reorder circular buffer
            start = self.window_pos % self.window_size
            window_k = torch.cat([
                self.window_cache[layer, 0, :, start:, :],
                self.window_cache[layer, 0, :, :start, :]
            ], dim=1)
            window_v = torch.cat([
                self.window_cache[layer, 1, :, start:, :],
                self.window_cache[layer, 1, :, :start, :]
            ], dim=1)
        
        return (
            torch.cat([sink_k, window_k], dim=1),
            torch.cat([sink_v, window_v], dim=1)
        )

Four sink tokens plus a 1K window is enough for unbounded streaming with stable perplexity. Many production streaming systems are built on this, and it underlies a lot of efficient long-context serving.

H2O: Heavy-Hitter Oracle

Not all KV cache entries are equal. H2O identifies "heavy hitter" tokens, the ones that accumulate high attention scores, and keeps them while evicting the rest.

The Observation

In practice, attention follows a power law: a small fraction of tokens receive most of the attention. H2O exploits this:

H2O keeps a running total of attention received by each token in the KV cache. As new tokens arrive it updates those totals and notes which tokens keep getting attention across many queries. When the cache goes over budget, it evicts the tokens with low accumulated attention while preserving the recent window and the top heavy hitters. Because it tracks the actual attention pattern during generation, it adapts in a way that fixed heuristics can't.

class H2OCache:
    def __init__(self, window_size: int = 256, heavy_hitter_budget: int = 256):
        self.window_size = window_size
        self.hh_budget = heavy_hitter_budget
        
        self.kv_cache = None
        self.attention_accumulator = None  # Track cumulative attention per token
    
    def update(self, new_k, new_v, attention_scores):
        """Update cache with new KV and attention information"""
        # Accumulate attention scores for existing tokens
        if self.attention_accumulator is not None:
            self.attention_accumulator += attention_scores.sum(dim=1)  # Sum over queries
        
        # Add new token
        self.kv_cache = torch.cat([self.kv_cache, new_k], dim=1)
        self.attention_accumulator = torch.cat([
            self.attention_accumulator, 
            torch.zeros(1)
        ])
        
        # Evict if over budget
        if len(self.kv_cache) > self.window_size + self.hh_budget:
            self._evict()
    
    def _evict(self):
        """Keep recent window + top heavy hitters"""
        cache_len = self.kv_cache.shape[1]
        
        # Always keep recent window
        recent_start = cache_len - self.window_size
        
        # Find heavy hitters in older tokens
        old_attention = self.attention_accumulator[:recent_start]
        hh_indices = old_attention.topk(self.hh_budget).indices
        
        # Combine: heavy hitters + recent
        keep_indices = torch.cat([hh_indices, torch.arange(recent_start, cache_len)])
        
        self.kv_cache = self.kv_cache[:, keep_indices]
        self.attention_accumulator = self.attention_accumulator[keep_indices]

Heavy hitters tend to be semantically important tokens (subjects, key entities), so keeping them preserves most of what matters.

Scissorhands: Persistence-Based Eviction

Similar to H2O, but the criterion is persistence of importance: tokens that stay important across many steps are kept.

The intuition: a token attended to heavily at step 100 but ignored from step 101-500 is less important than one consistently attended to.

SnapKV: Observation Window + Pooling

SnapKV discovers that heavy hitters can be identified from just a small "observation window" at the end of the prompt, then compresses the rest:

1. Run attention on last $w$ tokens to identify important positions
2. Select top-k positions per head based on pooled attention
3. Keep only those positions in KV cache for generation

def snapkv_compress(K, V, attention_scores, num_keep: int = 512, window: int = 64):
    """Compress KV cache using SnapKV strategy"""
    seq_len = K.shape[1]
    
    # Use attention from last 'window' queries to identify important keys
    observation_attn = attention_scores[:, :, -window:, :]  # [batch, heads, window, seq]
    
    # Pool attention across observation window
    importance = observation_attn.mean(dim=2)  # [batch, heads, seq]
    
    # Select top-k per head (excluding observation window itself)
    prefix_importance = importance[:, :, :-window]
    top_indices = prefix_importance.topk(num_keep, dim=-1).indices
    
    # Gather compressed KV
    K_compressed = K.gather(1, top_indices.unsqueeze(-1).expand(-1, -1, -1, K.shape[-1]))
    V_compressed = V.gather(1, top_indices.unsqueeze(-1).expand(-1, -1, -1, V.shape[-1]))
    
    # Append observation window (always kept)
    K_out = torch.cat([K_compressed, K[:, -window:]], dim=1)
    V_out = torch.cat([V_compressed, V[:, -window:]], dim=1)
    
    return K_out, V_out

One attention pass is enough to decide what to keep, which makes it cheap at prompt-processing time.

PyramidKV: Layer-Wise Budget Allocation

Different layers need different KV cache sizes. Lower layers capture local patterns (need less cache), higher layers capture global semantics (need more).

where $\alpha > 1$ gives more budget to later layers.

def pyramid_budgets(num_layers: int, total_budget: int, alpha: float = 1.5):
    """Allocate KV cache budget per layer (pyramid shape)"""
    # Later layers get more budget
    raw_budgets = [alpha ** (num_layers - l - 1) for l in range(num_layers)]
    total_raw = sum(raw_budgets)
    
    # Normalize to total budget
    budgets = [int(b / total_raw * total_budget) for b in raw_budgets]
    
    return budgets  # e.g., [64, 96, 144, 216, 324, 486, ...] for 7 layers

FastGen: Adaptive KV Compression

FastGen profiles the attention patterns during prompt processing and picks a compression strategy per head. It looks at attention entropy and variance across heads to set per-head compression ratios: heads with low entropy (focused attention) can be compressed aggressively, heads with high entropy (spread-out attention) need more of their context kept.

On top of that it applies token-type rules. Special tokens are always kept, since they draw high attention. Punctuation is usually dropped. Content tokens are pruned selectively based on their attention profile.

The point is that compression ratios differ per attention head, because some heads are naturally compressible and others aren't. That lets FastGen reach better compression-quality trade-offs than fixed heuristics that treat every token and head the same.

Comparison of Streaming Methods

The direction of travel is from fixed heuristics (Attention Sinks) toward learned or adaptive selection (SnapKV, FastGen) that copes with varied workloads.

Infini-Attention (Google's 1M+ Method)

The technique behind Gemini's very long context. It combines local attention with a compressive memory that summarizes the history that would otherwise be thrown away.

The Core Idea

Instead of discarding old KV pairs, compress them into a fixed-size memory:

where $M$ is a compressive memory updated incrementally.

Infini-Attention processes the sequence in segments. For each segment it runs two attention operations side by side: standard causal attention over the current segment's K/V pairs, and linear-attention retrieval from the compressive memory, which holds compressed information from all previous segments. A learnable gate $\beta$ decides how much to trust local attention versus memory. You get the precision of standard attention for recent context and the efficiency of compressed memory for distant history.

Infini-Attention architecture: hybrid of local causal attention and compressive memory

Memory as Associative Binding

Infini-attention uses linear attention to maintain memory. What makes it work is that the memory can be updated incrementally by accumulating compressed KV information:

As each segment is processed, its K/V pairs are compressed and added to the memory matrix. The memory $M$ has fixed size $O(d^2)$ regardless of how many segments have gone by. It is a compressed summary of everything so far.

When retrieving information, queries attend to the memory using linear attention:

where $z = z_{old} + \sum \sigma(K)$ is a normalization term that tracks the total "mass" of keys added to memory, and $\sigma$ is a non-linearity (e.g., ELU + 1) that keeps attention weights positive. Because this is linear attention, retrieval avoids the quadratic cost while still letting queries pull relevant information out of the compressed memory.

class InfiniAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int, segment_len: int = 2048):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.segment_len = segment_len
        
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        
        # Learnable gate for combining local + memory attention
        self.beta = nn.Parameter(torch.zeros(num_heads))
        
    def forward(self, x, memory_state=None):
        batch, seq_len, _ = x.shape
        
        Q = self.W_q(x).view(batch, seq_len, self.num_heads, self.head_dim)
        K = self.W_k(x).view(batch, seq_len, self.num_heads, self.head_dim)
        V = self.W_v(x).view(batch, seq_len, self.num_heads, self.head_dim)
        
        # Initialize memory if needed
        if memory_state is None:
            M = torch.zeros(batch, self.num_heads, self.head_dim, self.head_dim)
            z = torch.zeros(batch, self.num_heads, self.head_dim)
        else:
            M, z = memory_state
        
        outputs = []
        
        # Process in segments
        for seg_start in range(0, seq_len, self.segment_len):
            seg_end = min(seg_start + self.segment_len, seq_len)
            
            Q_seg = Q[:, seg_start:seg_end]
            K_seg = K[:, seg_start:seg_end]
            V_seg = V[:, seg_start:seg_end]
            
            # === Local causal attention ===
            local_out = F.scaled_dot_product_attention(
                Q_seg.transpose(1, 2), 
                K_seg.transpose(1, 2), 
                V_seg.transpose(1, 2),
                is_causal=True
            ).transpose(1, 2)
            
            # === Memory retrieval ===
            # σ(Q) @ M / (σ(Q) @ z)
            Q_norm = F.elu(Q_seg) + 1  # [batch, seg_len, heads, head_dim]
            
            # Retrieve from memory
            mem_out = torch.einsum('bshd,bhde->bshe', Q_norm, M)
            normalizer = torch.einsum('bshd,bhd->bsh', Q_norm, z).unsqueeze(-1) + 1e-6
            mem_out = mem_out / normalizer
            
            # === Combine with learnable gate ===
            beta = torch.sigmoid(self.beta).view(1, 1, -1, 1)
            combined = local_out + beta * mem_out
            outputs.append(combined)
            
            # === Update memory with this segment ===
            K_norm = F.elu(K_seg) + 1
            # M += σ(K)^T @ V
            M = M + torch.einsum('bshd,bshe->bhde', K_norm, V_seg)
            # z += sum(σ(K))
            z = z + K_norm.sum(dim=1)
        
        output = torch.cat(outputs, dim=1)
        output = output.reshape(batch, seq_len, self.d_model)
        
        return self.W_o(output), (M, z)

Why It Scales to Millions of Tokens

The compressive memory $M$ has fixed size $O(d^2)$ regardless of context length. Processing 1M tokens only requires linear compute $O(n \cdot d^2)$.

In the paper, a 1B parameter model with Infini-attention passes 1M-token passkey retrieval, achieves a 114x compression ratio over the baseline memory, and matches full attention on BookSum, which means summarizing books of 500K+ tokens.

LongRoPE: 2M Context Extension

Microsoft's approach to extreme length extension via progressive interpolation:

1. Search for optimal RoPE rescaling factors per dimension
2. Progressive extension: 256K → 512K → 1M → 2M in stages
3. Readjust for short contexts: Prevent degradation on original lengths

def longrope_scaling(dim: int, target_len: int, original_len: int = 4096):
    """LongRoPE non-uniform scaling factors"""
    # Different dimensions get different scaling
    # Low-frequency (high dim indices) scale more aggressively
    
    scale_factor = target_len / original_len
    
    # Searched optimal factors (simplified)
    lambda_factors = torch.ones(dim // 2)
    
    # High-frequency dimensions (local info): minimal scaling
    lambda_factors[:dim//4] = 1.0
    
    # Low-frequency dimensions (global info): aggressive scaling  
    lambda_factors[dim//4:] = scale_factor ** 0.5
    
    return lambda_factors

Part 3: Memory Systems (Beyond Context Windows)

What if we stop pretending everything fits in context?

Retrieval-Augmented Generation (RAG)

RAG has become super common these days, so I'll skip the basics. The core idea is simple: instead of putting everything in context, retrieve only what's relevant when you need it.

1. Index: Embed documents into vector database
2. Retrieve: Find top-k relevant chunks for query
3. Generate: Use retrieved chunks as context

class RAGSystem:
    def __init__(self, embedding_model, vector_db, llm):
        self.embedder = embedding_model
        self.db = vector_db
        self.llm = llm
    
    def query(self, question: str, k: int = 5):
        # Embed the question
        q_embedding = self.embedder.encode(question)
        
        # Retrieve relevant chunks
        chunks = self.db.search(q_embedding, top_k=k)
        
        # Build context
        context = "\n\n".join([c.text for c in chunks])
        
        # Generate answer
        prompt = f"Context:\n{context}\n\nQuestion: {question}\nAnswer:"
        return self.llm.generate(prompt)

The effective context is unlimited, since a database can hold terabytes. The limit moves to retrieval quality instead, and complex queries that span several documents are still hard.

Memorizing Transformers

Add an explicit kNN memory to attention:

where $M$ is an external memory of past (key, value) pairs.

Memorizing Transformer architecture: combining local attention with kNN retrieval from external memory

class MemorizingAttention(nn.Module):
    def __init__(self, d_model: int, memory_size: int = 65536):
        super().__init__()
        self.memory_keys = torch.zeros(memory_size, d_model)
        self.memory_values = torch.zeros(memory_size, d_model)
        self.memory_ptr = 0
        self.memory_size = memory_size
    
    def forward(self, Q, K, V):
        # Standard local attention
        local_attn = F.scaled_dot_product_attention(Q, K, V)
        
        # kNN lookup in memory
        # Find top-k most similar memory keys for each query
        similarities = Q @ self.memory_keys.T  # [batch, seq, memory_size]
        top_k_idx = similarities.topk(k=32, dim=-1).indices
        
        # Gather memory values
        memory_v = self.memory_values[top_k_idx]
        memory_attn = F.softmax(similarities.gather(-1, top_k_idx), dim=-1)
        memory_out = (memory_attn.unsqueeze(-1) * memory_v).sum(-2)
        
        # Combine
        return local_attn + 0.1 * memory_out
    
    def update_memory(self, K, V):
        """Add current K, V to memory"""
        batch_size = K.shape[0] * K.shape[1]
        end_ptr = (self.memory_ptr + batch_size) % self.memory_size
        
        # FIFO update
        self.memory_keys[self.memory_ptr:end_ptr] = K.flatten(0, 1)
        self.memory_values[self.memory_ptr:end_ptr] = V.flatten(0, 1)
        self.memory_ptr = end_ptr

Part 4: Context Compression

Instead of extending context, compress it.

Gisting / Prompt Compression

Learn special "gist" tokens that summarize long contexts:

class GistCompressor(nn.Module):
    def __init__(self, llm, num_gist_tokens: int = 10):
        super().__init__()
        self.llm = llm
        self.gist_tokens = nn.Parameter(torch.randn(num_gist_tokens, llm.d_model))
    
    def compress(self, long_context_ids: torch.Tensor):
        # Encode the long context
        hidden_states = self.llm.encode(long_context_ids)
        
        # Cross-attention: gist tokens attend to context
        gist_repr = self.cross_attention(
            query=self.gist_tokens,
            key=hidden_states,
            value=hidden_states
        )
        
        return gist_repr  # [num_gist_tokens, d_model]

A related idea, AutoCompressor, trains models to recursively summarize their own context.

ICAE: In-Context Autoencoder

Train an encoder to compress context, decoder to expand when needed:

Achieves ~30x compression with minimal quality loss.

LongLLMLingua: Selective Pruning

Not all tokens matter equally. Prune unimportant ones:

1. Score each token by perplexity contribution
2. Keep tokens with high information content
3. Compress 10K tokens to 2K with minimal information loss

def compress_context(context_ids: torch.Tensor, model, target_ratio: float = 0.3):
    """Compress context by keeping only important tokens"""
    # Get token importance scores
    with torch.no_grad():
        outputs = model(context_ids, output_attentions=True)
        
    # Aggregate attention across layers and heads
    importance = torch.stack(outputs.attentions).mean(dim=(0, 2))  # [seq, seq]
    token_importance = importance.sum(dim=0)  # How much each token is attended to
    
    # Keep top tokens
    num_keep = int(len(context_ids) * target_ratio)
    keep_idx = token_importance.topk(num_keep).indices.sort().values
    
    return context_ids[keep_idx]

Part 5: Recurrent and Streaming Approaches

Making recurrent methods perform better than transformers is another fascinating direction. The goal is to process unbounded streams while keeping a fixed-size state, which is a different approach from anything attention-based. I'll write these architectures up properly soon, but for now, here are the basics.

RWKV: Linear RNN with Attention-like Expressivity

RWKV (Receptance Weighted Key Value) replaces attention with a linear recurrence that can express similar patterns. It keeps a running weighted sum of values, where the weights decay exponentially over time.

The recurrence has two components: a decay mechanism that controls how much history to remember, and a weighted aggregation that combines past values:

This decay term $w_t$ determines how much weight to give to tokens at different positions. Recent tokens get more weight, but older tokens aren't forgotten outright, just weighted less.

The output combines this decay with key-value attention:

Here, $r_t$ is a "receptance" gate that controls how much of the aggregated information to use, $k_i$ are keys that determine relevance, and $v_i$ are the values being aggregated. The exponential terms $\exp(w_i + k_i)$ create attention-like weights, but the recurrence structure allows efficient computation.

The useful property is that this can run as an RNN at inference, $O(1)$ per token (just update the running sums), while training in parallel like standard attention. Cheap streaming inference and fast parallel training from one formulation, with the decay mechanism learning what history to keep.

Mamba / State Space Models

Maintain a fixed-size hidden state that summarizes all history:

Effective context is unlimited, since the state compresses everything that came before. The cost is that a fixed-size state is a lossy compression, and some information gets forgotten.

Transformer-XL: Segment-Level Recurrence

Process context in segments, passing hidden states between segments:

Segment 1: [tokens 0-512]     → hidden_1
Segment 2: [tokens 512-1024]  → hidden_2 (conditioned on hidden_1)
Segment 3: [tokens 1024-1536] → hidden_3 (conditioned on hidden_2)

class TransformerXL(nn.Module):
    def __init__(self, d_model: int, segment_len: int = 512):
        super().__init__()
        self.segment_len = segment_len
        self.layers = nn.ModuleList([...])
    
    def forward_segment(self, x, memory=None):
        """Process one segment with memory from previous segment"""
        for layer in self.layers:
            if memory is not None:
                # Concatenate memory for extended context
                k = torch.cat([memory, x], dim=1)
                v = torch.cat([memory, x], dim=1)
            else:
                k, v = x, x
            
            x = layer(x, k, v)
        
        return x
    
    def forward_stream(self, token_stream):
        """Process infinite stream of tokens"""
        memory = None
        
        for segment in chunk(token_stream, self.segment_len):
            output = self.forward_segment(segment, memory)
            memory = output.detach()  # Detach to prevent infinite backprop
            yield output

Comparison: What Works When?

* Local attention quality, global degraded
** Depends on retrieval quality

Some Final Thoughts

128K native context is more or less solved: YaRN, FlashAttention and GQA together make it practical. 1M is achievable, and Gemini is the proof, with Infini-attention and Ring attention doing the work. Streaming has several good answers, from the simple (Attention Sinks) to the adaptive (SnapKV, H2O), and the choice comes down to your latency and quality budget. Underneath most of the streaming methods is the same question, which tokens to keep in the KV cache, so that is where the real bottleneck sits. Anything you could honestly call infinite context is going to be a hybrid: native context plus compressive memory plus retrieval. And quality still degrades with length. Even at 10M tokens, models struggle with needle-in-a-haystack tasks, so retrieval within the context is still an open problem.

References

1. Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., & Liu, Y. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv preprint. arXiv:2104.09864

2. Press, O., Smith, N. A., & Lewis, M. (2022). Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. ICLR. arXiv:2108.12409

3. Chen, S., Wong, S., Chen, L., & Tian, Y. (2023). Extending Context Window of Large Language Models via Positional Interpolation. arXiv preprint. arXiv:2306.15595

4. Peng, B., Quesnelle, J., Fan, H., & Shippole, E. (2023). YaRN: Efficient Context Window Extension of Large Language Models. arXiv preprint. arXiv:2309.00071

5. Liu, H., Yan, W., Zaharia, M., & Abbeel, P. (2023). Ring Attention with Blockwise Transformers for Near-Infinite Context. arXiv preprint. arXiv:2310.01889

6. Mohtashami, A., & Jaggi, M. (2023). Landmark Attention: Random-Access Infinite Context Length for Transformers. arXiv preprint. arXiv:2305.16300

7. Wu, Y., Rabe, M. N., Hutchins, D., & Szegedy, C. (2022). Memorizing Transformers. ICLR. arXiv:2203.08913

8. Mu, J., Li, X., & Goodman, N. D. (2023). Learning to Compress Prompts with Gist Tokens. NeurIPS. arXiv:2304.08467

9. Ge, T., Hu, J., Wang, X., Chen, S., & Wei, F. (2024). In-context Autoencoder for Context Compression in a Large Language Model. ICLR. arXiv:2307.06945

10. Jiang, H., Wu, Q., Lin, C. Y., Yang, Y., & Qiu, L. (2023). LongLLMLingua: Accelerating and Enhancing LLMs in Long Context Scenarios via Prompt Compression. arXiv preprint. arXiv:2310.06839

11. Peng, B., Alcaide, E., Anthony, Q., et al. (2023). RWKV: Reinventing RNNs for the Transformer Era. EMNLP. arXiv:2305.13048

12. Gu, A., & Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv preprint. arXiv:2312.00752

13. Dai, Z., Yang, Z., Yang, Y., Carbonell, J., Le, Q. V., & Salakhutdinov, R. (2019). Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context. ACL. arXiv:1901.02860

14. Reid, M., Savinov, N., Teber, D., et al. (2024). Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint. arXiv:2403.05530

15. Bertsch, A., Alon, U., Neubig, G., & Gormley, M. R. (2024). Unlimiformer: Long-Range Transformers with Unlimited Length Input. NeurIPS. arXiv:2305.01625

16. Xiao, G., Tian, Y., Chen, B., Han, S., & Lewis, M. (2023). Efficient Streaming Language Models with Attention Sinks. ICLR. arXiv:2309.17453

17. Munkhdalai, T., Faruqui, M., & Gopal, S. (2024). Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention. arXiv preprint. arXiv:2404.07143

18. Ding, Y., Zhang, L., Shang, J., Xu, J., et al. (2024). LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens. arXiv preprint. arXiv:2402.13753

19. Han, C., Wang, Q., Xiong, W., et al. (2024). LM-Infinite: Simple On-the-Fly Length Generalization for Large Language Models. NAACL. arXiv:2308.16137

20. Zhang, Z., Sheng, Y., Zhou, T., Chen, T., et al. (2024). H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models. NeurIPS. arXiv:2306.14048

21. Liu, Z., Desai, A., Liao, F., Wang, W., Xie, V., Xu, Z., Kyrillidis, A., & Shrivastava, A. (2023). Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time. NeurIPS. arXiv:2305.17118

22. Li, Y., He, Y., Sun, Y., Tan, Z., Yan, G., et al. (2024). SnapKV: LLM Knows What You are Looking for Before Generation. arXiv preprint. arXiv:2404.14469

23. Cai, Z., Zhang, Y., Gao, B., Liu, Y., et al. (2024). PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling. arXiv preprint. arXiv:2406.02069

24. Ge, S., Zhang, Y., Liu, L., Zhang, M., Han, J., & Gao, J. (2024). Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs. ICLR. arXiv:2310.01801

25. Sheng, Y., Zheng, L., Yuan, B., Li, Z., et al. (2023). FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU. ICML. arXiv:2303.06865

26. Hooper, C., Kim, S., Mohammadzadeh, H., et al. (2024). KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization. arXiv preprint. arXiv:2401.18079