Evaluation at Long Context

The Real Problem: A 128k Number on a Spec Sheet

Every modern frontier model release is announced with a context-length number, usually a power of two. 128k. 200k. 1M. The number sits on the spec sheet next to the parameter count and the MMLU score, and it is the single most misunderstood metric in large-language-model training. A model that accepts a 128 000-token prompt without raising an error is technically a 128k model. A model that can reliably use the information in that prompt is a completely different beast — and the gap between the two is wide enough to bury an entire product launch in.

The mechanics from sections 12.1–12.3 explain why the gap exists. Position embeddings have to extrapolate or interpolate. Attention has to dilute mass across a longer key-set. The KV cache balloons and forces architectural compromises. Continued pretraining on synthetic long-context data only works for some patterns. By the time the model reaches release, you have a stack of small, mostly-invisible regressions that all show up in the same place: questions whose answer depends on a needle buried deep in a haystack.

Intuition: A 128k Model That Forgets at Token 8k

Picture an open-book exam. The student is handed a 600-page textbook five minutes before the test starts and told the answer to one of the ten questions is on page 312. The other nine questions can be answered from general knowledge. If you score the test only on the average across all ten questions, a student who never even opens the book can get a 90%. The one question that depends on the book is buried in the signal.

That is exactly what happens when you evaluate a long-context model on next-token loss. The vast majority of tokens in a long document are easy: they continue a sentence, they finish a code block, they complete a list. A model that has completely lost track of the start of the document still predicts most tokens just fine, because most tokens are predicted from a window of only a few hundred neighbours. The cross-entropy stays low. Perplexity drops. The headline number looks great.

The questions that actually require the model to remember the whole document are a tiny fraction of the token distribution. Cross-entropy averages over them. A long-context evaluation has to do the opposite: it has to construct prompts whose only possible answer requires reading the whole context, and then score the model on those prompts alone.

What “Long-Context Capability” Actually Means

Let the prompt be a sequence of tokens $x_{1:N}$ of length $N$, and let $y$ be the expected answer. A general capability measure for context length $N$ is the conditional accuracy $A(N) = \mathbb{E}_{(x_{1:N}, y) \sim D_N}\!\left[ \mathbf{1}[ M(x_{1:N}) = y ] \right]$, where $M$ is the model and $D_N$ is a distribution of (prompt, answer) pairs at length $N$. The whole game of long-context evaluation is choosing $D_N$ so that:

1. Every $y$ is uniquely determined by the prompt. No general-knowledge shortcuts.
2. The information required to answer $y$ depends on a position inside $x_{1:N}$ that we control. So we can sweep the “depth” axis.
3. The grader is deterministic and reproducible. Substring match, exact match, or a frozen LLM-as-judge with a fixed rubric — never a free-form human grader for a regression suite.

The Needle-in-a-Haystack (NIAH) construction is the simplest $D_N$ that satisfies all three. We embed a single sentence — “the needle” — at position $p \in [0, N]$ inside an otherwise unrelated document of length $N$, then ask a question whose only valid answer is the needle. Reporting $A(N, p)$ as a 2-D heatmap is the canonical long-context plot.

Two failure modes shape that heatmap. The first is the extrapolation cliff: for any $N$ beyond the model's trained window $N_{\text{train}}$ and any depth, accuracy crashes to near-zero, because positions past $N_{\text{train}}$ generate RoPE phases the model has never seen. This is what raw Llama-2 looks like at 16k. The second is the lost-in-the-middle dip: even inside the trained window, accuracy is highest near $p = 0$ and $p = N$ and dips in the middle. Liu et al. (2023) showed this U-shape holds for nearly every model trained with standard data, because attention mass naturally concentrates at the beginning (the system prompt) and the end (the question).

Manual Numerical Walkthrough: A 32-Token NIAH

Let us run NIAH by hand at a scale a human can actually verify. The haystack is six copies of the same boring sentence; the needle is one sentence injected at a chosen depth. We tokenize with whitespace splits to keep the arithmetic readable.

