Inference

Training happens once; inference happens every time anyone uses the model, so this is where the money goes. Generating text has two very different phases, one memory trick that makes it affordable (the KV cache), a few knobs that decide which token comes out (sampling), and a toolbox of speedups: quantization, speculative decoding, continuous batching and prompt caching. Every one of them follows from one fact:

Generating one token requires reading every weight of the model from memory, and memory is much slower than arithmetic.

Everyday picture You are handed a letter and asked to reply. Reading the letter is fast: your eyes take in whole lines at once. Writing the reply is slow: one word at a time, and before each word you must walk to a filing cabinet and flip through an entire reference binder. The walk, not the thinking, is what takes the time.

The model is the same. Prefill reads the whole prompt in one parallel pass. Decode then produces the answer one token at a time, and every single token requires streaming all the weights from GPU memory.

Tiny worked example An 8-billion-parameter model stored at 16 bits is 16 GB of weights. On a GPU that moves 3.35 TB/s from memory and does about 10¹⁵ 16-bit operations per second:

• Prefill of a 1,000-token prompt: 2 × 8×10⁹ × 1,000 = 1.6×10¹³ operations, about 16 ms.
• Decode: read 16 GB for every token, 16×10⁹ / 3.35×10¹² s = 4.78 ms per token, at most about 209 tokens per second for a single user.

sequenceDiagram
  participant U as User
  participant M as Model
  participant C as KV cache
  U->>M: Prompt of 1,000 tokens
  M->>C: Prefill: store K and V for all 1,000
  M-->>U: First token
  loop Each new token
    M->>C: Read cached K and V, add one entry
    M-->>U: Next token
  end

The reason the phases behave so differently is arithmetic intensity: how many operations you do for each byte you fetch from memory.

# 8 billion weights at 2 bytes each (16-bit)
P, b = 8e9, 2
# bytes read per second, operations per second
BW, FLOPS = 3.35e12, 1e15
def I(n):
    # operations per byte of weights read
    return 2 * n / b
# decode, then a 1,000-token prefill
I(1), I(1000)  # → (1.0, 1000.0)
# the break-even intensity
round(FLOPS / BW)  # → 299
# t_decode, in milliseconds per token
round(P * b / BW * 1000, 2)  # → 4.78
# t_prefill for 1,000 tokens, in milliseconds
round(2 * P * 1000 / FLOPS * 1000, 1)  # → 16.0

In code: arithmetic_intensity computes I, ridge_point finds the break-even, and bottleneck says which side a pass falls on; decode_seconds_per_token and prefill_seconds give the two time bounds.

Why it matters in practice. When a system feels slow, ask which phase dominates. Slow first token: the prompt is long (trim it, or cache it). Slow streaming: decode is memory-bound (quantize, batch, speculate).

Everyday picture Reading a mystery novel, you don't reread the whole book before each new sentence; you keep notes on every character and clue. For each new sentence you glance at your notes and add one line.

Attention needs every earlier token's key and value (see primer.ml.attention). Because of causal masking, an earlier token's keys and values never change when later tokens arrive, so they can be computed once and kept.

Tiny worked example An 8-token prompt, then 16 generated tokens.

• Without a cache, step k re-processes the whole sequence so far: 8 + 9 + … + 23 = 248 token positions.
• With a cache: prefill the 8 prompt tokens once, then 1 new position for each of the next 15 tokens = 23 token positions.

On TinyDecoder that is about 11× fewer operations for this short run, and the gap grows with length: without a cache the total work grows with the square of the length.

flowchart LR
  subgraph NOCACHE["No cache: every step starts over"]
    A1[step 1: tokens 1..8] --> A2[step 2: tokens 1..9] --> A3[step 3: tokens 1..10]
  end
  subgraph CACHE["KV cache: every step adds one"]
    B1[prefill: tokens 1..8<br/>store K,V] --> B2[token 9 only<br/>read K,V 1..8] --> B3[token 10 only<br/>read K,V 1..9]
  end

In code: TinyDecoder.forward_full processes a whole sequence (and fills a cache during prefill), TinyDecoder.forward_step processes one new token against the cache from TinyDecoder.new_cache, and TinyDecoder.generate runs either way, returning a Generation that holds the tokens and the work counted.

Why it matters in practice. The cache trades memory for speed, and that memory is what limits how many users one GPU can serve. Section 3 does the math.

Everyday picture Packing for a trip. The suitcase is GPU memory. The model's weights are the big fixed items that always go in. Every active conversation adds a bag of notes (its KV cache) whose size grows with the conversation's length. Once the suitcase is full, the next customer waits.

Tiny worked example

