Inference

How This Lesson Fits the Module & Volume

This is the Module 11.1 capstone. You now have the full stack: tokens → embeddings → hidden states → logits → softmax → decoding (temperature, top-k, top-p, beam) accelerated by a KV cache.

Inference ties them into one generation loop: tokenize the prompt, prefill, decode until stop, detokenize. Module 11.2 then shifts to encoder-style models starting with BERT.

Learning Objectives

By the end of this lesson, students should be able to:

Describe the end-to-end autoregressive inference pipeline.
Separate prefill from decode and name stop conditions.
Assemble temperature + top-k/p sampling inside a generate loop.
Explain where the KV cache plugs into that loop.
Contrast training forward passes with inference generation.
Choose decoding settings for a concrete product scenario.

Definition

Inference (generation / decoding) is the process of producing new tokens from a trained language model given a prompt—without updating weights. Autoregressive inference repeatedly maps context → next-token distribution → chosen token until an end condition (EOS, max length, or stop string).

The Generation Loop

Tokenize

Prompt → IDs

Prefill

Build KV cache

Decode

Sample / search

Detokenize

IDs → text

Stage	Inputs	Outputs	Module links
Tokenize	Text prompt	Token IDs	Tokenizer
Forward	IDs (+ cache)	Logits	Hidden State, Logits
Decode step	Logits + knobs	Next ID	Sampling family
Stop / emit	Sequence so far	Text stream	EOS, max_new_tokens

Training vs Inference

Training

Teacher forcing; parallel over T.
Loss on all positions.
Gradients update weights.

Sequential token decisions.
No grad; eval / no_grad mode.
KV cache for speed.

Shared

Same weights and architecture.
Causal attention semantics.
Same tokenizer vocabulary.

Choosing a Decode Strategy

Goal	Prefer	Knobs
Deterministic short answer	Greedy or low T	T≈0, or beam for structured MT
Chat / creative	Nucleus sampling	T≈0.7–1.0, top-p≈0.9
Code / formats	Low entropy	Low T, small top-k, stop sequences
Many alternatives	Independent samples	Vary seed; keep top-p

Code: Capstone Generate Loop

import torch
import torch.nn.functional as F

def top_k_top_p(logits, k=50, p=0.9):
    # logits: (1, V) — apply top-k then top-p style filters
    if k > 0 and k < logits.size(-1):
        vals, _ = torch.topk(logits, k)
        logits = torch.where(logits < vals[:, -1:], torch.full_like(logits, float("-inf")), logits)
    sorted_logits, sorted_idx = torch.sort(logits, descending=True)
    probs = F.softmax(sorted_logits, dim=-1)
    cum = torch.cumsum(probs, dim=-1)
    mask = cum - probs > p
    sorted_logits = sorted_logits.masked_fill(mask, float("-inf"))
    out = torch.full_like(logits, float("-inf"))
    out.scatter_(1, sorted_idx, sorted_logits)
    return out

@torch.no_grad()
def generate(model, input_ids, max_new_tokens=32, temperature=0.9, top_k=50, top_p=0.9, eos_id=None):
    """
    model(input_ids) -> logits (B, T, V)
    Pedagogical loop (no KV cache); production adds cache as in prior lecture.
    """
    ids = input_ids
    for _ in range(max_new_tokens):
        logits = model(ids)[:, -1, :]              # (B, V)
        logits = logits / max(temperature, 1e-5)
        logits = top_k_top_p(logits, k=top_k, p=top_p)
        probs = F.softmax(logits, dim=-1)
        next_id = torch.multinomial(probs, num_samples=1)  # (B, 1)
        ids = torch.cat([ids, next_id], dim=1)
        if eos_id is not None and (next_id == eos_id).all():
            break
    return ids

# Wire a real model + tokenizer in class; here shapes only:
# prompt = tokenizer.encode("Hello", return_tensors="pt")
# out = generate(model, prompt)
# print(tokenizer.decode(out[0]))

Production Checklist

Must-haves

model.eval() and torch.no_grad().
Stop rules: EOS, max_new_tokens, optional stop strings.
KV cache (or engine) for interactive latency.
Documented decode defaults per feature.

Common failures

Leaving dropout on at inference.
Forgetting to append the sampled token.
Unbounded generation without max length.
Mixing train-time teacher forcing with decode bugs.

Common Misconception

“Inference means running one forward pass and reading the answer.” For causal LMs, inference is a loop: each new token changes the context for the next forward (or cached) step. A single forward over the prompt only scores the immediate next token unless you keep decoding.

Knowledge Check

Short Answer: Name the four high-level stages of generation. Answer: Tokenize, prefill, decode loop, detokenize.
True/False: Autoregressive inference updates model weights. Answer: False.
Multiple Choice: Prefill’s main job is to: (a) fill KV cache for the prompt, (b) shuffle the vocab, (c) train embeddings. Answer: (a).
Short Answer: Give two stop conditions. Answer: EOS token and max_new_tokens (also stop strings).
True/False: One forward over the prompt alone generates a full paragraph. Answer: False—you must loop.
Multiple Choice: Chat defaults often use: (a) top-p sampling, (b) only random bytes, (c) no softmax. Answer: (a).
Short Answer: Why call torch.no_grad()? Answer: Disable autograd to save memory and compute during generation.
True/False: KV cache and top-p solve the same problem. Answer: False—speed vs distribution shaping.
Multiple Choice: Module 11.2 begins with: (a) BERT, (b) CNNs, (c) k-means. Answer: (a).
Short Answer: Where do logits come from each decode step? Answer: LM head on the latest hidden state (last position).

Key Takeaways

Inference is tokenize → prefill → decode loop → detokenize.
Decoding knobs shape each step; KV cache makes steps fast.
Training is parallel teacher forcing; inference is sequential.
Always set stop conditions and run in eval / no_grad.
Next module: 11.2 BERT.

Trainer’s Guide

Hands-on idea: Have teams implement generate on a tiny character LM, then A/B greedy vs top-p=0.9 on the same prompt and read outputs aloud.

Discussion prompt: Which inference failure is worse for your product: latency spikes, repetitive loops, or factual drift—and which knob addresses each?

Recap: Inference assembles Module 11.1 into a full generation loop. Continue with BERT in Module 11.2.