Filler sentence (4 tokens): the cat sat down
Needle (7 tokens): the magic number is forty two today
Question: What is the magic number?

Target length is 32 tokens. Six filler copies give 24 tokens; the needle adds 7 (we drop one filler copy to keep the total close to target). We sweep three depths: 0% (front), 50% (middle), 100% (back). The harness produces three prompts:

Now imagine three different models answering all three prompts. Model A is a tiny RNN that only remembers the last 8 tokens; it sees only the question and the surrounding filler, so it answers all three with “I don't know.” Score: 0 / 3.

Model B is a transformer trained on sequences up to length 16. It can see the whole 32-token prompt syntactically, but its position embeddings only ever saw positions 0–15 during training. For the depth = 0% prompt, the needle sits inside the trained position range, so it retrieves and answers “forty two.” For the depth = 50% prompt the needle starts at position 13 (still inside the trained range) and the question is appended past position 16 (so RoPE phases 13–32 are partially extrapolated, attention is brittle, but the retrieval just barely succeeds half the time). For depth = 100%, the needle lives in positions the model never saw any training data for — complete miss. Score: ~1.5 / 3.

Model C is the same transformer with YaRN applied to RoPE so all positions up to 32 behave like they were trained. All three prompts are inside the post-extension trained range. Retrieval succeeds in every case. Score: 3 / 3.

Now repeat this experiment at length $N = 32{,}000$ with 11 depths instead of 3, and you have one column of a real NIAH heatmap. Now repeat for eight different lengths and you have the full heatmap. The principle is unchanged from this hand example — only the scale and the model.

Visualizing NIAH Failure Modes

The interactive heatmap below switches between five archetypal long-context profiles. Click each profile to see the characteristic shape — the extrapolation cliff, the lost-in-the-middle dip, the smooth PI decay, the YaRN-preserved high-frequency floor. These are idealised curves, but every one of them maps onto a real published figure.

Two things are worth pausing on. First, the cliff profile is the one most release candidates produce when you naively push their max context past the trained window. It is also the profile a careless eval can miss entirely if it never tests past 8k. Second, the lost-in-the-middle profile is real and persistent — every long-context paper since Liu et al. 2023 has had to address it, and it is the main reason agentic frameworks aggressively re-rank and summarise their contexts even when the model technically accepts the full document.

The Perplexity Trap

The most common temptation when extending a model's context is to evaluate using perplexity on long documents. It is fast, deterministic, and produces a single scalar per model. It also lies. The chart below shows what happens if you let perplexity be your only signal: PPL stays in the healthy 3–4 range out to the full 128k window, while retrieval accuracy collapses to zero the moment the prompt exits the trained region.

The mechanism is simple. Cross-entropy on natural text is dominated by short-range structure: the local n-gram, the local syntax, the ongoing topic. A model that has completely failed to integrate information from 30 000 tokens ago still predicts the next sentence from the previous sentence. The errors that matter for retrieval — sharp, localised, on a single named entity — are washed out by the thousands of cheap next-token predictions around them.

This is not a new observation. It is one of the lessons every long-context team relearns the hard way. The first generation of long-context papers (early 2023) leaned heavily on PPL because no good alternative existed. NIAH (Kamradt, mid-2023) and then RULER (Hsieh et al., 2024) gave the field something measurable, and within a year nobody serious was publishing PPL-only long-context numbers anymore.

Rule of thumb: if a long-context release announces a new max context without showing a NIAH heatmap or RULER table, treat the context number as marketing, not capability.

Beyond NIAH: RULER, LongBench, InfiniteBench

NIAH is the right first eval, but it is not the only eval, and it rewards a very specific capability — single-key retrieval — that not all long-context tasks need. Real applications ask the model to reason across multiple documents, aggregate frequencies, track variables across an argument, and answer questions whose evidence is scattered. The community has converged on three benchmarks that cover these dimensions.

