← Inference From First Principles

Long-Context Inference

A context length is five claims, not one: that the model accepts the tokens, behaves at those positions, fits in memory, answers in time, and uses what it is given. Each measured on a model that advertises 128K, on a card that can hold a quarter of it.

Llama-3.2-1B advertises a context of 131,072 tokens. On the 12 GiB card of this book, asked to prefill 4,096 of them with the library’s default attention, it used 4.9 GiB of extra memory, nearly four times what 2,048 needed. A prefill of 8,192, tried once outside the final measurements, wanted nearly 19 GiB and took two and a half minutes because the driver quietly spilled the excess into system memory. The model accepted the tokens in every case. Nothing in “accepts 131,072 tokens” said that.

A context length bundles five claims, and a system can honour some of them while failing others. The tokenizer and the API accept that many tokens. The positional scheme and the trained model behave sensibly at those positions. The memory can hold the state, including the state of the computation that builds it. The latency is tolerable. And the model uses what it is given. This chapter measures each in turn on models cached here, finds that the stack’s failure on the third claim is not the one the cache arithmetic of earlier chapters would predict, and ends with a retrieval test whose result depends more on where a fact sits than on how long the passage is.

Where this sits

Chapter 10 gave the cache’s size, Chapters 14 and 16 managed it, Chapter 21 shrank it with fewer key/value heads and Chapter 22 with a window. This chapter asks what happens when none of those is enough and the context simply gets long. Chapter 8 measured attention cost and Chapter 13 built the tiled algorithm that avoids its memory; Chapter 9 introduced chunked prefill, which turns out to be the chapter’s most useful tool. The kernel finding of Chapter 21 (the library’s grouped-attention path is slow on this build) returns here with larger consequences.

The Inference Lab piece added here is the long-context harness: cache and FLOP arithmetic, a fused attention override that restores linear prefill memory, chunked prefill, per-position loss with a growing cache, a needle-at-depth generator, truncation and cache-reuse helpers, and the experiments. Chapter 24 asks which of the things measured in Chapters 17 to 23 belong to the model and which to the runtime.

For the Pi series, the lesson is a reporting discipline. A claim that a system supports a given context should say which of the five it tested. The answers were different for each, here.

Claim one: it accepts the tokens

The first claim is a configuration. These are the trained positions and RoPE settings in the cached configs, and the cache they imply at full length in fp16:

Model Trained positions RoPE KV cache at full length
GPT-2 1,024 learned absolute positions 0.04 GiB
TinyLlama-1.1B 2,048 plain 0.04 GiB
Phi-3-mini-4k 4,096 plain 1.50 GiB
Gemma-2-2B 8,192 plain 0.81 GiB (counting every layer at full length)
Mistral-7B-Instruct-v0.3 32,768 plain 4.00 GiB
Llama-3.2-1B 131,072 llama3, factor 32, from 8,192 4.00 GiB
Llama-3.2-3B 131,072 llama3, factor 32, from 8,192 14.0 GiB
Llama-3.1-8B 131,072 llama3, factor 8, from 8,192 16.0 GiB

The Llama-3.2-3B row is the first failure of a different kind from the others. Its cache at full length is larger than the whole card, before counting its 6.4 GB of weights. The model accepts 131,072 tokens in the sense that the configuration permits it; no single sequence of that length can be held here. The maximum a given card can hold is a function of the card, the weights and the cache per token, and for this 3B model the arithmetic gives 4 GiB of cache for about 37,000 tokens.

Claim two: it behaves at those positions

A RoPE model rotates query and key components by angles proportional to position. A model trained on positions up to \(N\) has not seen the angles beyond it. The Position Interpolation paper’s remedy, and the idea behind much of what followed, is to scale the position indices down so that a longer sequence maps onto the trained range; its abstract reports extending LLaMA models to 32,768 tokens with fine-tuning of at most 1,000 steps, and attributes the failure of plain extrapolation to extremely high attention scores. YaRN and LongRoPE extend the idea (the abstracts report needing 10× fewer tokens and 2.5× fewer steps, and reaching 2 million tokens, respectively). I ran none of those. What I could run, because the installed library takes the variant as a configuration, is the effect of changing the positional scheme under a checkpoint, and that is what makes the point: the scheme is part of the weights’ meaning.

