Managing LLM Context: Beyond Simple Truncation

A Lecture Module for Graduate AI Agents Course

Prerequisite Review: Attention and the KV Cache

This section takes approximately 10–15 minutes. It can be skipped if students are already comfortable with these mechanics, but it establishes the vocabulary and notation used throughout the rest of the lecture.

How Attention Works in a Decoder-Only LLM

A decoder-only transformer (GPT-style, LLaMA-style) generates tokens one at a time, left to right. At each step it takes the full sequence of tokens produced so far and predicts the next one. The mechanism that lets each token "look back" at previous tokens is causal self-attention.

For each token at position i, the model learns three projection matrices — W_Q, W_K, W_V — that transform the token's embedding x_i into a query vector q_i, a key vector k_i, and a value vector v_i:

q_i = x_i @ W_Q      # "What am I looking for?"
k_i = x_i @ W_K      # "What do I offer to others looking?"
v_i = x_i @ W_V      # "What information do I carry?"

To compute the output for token i, the model scores how well its query matches every preceding key (causal masking prevents looking forward), converts those scores to weights via softmax, then takes a weighted sum of the values:

# For token i attending to positions 1..i:
scores[j] = dot(q_i, k_j) / sqrt(d_head)    # for j = 1..i
weights    = softmax(scores)                  # sum to 1
output_i   = sum(weights[j] * v_j for j in 1..i)

The output is then passed through the rest of the transformer layer (feed-forward network, layer norm, residual connection) and eventually projects to a probability distribution over the vocabulary.

Two things to notice:

1. Every token attends to every previous token. For a sequence of length n, computing attention requires n² dot products per layer per head. This is where the O(n²) cost comes from.

2. The queries, keys, and values depend only on the token's embedding and the learned weight matrices — not on anything else. This makes the KV cache possible.

What the KV Cache Does and Why It Matters

At generation time, the model produces one token per step. At step t, it has tokens 1…t in context and must compute attention for token t over all of tokens 1…t.

Without any optimization, you would recompute k_j and v_j for every previous token at every step — even though those values haven't changed:

# Naive generation (wasteful):
for step t = 1, 2, 3, ...:
  for j = 1..t:
    k_j = x_j @ W_K     # RECOMPUTED every step, even though x_j hasn't changed
    v_j = x_j @ W_V     # RECOMPUTED every step

  scores = [dot(q_t, k_j) / sqrt(d) for j in 1..t]
  output_t = weighted_sum(softmax(scores), values)

The KV cache eliminates this redundancy. After each token is processed, its key and value vectors are stored. On the next step, you only compute the new token's K and V, then append them:

# KV-cached generation (efficient):
kv_cache = []   # list of (k_j, v_j) for all past tokens

for step t = 1, 2, 3, ...:
  k_t = x_t @ W_K
  v_t = x_t @ W_V
  kv_cache.append((k_t, v_t))

  # Attend: only q_t is new; all keys and values come from cache
  scores  = [dot(q_t, k_j) / sqrt(d) for (k_j, v_j) in kv_cache]
  output_t = weighted_sum(softmax(scores), [v_j for (k_j, v_j) in kv_cache])

This reduces the per-step computation from O(n) matrix multiplications to O(1) — one new K/V pair per step, plus a single attention pass over the cached K's. It is why modern LLM inference is fast for generation despite the apparent complexity of attention.

The cost: memory. Every token that has ever been processed leaves a K and V tensor in the cache — one per layer, one per attention head. For a model with L layers, H heads, and head dimension d, the memory required for a sequence of n tokens is:

kv_cache_size = n × L × H × d × 2 × dtype_bytes
               ^   ^   ^   ^   ^
               |   |   |   |   two tensors (K and V)
               |   |   |   head dimension
               |   |   number of heads
               |   number of layers
               sequence length

For Llama-3.1-8B (32 layers, 32 KV heads, head dimension 128, bfloat16): each token occupies 32 × 32 × 128 × 2 × 2 bytes = 524 KB. A 128K-token context therefore requires roughly 64 GB for the KV cache alone — more than the model weights.

This is the central tension the rest of the lecture addresses: the KV cache makes generation tractable, but it grows linearly with context length and can dominate GPU memory. Everything that follows is, at some level, an attempt to manage or reduce that cost intelligently.

Part 1: Motivation — Why Context Management is Hard

The naive answer to "what do we do when we run out of context?" is: cut the oldest tokens and keep the newest. This works — barely — but it throws away information arbitrarily, ignores what the model actually attends to, and leaves you with no guarantees about what was lost.