The biggest practical upgrade NIAH → RULER buys you is the multi-task degradation profile. NIAH hides the fact that single-needle retrieval is the easiest long-context capability. Variable tracking, multi-key retrieval, and frequency aggregation all degrade much faster with length. The grid below switches between three model archetypes and shows how the cells change task by task.

Notice the pattern that holds across all three model classes: as length grows, the bottom rows (aggregation, multi-hop QA) lose accuracy before the top rows (single-needle retrieval). A model can score 95% on NIAH at 128k and 25% on RULER's aggregation tasks at the same length. Both numbers are real. Only one of them gets put on the press release.

Plain Python: NIAH from Scratch

The harness is small — fewer than a hundred lines of Python with no ML dependencies. We build the prompt, we score the answer, we loop over a grid. Every long-context evaluation harness in production (Anthropic's, Google's, the open-source lm-evaluation-harness suite) has this same skeleton, with more polished prompts and a bigger task library bolted on.

NEEDLE = 'The best thing to do in San Francisco is eat a sandwich at Dolores Park on a sunny day.'

QUESTION = 'What is the best thing to do in San Francisco?'

ANSWER_KEY = 'eat a sandwich at dolores park'

~32-token sentence about Rome

build_prompt(32_000, 0.5) → 32 000-token prompt with needle in the middle

context_tokens=8000 → n_sentences=333; context_tokens=128000 → n_sentences=5333

n_sentences=333: depth 0.0 → idx 0; depth 0.5 → idx 166; depth 1.0 → idx 332

haystack = concatenated string ~ context_tokens tokens long

score('the answer is eat a sandwich at Dolores Park.') → 1.0

Result(length=32_000, depth_pct=0.5, accuracy=0.66)

niah_sweep(model, [1_000,4_000,8_000,16_000,32_000], [0,0.25,0.5,0.75,1.0], 5) → 125 generations

fake_model_factory(8000) — about Llama-2-7B's native ctx

n_tokens = len(prompt.split())  # rough whitespace token count

lengths=[1k,4k,8k,16k,32k], depths=[0,0.25,0.5,0.75,1.0]