TinyLlama, trained to 2,048 positions. Six passages of 4,096 tokens from this book (chapters 17 to 22), mean negative log-likelihood in nats per token by position:

Positions Default RoPE Linear scaling × 2, no training “Dynamic” scaling × 2, no training
0–512 2.92 3.62 2.92
512–1,024 2.56 3.39 2.56
1,024–2,048 2.42 3.45 2.42
2,048–3,072 7.20 3.54 7.86
3,072–4,095 8.43 3.57 7.18

Past its trained length the default model’s loss triples. Halving the position indices keeps the loss flat beyond 2,048 at about 3.5 nats, but it costs the positions it was trained on between 0.7 and 1.0 nats, since every position is now read at half the angle it learned. The library’s dynamic variant leaves the trained range untouched (identical to the default) and did not rescue the range beyond it in this evaluation (7.9 and 7.2). I did not investigate whether a different way of calling it, other than chunks of 1,024 tokens against a cache, would behave differently. None of the scaled variants was trained afterwards; the papers’ recipes all include some fine-tuning, and what is shown here is the cost of skipping it.

Llama-3.2-1B. Its checkpoint ships with the llama3 scaling in its config, and has presumably been trained with it. Two passages of 16,384 tokens, shipped against the same weights with the scaling replaced by plain RoPE at the same base:

Positions As shipped Scaling removed
0–1,024 3.29 3.30
1,024–2,048 3.15 3.19
2,048–4,096 2.48 2.54
4,096–8,192 2.76 2.94
8,192–12,288 2.84 3.65
12,288–16,383 2.73 5.68

As shipped, the model’s loss does not rise between 2,048 and 16,384 tokens (it is 2.7 to 2.8 from 4,096 to the end); with the scaling removed it stays within 0.18 nats of the shipped model inside the original 8,192 positions and degrades sharply after them. This is the same model with the same weights, evaluated under a different positional setting. A serving system that loads a checkpoint without honouring its RoPE parameters, or that applies a context-extension setting the checkpoint was not trained for, has changed the function. It is Chapter 2’s lesson about identity once more, with the identity extended to the positional scheme.

Claim three: it fits

The cache arithmetic is the easy part. The hard part, in this stack, was the computation that builds the cache.

Prefill peak memory for Llama-3.2-1B, in GiB beyond the weights and the cache, by the way attention is computed:

Tokens Cache Default attention, one shot Default, chunked by 1,024 Fused attention, one shot Fused, chunked by 2,048
1,024 0.03 0.38 0.37 0.09 0.09
2,048 0.06 1.31 0.37 0.19 0.19
4,096 0.12 4.89 0.37 0.38 0.26
8,192 0.25 (≈ 19, probe) 0.37 0.75 0.40
16,384 0.50 — 0.69 1.50 0.76
32,768 1.00 — 1.36 3.01 1.48

The default path’s extra memory grows by 3.4 then 3.7 times each time the context doubles. It is quadratic: the attention scores are being materialised, as the Chapter 21 finding that the library’s grouped-attention path takes a slow route on this build predicted. A 32,768-token prefill that way would, by extrapolating the square, need hundreds of gibibytes. Chunking the prefill into 1,024-token pieces (Chapter 9’s idea) bounds the extra memory by the chunk, at 0.37 GiB flat up to 8,192 tokens, and was 100 times faster than the one-shot attempt at 8,192 in the probe, which had spilled. A fused attention (the library’s sdpa call, with the shared key/value heads repeated for prefill and the queries grouped for decode) is linear in memory by itself and also fast. I wrote it as a drop-in override of the library’s attention function and tested it against the default path on a small model: logits agree to \(10^{-4}\), chunked prefill agrees with one-shot, and the original function is restored afterwards.

