From 774f899ecfff54a1da0a6e4792ab616a1570652a Mon Sep 17 00:00:00 2001 From: xuanyili Date: Sat, 13 Jun 2026 00:21:46 -0700 Subject: [PATCH 1/2] Add speculative decoding lesson --- README.md | 2 +- book/src/SUMMARY.md | 1 + book/src/week3-04-speculative-decoding.md | 240 ++++++++++++++++++++++ src/tiny_llm_ref/generate.py | 84 ++++---- tests_refsol/test_week_3_day_4.py | 225 ++++++++++++++++++++ 5 files changed, 509 insertions(+), 43 deletions(-) create mode 100644 book/src/week3-04-speculative-decoding.md create mode 100644 tests_refsol/test_week_3_day_4.py diff --git a/README.md b/README.md index 73164461..1031bbd1 100644 --- a/README.md +++ b/README.md @@ -47,7 +47,7 @@ Week 1 and 2 is complete. Week 3 is in progress. | 3.1 | Paged Attention - Part 1 | ✅ | ✅ | 🚧 | | 3.2 | Paged Attention - Part 2 | ✅ | ✅ | 🚧 | | 3.3 | MoE (Mixture of Experts) | ✅ | ✅ | 🚧 | -| 3.4 | Speculative Decoding | 🚧 | ✅ | 🚧 | +| 3.4 | Speculative Decoding | ✅ | ✅ | ✅ | | 3.5 | RAG Pipeline | 🚧 | 🚧 | 🚧 | | 3.6 | AI Agent / Tool Calling | 🚧 | 🚧 | 🚧 | | 3.7 | Long Context | 🚧 | 🚧 | 🚧 | diff --git a/book/src/SUMMARY.md b/book/src/SUMMARY.md index 590dbc9d..5eb0c954 100644 --- a/book/src/SUMMARY.md +++ b/book/src/SUMMARY.md @@ -22,6 +22,7 @@ - [Paged Attention, Part 1](./week3-01-paged-attention-part1.md) - [Paged Attention, Part 2](./week3-02-paged-attention-part2.md) - [Mixture of Experts](./week3-03-moe.md) + - [Speculative Decoding](./week3-04-speculative-decoding.md) --- diff --git a/book/src/week3-04-speculative-decoding.md b/book/src/week3-04-speculative-decoding.md new file mode 100644 index 00000000..99fdbec8 --- /dev/null +++ b/book/src/week3-04-speculative-decoding.md @@ -0,0 +1,240 @@ +# Week 3 Day 4: Speculative Decoding + +In this chapter, we will implement **speculative decoding** for the tiny-llm +generation loop. + +By this point, the model can already decode with a KV cache. That is correct, +but it still uses the large target model once per output token: + +```plain +large model -> token 1 +large model -> token 2 +large model -> token 3 +... +``` + +Speculative decoding changes the shape of that loop. A smaller draft model +proposes several tokens, and the larger target model verifies them in one +forward pass: + +```plain +draft model: propose token 1, token 2, token 3, token 4 +target model: verify those tokens together +``` + +If the target model agrees with the draft tokens, we can emit several tokens +after one target-model call. If it disagrees, we keep the accepted prefix, +replace the first wrong draft token with the target model's token, and rewind +the cached K/V for the rejected suffix. + +This chapter uses a deterministic greedy version of speculative decoding. The +sampled version adds probabilistic acceptance, but the cache and verification +structure are the same. + +**Readings** + +- [Fast Inference from Transformers via Speculative Decoding](https://arxiv.org/abs/2211.17192) +- [Accelerating Large Language Model Decoding with Speculative Sampling](https://arxiv.org/abs/2302.01318) + +## The Core Idea + +Suppose the accepted sequence currently ends with token `A`. + +The draft model predicts four more tokens: + +```plain +A -> b1 -> b2 -> b3 -> b4 +``` + +Then the target model runs once on: + +```plain +[A, b1, b2, b3, b4] +``` + +The target model's logits tell us: + +```plain +target(A) = t1 +target(b1) = t2 +target(b2) = t3 +target(b3) = t4 +target(b4) = t5 +``` + +To compare predictions with drafted tokens, shift the target predictions one +position: + +```plain +[A, t1, t2, t3, t4] +``` + +Now compare against: + +```plain +[A, b1, b2, b3, b4] +``` + +The first token always matches because it is the already accepted seed token. +After that: + +- if `t1 == b1`, then `b1` is accepted, +- if `t2 == b2`, then `b2` is accepted, +- if `t3 != b3`, then `b3` and everything after it is rejected. + +When all drafted tokens match, the target model has also produced one extra +token, `t5`. That token becomes the seed for the next speculative round. + +## Why Rewind Is Needed + +Both models use KV cache. During a speculative round: + +1. the draft model writes K/V for the tokens it used while drafting, +2. the target model writes K/V for the full verification sequence, +3. some suffix of those tokens may be rejected. + +The cache must represent only the accepted prefix. If the target rejects a draft +token, both caches need to forget the rejected suffix: + +```plain +accepted: [A, b1] +rejected: [b2, b3, b4] + +target cache: rewind rejected verification tokens +draft cache: rewind drafted tokens that must be regenerated +``` + +The exact rewind lengths differ because the draft model has not cached the final +draft token yet. In the reference implementation, the draft model is caught up +by one token only when the whole draft window is accepted. + +Paged attention from Week 3 already introduced `rewind(n)` for this reason. +Dense Week 2 caches also need the same lifecycle hook. + +## Task 1: Shared Greedy Step + +``` +src/tiny_llm/generate.py +``` + +Implement a helper inside `speculative_generate` that calls a Week 2 or Week 3 +model and returns greedy next tokens: + +```plain +logits = model(tokens[None], offset, kv_cache) +logits = logits[:, -n_tokens:, :] +next_tokens = argmax(log_softmax(logits), axis=-1) +``` + +For normal one-token decode, `n_tokens` is `1`. For target verification, +`n_tokens` is the size of the verification sequence. + +## Task 2: Prefill Both Models + +Create one KV cache for the draft model and one for the target model: + +```python +draft_kv_cache = draft_model.create_kv_cache() +kv_cache = model.create_kv_cache() +``` + +Then prefill both models with the same prompt. Each prefill should return: + +- the first generated token, +- the offset after the prompt tokens have been written into the cache. + +The prompt should be encoded with the tokenizer passed to that model. + +## Task 3: Draft Tokens + +Implement a small helper that runs the draft model for a fixed number of tokens: + +```plain +last accepted token -> draft token 1 +draft token 1 -> draft token 2 +draft token 2 -> draft token 3 +draft token 3 -> draft token 4 +``` + +Use `num_drafts = 4` for this teaching implementation. + +The helper should update the draft cache and advance the draft offset by one +for each drafted step. + +## Task 4: Verify and Rewind + +Build the target verification sequence: + +```python +draft_tokens = mx.concat([token, mx.array(draft_tokens)]) +``` + +Run the target model once with `n_tokens = num_drafts + 1`. + +Then shift the target predictions so they line up with the drafted tokens: + +```plain +target predictions: [t1, t2, t3, t4, t5] +shifted compare: [token, t1, t2, t3, t4] +draft sequence: [token, b1, b2, b3, b4] +``` + +Walk from left to right: + +- emit every matching token, +- on the first mismatch, rewind both caches to the accepted prefix, +- set `token` to the target model's replacement token, +- continue the next speculative round from that replacement token. + +When all tokens match, emit the whole accepted draft window. Then run the draft +model one more step on the final draft token so the draft cache catches up with +the target cache, and use the target model's extra prediction as the next seed. + +## Task 5: Stop and Release Caches + +Stop when the token to emit is `tokenizer.eos_token_id`. Do not add EOS to the +detokenizer output. + +Always release both KV caches in a `finally` block: + +```python +finally: + _release_kv_cache(draft_kv_cache) + _release_kv_cache(kv_cache) +``` + +This keeps paged-cache pages reusable after generation finishes or exits early. + +## Verify + +Run the fast speculative decoding tests: + +```bash +pdm run test --week 3 --day 4 +``` + +The tests use deterministic fake models and fake tokenizers. They check the +algorithmic behavior directly: + +- accepting a full draft window, +- carrying the target model's extra token into the next round, +- rejecting a bad draft suffix, +- rewinding draft and target caches by the right lengths, +- stopping on EOS without emitting it, +- releasing both caches. + +You can also run the reference solution tests: + +```bash +pdm run test-refsol --week 3 --day 4 +``` + +To try the real model path after completing earlier chapters, pass both a target +model and a draft model: + +```bash +pdm run main --solution tiny_llm --loader week2 \ + --model qwen3-4b --draft-model qwen3-0.6b +``` + +{{#include copyright.md}} diff --git a/src/tiny_llm_ref/generate.py b/src/tiny_llm_ref/generate.py index 5f3324ab..d5c3cc18 100644 --- a/src/tiny_llm_ref/generate.py +++ b/src/tiny_llm_ref/generate.py @@ -88,6 +88,8 @@ def speculative_generate( ) -> str: draft_kv_cache = draft_model.create_kv_cache() kv_cache = model.create_kv_cache() + detokenizer = tokenizer.detokenizer + detokenizer.reset() def _step(model, y, offset, kv_cache, n_tokens=1): logits = model(y[None], offset, kv_cache) @@ -100,54 +102,52 @@ def _step(model, y, offset, kv_cache, n_tokens=1): y = sampler(logprobs) return y, logprobs.squeeze(0) - # prefill with the prompt, using the large model def _prefill(model, tokenizer, prompt, kv_cache): prefill_tokens = mx.array(tokenizer.encode(prompt, add_special_tokens=False)) - offset = 0 - token, _ = _step(model, prefill_tokens, offset, kv_cache) + token, _ = _step(model, prefill_tokens, 0, kv_cache) mx.eval(token) if token.item() == tokenizer.eos_token_id: - return - offset = prefill_tokens.size - return token, offset + return token, prefill_tokens.size + return token, prefill_tokens.size + + def _decode_one(token): + if token.item() == tokenizer.eos_token_id: + return False + detokenizer.add_token(token.item()) + return True + + def _draft_generate(model, last_token, offset, kv_cache, num_drafts): + tokens = [] + current_offset = offset + for _ in range(num_drafts): + token, _ = _step(model, last_token, current_offset, kv_cache) + mx.eval(token) + tokens.append(token.item()) + last_token = token + current_offset += 1 + return tokens + + def _rewind_cache(kv_cache, n_tokens): + for layer in kv_cache: + layer.rewind(n_tokens) + + def _print_progress(progress): + newline = "\n" + print(f"+{progress} {detokenizer.text.replace(newline, ' ')[-80:]}") try: - draft_token, draft_offset = _prefill( + _draft_token, draft_offset = _prefill( draft_model, draft_tokenizer, prompt, draft_kv_cache ) token, offset = _prefill(model, tokenizer, prompt, kv_cache) - - def _decode_one(token, tokenizer): - if token.item() == tokenizer.eos_token_id: - return False - detokenizer = tokenizer.detokenizer - detokenizer.add_token(token.item()) - return True - - def draft_generate(model, last_token, offset, kv_cache, num_drafts): - tokens = [] - current_offset = offset - for _ in range(num_drafts): - token, _ = _step(model, last_token, current_offset, kv_cache) - mx.eval(token) - tokens.append(token.item()) - last_token = token - current_offset += 1 - return tokens + if token.item() == tokenizer.eos_token_id: + return detokenizer.text num_drafts = 4 - def _rewind_cache(kv_cache, revert_len): - for layer in kv_cache: - layer.rewind(revert_len) - - def _print_text(text, progress): - newline = "\n" - print(f"+{progress} {text.replace(newline, ' ')[-80:]}") - # speculative decode while True: - draft_tokens = draft_generate( + draft_tokens = _draft_generate( draft_model, token, draft_offset, draft_kv_cache, num_drafts ) draft_offset += num_drafts @@ -162,8 +162,8 @@ def _print_text(text, progress): accept_all = True for i in range(len(new_tokens)): if new_tokens[i] != draft_tokens[i]: - # revert the full draft generation; re-generate next time - # or we matched full, then no rewind and use the last token + # The accepted prefix is already emitted. Rewind both caches + # so the next round starts from the target replacement token. assert i >= 1 # first token is always the same revert_len = len(draft_tokens) - i _rewind_cache(draft_kv_cache, revert_len - 1) @@ -173,15 +173,15 @@ def _print_text(text, progress): offset -= revert_len assert offset == draft_offset assert offset == kv_cache[0].offset - _print_text(tokenizer._detokenizer.text, i) + _print_progress(i) accept_all = False break - if not _decode_one(new_tokens[i], tokenizer): - print(tokenizer._detokenizer.text) - return tokenizer._detokenizer.text + if not _decode_one(new_tokens[i]): + print(detokenizer.text) + return detokenizer.text if accept_all: - _print_text(tokenizer._detokenizer.text, len(new_tokens)) - draft_generate( + _print_progress(len(new_tokens)) + _draft_generate( draft_model, mx.array(draft_tokens[-1:]), draft_offset, diff --git a/tests_refsol/test_week_3_day_4.py b/tests_refsol/test_week_3_day_4.py new file mode 100644 index 00000000..7b60ad83 --- /dev/null +++ b/tests_refsol/test_week_3_day_4.py @@ -0,0 +1,225 @@ +import mlx.core as mx + +from .tiny_llm_base import speculative_generate + + +EOS = 0 +PROMPT = "prompt" +PROMPT_TOKENS = [1, 2] +PIECES = { + 10: "A", + 11: "B", + 12: "C", + 13: "D", + 14: "E", + 15: "F", + 21: "X", +} + + +class FakeLayerCache: + def __init__(self, name: str): + self.name = name + self.offset = 0 + self.updates = [] + self.rewinds = [] + self.release_count = 0 + self.released_offsets = [] + + def append(self, offset: int, tokens: list[int]): + assert self.offset == offset, ( + f"{self.name} expected offset {self.offset}, got {offset}" + ) + self.updates.append((offset, tuple(tokens))) + self.offset += len(tokens) + + def rewind(self, n: int): + assert 0 <= n <= self.offset + self.rewinds.append(n) + self.offset -= n + + def release(self): + self.release_count += 1 + self.released_offsets.append(self.offset) + self.offset = 0 + + +class ScriptedModel: + def __init__( + self, + transitions: dict[int, int], + name: str, + vocab_size: int = 128, + ): + self.transitions = {EOS: EOS, **transitions} + self.name = name + self.vocab_size = vocab_size + self.created_caches = [] + + def create_kv_cache(self): + cache = [FakeLayerCache(self.name)] + self.created_caches.append(cache) + return cache + + def __call__(self, inputs: mx.array, offset: int, cache: list[FakeLayerCache]): + tokens = [int(token) for token in inputs.tolist()[0]] + for layer in cache: + layer.append(offset, tokens) + + logits = [] + for token in tokens: + next_token = self.transitions.get(token, EOS) + row = [-1000.0] * self.vocab_size + row[next_token] = 1000.0 + logits.append(row) + return mx.array([logits], dtype=mx.float32) + + +class FakeDetokenizer: + def __init__(self, pieces: dict[int, str]): + self.pieces = pieces + self.reset() + + def reset(self): + self.text = "" + self.last_segment = "" + + def add_token(self, token: int): + assert token != EOS, "EOS should stop generation, not be detokenized" + self.last_segment = self.pieces[token] + self.text += self.last_segment + + +class FakeTokenizer: + eos_token_id = EOS + + def __init__(self, pieces: dict[int, str] | None = None): + self._detokenizer = FakeDetokenizer(pieces or PIECES) + + @property + def detokenizer(self): + return self._detokenizer + + def encode(self, prompt: str, add_special_tokens: bool = False): + assert prompt == PROMPT + assert add_special_tokens is False + return PROMPT_TOKENS + + +def run_speculative( + target_transitions: dict[int, int], + draft_transitions: dict[int, int], +): + draft_model = ScriptedModel(draft_transitions, name="draft") + target_model = ScriptedModel(target_transitions, name="target") + draft_tokenizer = FakeTokenizer() + tokenizer = FakeTokenizer() + + text = speculative_generate( + draft_model, + target_model, + draft_tokenizer, + tokenizer, + PROMPT, + ) + + draft_layer = draft_model.created_caches[0][0] + target_layer = target_model.created_caches[0][0] + return text, draft_layer, target_layer + + +def test_task_1_accepts_full_draft_window_and_carries_extra_token(): + target = { + 2: 10, + 10: 11, + 11: 12, + 12: 13, + 13: 14, + 14: 15, + 15: EOS, + } + draft = { + 2: 10, + 10: 11, + 11: 12, + 12: 13, + 13: 14, + 14: 21, + 15: 21, + 21: 21, + } + + text, draft_layer, target_layer = run_speculative(target, draft) + + assert text == "ABCDEF" + assert target_layer.updates[1] == (2, (10, 11, 12, 13, 14)) + assert (6, (14,)) in draft_layer.updates + assert draft_layer.release_count == 1 + assert target_layer.release_count == 1 + + +def test_task_2_rejects_bad_suffix_and_rewinds_to_accepted_prefix(): + target = { + 2: 10, + 10: 11, + 11: 21, + 21: EOS, + } + draft = { + 2: 10, + 10: 11, + 11: 12, + 12: 13, + 13: 14, + 21: 14, + 14: 14, + } + + text, draft_layer, target_layer = run_speculative(target, draft) + + assert text == "ABX" + assert draft_layer.rewinds[0] == 2 + assert target_layer.rewinds[0] == 3 + assert target_layer.updates[1] == (2, (10, 11, 12, 13, 14)) + + +def test_task_2_rejects_at_first_draft_position(): + target = { + 2: 10, + 10: 21, + 21: EOS, + } + draft = { + 2: 10, + 10: 11, + 11: 12, + 12: 13, + 13: 14, + 21: 14, + 14: 14, + } + + text, draft_layer, target_layer = run_speculative(target, draft) + + assert text == "AX" + assert draft_layer.rewinds[0] == 3 + assert target_layer.rewinds[0] == 4 + + +def test_task_3_stops_on_eos_and_releases_caches(): + target = { + 2: 10, + 10: EOS, + } + draft = { + 2: 10, + 10: EOS, + } + + text, draft_layer, target_layer = run_speculative(target, draft) + + assert text == "A" + assert draft_layer.release_count == 1 + assert target_layer.release_count == 1 + assert draft_layer.rewinds == [] + assert target_layer.rewinds == [] From 51ca11b7977cd9dcfd9c87076b6e3b8e925acc44 Mon Sep 17 00:00:00 2001 From: xuanyili Date: Sat, 13 Jun 2026 00:37:24 -0700 Subject: [PATCH 2/2] Drop speculative reference refactor --- src/tiny_llm_ref/generate.py | 84 ++++++++++++++++++------------------ 1 file changed, 42 insertions(+), 42 deletions(-) diff --git a/src/tiny_llm_ref/generate.py b/src/tiny_llm_ref/generate.py index d5c3cc18..5f3324ab 100644 --- a/src/tiny_llm_ref/generate.py +++ b/src/tiny_llm_ref/generate.py @@ -88,8 +88,6 @@ def speculative_generate( ) -> str: draft_kv_cache = draft_model.create_kv_cache() kv_cache = model.create_kv_cache() - detokenizer = tokenizer.detokenizer - detokenizer.reset() def _step(model, y, offset, kv_cache, n_tokens=1): logits = model(y[None], offset, kv_cache) @@ -102,52 +100,54 @@ def _step(model, y, offset, kv_cache, n_tokens=1): y = sampler(logprobs) return y, logprobs.squeeze(0) + # prefill with the prompt, using the large model def _prefill(model, tokenizer, prompt, kv_cache): prefill_tokens = mx.array(tokenizer.encode(prompt, add_special_tokens=False)) - token, _ = _step(model, prefill_tokens, 0, kv_cache) + offset = 0 + token, _ = _step(model, prefill_tokens, offset, kv_cache) mx.eval(token) if token.item() == tokenizer.eos_token_id: - return token, prefill_tokens.size - return token, prefill_tokens.size - - def _decode_one(token): - if token.item() == tokenizer.eos_token_id: - return False - detokenizer.add_token(token.item()) - return True - - def _draft_generate(model, last_token, offset, kv_cache, num_drafts): - tokens = [] - current_offset = offset - for _ in range(num_drafts): - token, _ = _step(model, last_token, current_offset, kv_cache) - mx.eval(token) - tokens.append(token.item()) - last_token = token - current_offset += 1 - return tokens - - def _rewind_cache(kv_cache, n_tokens): - for layer in kv_cache: - layer.rewind(n_tokens) - - def _print_progress(progress): - newline = "\n" - print(f"+{progress} {detokenizer.text.replace(newline, ' ')[-80:]}") + return + offset = prefill_tokens.size + return token, offset try: - _draft_token, draft_offset = _prefill( + draft_token, draft_offset = _prefill( draft_model, draft_tokenizer, prompt, draft_kv_cache ) token, offset = _prefill(model, tokenizer, prompt, kv_cache) - if token.item() == tokenizer.eos_token_id: - return detokenizer.text + + def _decode_one(token, tokenizer): + if token.item() == tokenizer.eos_token_id: + return False + detokenizer = tokenizer.detokenizer + detokenizer.add_token(token.item()) + return True + + def draft_generate(model, last_token, offset, kv_cache, num_drafts): + tokens = [] + current_offset = offset + for _ in range(num_drafts): + token, _ = _step(model, last_token, current_offset, kv_cache) + mx.eval(token) + tokens.append(token.item()) + last_token = token + current_offset += 1 + return tokens num_drafts = 4 + def _rewind_cache(kv_cache, revert_len): + for layer in kv_cache: + layer.rewind(revert_len) + + def _print_text(text, progress): + newline = "\n" + print(f"+{progress} {text.replace(newline, ' ')[-80:]}") + # speculative decode while True: - draft_tokens = _draft_generate( + draft_tokens = draft_generate( draft_model, token, draft_offset, draft_kv_cache, num_drafts ) draft_offset += num_drafts @@ -162,8 +162,8 @@ def _print_progress(progress): accept_all = True for i in range(len(new_tokens)): if new_tokens[i] != draft_tokens[i]: - # The accepted prefix is already emitted. Rewind both caches - # so the next round starts from the target replacement token. + # revert the full draft generation; re-generate next time + # or we matched full, then no rewind and use the last token assert i >= 1 # first token is always the same revert_len = len(draft_tokens) - i _rewind_cache(draft_kv_cache, revert_len - 1) @@ -173,15 +173,15 @@ def _print_progress(progress): offset -= revert_len assert offset == draft_offset assert offset == kv_cache[0].offset - _print_progress(i) + _print_text(tokenizer._detokenizer.text, i) accept_all = False break - if not _decode_one(new_tokens[i]): - print(detokenizer.text) - return detokenizer.text + if not _decode_one(new_tokens[i], tokenizer): + print(tokenizer._detokenizer.text) + return tokenizer._detokenizer.text if accept_all: - _print_progress(len(new_tokens)) - _draft_generate( + _print_text(tokenizer._detokenizer.text, len(new_tokens)) + draft_generate( draft_model, mx.array(draft_tokens[-1:]), draft_offset,