1"""
2Needle-in-a-Haystack (NIAH) from scratch.
3
4We build a haystack of "filler" text, insert a single sentence ("the needle")
5at a chosen depth, then ask the model a question whose only honest answer is
6the needle. Repeating this over many (length, depth) pairs gives the
7characteristic 2-D heatmap that long-context releases live and die by.
8
9The harness is intentionally framework-agnostic. The only thing we need from
10the model is a callable that returns a string answer given a prompt.
11"""
12
13from dataclasses import dataclass
14from typing import Callable, List, Tuple
15
16# ---------------------------------------------------------------------------
17# 1. Build one NIAH prompt.
18# ---------------------------------------------------------------------------
19
20NEEDLE = "The best thing to do in San Francisco is eat a sandwich at Dolores Park on a sunny day."
21QUESTION = "What is the best thing to do in San Francisco?"
22ANSWER_KEY = "eat a sandwich at dolores park"   # substring we will look for
23
24# A long, repetitive haystack. In real evals you use Paul Graham essays
25# (the "PG haystack" from Greg Kamradt's original NIAH), but for a toy run
26# any unrelated text whose total length we can control will do.
27HAYSTACK_SENTENCE = (
28    "The Roman Empire spanned three continents and lasted for over a "
29    "thousand years, shaping the laws, languages, and architecture of "
30    "modern Europe. "
31)
32
33
34def build_prompt(context_tokens: int, depth_pct: float,
35                 tokens_per_sentence: int = 24) -> str:
36    """Construct the (very long) prompt for one NIAH datapoint."""
37    n_sentences = max(1, context_tokens // tokens_per_sentence)
38    filler = [HAYSTACK_SENTENCE] * n_sentences
39
40    insert_idx = int(round(depth_pct * (n_sentences - 1)))
41    filler[insert_idx] = NEEDLE + " " + filler[insert_idx]
42
43    haystack = "".join(filler)
44    return (
45        "Read the following document carefully and then answer the question "
46        "using only information from the document.\n\n"
47        f"<document>\n{haystack}\n</document>\n\n"
48        f"Question: {QUESTION}\nAnswer:"
49    )
50
51
52# ---------------------------------------------------------------------------
53# 2. Score one model answer.
54# ---------------------------------------------------------------------------
55
56def score(answer: str) -> float:
57    """1.0 if the needle appears (case-insensitive substring), else 0.0."""
58    return 1.0 if ANSWER_KEY in answer.lower() else 0.0
59
60
61# ---------------------------------------------------------------------------
62# 3. Sweep over a grid of (length, depth) pairs.
63# ---------------------------------------------------------------------------
64
65@dataclass
66class Result:
67    length: int        # context length in tokens
68    depth_pct: float   # where the needle was placed
69    accuracy: float    # 0.0 or 1.0 for a single sample; mean over k samples
70
71def niah_sweep(
72    generate: Callable[[str], str],
73    lengths: List[int],
74    depths: List[float],
75    samples_per_cell: int = 5,
76) -> List[Result]:
77    """Run a full NIAH grid. Returns one Result per cell."""
78    out: List[Result] = []
79    for L in lengths:
80        for d in depths:
81            scores = []
82            for _ in range(samples_per_cell):
83                prompt = build_prompt(L, d)
84                ans = generate(prompt)
85                scores.append(score(ans))
86            out.append(Result(L, d, sum(scores) / len(scores)))
87    return out
88
89
90# ---------------------------------------------------------------------------
91# 4. A fake "model" so we can run the harness end-to-end without a GPU.
92#    The fake reads the needle at depth d only if the prompt is shorter
93#    than 8k tokens (about the trained window for a stock Llama-2-7B).
94# ---------------------------------------------------------------------------
95
96def fake_model_factory(trained_tokens: int = 8000):
97    def generate(prompt: str) -> str:
98        n_tokens = len(prompt.split())
99        if n_tokens <= trained_tokens and NEEDLE.lower() in prompt.lower():
100            return ANSWER_KEY
101        return "I'm not sure."
102    return generate
103
104
105if __name__ == "__main__":
106    lengths = [1_000, 4_000, 8_000, 16_000, 32_000]
107    depths  = [0.0, 0.25, 0.5, 0.75, 1.0]
108
109    results = niah_sweep(fake_model_factory(8_000), lengths, depths,
110                         samples_per_cell=3)
111
112    for r in results:
113        print(f"len={r.length:>6}  depth={r.depth_pct:>4.0%}  "
114              f"acc={r.accuracy:.2f}")

PyTorch: Running NIAH on a Real Model

Swapping the fake generator for a real model touches three things: the tokenizer (use the model's own), the dtype and attention implementation (bfloat16 + FlashAttention 2 at this scale), and the decoding settings (greedy, low temperature). The harness itself does not move.

MODEL_ID = 'meta-llama/Meta-Llama-3.1-8B-Instruct'

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

model loads at ~16 GB; with flash-attn 2, 128k prompt prefill fits in ~62 GB total

generator = make_hf_generator(max_new_tokens=64)

inputs['input_ids'] → torch.int64 of shape [1, S] where S is the real token count

n_tokens = inputs['input_ids'].shape[1]  # e.g. 31 740

Llama-3.1-8B: model.config.max_position_embeddings = 131072

out = model.generate(input_ids=[1, 31740], max_new_tokens=64, do_sample=False)

answer_ids = out[0, 31740:]  # shape [≤ 64]

from niah_plain import build_prompt, score, niah_sweep, Result

lengths=[1_000,2_000,4_000,8_000,16_000,32_000,64_000,128_000], depths=[0,0.1,...,1.0]

pathlib.Path('niah_llama31_8b.json').write_text(json.dumps(...))

1"""
2NIAH on a real PyTorch model.
3
4We swap the fake generator from the Python script for a Hugging Face
5transformers call. The harness is unchanged — niah_sweep does not know
6whether it is calling a 7B Llama or a 405B frontier model.
7
8Key engineering details:
9  - load in bfloat16 on a single H100 (80 GB) — fits up to ~70B at 128k
10    with FlashAttention 2 and reasonable KV-cache sharding;
11  - use the model's *own* tokenizer for accurate length control;
12  - cap output to ~64 tokens — the answer is short, longer outputs only
13    add cost and let the model wander.
14"""
15
16from typing import Callable
17import torch
18from transformers import AutoModelForCausalLM, AutoTokenizer
19
20MODEL_ID = "meta-llama/Meta-Llama-3.1-8B-Instruct"
21
22tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
23model = AutoModelForCausalLM.from_pretrained(
24    MODEL_ID,
25    torch_dtype=torch.bfloat16,
26    device_map="auto",
27    attn_implementation="flash_attention_2",
28)
29model.eval()
30
31
32def make_hf_generator(max_new_tokens: int = 64) -> Callable[[str], str]:
33    """Return a generate(prompt) -> str closure suitable for niah_sweep."""
34
35    def generate(prompt: str) -> str:
36        # 1. Tokenize WITH attention to the model's real context limit.
37        inputs = tokenizer(
38            prompt,
39            return_tensors="pt",
40            truncation=False,                # we want NIAH to actually fail
41        ).to(model.device)
42
43        n_tokens = inputs["input_ids"].shape[1]
44        if n_tokens > model.config.max_position_embeddings:
45            # Past the trained ctx — this is the failure regime we are
46            # *measuring*, so we still hand it to the model.
47            pass
48
49        # 2. Greedy decode. NIAH is a retrieval test, sampling would
50        #    add noise without buying any quality signal.
51        with torch.inference_mode():
52            out = model.generate(
53                **inputs,
54                max_new_tokens=max_new_tokens,
55                do_sample=False,
56                temperature=0.0,
57                pad_token_id=tokenizer.eos_token_id,
58            )
59
60        # 3. Slice off the prompt — generate() returns prompt+answer.
61        answer_ids = out[0, inputs["input_ids"].shape[1]:]
62        return tokenizer.decode(answer_ids, skip_special_tokens=True)
63
64    return generate
65
66
67# Reuse build_prompt, score, niah_sweep from the plain-Python file above.
68from niah_plain import build_prompt, score, niah_sweep, Result      # noqa: E402
69
70if __name__ == "__main__":
71    generator = make_hf_generator(max_new_tokens=64)
72
73    # A realistic 11-depth x 8-length sweep — 88 cells, samples_per_cell=3.
74    lengths = [1_000, 2_000, 4_000, 8_000, 16_000, 32_000, 64_000, 128_000]
75    depths  = [i / 10 for i in range(11)]
76
77    results = niah_sweep(generator, lengths, depths, samples_per_cell=3)
78
79    # Write to disk so the heatmap script can plot it offline.
80    import json, pathlib
81    pathlib.Path("niah_llama31_8b.json").write_text(json.dumps(
82        [r.__dict__ for r in results], indent=2,
83    ))

From toy to production

The harness above runs in tens of GPU-minutes on an 8B model. When the same idea is run inside a frontier lab, three things change:

1. Scale. The eval is run on every checkpoint of every continued-pretraining run, often with thousands of cells (multiple needles per cell, multiple needle positions, multiple haystack themes). Each run is a non-trivial GPU bill, so the eval cluster is a budgeted resource.
2. Distribution shift. The haystack is no longer a single domain (Paul Graham essays). It is a sampled mixture of domains the model is expected to handle in production — legal, medical, codebases, technical reports — because long-context performance varies sharply with domain.
3. Adversarial needles. The needle is no longer a single distinctive sentence. It is paraphrased, embedded in similar-looking distractor sentences, or split across multiple positions. This catches models that learned to pattern-match the NIAH format rather than actually retrieve.

At Massive Scale: Evaluating 1M-Context Models

Once the trained context reaches 1M tokens, every previous bottleneck gets worse by an order of magnitude.

The compute cost of the eval itself

One NIAH cell at 1M tokens needs the full prefill of a 1M-token KV cache. With FlashAttention 2 the cost is roughly $O(N \cdot d_{\text{model}})$ per layer for attention plus the usual $O(N \cdot d_{\text{model}}^2)$ for the FFN — at $N = 10^6$ this is roughly $10^{12}$ FLOPs per layer per forward, or $10^{14}$ for a 100-layer model. A full 88-cell sweep at 1M is therefore ~$10^{16}$ FLOPs — about three minutes of an 8-H100 node at fp8 peak, which is fine; but the KV cache alone takes ~10 GB at hidden size 7 168 and head dimension 128 (two cache types) — and that has to fit on the device while the model weights are loaded.

The KV cache becomes the eval bottleneck

For a Llama-3-class 70B model at 1M context, the KV cache is bigger than the model itself. A single eval prompt requires sharding the KV cache across multiple devices, exactly the same tensor-parallel choreography as serving — except for an eval workload that is mostly prefill, not decode. This is why long-context eval infrastructure at a frontier lab is co-designed with the inference stack, not built as a side project.

The grading function has to keep up

At 1M context the model's answer is no longer a 64-token substring match. Tasks like “summarise this novel” require LLM-as-judge with a careful rubric. The judge model is often the same family as the model under test, which raises the question of grading bias — frontier labs typically use a held-out judge of comparable but different lineage (e.g. grading Claude with GPT-4 and vice versa) and report inter-judge agreement.

The training data has to teach the eval

Naively training on long documents does not produce a 1M-context model that scores well on RULER. The training mix needs synthetic long-context tasks — book-length question answering, paraphrased needles at random depths, multi-document aggregation — because those are the capabilities we are about to measure. This is the deepest connection between training and evaluation in long-context: the evaluation set effectively defines what “1M context” means, and the training set has to be engineered to match.

Engineering Reality: What Actually Fails

Two years of long-context releases have produced a small catalogue of recurring failure modes. The good ones come up in code reviews before the release ships; the bad ones get caught in the wild and named after the team that hit them first.

• Tokenizer mismatch. The eval harness uses a different tokenizer from the model. The horizontal axis of the heatmap is off by ±15%, the cliff is in the wrong place, and the team spends a week debugging a regression that does not exist. Always tokenize with the model's own vocabulary.
• Truncation silently kicks in. The HuggingFace tokenizer truncates by default at the model's configured max position. Pass truncation=False when running NIAH or the model will pass cells it should fail.
• The prompt template is the system prompt. Some chat-tuned models prepend a default system prompt unless you explicitly pass one. That extra ~500 tokens at the front shifts every depth percentage and is invisible unless you log token counts.
• The needle is too memorable. “The magic number is 42” appears in pretraining data. Use a needle the model could not have memorised — a randomly-generated string of words, a date that does not exist, a fictional named entity.
• Greedy decoding silently goes off the rails. At very long context, the argmax of the first answer token is sometimes a stop token. Setting min_new_tokens=1 forces at least one real token, which is enough to catch the failure.
• The KV cache is paged out of GPU memory. Tools like vLLM page the KV cache to CPU when it overflows. At 128k this is fine for throughput but is a 50× slowdown per token, which makes the eval take long enough that intermediate fails (NaNs, OOMs) become disproportionately likely.
• Single-trial heatmaps overfit to one needle. Always run $k \geq 3$ samples per cell with different needles. The published Kamradt heatmaps are $k = 1$ and have ~10 percentage points of per-cell noise that nobody talks about.

Once these are fixed, what is left is the actual capability of the model. That capability is what the next chapter — on inference and serving — has to deliver to the user at runtime.

The mental model that unifies this section: long-context capability is not a single number. It is a 2-D surface (length × depth) across a panel of tasks (single needle, multi-hop, aggregation, QA). A “128k model” is a model with a particular shape on that surface — and the shape is what long-context training is trying to flatten.