So on this build the practical limit on context is not the 4 GiB of cache for a 1B model at 128K but the 19 GiB a default prefill wants at 8K. Chapter 21’s kernel finding and this one are the same fact seen at two scales. What the cache arithmetic says is the limit is only a limit for a stack that computes attention without materialising it.

Two further notes on memory, both about how a system fails rather than how much it needs.

The logits. A language model produces a vector over its vocabulary at every position. Llama’s vocabulary is 128,256 entries, so logits for a 4,096-token prefill are a 1 GiB tensor, and for 32,768 tokens 7.8 GiB, for 131,072 tokens 31 GiB. A prefill needs the logits of its last position only; asking the library for just that (logits_to_keep) cut the extra memory at 2,048 and 4,096 tokens from 0.56 to 0.19 GiB and from 1.12 to 0.38 GiB. A server that forgets this will fail on long prompts for reasons that have nothing to do with the cache.

Oversubscription is silent. On Windows the driver lets a process allocate beyond the card’s memory and falls back to system memory. A probe run with a 5.4 GB model server running alongside did exactly that: a 32,768-token prefill took 25.5 seconds, against 6.6 seconds when the same run fitted in memory, with no error. Every long run in this chapter checks the free memory first and records a skip instead of proceeding, and I discarded the contaminated ones. A latency figure measured while paging is a measurement of the paging.

Claim four: it answers in time

Time to first token for Llama-3.2-1B (best of three, fused attention):

Tokens One shot Chunked by 4,096 Attention’s share of the FLOPs
1,024 0.093 s 0.093 s 3%
4,096 0.40 s 0.40 s 10%
8,192 0.92 s 1.12 s 18%
16,384 2.34 s 3.18 s 30%
32,768 6.58 s 9.91 s 46%
49,152 not run 20.2 s 57%

Prefill FLOPs have two terms: \(2NT\) for the dense layers and \(2 L\, d_{attn}\, T^2\) for attention (causality halves each of the two matrix products in \(4 L d_{attn} T^2\)). They are equal at \(T = N / (L\, d_{attn})\), which for this 1.24-billion-parameter, 16-layer, 2,048-wide-attention model is 37,700 tokens. Below that, prefill time grows roughly linearly in the prompt length; above it, quadratically. Both regimes show in the table: 0.093 s per thousand tokens at 1K and 0.20 at 32K.

A single throughput, fitted on the short prompts (28 TFLOP/s), underestimates the 32K prefill by 18%, because attention runs less efficiently than the dense layers on this card, whose fused kernels are not the flash-style ones (Chapter 8). A model with two throughputs, fitted on only the 4K, 8K and 16K runs, gave 30.0 TFLOP/s for the dense layers and 17.8 for attention, and predicted the held-out 32K time as 6.65 seconds against 6.58 measured. That is a two-parameter fit to three points with one degree of freedom, checked once; it is evidence that the shape of the cost is right, not a calibrated predictor.

Chunking costs time: 9.9 seconds against 6.6 at 32K. My guess is that each chunk’s attention repeats the keys and values of the whole cache so far, so that the cost lands in the part that grows; I did not isolate it.

Decode has its own story. A decode step reads the weights (2.5 GB for this model) and the whole cache, so by the bytes alone the step should rise from 8.2 ms at 1K cached tokens to 13.4 ms at 49K. Measured, at batch 1:

Cached tokens 1K 4K 8K 16K 32K 49K
Step, ms 14.7 17.3 23.0 33.1 56.6 85.2
Bytes alone, ms 8.2 8.6 9.0 9.9 11.7 13.4

The step grows by 1.47 ms for every thousand tokens of cache, 13 times what reading those bytes would cost. The cause is in the kernel. For a single sequence the attention of a decode step is one query against a long cache. The timings are what a kernel that parallelises over heads and queries, and not over the length of the cache, would show: with 8 key/value heads and one query token most of the card would sit idle. I did not inspect the kernel to confirm that. In isolation, one layer of that attention against 32,768 cached tokens took 2.02 ms where the bytes alone would take 0.22; over 16 layers that is 32 ms, which accounts for 77% of the 42 ms by which the step grew between 1K and 32K tokens. Splitting the cache into 8 parts and merging the partial softmaxes (the idea of Chapter 13’s split-KV decode) halved the kernel to 0.98 ms and still sat 4.5 times above its floor. I did not put the split-KV kernel inside the model, so I make no end-to-end claim about it. And I did not isolate the remaining 23% of the growth.