• Weights: 70×10⁹ parameters × 2 bytes (16-bit) = 140 GB: more than one 80 GB GPU. At 4 bits (half a byte) it is 35 GB and fits on one.
• KV cache per token for a Llama-3-8B-shaped model (32 layers, 8 KV heads, 128 dimensions per head, 16-bit): 2 × 32 × 8 × 128 × 2 = 131,072 bytes, about 128 KB.
• A 32,000-token conversation: 32,000 × 131,072 ≈ 4.2 GB of cache.
• An 80 GB GPU holding 16 GB of weights has 64 GB left: room for 15 such conversations at once.

P = 7e10
# weight GB at 16 bits, then at 4 bits
P * 16 / 8 / 1e9, P * 4 / 8 / 1e9  # → (140.0, 35.0)
L, H_kv, d_h, b = 32, 8, 128, 2
# KV bytes per token: a key and a value, per layer, per KV head
2 * L * H_kv * d_h * b  # → 131072
# GB of cache for a 32,000-token conversation
round(32_000 * 131_072 / 1e9, 1)  # → 4.2

In code: weight_bytes and kv_cache_bytes_per_token are the two formulas, kv_cache_bytes scales the cache to a context and batch, and max_concurrent_requests counts how many conversations fit beside the weights.

Everyday picture Choosing where to eat. Greedy always picks the top-rated place. Sampling holds a lottery weighted by rating. Temperature is how adventurous you feel: low means you nearly always pick the favourite, high means long shots get a real chance. Top-k says "only consider the top 3". Top-p says "consider just enough places to cover 90% of my enthusiasm": one place if you have a clear favourite, several if you're torn.

Tiny worked example Scores (logits) 2, 1, 0 for three tokens:

With probabilities 0.5, 0.3, 0.15, 0.05: top-k = 2 keeps 0.5 and 0.3, renormalised to 0.625 and 0.375. Top-p = 0.9 keeps 0.5, 0.3 and 0.15 (the first set whose total reaches 0.9) and cuts the 0.05 tail.

flowchart LR
  L[Logits, one per<br/>vocabulary token] --> T[Divide by<br/>temperature T]
  T --> S[Softmax<br/>probabilities]
  S --> K[Top-k: keep the<br/>k likeliest]
  K --> P[Top-p: keep the smallest set<br/>reaching probability p]
  P --> R[Renormalise<br/>and draw one token]

import math
z, T = [2.0, 1.0, 0.0], 0.5
# e^(z_i / T)
exps = [math.exp(z_i / T) for z_i in z]
# Σ_j e^(z_j / T)
round(sum(exps), 1)  # → 63.0
# p_i
[round(e / sum(exps), 3) for e in exps]  # → [0.867, 0.117, 0.016]

Temperature 0 is not a determinism guarantee. Floating-point addition isn't associative: in 32-bit floats, (10⁸ + 1) − 10⁸ = 0 but (10⁸ − 10⁸) + 1 = 1. On a GPU the order of additions can depend on the kernel chosen and on what else is in the batch, so two nearly tied tokens can swap places between otherwise identical requests.

In code: temperature_probs applies the formula, top_k_filter and top_p_filter cut the tail, and sample_next chains them into one draw. float32_sum and greedy_pick_with_summation_order show a near tie flipping with the order of additions.

Why it matters in practice. Use low temperature for extraction, classification and tool calls; higher for brainstorming and creative text.

Everyday picture A junior writer drafts the next few sentences quickly. A senior editor reads the whole draft at once, keeps every sentence they would have written themselves, rewrites the first one they wouldn't, and throws away the rest. The senior's reading is fast; their writing is slow. So the team moves at the junior's speed but produces the senior's text.

This works because decode is memory-bound: checking 5 draft tokens in one pass of the big model costs about the same as generating 1.

Tiny worked example The big model's next-token probabilities are (0.5, 0.3, 0.2); the small model's are (0.3, 0.3, 0.4). A draft token survives with probability Σ min = 0.3 + 0.3 + 0.2 = 0.8. Drafting 4 tokens per round yields on average (1 − 0.8⁵) / (1 − 0.8) = 3.36 tokens per pass of the big model instead of 1.

flowchart TD
  D[Small model drafts γ tokens<br/>one by one, cheap] --> V[Big model scores all γ positions<br/>in ONE parallel pass]
  V --> C{For each draft x in order:<br/>keep with probability min 1, p/q}
  C -->|kept| N[Next draft]
  N --> C
  C -->|rejected| F[Replace x with a draw from<br/>max 0, p − q, renormalised. Stop.]
  C -->|all kept| B[Bonus: draw one more<br/>token from the big model]