The real answer requires understanding what the context window is at a mechanical level, and why simply raising the limit doesn't solve the problem.

The Cost of Full Attention

Self-attention in a transformer computes a score between every pair of tokens. For a sequence of length n, this is O(n²) in both time and memory. Doubling the context window quadruples the attention cost. A 128K context window is not just 32× harder than a 4K window — with naïve attention, it is 1,024× harder in the attention computation alone.

The KV cache compounds this: at each autoregressive decoding step, the model needs the key and value tensors for every previous token. A KV cache for a 128K-token context of Llama-3.1-8B can consume tens of gigabytes of GPU memory — often more than the model weights themselves.

So context management is a two-sided problem:

1. Fitting more context into the window — either by making the window bigger (architectural/training approaches) or by compressing what goes into it

2. Using fixed-size context effectively — intelligent selection of what lives in the window at any given moment

This lecture traces the major lines of attack, roughly from the model level down to the application level.

Part 2: Architectural Approaches — Sparse Attention

The most principled way to handle long sequences is to change what the attention mechanism computes. Rather than attending to every pair of tokens, use a sparse pattern that is theoretically expressive but computationally tractable.

2.1 Longformer: Local + Global Attention (2020)

Citation: Beltagy, Peters & Cohan. Longformer: The Long-Document Transformer. arXiv:2004.05150, 2020.

Longformer replaces the full O(n²) attention matrix with a pattern that is O(n):

• Sliding window attention: Each token attends only to a window of w neighbors (both left and right). This captures local context — the dominant pattern for most positions in natural language.

• Global attention: Certain designated tokens (e.g., [CLS] for classification, or all question tokens in QA tasks) attend to every position and are attended to by every position.

For each token i:
  local_neighbors = tokens in [i - w/2, i + w/2]
  if token i is a global token:
    attend to ALL tokens
  else:
    attend to local_neighbors UNION global_tokens

The sliding window means that information propagates across the full sequence through successive layers: a token at position 100 reaches position 500 in about 400/w layers. This "local attention + depth" tradeoff is central to Longformer's design. It supports sequences up to 4,096 tokens at roughly the same compute as BERT-style full attention on 512 tokens.

2.2 BigBird: Random + Window + Global (2020)

Citation: Zaheer et al. Big Bird: Transformers for Longer Sequences. NeurIPS 2020. arXiv:2007.14062.

BigBird generalizes Longformer by adding a third attention component: random attention. Each token attends to r randomly sampled positions in addition to its local window and any global tokens.

For each token i:
  random_keys = r randomly selected tokens
  local_keys  = w/2 tokens to the left, w/2 to the right
  global_keys = g designated global tokens
  attend to: random_keys UNION local_keys UNION global_keys

Why random? The paper frames this as a graph sparsification problem. A complete graph (full attention) can be approximated by a sparse graph — and random graphs are expanders, meaning information flows efficiently even with few edges. Critically, BigBird proves that this sparse attention mechanism is:

• A universal approximator of sequence-to-sequence functions (as powerful as full attention in a formal sense)

• Turing complete (preserving the computational power of standard transformers)

BigBird handles sequences up to 8x longer than was previously possible on the same hardware.

2.3 What These Approaches Give Up

Both Longformer and BigBird require the sparse attention pattern to be baked into training. You cannot take a pretrained LLaMA or GPT-4 model and retroactively apply Longformer-style attention without retraining. This limits their practical relevance to the current ecosystem, where nearly all production models use dense causal attention. They remain architecturally important as the theoretical backbone for understanding what later, more pragmatic techniques approximate.

Part 3: Extending the Positional Range — Context Window Scaling

Rather than changing the attention pattern, a second family of approaches extends how far a model can count — i.e., how it assigns position encodings to tokens beyond its training length.

3.1 Why Position Encodings Break at Long Range

Modern LLMs almost universally use Rotary Position Embeddings (RoPE), introduced by Su et al. (2021). RoPE encodes position by rotating key and query vectors by angles that are a function of position. The inner product of a query at position m and a key at position n depends on their relative distance (m - n) — a desirable property.

The problem: during pretraining, the model only sees positions up to some maximum L. The rotation angles at positions L+1, L+2, ... are extrapolations the model has never encountered. This causes catastrophically high attention scores and incoherent outputs.