The conclusion for a long-context service is the batch-one version of Chapter 21’s: the cache’s bytes are not what limits decode at long context; the kernel’s ability to use the card is. A batch of sequences gives the kernel the parallelism a single one lacks (Chapter 21’s batch-16 measurements stayed near the floor), which is one more reason long-context serving at batch one is unusually expensive per token.

Claim five: it uses what it is given

Accepting, fitting and answering say nothing about whether the model draws on the tokens it has been given. Two measurements.

Does distant text lower the loss? I took two 32,768-token passages from chapters 1 to 14, split each into eight segments of 4,096 tokens, and compared the loss on the second half of each segment (the same target tokens) when the model had the whole text so far with when it had only its own segment. In segment 0 the two are the same computation and differed by 0.000, which is the check. In segments 1 to 7 the whole prefix gave a loss lower by 0.33 nats on average (range 0.06 to 0.50), which is a perplexity 1.39 times lower. The model uses text from thousands of tokens back. The variation across segments is large and the sample is two passages of a text that, being a book about one subject with repeated section structure, rewards copying; the figure says the effect exists, not how large it is for other text.

Can it find a fact? The needle test places one short fact with a number in it (such as “The vault access code is 7391.”) at a chosen depth in a passage of book text of a chosen length, asks a question about it at the end, and checks the answer for the number. I used 12 different facts and, for each, a different stretch of the book as filler; 5 depths (0, 25, 50, 75 and 100 per cent) and 5 lengths (2K to 32K tokens); 300 prompts, greedy answers of up to 12 tokens, the model’s chat template, no fine-tuning. As a control, the same questions with the fact removed scored 0 of 12 at both 2K and 8K, so the answers cannot be guessed. Correct answers per 12:

Length Depth 0% 25% 50% 75% 100%
2K 12 9 10 7 11
4K 12 10 11 11 12
8K 12 10 8 9 12
16K 11 8 6 9 12
32K 11 7 12 8 12

Two patterns, one clear and one not. Position matters. Pooled over lengths, the start and the end of the passage scored 58 and 59 of 60; the three middle depths 44, 47 and 44. The facts at the ends were found 117 times in 120 (97.5%), those in the middle 135 times in 180 (75%), a difference with an exact-test probability of \(3 \times 10^{-8}\) if position did not matter. This is the pattern Liu et al. report under the title “Lost in the Middle,” on multi-document question answering and key-value retrieval, for models including ones built for long inputs; here it appears in a small instruction-tuned model on a different task. Length matters much less than one would expect. Pooled over depths, correct answers by length were 49, 56, 51, 46 and 50 of 60: not monotone, and 32K is no worse than 2K. The 2K row is the second worst, and its weakest cell (7 of 12) is as weak as the 32K row’s. With twelve facts per cell and sampling-noise-sized fluctuations, I would not read anything into the ordering of lengths.

The failures were not spread evenly across facts. Three of the twelve were found every time (pier 48, the key number 913, the card number 40517), and two were missed often (the ship’s bell, 14 times, 13 of 25 correct; the lighthouse steps, 117, 15 of 25). Of the 48 wrong answers, by a rough classification, 27 were a bare wrong number (“12” for 14), 13 restated the sentence with another number, and 8 contained no number, as when the model summarised the chapter instead of answering. Failing on a number it has just been shown is an unreliable-copying failure and not a failure to find the passage, and the test as built cannot tell the two apart. A model larger than 1B parameters may behave differently; the 3B model did not fit beside the other program’s memory, and I did not run it.

Truncating, and what a cache survives

