Lab 02 — Speculative Decoding

Phase: 08 — Inference Optimization at Depth | Difficulty: ⭐⭐⭐⭐⭐ | Time: 4–5 hours

Speculative decoding (Chen et al., 2023) uses a small draft model to generate K tokens, then verifies them all in one target model pass — achieving near-lossless 2–4× throughput gains.

What you build

  • TinyLM — small draft language model
  • speculative_step — draft K tokens, target verifies in one forward pass, accept/reject with min(1, p_t/p_d), sample corrected distribution on rejection, add bonus token
  • speculative_generate — full generation loop
  • compare_distributions — total variation distance between distributions

Key concepts

ConceptWhat to understand
Accept/rejectToken t accepted with prob min(1, p_target(t) / p_draft(t))
Corrected distributionOn rejection: sample from max(0, p_t - p_d) / Z
Bonus tokenWhen all K accepted, target generates one extra token "for free"
Throughput gainWith high acceptance rate α, gain ≈ K·α + 1 tokens per target call
Draft-target alignmentMore similar models → higher α → more speedup

Files

FilePurpose
lab.pyStubs with TODO markers
solution.pyComplete working implementation
test_lab.pypytest suite
requirements.txtDependencies

Run

pip install -r requirements.txt
python lab.py
pytest test_lab.py -v