Building GPT from Scratch — Part 6: Inference
TL;DR
Naive generation recomputes all previous tokens every step — O(n²) total. KV cache stores past Keys and Values so each new token only computes one query against the cached context — O(n) total. This is how every production LLM serves responses fast. For sampling: temperature controls randomness (low=focused, high=creative), top-k limits candidates to the k most likely, top-p adapts the candidate set to the model's confidence. Good defaults: temperature=0.8, top_k=40, top_p=0.9.
The Problem
The model is trained. Now we need to generate text. This seems simple — run the model, get the next token, repeat — but there are two critical challenges:
Speed: the naive approach is O(n²) per generated token. For long sequences, this is unacceptably slow.
Quality: the model outputs probabilities, not words. How you convert probabilities to tokens dramatically affects output quality.
How Autoregressive Generation Works
GPT generates text one token at a time:
Step 1: Input "The" → Model predicts next token → "cat"
Step 2: Input "The cat" → Model predicts next token → "sat"
Step 3: Input "The cat sat" → Model predicts next token → "on"
...Each step, we:
Feed the entire sequence so far into the model
Take the logits at the last position (the model's prediction for the next token)
Sample a token from those logits
Append it to the sequence
Repeat
The code (naive version, no optimization):
for i in range(max_new_tokens):
# Feed ENTIRE sequence every time
x_cond = x[:, -context_length:] # truncate if too long
logits, _ = model(x_cond) # full forward pass
logits = logits[:, -1, :] # only the last position matters
next_token = sample(logits) # pick a token
x = torch.cat([x, next_token], dim=1) # appendThe problem: at step 100, we process 100 tokens through all layers. At step 200, we process 200 tokens. Most of this computation is redundant — we're recomputing the same attention outputs for tokens we've already processed.
The KV Cache: The Key Optimization
The Insight
In causal self-attention, a fundamental property holds: past tokens' Keys and Values never change.
When generating token 50, the Key and Value vectors for tokens 0–49 are identical to what they were when we generated token 49. The causal mask guarantees this — token 50 doesn't affect the computation of earlier tokens.
So why recompute them?
How It Works
The KV cache stores the Key and Value tensors for all previously processed tokens. On each new step, we only compute Q, K, V for the new token, then attend against the cached K, V.
Quadratic → linear. For generating 100 tokens with context 100, that's ~100× less compute.
The Two Phases
Generation with KV cache has two distinct phases:
Phase 1 — Prefill: Process the entire prompt at once. This populates the cache.
model.clear_cache()
logits, _ = model(prompt_tokens, use_cache=True, pos_offset=0)
# Cache now holds K,V for all prompt tokensPhase 2 — Decode: Generate tokens one at a time, using the cache.
for i in range(max_new_tokens):
pos_offset = current_sequence_length - 1
logits, _ = model(last_token_only, use_cache=True, pos_offset=pos_offset)
next_token = sample(logits)Implementation in the Attention Layer
class CausalSelfAttention(nn.Module):
def __init__(self, config):
...
self.cache_k = None
self.cache_v = None
def clear_cache(self):
self.cache_k = None
self.cache_v = None
def forward(self, x, use_cache=False):
B, T, C = x.shape
# Compute Q, K, V for the NEW tokens only
qkv = self.qkv(x)
q, k, v = qkv.split(C, dim=-1)
# reshape to [B, n_heads, T, head_dim]...
if use_cache:
if self.cache_k is not None:
# Append new K,V to cached K,V
k = torch.cat([self.cache_k, k], dim=2)
v = torch.cat([self.cache_v, v], dim=2)
# Update cache
self.cache_k = k
self.cache_v = v
# Attention: new Q attends to ALL K (cached + new)
# q: [B, heads, T_new, head_dim] (usually T_new = 1 during decode)
# k: [B, heads, T_total, head_dim] (all cached + new)
attn = (q @ k.transpose(-2, -1)) * (head_dim ** -0.5)
...During decode, T_new = 1 (one new token) and T_total grows by 1 each step. The attention computation is just one query attending to N keys — a vector-matrix multiply instead of a matrix-matrix multiply.
Position Encoding with Cache
With KV cache, we only process one token at a time during decode. But it still needs the correct positional embedding:
def forward(self, idx, use_cache=False, pos_offset=0):
B, T = idx.shape
# Position is NOT always 0,1,2,...T
# During decode, it's the absolute position: e.g., 57 for the 58th token
pos = torch.arange(pos_offset, pos_offset + T, device=idx.device)
x = self.token_emb(idx) + self.pos_emb(pos)Without pos_offset, every decode step would use position embedding 0 — the model would think every token is the first token. This is a subtle but critical bug.
When Cache Doesn't Help
KV cache adds overhead:
Memory: storing K,V for every layer × every head × every position
torch.catoperations each step (tensor concatenation)
For our tiny model (3.3M params, context_length=128), the overhead exceeds the savings. The cache starts paying off with:
Larger models (more computation per layer)
Longer sequences (more redundant recomputation to avoid)
Production serving (where latency matters)
At GPT-4 scale, KV cache isn't optional — it's the difference between 2 seconds and 2 minutes per response.
Sampling: Turning Probabilities into Text
The model outputs logits: one number per vocabulary token. Higher logit = more likely token. But we need to pick ONE token. How?
Greedy Decoding (Always Pick the Most Likely)
next_token = logits.argmax(dim=-1)Pros: deterministic, "safest" choice. Cons: repetitive, boring, often loops ("the the the the...").
Why it loops: if "the" is the most likely token in some context, greedy always picks it, creating a context where "the" is again most likely. Greedy decoding is almost never used for open-ended generation.
Temperature: Controlling Randomness
logits = logits / temperature
probs = softmax(logits)
next_token = multinomial(probs, num_samples=1)Temperature reshapes the probability distribution:
Original logits: [2.0, 1.0, 0.5, -1.0]
Softmax: [0.47, 0.17, 0.10, 0.02] (roughly)
Temperature = 0.5 (sharper):
Scaled logits: [4.0, 2.0, 1.0, -2.0]
Softmax: [0.84, 0.11, 0.04, 0.00] ← top token dominates
Temperature = 2.0 (flatter):
Scaled logits: [1.0, 0.5, 0.25, -0.5]
Softmax: [0.34, 0.21, 0.16, 0.08] ← more uniformLow temperature (0.2–0.5): model is very confident, picks top tokens almost always → coherent but repetitive
Temperature = 1.0: model's original distribution → natural variety
High temperature (1.5–2.0): all tokens become more likely → creative but potentially incoherent
Top-k Sampling: Limiting Choices
# Only consider the k most probable tokens
top_k_values, _ = torch.topk(logits, k=40)
threshold = top_k_values[:, -1:] # value of the 40th token
logits[logits < threshold] = float("-inf") # zero out everything else
probs = softmax(logits)
next_token = multinomial(probs)The problem top-k solves: even with reasonable temperature, low-probability tokens sometimes get sampled. Token 500 with 0.001% probability is almost certainly wrong — sampling it produces garbage.
Top-k says: only consider the top 40 candidates. Everything else is impossible.
The limitation: k=40 is fixed. Sometimes the model is very confident (2 tokens have 90% combined probability) and sometimes uncertain (100 tokens each have 1%). A fixed k doesn't adapt.
Top-p (Nucleus) Sampling: Adaptive Filtering
sorted_logits, sorted_idx = torch.sort(logits, descending=True)
cumulative_probs = torch.cumsum(softmax(sorted_logits), dim=-1)
# Keep tokens until cumulative probability exceeds p
mask = cumulative_probs - softmax(sorted_logits) >= top_p
sorted_logits[mask] = float("-inf")Top-p keeps the smallest set of tokens whose combined probability exceeds p (e.g., 0.9).
If the model is confident: "Paris" has 85% → only 2-3 tokens in the nucleus
If the model is uncertain: many tokens at 2-3% each → 30-40 tokens in the nucleus
This adapts to the model's confidence. When the answer is clear, it narrows the choices. When it's ambiguous, it allows variety.
Combining Them
In practice, we use all three together:
logits = logits / temperature # reshape distribution
logits = top_k_filter(logits, k) # remove long tail
logits = top_p_filter(logits, p) # adaptive narrowing
probs = softmax(logits)
next_token = multinomial(probs)Good defaults: temperature=0.8, top_k=40, top_p=0.9
This means: slightly sharpen the distribution, limit to top 40 candidates, then further narrow to 90% cumulative probability, then sample.
The Full Generation Function
def generate(model, tokenizer, prompt, max_new_tokens=200,
temperature=0.8, top_k=40, top_p=0.9, use_cache=True):
model.eval()
token_ids = tokenizer.encode(prompt)
x = torch.tensor([token_ids], device=device)
with torch.no_grad(): # no gradients needed for inference
if use_cache:
model.clear_cache()
# PREFILL: process prompt, populate cache
logits, _ = model(x, use_cache=True, pos_offset=0)
logits = logits[:, -1, :] / temperature
next_token = sample(logits, top_k, top_p)
x = torch.cat([x, next_token], dim=1)
# DECODE: one token at a time using cache
for i in range(1, max_new_tokens):
pos_offset = x.shape[1] - 1
if pos_offset >= model.config.context_length:
break # can't exceed context window
logits, _ = model(x[:, -1:], use_cache=True, pos_offset=pos_offset)
logits = logits[:, -1, :] / temperature
next_token = sample(logits, top_k, top_p)
x = torch.cat([x, next_token], dim=1)
model.clear_cache()
return tokenizer.decode(x[0].tolist())Key details:
model.eval(): disables dropout (we want deterministic behavior during inference)torch.no_grad(): disables gradient tracking (saves memory, not training)pos_offset: ensures correct positional encoding for cached generationContext length check: stops if we hit the model's maximum sequence length
What Happens During Generation (Traced)
Let's trace one step of generation with verbose output:
Step 4 of generation:
Input: token 395 ("qu")
Position: 6 (absolute position in sequence)
→ Token embedding: [0.02, -0.15, 0.31, ...] (256 dims)
→ + Position embedding for pos 6
→ Through 4 transformer blocks:
Block 0: attention to all 6 previous tokens, FFN
Block 1: deeper patterns
Block 2: even deeper
Block 3: final representation
→ LayerNorm → Linear head → 512 logits
Top-5 predictions:
token 354 ("ro"): logit 3.2, prob 14.9%
token 106 ("j"): logit 2.7, prob 9.4%
token 276 ("en"): logit 2.5, prob 7.8%
token 101 ("e"): logit 2.3, prob 6.2%
token 349 ("it"): logit 2.2, prob 5.8%
After top-k (40) + top-p (0.9) + temperature (0.8):
Sampled: token 276 ("en")
→ Append to sequence, next step uses token 276 as inputEach step: one token in, one token out. The model's entire "intelligence" is compressed into this: look at the context, predict what comes next.
Performance Considerations
Memory
KV cache memory per layer = 2 × batch_size × n_heads × seq_len × head_dim × sizeof(float)
For our model: 2 × 1 × 4 × 128 × 64 × 4 bytes = 262 KB per layer × 4 layers = ~1MB
For GPT-4 scale (128 layers, 96 heads, 8K context, 128 head_dim): 2 × 1 × 96 × 8192 × 128 × 2 bytes (fp16) = ~3GB per layer × 128 layers = ~384GB
This is why serving large models is expensive — the KV cache alone can exceed GPU memory.
Latency
Two metrics matter:
Time to First Token (TTFT): how long the prefill takes (proportional to prompt length)
Time per Output Token (TPOT): how fast decode generates (one forward pass of 1 token)
For our model on MacBook MPS:
Prefill: ~5ms for 128 tokens
Decode: ~15ms per token (~65 tokens/second)
Stopping
When does generation stop?
Hit
max_new_tokens→ hard limitHit
context_length→ model can't process moreGenerate an end-of-sequence token → (we haven't implemented this, but production models have a special
<EOS>token)
Running the Full Pipeline
# Complete pipeline: download → tokenize → train → finetune → generate
python run.py
# Or just interactive generation with existing checkpoint
python generate.py
# Or run the test suite
python inference.pyThe interactive mode:
>>> The meaning of life is
The meaning of life is the great force of those
considerations which the character, and the more
profoundly interested in the subject...
>>> ### Instruction:\nExplain gravity.\n\n### Response:\n
Gravity is the force that attracts objects...Key Takeaways
KV cache is the fundamental inference optimization. It reduces generation from O(n²) to O(n). Every production LLM uses it. Without it, ChatGPT responses would take minutes instead of seconds.
Prefill and decode are different computational profiles. Prefill is compute-bound (processing many tokens in parallel). Decode is memory-bound (reading cached KV from memory for one new token).
Sampling strategy controls the creativity/coherence tradeoff. Temperature, top-k, and top-p work together. There's no single "right" setting — it depends on the task (code generation wants low temperature; creative writing wants high).
The model has a fixed context window. It cannot process or generate beyond
context_lengthtokens. This is a fundamental architectural constraint. (Solutions: sliding window, RoPE scaling, but those are advanced topics.)Generation is embarrassingly simple. Predict next token, append, repeat. All the complexity is in the model itself (trained during pretraining/finetuning). Inference is just running the function.
Series Complete
We've built a complete GPT from scratch:
download_data.py → Get text corpus
tokenizer.py → BPE tokenization (text → numbers)
model.py → Transformer architecture
train.py → Pretraining (learn language)
finetune.py → Fine-tuning (learn behavior)
generate.py → Inference (produce text with KV cache)~500 lines of Python. Same architecture as GPT-4. Running on a MacBook.
The gap between our 3.3M model and GPT-4's ~1.8T parameters is scale, not mechanism. More layers, more heads, more data, more compute — but the same attention, the same next-token prediction, the same sampling loop.