Lab 01 — Sampling & Constrained-Decoding Engine

Phase: 08 — Sampling, Decoding & Constrained Generation Difficulty: ⭐⭐⭐☆☆ (the code is small; the order and the invariants are the hard part) Time: 3–4 hours

The model gives you a vector of logits; you turn it into the next token. This lab builds the part of model.generate() nobody shows you: a numerically-stable softmax with the temperature knob, greedy argmax, the truncation filters (top-k, top-p/nucleus, min-p), the repetition penalty, the full ordered sample() pipeline, and a grammar/FSM constrained decoder with logit masking. Two facts are the soul of it: greedy == sample(temperature→0) == sample(top_k=1), and a grammar makes invalid output impossible — "please return JSON" is best-effort; a grammar is a guarantee.

What you build

• softmax(logits, temperature) — the stable, max-subtracted map from logits to a probability distribution; temperature divides the logits first (T>1 flattens, T<1 sharpens, T→0 is greedy). The max-subtraction is the lesson, not a detail — it is what keeps a 1000-magnitude logit from overflowing.
• greedy(logits) — argmax with a deterministic lowest-index tie-break.
• apply_temperature / top_k_filter / top_p_filter / min_p_filter — the truncation filters, all operating in logit space (rejected tokens → -inf, which softmax maps to probability 0). Top-p is the smallest set whose cumulative prob ≥ p; min-p is a cutoff relative to the max prob.
• repetition_penalty(logits, generated, penalty) — the HF convention: divide positive logits / multiply negative logits of already-seen tokens, both pushing them downward.
• sample(...) — the whole pipeline in the canonical order (penalty → temperature → top-k → top-p → min-p → softmax → draw), collapsing to greedy at the temperature→0 / top_k=1 limit.
• FSM + constrained_logits_mask + constrained_decode — a finite-state grammar, the mask that lists the only legal next token IDs, and a decoder that masks logits so the produced sequence is always accepted by the FSM. The concrete grammar is a tiny fixed-format JSON object {"ok": <true|false>}.

Key concepts

Files

Run

pip install -r requirements.txt
pytest test_lab.py -v                       # against your lab.py
LAB_MODULE=solution pytest test_lab.py -v   # against the reference
python solution.py                          # the worked example

Success criteria

• All tests pass against your implementation (and LAB_MODULE=solution).
• You can explain why softmax must subtract the max and prove it leaves the result unchanged.
• You can explain why test_soul_greedy_equals_sample_zero_temperature_and_top_k_one is the whole point — the sampler is a generalization of greedy, not a separate algorithm.
• You can state, from memory, the order the filters apply and why softmax runs once at the end.
• You can explain the difference between top-k (fixed size) and top-p (adaptive size) and what failure mode each one fixes.
• You can explain why test_constrained_soul_never_emits_invalid_token_and_is_accepted holds for every seed even when the logits scream for an invalid token — and why that is a stronger property than any prompt.
• Given any sampling-config question you can name the knob, its effect on entropy, and its failure mode.

How this maps to the real stack

Limits of the miniature (say these in the interview): a real vocabulary is 100k+ tokens, so top-p/top-k use partial sorts and the mask is built with a vocab→state index, not a dict scan; a real grammar compiles a regex/CFG to a DFA and precomputes the allowed-token mask per state so masking is O(1) per step, not O(vocab); subword tokenization means one grammar terminal can span many tokens (a token can be a partial match), which is the genuinely hard part Outlines solves; and beam search, contrastive search, and speculative decoding (P09) are whole other decoding strategies this lab only previews.

Extensions (build these for real depth)

• Add presence and frequency penalties (the OpenAI variants: presence subtracts a flat amount once a token appears; frequency subtracts proportional to count) and contrast them with the multiplicative HF repetition_penalty.
• Add beam search (beam_width) and show it on a translation-style toy task — then show it collapses to repetitive, generic text on open-ended generation (the Holtzman result).
• Replace the dict-scan mask with a precompiled per-state allowed-id array and benchmark the speedup as the vocab grows.
• Build a tiny regex-to-DFA compiler (e.g. for \d{4}-\d{2}-\d{2}) and feed its DFA into constrained_decode to emit guaranteed-valid dates.
• Handle multi-token terminals: let a grammar terminal match a prefix of a token and track partial progress (the subword problem real constrained decoders must solve).

Interview / resume

• Talking points: "Why must softmax subtract the max?" "Greedy vs sampling — when and why?" "Top-k vs top-p vs min-p — what does each fix?" "In what order are sampling filters applied, and why does softmax run once at the end?" "Why is 'please return JSON' weaker than a grammar, and how does constrained decoding guarantee validity?"
• Resume bullet: Built a from-scratch LLM decoding engine — numerically-stable temperature softmax, greedy/top-k/top-p/min-p sampling with the canonical filter ordering, HF-style repetition penalty, and an FSM/grammar constrained decoder with logit masking that guarantees schema-valid output — mirroring the samplers in vLLM/HF and the structured-generation guarantees of Outlines and OpenAI structured outputs.