When a prompt exceeds what the system will take, something is dropped. The policy matters twice. A policy that drops the oldest tokens after a preserved system prompt keeps the prompt’s instructions and the most recent turns, which truncate_keep_prefix implements and tests. And it decides what the cache can reuse: a cache built for the long prompt is valid, exactly, only for the tokens before the first one that changed (the common prefix), because every later token’s keys and values at every layer after the first were computed having attended to what has now been removed (and, if positions are renumbered after the cut, their rotary phases change as well). For a truncation that drops tokens from the middle the reusable part is the preserved prefix alone, which the test shows for a 100-token prompt cut to 40 (reusable: the first 10). Systems that keep more than that are keeping approximate state, and what approximation costs I did not measure.

This is Chapter 16’s identity rule met from the other side: a prefix cache hits when the prefix is the same, and truncation is the operation that most reliably makes it different.

What the numbers do not show

  • One 1B model and a smaller one. Llama-3.2-1B and TinyLlama. A 3B or 8B model at long context did not fit here, and may use its context differently.
  • A passage of one book. Both the loss and needle tests use this book’s text as filler; its repetition favours copying, and a different corpus would give different numbers.
  • Needle retrieval is the easiest long-context task. It tests finding a short, distinctive string, not combining several facts, reasoning across them or summarising.
  • Twelve facts per cell. Individual cells are noisy; the pooled position effect is strong, the length ordering is not established.
  • No fine-tuned extension. Position Interpolation, YaRN and LongRoPE were described from their abstracts and not run.
  • A hand-written attention override. It is tested against the default and fixes the memory problem on this build, but it is mine, and the quadratic path I measured is a property of this build and this card (as Chapter 21 found).
  • A shared GPU. Other programs’ memory forced skips and contaminated early probes. The tabulated numbers are from runs that checked free memory first.
  • Prefill throughputs are for this card. The two-constant fit has one degree of freedom.

Inference Lab increment

The Lab now has:

  • kv_bytes, max_context_for_budget, prefill_flops, attention_crossover_tokens, decode_bytes_per_step: the arithmetic of what a context costs.
  • fused_attention: a context manager replacing the library’s attention with a fused repeated-heads form for prefill and a grouped-query form for decode, restoring the original afterwards.
  • prefill_chunked, token_nll: chunked cache filling and per-position loss against a growing cache.
  • insert_at_depth, contains_answer, a set of twelve facts, and a chat-templated prompt builder for the needle test.
  • truncate_keep_prefix, reusable_prefix_length: truncation that preserves a prefix, and the exact amount of cache that survives it.
  • limits, latency, positions, use, decode_kernel and needle commands.

The rule this adds: a context length is five claims, and a system supports it only to the extent that it has measured each. Configuration says what is accepted; the positional scheme says where the model is trustworthy; the memory arithmetic says what fits only if the computation does not materialise more than the cache; the latency has a quadratic term and a kernel term; and what the model uses depends on where the evidence is.

What you should now be able to do

  • Compute the cache of a model at its full trained context and the context that fits a budget, and notice when a model’s advertised length cannot be held on the hardware at hand.
  • Explain what a RoPE scaling setting is part of in a checkpoint’s identity, and what skipping the training that goes with it costs.
  • Measure prefill memory against context for a given stack, and recognise materialised attention scores, a full-vocabulary logits tensor and silent oversubscription as separate failure modes.
  • Split a prefill’s cost into dense and attention terms, find where attention dominates, and say why decode at long context can be limited by a kernel rather than by bytes.
  • Design a needle test with a control, a depth sweep and enough items per cell to separate position from noise.
  • State exactly how much of a cache survives a truncation.

Key terms

  • Context length: the number of tokens a system accepts; here separated into five claims.
  • RoPE scaling: a change to the rotary angles that lets a checkpoint run at positions beyond its original training length.
  • Position interpolation: scaling position indices down to the trained range.
  • Chunked prefill: filling the cache in pieces so that peak memory is bounded by the chunk.
  • Logits to keep: the number of final positions for which logits are computed.
  • Needle in a haystack: a retrieval test with a planted fact.
  • Lost in the middle: lower accuracy for evidence in the middle of a long input.
  • Reusable prefix: the part of a cache still valid after the prompt changes.

