How LLMs Actually Process Context Windows: A Deep Dive

The transformer architecture introduced in "Attention Is All You Need" gave us something remarkable: a model that can attend to any token in its input regardless of position. But "attending to any token" is not the same as "equally remembering any token." The difference between those two things is responsible for most of the surprising behaviors we see in production LLM deployments.

What attention actually computes

At its core, self-attention computes a weighted sum over value vectors, where the weights come from comparing query vectors against key vectors. Written out:

import torch
import math

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q, K, V: (batch, heads, seq_len, d_k)
    """
    d_k = Q.shape[-1]

    # Similarity scores
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)

    # Causal mask: prevent attending to future tokens
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    # Softmax over the key dimension
    weights = torch.softmax(scores, dim=-1)

    return torch.matmul(weights, V), weights

The sqrt(d_k) scaling is critical. Without it, for large d_k, the dot products grow large in magnitude, pushing softmax into regions with vanishingly small gradients. This was a known problem in early transformer implementations that caused training instability.

The crucial thing to notice: attention weights are computed over all tokens in the context, every time, for every layer. This is the O(n²) complexity that makes long contexts expensive.

The KV cache: why inference is different from training

During training, you process all tokens in a sequence simultaneously (with the causal mask preventing future-token leakage). During autoregressive inference — where you generate one token at a time — something more efficient is possible: you only need to compute Q, K, V for the new token, then attend against all previously computed K/V pairs. This is the KV cache.

class KVCache:
    def __init__(self, max_batch, max_seq, n_heads, d_head, dtype=torch.float16):
        self.k = torch.zeros(max_batch, n_heads, max_seq, d_head, dtype=dtype)
        self.v = torch.zeros(max_batch, n_heads, max_seq, d_head, dtype=dtype)
        self.pos = 0

    def update(self, k_new, v_new):
        """Append new key/value pair, return full cache."""
        seq_len = k_new.shape[2]
        self.k[:, :, self.pos:self.pos + seq_len] = k_new
        self.v[:, :, self.pos:self.pos + seq_len] = v_new
        self.pos += seq_len
        return self.k[:, :, :self.pos], self.v[:, :, :self.pos]

The memory cost is significant. For a model like Llama-3-70B with 80 attention layers, 64 heads, and d_head=128, a single 128K-token sequence requires:

80 layers × 2 (K+V) × 64 heads × 128 d_head × 128K tokens × 2 bytes (fp16)
= 80 × 2 × 64 × 128 × 131072 × 2
≈ 27.5 GB just for the KV cache

This is why providers charge so much more per input token for long-context requests, and why techniques like GQA (Grouped-Query Attention) and MQA (Multi-Query Attention) exist — they share key/value heads across query heads to reduce this footprint.

The lost-in-the-middle problem

A 2023 paper by Liu et al. documented something practitioners had already noticed empirically: LLMs are significantly better at recalling information placed at the beginning or end of a long context, versus information buried in the middle. They called this the "lost-in-the-middle" problem.

The mechanism isn't fully understood, but there are two plausible explanations. First, during pretraining, documents are often truncated to fit context windows, so the model saw many examples of "important information at start and end" simply due to how text is structured. Second, positional embeddings (particularly RoPE-based ones) may create implicit recency bias that favors recent tokens.

Attention sinks and StreamingLLM

Another non-obvious behavior: transformer models tend to assign disproportionate attention weight to the first few tokens in a sequence, regardless of their semantic content. These are called attention sinks.

The StreamingLLM paper (Xiao et al., 2023) exploited this: they showed you could run inference on infinite-length streams by always keeping the first four tokens in the KV cache (to serve as attention sinks) plus a sliding window of recent tokens. The model stays coherent even though most of its past context has been evicted.

def streaming_kv_cache(new_k, new_v, cache_k, cache_v,
                        sink_size=4, window_size=1024):
    """
    Maintain attention sinks + sliding window.
    sink_size: number of initial tokens to always keep
    window_size: recent tokens to keep
    """
    sinks_k, sinks_v = cache_k[:sink_size], cache_v[:sink_size]
    window_k = cache_k[sink_size:][-window_size + 1:]
    window_v = cache_v[sink_size:][-window_size + 1:]

    return (
        torch.cat([sinks_k, window_k, new_k], dim=0),
        torch.cat([sinks_v, window_v, new_v], dim=0),
    )

What this means in practice

If you're building applications on top of LLMs, these mechanics have concrete implications:

Put critical information first or last. If you're stuffing retrieved documents into a context window for RAG, don't assume the model will find information in position 12 out of 20 equally well as position 1. Reranking and structural placement matter.

Long context ≠ long memory. A model with a 128K context window will still perform worse on retrieval tasks with 100K tokens of context versus a well-structured 8K prompt. Context length is a ceiling, not a performance guarantee.

KV cache invalidation is expensive. If you're building a chat application and you modify an earlier system prompt, you've invalidated everything that was cached after it. Anthropic, OpenAI, and Google all have explicit prompt caching APIs now — structure your prompts so the stable prefix is long and comes first.

Watch out for the O(n²) wall. Even with flash attention and paged attention optimizations, inference latency is superlinear in context length. At 128K tokens with a fast A100, prefill alone can take 10-30 seconds. Design your UX accordingly.

What's changing

Linear attention variants (Mamba, RWKV, RetNet) try to escape the O(n²) wall by replacing softmax attention with recurrent or linear mechanisms. They're genuinely promising for long-context tasks, but as of early 2026 they still lag behind transformer baselines on general reasoning benchmarks.

Mixture-of-Experts (MoE) architectures help with compute cost but don't fundamentally change the attention complexity. What does change it is speculative decoding and multi-token prediction — both of which improve decode throughput without changing the prefill story.

For now, the pragmatic advice is: understand the mechanics, test your specific use case with real context lengths, and don't assume that "128K context" means your application will work the same at 100K tokens as at 4K. Benchmark at the extremes.

If you found this useful, the Lost in the Middle paper and the StreamingLLM paper are worth reading in full. The HuggingFace transformers source is also surprisingly readable for understanding KV cache implementation details.