A simple fix — just fine-tune on longer sequences — works poorly in practice. Training on 10,000 steps with 32K context causes the effective window to grow from 2,048 to only ~2,560 tokens. The model resists.

3.2 Position Interpolation (2023)

Citation: Chen, Wong, Chen & Tian (Meta). Extending Context Window of Large Language Models via Positional Interpolation. arXiv:2306.15595, 2023.

The key insight: instead of extrapolating to unseen positions, compress the existing position range so that a longer sequence maps back into the trained range.

For a model trained on length L that you want to extend to length L' (where L' > L), scale every position index m by a factor of L/L':

# Standard RoPE
position_index = m  (where m can be up to L)

# Position Interpolation: squeeze L' positions into [0, L]
position_index = m * (L / L')  (values always in [0, L])

The rotation angles at any position now fall within a range the model has already learned. A brief fine-tuning (just ~1,000 gradient steps) allows the model to adapt to the new spacing. The authors extend LLaMA 7B–65B to context windows of up to 32,768 tokens this way.

Theoretical justification: the upper bound on interpolated attention score inflation is ~600x smaller than for extrapolation. This is why interpolation is stable and extrapolation is not.

Limitation: Linear interpolation "crowds" nearby tokens together. At very high extension ratios (e.g., 16x or more), distinguishing closely-adjacent tokens becomes difficult, and performance degrades. Later methods like YaRN (Peng et al., 2023, arXiv:2309.00071) address this by applying different scaling factors to different frequency dimensions of RoPE, trading off precision at different ranges.

Part 4: Inference-Time KV Cache Management

The approaches in Parts 2 and 3 require architectural choices made at training time. The methods in this section operate entirely at inference time, on top of whatever model you have. They answer the question: given a fixed model with a fixed context window, which tokens should live in the KV cache right now?

4.1 The KV Cache Eviction Problem

At each decoding step, the model needs KV tensors for every token it attends to. If the conversation exceeds the context window, something has to go. The question is what.

Key observations that motivate intelligent eviction:

• Attention matrices in practice are highly sparse — each token tends to strongly attend to only a small subset of predecessors

• This sparsity follows a power law: a small number of "heavy hitter" tokens receive disproportionately high attention

• Evicting uniformly at random destroys important structure; evicting the wrong tokens causes severe degradation

4.2 StreamingLLM: Attention Sinks (2023)

Citation: Xiao, Tian, Chen, Han & Lewis (MIT/Meta). Efficient Streaming Language Models with Attention Sinks. ICLR 2024. arXiv:2309.17453.

StreamingLLM discovered an important phenomenon: initial tokens receive disproportionately high attention scores even when they are not semantically important. The authors call these "attention sinks."

Why do attention sinks exist? The softmax in attention must sum to 1. When no token is strongly relevant, the model dumps probability mass somewhere — and the first token, which is always visible from any position in causal attention, is a natural destination. The behavior is similar to outlier channels in quantization: it emerges from training dynamics, not from semantic content.

The implication: a naive sliding window (keep the last w tokens, evict everything older) fails because it discards the attention sink tokens, causing the attention distribution to destabilize. Perplexity explodes.

StreamingLLM's fix:

xxxxxxxxxx
KV_cache = attention_sinks UNION recent_window
# attention_sinks: first 4 tokens (constant cost)
# recent_window:   last W tokens (sliding)
# everything in between: evicted

At each new token:
  if len(KV_cache) > budget:
    evict oldest non-sink tokens
  append new token's KV to cache

With just 4 initial tokens retained as sinks, StreamingLLM enables stable language modeling at up to 4 million tokens — with up to 22x speedup over the recomputation baseline. It does not extend the effective context window (the model only truly sees the recent window), but it enables infinite-length streaming without cache resets.

4.3 H2O: Heavy Hitter Oracle (2023)

Citation: Zhang, Sheng et al. (UT Austin/Stanford/CMU/Berkeley). H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models. NeurIPS 2023. arXiv:2306.14048.

H2O formalizes the eviction problem differently. Rather than retaining a fixed structural pattern (sinks + recent), it asks: which tokens have accumulated the most attention over the course of generation?

Key observation: Attention scores over all tokens follow a power-law distribution. A small set of "heavy hitter" (H2) tokens account for a disproportionate share of total attention. These tokens are content-dependent, not position-dependent.

H2O's eviction policy:

xxxxxxxxxx
score[i] = accumulated_attention_to_token_i  (sum over all decoding steps)

At each step:
  if len(KV_cache) > budget:
    evict token i with lowest score[i] that is not in the recent window

  # Update scores for next step
  score[new_token] = 0
  for each retained token i:
    score[i] += attention_weight(current_query -> key_i)

H2O retains a balance of recent tokens (which may become heavy hitters) and established heavy hitters (which are important but not necessarily recent). It is formalized as a dynamic submodular problem, and the greedy algorithm has a provable near-optimality guarantee under mild assumptions.

Result: H2O with 20% of tokens retained as heavy hitters improves throughput by up to 29x on OPT-6.7B and OPT-30B over leading inference systems, with 1.9x latency reduction.

4.4 SnapKV: Query-Aware Prefill Compression (2024)

Citation: Li, Huang, Yang et al. SnapKV: LLM Knows What You are Looking for Before Generation. NeurIPS 2024. arXiv:2404.14469.

H2O operates during generation. SnapKV targets a different moment: the prefill phase, where the model processes a long input prompt before generating any tokens.

Key insight: Each attention head consistently focuses on the same regions of the prompt during generation, regardless of what output is produced. And you can predict which regions those are from the attention patterns in a small observation window at the end of the prompt (i.e., the instruction or question itself).

xxxxxxxxxx
observation_window = last O tokens of prompt (the instruction / query)

# During prefill, before generation begins:
for each attention head h:
  # Compute attention from observation window tokens to the full prefix
  attn_scores_h = softmax(Q_obs @ K_prefix.T / sqrt(d))
  
  # Pool scores over observation window (vote for each prefix position)
  importance_h[i] = max(attn_scores_h[:, i])  for each prefix position i
  
  # Select top-k positions (k = per-head budget)
  selected_h = top_k(importance_h, k=budget)

# Build compressed KV cache:
KV_cache = union_over_heads(KV[selected_h]) UNION KV[observation_window]
# observation window is always fully retained

This compression is done once before generation begins and requires no changes during decoding. The selected KV pairs are chosen per head, not globally, so different heads can prioritize different parts of the prompt.

Results: 3.6x generation speedup and 8.2x memory efficiency improvement on 16K-token inputs, with no meaningful accuracy drop across 16 long-sequence benchmarks. On a single A100-80GB, SnapKV can process up to 380K tokens.

4.5 Comparison of Eviction Strategies

A practical note: recent work (e.g., Hold Onto That Thought, 2025) shows that for reasoning models specifically, attention-based methods (H2O, SnapKV) significantly outperform positional methods (StreamingLLM) because reasoning chains involve long-range dependencies that simple recency cannot capture.

Part 5: Application-Level Context Management — Virtual Memory for LLMs

All of the approaches above operate inside a single inference call. The final family of techniques lifts context management to the application level: the LLM itself decides what to keep, store, retrieve, and summarize across its own context window.

5.1 The "Lost in the Middle" Problem

Before introducing application-level memory, it is worth noting a failure mode of simply providing more context. Liu et al. (2023, arXiv:2307.03172) showed that LLMs perform best when relevant information appears at the beginning or end of a long context, and significantly worse when it is buried in the middle — even if the model nominally "sees" the whole input.

This means a large context window is not a free lunch. Models can struggle to utilize information that is present but not salient. Application-level management can mitigate this by controlling what information appears and where.

5.2 Recursive Summarization: The Simple Baseline

The most common application-level approach in practice — and worth understanding as a baseline — is recursive summarization:

xxxxxxxxxx
messages = conversation_history

while token_count(messages) > limit:
  oldest_chunk = messages[:chunk_size]
  summary = LLM.summarize(oldest_chunk)
  messages = [summary] + messages[chunk_size:]

response = LLM.generate(messages)

This is simple to implement and works reasonably well in practice. The failure modes are predictable: information in summarized segments is lossy, specific details (names, numbers, quotes) tend not to survive multiple rounds of summarization, and the LLM has no agency over what gets compressed. Summarization also happens upfront — the model discards information before knowing what the user will ask.

5.3 MemGPT: LLMs as Operating Systems (2023)

Citation: Packer, Wooders, Lin, Fang, Patil, Stoica & Gonzalez (UC Berkeley). MemGPT: Towards LLMs as Operating Systems. arXiv:2310.08560, 2023.

MemGPT draws an explicit analogy to virtual memory in operating systems. An OS gives applications the illusion of a large, flat address space by moving data between fast RAM and slow disk storage as needed. MemGPT does the same for LLMs.

Memory tiers:

• Main context (analogous to RAM): the fixed-size context window currently visible to the LLM — fast, limited

• External context (analogous to disk): a database of conversation history, documents, and notes — too large for the context window but queryable

The LLM is given a set of memory management functions as tools:

xxxxxxxxxx
# Tools available to the LLM as function calls:
recall_memory.search(query, n_results)      # semantic search over conversation history
archival_memory.insert(content)            # write content to external storage
archival_memory.search(query, n_results)   # retrieve from external storage
core_memory.append(field, content)         # append to always-in-context notes
core_memory.replace(field, old, new)       # edit always-in-context notes

# System interrupt: triggered when main context is near full
on_context_limit_warning():
  # LLM decides what to summarize, archive, or retrieve
  summary = summarize(recent_messages)
  archival_memory.insert(summary)
  retrieve_relevant = recall_memory.search(current_topic)
  # update main context accordingly

The LLM treats these as first-class actions in its reasoning loop — deciding, based on its goals and the current task, what needs to stay in working memory and what can be paged out to storage.

Key design features:

• Core memory: a small, always-in-context block for persistent facts (persona, user preferences, current goal)

• Recall memory: full conversation history, stored externally, retrieved by semantic or time-based search

• Archival memory: general external storage for documents and notes

• Interrupt mechanism: the system alerts the LLM when context pressure is high, triggering explicit memory management

MemGPT was evaluated on document QA (documents far exceeding the LLM's context window) and multi-session chat (conversations across multiple sessions). In both cases it significantly outperforms truncation and recursive summarization baselines. As of 2024 it is available as an open-source framework called Letta.

Contrast with recursive summarization: MemGPT retrieves information on demand rather than compressing it upfront. The LLM decides what it needs, not what to throw away.

Part 6: Synthesis — Choosing the Right Strategy

Different approaches suit different deployment contexts. Here is a decision framework:

If you are training a new model from scratch for long-document use: Sparse attention architectures (Longformer, BigBird) provide theoretically grounded O(n) complexity. This remains the state-of-the-art for encoder-heavy tasks like document classification and extractive QA.

If you have a pretrained dense-attention model and need to extend its range: Position interpolation (RoPE scaling, YaRN) with light fine-tuning can extend context 8–32x. This is how most production models achieved their large context windows — it requires some retraining but far less than pretraining from scratch.

If you need to serve long contexts at inference time efficiently: KV cache eviction methods (StreamingLLM, H2O, SnapKV) can reduce memory by 2–8x with minimal accuracy loss. These are orthogonal to the context window size and can be composed with position interpolation.

If your application involves persistent memory across sessions or document corpora larger than any context window: Application-level virtual context management (MemGPT/Letta, or simpler RAG + summarization pipelines) is the right level of abstraction. No attention-level technique handles truly unbounded external state.

A Note on Composability

These approaches are not mutually exclusive. A production deployment might use:

• A model trained with position interpolation (extended context window)

• SnapKV during prefill (compress long system prompts before generation)

• StreamingLLM's sliding window during generation (streaming chat at low memory cost)

• A MemGPT-style external memory system (persistent user state across sessions)

Each technique addresses a different bottleneck at a different level of the stack. The prerequisite review at the top of this document — understanding exactly what the KV cache is and why it is expensive — is what makes all of these tradeoffs legible.

Discussion Questions

1. StreamingLLM keeps "attention sink" tokens not because they contain important information, but because removing them destabilizes attention. What does this tell you about the relationship between model behavior and model intent?

2. H2O defines "important" tokens as those receiving high accumulated attention. Can you construct a scenario where this heuristic fails? What kind of information would be systematically discarded?

3. SnapKV uses the observation window (the instruction portion of the prompt) to predict which prefix tokens matter for generation. What assumptions does this make? When would those assumptions break?

4. MemGPT's memory architecture is inspired by OS virtual memory, but it asks the LLM itself to manage memory rather than the OS. What are the risks of this design? What happens if the LLM's memory management decisions are wrong?

5. The "lost in the middle" phenomenon suggests that long contexts are not uniformly utilized. How does this interact with KV cache eviction? If a model already does not attend well to middle tokens, does it matter whether we evict them?

6. All of the inference-time eviction methods (StreamingLLM, H2O, SnapKV) make a one-way decision: once a token is evicted, it is gone. What would a system look like that could retrieve evicted tokens on demand? What tradeoffs would it face?

References

End of Module — Managing LLM Context: Beyond Simple Truncation