Reader exercises

  1. Run the needle test on Llama-3.2-3B when the card is free. Is the middle dip larger or smaller?
  2. Put the split-KV decode kernel inside the model, and measure the step at 32K cached tokens. How much of the 77% does it recover?
  3. Fine-tune TinyLlama with linear position scaling for 500 steps and repeat the position table. How much of the 0.7 to 1.0 nats lost at short positions comes back?
  4. Add a question that needs two facts planted at different depths. How does accuracy change with the distance between them?
  5. Measure the quality of approximate cache reuse after a mid-prompt truncation: keep the stale keys and values and compare the next 128 tokens’ loss with recomputation.
  6. Compare the needle result with retrieving the paragraph by embedding similarity and putting only it in a 1K prompt. How do accuracy and cost compare?
  7. Repeat the prefill-memory table for a 3B model. At what length does the default path exhaust the card?
  8. Use the two-constant model to predict the 49K prefill from the 4K to 16K fit. How far off is it?

References

Evidence in this chapter

Claim Category
Position Interpolation, YaRN and LongRoPE reach the lengths their abstracts state Reported by the authors; not run
“Lost in the middle” pattern in long-context models Established background (the paper); reproduced here for one small model
A model can be configured for 131,072 positions and have a cache that cannot be held (Llama-3.2-3B: 14 GiB) Arithmetic from the cached config
TinyLlama’s loss triples past its trained length; linear scaling flattens it at a cost of 0.7 to 1.0 nats at short positions Book result, 6 passages
Llama-3.2-1B degrades after 8,192 positions when its shipped scaling is removed (2.73 to 5.68 nats at 12K to 16K) Book result, 2 passages
The library’s default attention needs memory quadratic in the context on this build (0.38, 1.31, 4.89 GiB at 1K, 2K, 4K); chunking or a fused override makes it linear or constant Book result for this build
Logits for every position cost 0.56 to 1.12 GiB at 2K to 4K tokens against 0.19 to 0.38 for the last position Book result
Oversubscription of the card is silent and multiplies the time (25.5 s against 6.6 s) Book result from a probe
A two-throughput model fitted at 4K to 16K predicts the 32K prefill within 1% Book result, one degree of freedom
Decode at batch 1 grows 1.47 ms per 1,000 cached tokens, 13× the bytes; the isolated kernel accounts for 77% of the growth Book result for this build
Whole-prefix context lowers loss by 0.33 nats (perplexity 1.39× lower) on book text Book result, 2 passages
Needle retrieval is 97.5% at the ends and 75% in the middle (p = 3e-8), and flat in length from 2K to 32K Book result, one 1B model, 12 facts
Models above 1B, fine-tuned extensions, multi-fact tasks, approximate cache reuse Not measured

Implementation status

  • Manuscript: content/books/inference/23-chapter.md
  • Runnable examples: examples/ch23/long_context.py (original toy; superseded), examples/ch23/long_context_lab.py (arithmetic, fused attention override, chunked prefill and the experiments)
  • Tests: tests/test_ch23_long_context.py, tests/test_ch23_long_context_lab.py (cache and FLOP arithmetic, the attention crossover, prefix reuse and truncation, needle insertion and exact-number matching, the fused override equal to the default and restored afterwards, chunked prefill equal to one-shot and continued by decoding, per-position loss independent of chunk size, the effect of a RoPE change; all on CPU with a small random Llama)
  • Evidence: research/ch23/sources.md, research/ch23/results.json, research/evidence-register.md
  • Verification command: python -m unittest discover -s tests -v

The chapter is a draft with draft = false. It makes no claim about models above 1B parameters at long context, about fine-tuned context extension, or about long-context reasoning. Chapter 24 steps back from the individual techniques of Chapters 17 to 23 and asks which of them belong to the model and which to the runtime.