# big model
p = [0.5, 0.3, 0.2]
# small model
q = [0.3, 0.3, 0.4]
# P(keep x) for each token
[round(min(1.0, p_x / q_x), 2) for p_x, q_x in zip(p, q)]  # → [1.0, 1.0, 0.5]
alpha = sum(min(p_x, q_x) for p_x, q_x in zip(p, q))
round(alpha, 2)  # → 0.8
gamma = 4
# expected tokens per big-model pass
round((1 - alpha ** (gamma + 1)) / (1 - alpha), 2)  # → 3.36

In code: acceptance_rate computes α and expected_tokens_per_round the geometric series; speculative_round runs the draft, verify and replace loop once, and speculative_generate repeats it until enough tokens exist.

Everyday picture Writing prices to the nearest dollar instead of the nearest cent: shorter to store, slightly less exact. Using one ruler per row of a spreadsheet, instead of one for the whole sheet, keeps a single huge number in one row from making every other row coarse.

Tiny worked example A row of weights (0.5, −1.27, 0.02).

• int8: the largest magnitude, 1.27, maps to 127, so the step size is 0.01. Codes: 50, −127, 2. Decoding gives back exactly (0.5, −1.27, 0.02).
• int4: only 15 levels (−7 … 7), step size 1.27/7 = 0.181. Codes: 3, −7, 0. Decoding gives (0.544, −1.27, 0): the small weight vanished.

w = [0.5, -1.27, 0.02]
def quantize(w, bits):
    # the step size
    s = max(abs(w_j) for w_j in w) / (2 ** (bits - 1) - 1)
    # c_j: whole steps
    return s, [round(w_j / s) for w_j in w]
s, c = quantize(w, bits=8)
round(s, 3), c  # → (0.01, [50, -127, 2])
s, c = quantize(w, bits=4)
round(s, 3), c  # → (0.181, [3, -7, 0])
# ŵ_j = s · c_j: the 0.02 is gone
[round(s * c_j, 3) for c_j in c]  # → [0.544, -1.27, 0.0]

In code: quantize returns the integer codes and one scale per row, dequantize multiplies them back, and quantization_error measures how far the round trip lands from the original weights.

Why it matters in practice. 8-bit weights are nearly lossless; 4-bit methods with smarter rounding (GPTQ, AWQ) keep most quality at a quarter of the memory. Fewer bytes per weight also means faster memory-bound decode.

Everyday picture A restaurant with two tables. Static batching seats two parties and won't seat anyone new until both have left, so a table sits empty while one slow diner lingers. Continuous batching seats the next party the moment any table frees.

Tiny worked example Four requests needing 4, 1, 1 and 1 decode steps, two slots. Static: {4, 1} runs 4 steps with one slot idle for 3, then {1, 1} runs 1: 5 steps, 70% of slot-steps busy. Continuous: the short requests slide into the freed slot while the long one runs: 4 steps, 87.5% busy.

# decode steps the four requests need
useful = 4 + 1 + 1 + 1
slots = 2
# static takes 5 steps, continuous 4
useful / (5 * slots), useful / (4 * slots)  # → (0.7, 0.875)

In code: simulate_static_batching and simulate_continuous_batching play out the two policies step by step, each returning a ServingRun that holds the slot timeline and its utilisation U.

Why it matters in practice. Continuous batching (together with paged KV-cache memory, as in vLLM) is a large part of why modern inference servers reach high throughput.

Everyday picture A kitchen that pre-chops the ingredients every order uses, instead of chopping them again for each plate.

The KV cache normally lives for one request. Prompt caching keeps it across requests: if many calls start with the same long system prompt, tool definitions or reference document, the provider stores that prefix's keys and values and skips its prefill next time. It is only valid up to the first differing token, because every token's keys depend on everything before it, so put stable content first and volatile content last.

Tiny worked example 100 requests, each a 10,000-token shared prefix plus a 500-token question, at $3 per million input tokens, with cache writes at 1.25× and cache reads at 0.1× (typical of providers; check current pricing):

• Without caching: 100 × 10,500 × $3/10⁶ = $3.15.
• With caching: the first call writes the cache ($0.039); each of the other 99 costs $0.0045. Total $0.48, an 85% saving, and each cached call also skips 10,000 tokens of prefill, so its first token arrives sooner.

flowchart LR
  subgraph Prompt["One request's prompt, in order"]
    S[System prompt<br/>stable] --> T[Tool definitions<br/>stable] --> D[Reference docs<br/>stable] --> Q[User question<br/>changes every call]
  end
  S & T & D -.->|cached after the first call| C[(Prefix KV cache)]

In code: reusable_prefix_tokens counts how many leading tokens two prompts share, and prompt_cache_cost prices a run of requests with and without the cache.

Answer each one out loud or on paper before you open it. If you can explain it, you know it.