Add speculative decoding with exact-distribution tests; expand README

Draft proposes gamma tokens, target verifies them in one pass using rejection sampling with residual resampling, and both KV caches roll back to the accepted prefix. Tests check greedy equivalence and compare sampled output against the brute-force target distribution; mutation checks confirm the test fails on deliberately broken variants.
2026-11-07 20:33:28 -07:00
parent 73ff04d77f
commit 046b09c88f
6 changed files with 546 additions and 51 deletions
@@ -1,55 +1,78 @@
 # microgpt
-A GPT-style language model written from scratch in PyTorch, with the pieces that matter in
+A GPT-style language model written from scratch in PyTorch, plus the inference techniques that
-production LLM systems: rotary position embeddings, RMSNorm, SwiGLU, an incremental-decoding
+matter in production LLM serving: a **KV cache**, **grouped-query attention**, **continuous
-**KV cache**, nucleus sampling, and a test suite that proves each piece is correct.
+batching** over variable-length requests, and **speculative decoding**. Every one of them is
 backed by tests that check correctness against a reference, not just "it runs".
-No `nn.Transformer`, no Hugging Face. Attention, masking, caching and sampling are all
+No `nn.Transformer`, no Hugging Face. Attention, masking, caching, batching and sampling are all
-written out explicitly (~220 lines of model code) so every line can be reasoned about.
+written out explicitly (~900 lines of library code) so every line can be reasoned about.
 ## What's inside
-| Piece | Where | Why it's here |
+| Piece | Where | The interesting part |
 |---|---|---|
-| Causal self-attention with explicit mask | `model.py` | Mask is offset by the cache length, which is the classic place incremental decoding goes wrong |
+| Causal attention, RoPE, RMSNorm, SwiGLU, tied embeddings | `model.py` | The Llama recipe at toy scale; mask offset by cache length |
-| RoPE | `model.py` | Relative position without a learned table; cached keys keep their absolute rotation |
+| Grouped-query / multi-query attention | `model.py` | Cache stores only the KV heads; q is regrouped so matmul broadcasts instead of copying K/V |
 | RMSNorm + SwiGLU, tied embeddings | `model.py` | The Llama-style recipe at toy scale |
 | KV cache + sliding window | `model.py` | `O(T)` work per generated token instead of `O(T^2)` |
-| Temperature / top-k / top-p | `sampling.py` | Written from first principles, tested against analytic distributions |
+| Continuous batching | `batching.py` | Left-padded shared cache, per-row RoPE positions, key mask, mid-flight admission, column trimming |
-| Trainer | `train.py` | AdamW with decoupled weight decay (no decay on norms), warmup + cosine LR, grad clipping |
+| Speculative decoding | `speculative.py` | Exact target distribution (rejection sampling + residual), cache rollback |
-| Benchmark | `scripts/benchmark.py` | Measures the cache speedup instead of asserting it |
+| Sampling: temperature / top-k / top-p | `sampling.py` | `logits_to_probs` exposes the exact sampling distribution (greedy = one-hot) |
 | Trainer | `train.py` | AdamW with decoupled decay, warmup + cosine LR, grad clipping |
 ## Quickstart
 ```bash
 python3 -m venv .venv && source .venv/bin/activate
 pip install -e ".[dev]"
-pytest                                  # 21 tests, ~3s on CPU
+pytest                                   # 49 tests, ~15s on CPU
 mkdir -p data && curl -sSL -o data/shakespeare.txt \
  https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
 python -m microgpt.train --data data/shakespeare.txt --steps 1200 --batch-size 24
 python -m microgpt.train --data data/shakespeare.txt --n-layer 1 --n-embd 64 --n-head 2 \
    --steps 2000 --out checkpoints/draft.pt            # small draft model
 python -m microgpt.generate --prompt "ROMEO:" --tokens 200 --top-k 40
-python scripts/benchmark.py
+python -m microgpt.generate --prompt "ROMEO:" --tokens 200 --draft checkpoints/draft.pt  # speculative
 python scripts/benchmark.py              # KV cache vs no cache
 python scripts/bench_gqa.py              # cache memory: MHA vs GQA vs MQA
 python scripts/bench_batching.py         # continuous batching throughput
 python scripts/bench_speculative.py      # speculative decoding speedup
 ```
-## Correctness: what the tests actually prove
+## How each feature is verified
-* **No future leakage** — perturbing tokens at positions >= 8 leaves logits at positions < 8 unchanged.
+**KV cache.** Prefill-then-decode reproduces the full forward pass to `1e-5`, including
-* **KV cache == full forward** — prefill then one-token-at-a-time decoding reproduces the full
+multi-token continuation (which exercises the offset causal mask). Perturbing future tokens never
-  forward pass to `1e-5`; so does multi-token continuation (which exercises the offset causal mask).
+changes earlier logits. RoPE preserves norms and its attention scores depend only on `m - n`.
-* **Greedy generation is identical with and without the cache** (while within the context window).
+
-* **RoPE** preserves vector norms, and `<rope(q, m), rope(k, n)>` depends only on `m - n`.
+**Grouped-query attention.** A GQA model must equal a multi-head model whose K/V heads are
-* **Sampling** — top-k and top-p keep exactly the intended tokens; empirical frequencies match
+*copies* of the GQA ones (weights are copied across and logits compared, both full-forward and
-  the softmax distribution; seeded runs are reproducible.
+token-by-token). Cache bytes are asserted to scale exactly with the number of KV heads.
-* **Learning** — the model can overfit a tiny batch (loss drops by more than 70%).
+
 **Continuous batching.** Requests of different lengths share one cache. The tests require that
 greedy output for every request is *identical* to running it alone, across: queueing behind a
 full batch, requests joining mid-flight with longer *and* shorter prompts than the current
 batch, EOS stopping one row without disturbing others, and the `block_size` limit. A padded row
 must not produce NaN (more on that below).
 **Speculative decoding.** At temperature 0 the output equals greedy target decoding
 token-for-token. For sampling, the test computes the target's exact distribution over the first
 two generated tokens and the marginal of the third by brute force on a 5-token vocabulary, then
 checks 2,500 speculative samples against it. A power check asserts the draft really is a
 different distribution, and I confirmed the test fails when the algorithm is deliberately broken
 (always accepting proposals, or resampling from `p` instead of the residual `max(0, p - q)`).
 ## Results
-Measured on an Intel iMac (6 CPU threads, no GPU). Numbers are from this repo's own scripts.
+All numbers were measured with this repo's scripts on an Intel iMac (6 CPU threads, no GPU).
 CPU timings are noisy; treat small differences as approximate.
-**Training** — 0.81M parameters, 4 layers, 128-dim, character-level Tiny Shakespeare,
+### Training (character-level Tiny Shakespeare)
-1200 steps (~5.6 min on CPU):
+
 0.81M parameters (4 layers, 128-dim), 1200 steps, about 5.6 minutes:
 | step | train loss | val loss |
 |---:|---:|---:|
@@ -58,36 +81,97 @@ Measured on an Intel iMac (6 CPU threads, no GPU). Numbers are from this repo's
 | 800 | 1.444 | 1.655 |
 | 1200 | 1.348 | 1.596 |
-(Initial loss 4.24 ≈ ln 65, i.e. uniform over the 65-character vocabulary, as expected.)
+Initial loss 4.24 is close to ln 65 = 4.17 (uniform over the 65-character vocabulary), as expected.
-**KV-cache decoding** — 2.68M-parameter model, greedy, 16-token prompt:
+### KV cache
 | new tokens | no cache (tok/s) | KV cache (tok/s) | speedup | identical output |
 |-----------:|-----------------:|-----------------:|--------:|:-----------------:|
 | 32 | 158.3 | 194.2 | 1.2x | yes |
 | 64 | 166.2 | 259.3 | 1.6x | yes |
 | 128 | 98.2 | 174.4 | 1.8x | yes |
 | 256 | 74.5 | 198.8 | 2.7x | yes |
 | 480 | 34.7 | 167.4 | 4.8x | yes |
-Cached throughput stays roughly flat as the sequence grows while uncached throughput decays,
+Cached throughput stays roughly flat as the sequence grows while uncached throughput decays.
 which is the `O(T)` vs `O(T^2)` difference showing up in practice. Timings on a shared CPU are
 noisy, so treat the small-length rows as approximate.
-## Design notes and known limitations
+### Grouped-query attention
-* **Sliding window past `block_size`.** The uncached path always sees the last `block_size`
+8 query heads, 6 layers, 256-dim, cache holding 496 tokens:
-  tokens. The cached path drops its cache when full and rebuilds it from the latest half-window.
+
-  Outputs are therefore identical only while the sequence fits in the window; beyond that they
+| KV heads | variant | params | KV cache | vs MHA |
-  differ by design. (RoPE makes true cache eviction possible, but keys would need re-rotation.)
+|---:|:--|---:|---:|---:|
-* **Attention is `O(T^2)` memory** per layer: no FlashAttention, no grouped-query attention.
+| 8 | MHA | 4.84M | 5.8 MiB | 1.00x |
-  Both are natural next steps and would also shrink the KV cache.
+| 4 | GQA | 4.45M | 2.9 MiB | 0.50x |
-* **Character-level, small scale.** The point is correctness and clarity, not sample quality.
+| 2 | GQA | 4.25M | 1.5 MiB | 0.25x |
-  Samples are Shakespeare-flavoured but not coherent.
+| 1 | MQA | 4.16M | 0.7 MiB | 0.12x |
 Memory is exact. Decode speed on this CPU (126 / 176 / 136 / 133 tok/s for 8 / 4 / 2 / 1 KV
 heads) is not monotonic, which suggests noise, but I did not measure run-to-run variance, so no
 speed claim is made. The benefit of GQA shown here is cache size, which matters when serving
 long contexts or many concurrent sequences.
 ### Continuous batching
 16 requests (prompts of 8-47 tokens, 64 new tokens each), 6-layer 192-dim model:
 | strategy | tokens/s | speedup | mean batch |
 |:--|--:|--:|--:|
 | one at a time | 240 | 1.0x | 1.0 |
 | continuous, max_batch=2 | 303 | 1.3x | 2.0 |
 | continuous, max_batch=4 | 575 | 2.4x | 4.0 |
 | continuous, max_batch=8 | 1144 | 4.8x | 8.0 |
 | continuous, max_batch=16 | 1767 | 7.4x | 16.0 |
 Caveat: every request here asks for the same number of tokens, so the batch stays full. Real
 traffic with mixed lengths would see lower occupancy; the scheduler refills slots as rows finish.
 ### Speculative decoding
 Target: the 0.81M model (val loss 1.60). Draft: 57.6K parameters, 1 layer (val loss 1.81).
 Temperature 0.8, 100 new tokens, 12 runs per row:
 | method | s / run | speedup | acceptance | tokens / target pass |
 |:--|--:|--:|--:|--:|
 | target only | 0.36 | 1.00x | - | 1.0 |
 | speculative, gamma=2 | 0.26 | 1.38x | 64% | 2.26 |
 | speculative, gamma=4 | 0.22 | 1.65x | 50% | 2.96 |
 | speculative, gamma=6 | 0.22 | 1.68x | 44% | 3.59 |
 | speculative, gamma=8 | 0.25 | 1.42x | 32% | 3.50 |
 The speedup peaks at moderate `gamma` and then falls: more proposals per round means each is less
 likely to be accepted, so the draft's extra work stops paying for itself. The output
 distribution is the target's in every row (see verification above), so this is a pure speed
 trade-off. Wall-clock gains depend on the draft being much cheaper than the target; with a
 GPU-scale target the same acceptance rates would give a larger speedup.
 ## Pitfalls hit (and avoided) while building this
 * **Sliding window.** An early version claimed cached and uncached generation were identical.
  They are, but only while the text fits in `block_size`; past it the two paths slide the window
  differently. Caught by seeded samples diverging at token 135; the docstring and README now say so.
 * **NaN from padding.** In a left-padded batch, a padding query has no visible key, so softmax over
  all `-inf` gives NaN, and NaN * 0 in the value product poisons real tokens. Fix: every query may
  always attend to itself (its output is discarded anyway). Covered by a finite-logits assertion.
 * **Ties in parity tests (avoided).** The default init gives near-uniform logits, where greedy
  argmax could flip on float noise between code paths. Parity tests therefore use a wider init
  (`tests/helpers.py`) so they compare confident outputs rather than ties.
 * **Wrong test, not wrong code.** One batching test failed because it sampled cache width *after*
  a row had already been retired in the same step; the engine was right, the assertion was not.
 ## Limitations
 * Past `block_size`, the single-sequence cache path drops and rebuilds its cache from the latest
  half-window; the batching engine and speculative decoding instead require sequences to fit in
  `block_size`.
 * Speculative decoding handles one sequence at a time and does not combine with the batch engine.
 * The cache grows by `torch.cat` each step (simple and correct, but it reallocates); a production
  system would use preallocated or paged blocks.
 * Attention is the plain `O(T^2)`-memory formulation: no FlashAttention.
 * Character-level and small-scale: the point is correctness and clarity, not sample quality.
 ## Roadmap
- [ ] Grouped-query attention (smaller KV cache) with a parity test against full multi-head
+- [ ] BPE tokenizer trained from scratch, with round-trip tests
- [ ] Batched generation with per-sequence cache lengths
+- [ ] Paged KV cache (fixed-size blocks, block table) to remove the `torch.cat` reallocation
- [ ] Speculative decoding with a draft model, verified to preserve the target distribution
+- [ ] Speculative decoding inside the batch engine
- [ ] BPE tokenizer
+- [ ] Load a public checkpoint into this architecture and match Hugging Face logits
@@ -0,0 +1,79 @@
 """Speculative decoding vs plain target decoding, on the trained Shakespeare models.
    python scripts/bench_speculative.py --target checkpoints/model.pt --draft checkpoints/draft.pt
 Reports hardware-independent numbers (acceptance rate, tokens per target forward pass) and
 wall-clock speedup, which depends on the draft being much cheaper than the target.
 """
 from __future__ import annotations
 import argparse
 import time
 import torch
 from microgpt.generate import load_checkpoint
 from microgpt.speculative import speculative_generate
 def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--target", default="checkpoints/model.pt")
    p.add_argument("--draft", default="checkpoints/draft.pt")
    p.add_argument("--tokens", type=int, default=100)
    p.add_argument("--trials", type=int, default=12)
    p.add_argument("--temperature", type=float, default=0.8)
    args = p.parse_args()
    target, tok = load_checkpoint(args.target)
    draft, _ = load_checkpoint(args.draft)
    prompts = ["ROMEO:", "First Citizen:\n", "KING RICHARD:\n", "The ", "MENENIUS:\n", "What "]
    print(
        f"target {target.num_params() / 1e6:.2f}M params, draft {draft.num_params() / 1e6:.2f}M, "
        f"temperature {args.temperature}, {args.tokens} new tokens x {args.trials} runs\n"
    )
    def run(fn) -> float:
        fn(0)  # warm up
        t0 = time.perf_counter()
        for i in range(args.trials):
            fn(i)
        return (time.perf_counter() - t0) / args.trials
    def encode(i: int) -> torch.Tensor:
        return torch.tensor([tok.encode(prompts[i % len(prompts)])])
    base = run(
        lambda i: target.generate(
            encode(i), args.tokens, temperature=args.temperature,
            generator=torch.Generator().manual_seed(i),
        )  # fmt: skip
    )
    print("| method | s / run | speedup | acceptance | tokens / target pass |")
    print("|:--|--:|--:|--:|--:|")
    print(f"| target only | {base:.2f} | 1.00x | - | 1.0 |")
    stats = []
    for gamma in (2, 4, 6, 8):
        stats.clear()
        def spec(i: int, gamma: int = gamma) -> None:
            _, s = speculative_generate(
                target, draft, encode(i), args.tokens, gamma=gamma,
                temperature=args.temperature, generator=torch.Generator().manual_seed(i),
            )  # fmt: skip
            stats.append(s)
        t = run(spec)
        tail = stats[1:]  # drop the warm-up run
        acc = sum(s.accepted for s in tail) / sum(s.proposed for s in tail)
        per_pass = sum(s.new_tokens for s in tail) / sum(s.target_calls for s in tail)
        print(
            f"| speculative, gamma={gamma} | {t:.2f} | {base / t:.2f}x | {acc:.0%} "
            f"| {per_pass:.2f} |"
        )
 if __name__ == "__main__":
    main()
@@ -3,11 +3,13 @@
 from __future__ import annotations
 import argparse
 import sys
 import torch
 from microgpt.data import CharTokenizer
 from microgpt.model import GPT, GPTConfig
 from microgpt.speculative import speculative_generate
 def load_checkpoint(path: str):
@@ -28,11 +30,26 @@ def main() -> None:
    p.add_argument("--top-p", type=float, default=None)
    p.add_argument("--no-cache", action="store_true", help="disable the KV cache")
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--draft", default=None, help="draft checkpoint: enables speculative decoding")
    p.add_argument("--gamma", type=int, default=4, help="tokens proposed per speculative round")
    args = p.parse_args()
    model, tok = load_checkpoint(args.ckpt)
    g = torch.Generator().manual_seed(args.seed)
    idx = torch.tensor([tok.encode(args.prompt)], dtype=torch.long)
    if args.draft:
        draft, _ = load_checkpoint(args.draft)
        out, stats = speculative_generate(
            model, draft, idx, args.tokens, gamma=args.gamma, temperature=args.temperature,
            top_k=args.top_k, top_p=args.top_p, generator=g,
        )  # fmt: skip
        print(tok.decode(out[0].tolist()))
        print(
            f"\n[speculative: {stats.acceptance_rate:.0%} accepted, "
            f"{stats.tokens_per_target_call:.2f} tokens per target pass]",
            file=sys.stderr,
        )
        return
    out = model.generate(
        idx,
        args.tokens,
@@ -25,6 +25,29 @@ def filter_top_p(logits: torch.Tensor, p: float) -> torch.Tensor:
    return torch.full_like(logits, float("-inf")).scatter(-1, sorted_idx, sorted_logits)
 def logits_to_probs(
    logits: torch.Tensor,
    temperature: float = 1.0,
    top_k: Optional[int] = None,
    top_p: Optional[float] = None,
 ) -> torch.Tensor:
    """The exact distribution a sampler draws from: temperature, then top-k, then top-p.
    temperature == 0 gives a one-hot at the argmax, so greedy decoding is just the
    degenerate case of sampling (speculative decoding relies on this uniform view).
    """
    if temperature < 0:
        raise ValueError("temperature must be >= 0")
    if temperature == 0:
        return F.one_hot(logits.argmax(dim=-1), logits.size(-1)).to(logits.dtype)
    logits = logits / temperature
    if top_k is not None:
        logits = filter_top_k(logits, top_k)
    if top_p is not None:
        logits = filter_top_p(logits, top_p)
    return F.softmax(logits, dim=-1)
 def sample_next(
    logits: torch.Tensor,
    temperature: float = 1.0,
@@ -40,10 +63,5 @@ def sample_next(
        raise ValueError("temperature must be >= 0")
    if temperature == 0:
        return logits.argmax(dim=-1, keepdim=True)
-    logits = logits / temperature
+    probs = logits_to_probs(logits, temperature, top_k, top_p)
    if top_k is not None:
        logits = filter_top_k(logits, top_k)
    if top_p is not None:
        logits = filter_top_p(logits, top_p)
    probs = F.softmax(logits, dim=-1)
    return torch.multinomial(probs, num_samples=1, generator=generator)
@@ -0,0 +1,151 @@
 """Speculative decoding (Leviathan et al. 2023; Chen et al. 2023).
 A cheap *draft* model proposes `gamma` tokens; the expensive *target* model scores all of them
 in a single forward pass and keeps the longest prefix it agrees with. Output is distributed
 **exactly** as if sampled from the target alone, but the target runs once per round instead of
 once per token.
 Acceptance rule, for proposal x_i drawn from the draft distribution q_i, with target p_i:
    accept with probability min(1, p_i(x_i) / q_i(x_i));
    on the first rejection, resample from the residual  norm(max(0, p_i - q_i))  and stop;
    if every proposal is accepted, draw one bonus token from the target's next distribution.
 Both p and q are the *processed* distributions (after temperature / top-k / top-p), so the
 guarantee holds for whatever sampler is configured. Temperature 0 reduces to exact greedy.
 Cache bookkeeping: both KV caches are rolled back (truncated) to the accepted prefix after each
 round. The draft's cache may lag by one token (the last proposal is never fed), so each round the
 draft simply consumes `seq[d_len:]`, which makes catching up automatic.
 Single sequence only, and the whole run must fit in `block_size` (no sliding window).
 """
 from __future__ import annotations
 from dataclasses import dataclass
 from typing import Optional, Tuple
 import torch
 from microgpt.model import GPT, KVCache
 from microgpt.sampling import logits_to_probs
@dataclass
 class SpecStats:
    rounds: int = 0
    proposed: int = 0
    accepted: int = 0
    target_calls: int = 0
    draft_calls: int = 0
    new_tokens: int = 0
    @property
    def acceptance_rate(self) -> float:
        return self.accepted / self.proposed if self.proposed else 0.0
    @property
    def tokens_per_target_call(self) -> float:
        return self.new_tokens / self.target_calls if self.target_calls else 0.0
 def _truncate(cache: KVCache, n: int) -> KVCache:
    return [(k[:, :, :n], v[:, :, :n]) for k, v in cache]
@torch.no_grad()
 def speculative_generate(
    target: GPT,
    draft: GPT,
    idx: torch.Tensor,
    max_new_tokens: int,
    gamma: int = 4,
    temperature: float = 1.0,
    top_k: Optional[int] = None,
    top_p: Optional[float] = None,
    generator: Optional[torch.Generator] = None,
 ) -> Tuple[torch.Tensor, SpecStats]:
    """Extend idx (1, T) by max_new_tokens tokens. Returns (tokens, stats)."""
    if idx.dim() != 2 or idx.size(0) != 1:
        raise ValueError("speculative_generate supports a single sequence: idx must be (1, T)")
    if target.cfg.vocab_size != draft.cfg.vocab_size:
        raise ValueError("target and draft must share a vocabulary")
    if gamma < 1:
        raise ValueError("gamma must be >= 1")
    limit = min(target.cfg.block_size, draft.cfg.block_size)
    if idx.size(1) + max_new_tokens > limit:
        raise ValueError(f"prompt + max_new_tokens exceeds the shorter block_size ({limit})")
    modes = (target.training, draft.training)
    target.eval()
    draft.eval()
    def probs(logits: torch.Tensor) -> torch.Tensor:
        return logits_to_probs(logits, temperature, top_k, top_p)
    seq = idx.clone()
    goal = idx.size(1) + max_new_tokens
    t_cache: Optional[KVCache] = None
    d_cache: Optional[KVCache] = None
    t_len = d_len = 0  # tokens of `seq` already inside each cache
    stats = SpecStats()
    while seq.size(1) < goal:
        n_old = seq.size(1)
        g = min(gamma, goal - n_old - 1)  # always leave room for the guaranteed extra token
        # 1) Draft proposes g tokens autoregressively, remembering its distributions q_i.
        proposals, q_dists = [], []
        if g > 0:
            feed = seq[:, d_len:]
            for _ in range(g):
                logits, _, d_cache = draft(feed, cache=d_cache)
                d_len += feed.size(1)
                stats.draft_calls += 1
                q = probs(logits[:, -1])[0]
                x = torch.multinomial(q, 1, generator=generator).view(1, 1)
                proposals.append(x)
                q_dists.append(q)
                feed = x
        # 2) Target scores the uncached suffix plus every proposal in ONE forward pass.
        feed = seq[:, t_len:]
        if g > 0:
            feed = torch.cat([feed] + proposals, dim=1)
        logits, _, t_cache = target(feed, cache=t_cache)
        t_len += feed.size(1)
        stats.target_calls += 1
        p_dists = probs(logits[0, -(g + 1) :])  # rows 0..g-1 judge x_1..x_g; row g is the bonus
        # 3) Accept / reject left to right.
        n_acc, extra = 0, None
        for i in range(g):
            x = int(proposals[i])
            p, q = p_dists[i], q_dists[i]
            if torch.rand((), generator=generator) < p[x] / q[x]:  # q[x] > 0: x was drawn from q
                n_acc += 1
                continue
            residual = (p - q).clamp_min(0)
            total = residual.sum()
            residual = residual / total if total > 0 else p
            extra = torch.multinomial(residual, 1, generator=generator)
            break
        if extra is None:  # all g accepted (or g == 0): free token from the target
            extra = torch.multinomial(p_dists[g], 1, generator=generator)
        # 4) Commit accepted proposals + the extra token, then roll both caches back.
        new = torch.cat(proposals[:n_acc] + [extra.view(1, 1)], dim=1)
        seq = torch.cat([seq, new], dim=1)
        t_len = n_old + n_acc
        t_cache = _truncate(t_cache, t_len)
        if g > 0 and d_len > n_old + n_acc:
            d_len = n_old + n_acc
            d_cache = _truncate(d_cache, d_len)
        stats.rounds += 1
        stats.proposed += g
        stats.accepted += n_acc
        stats.new_tokens += n_acc + 1
    target.train(modes[0])
    draft.train(modes[1])
    return seq, stats
@@ -0,0 +1,146 @@
 import itertools
 import pytest
 import torch
 from helpers import make_model
 from microgpt.model import GPTConfig
 from microgpt.sampling import logits_to_probs
 from microgpt.speculative import speculative_generate
 V = 50
 TARGET_CFG = GPTConfig(vocab_size=V, block_size=64, n_layer=2, n_head=4, n_embd=32)
 DRAFT_CFG = GPTConfig(vocab_size=V, block_size=64, n_layer=1, n_head=2, n_embd=16)
@pytest.fixture(scope="module")
 def target():
    return make_model(TARGET_CFG, seed=1)
@pytest.fixture(scope="module")
 def draft():
    return make_model(DRAFT_CFG, seed=2)
 def prompt_of(n, seed=0):
    g = torch.Generator().manual_seed(seed)
    return torch.randint(0, V, (1, n), generator=g)
@pytest.mark.parametrize("gamma", [1, 3, 6])
 def test_greedy_speculation_equals_greedy_target(target, draft, gamma):
    """At temperature 0 the output must be token-for-token the target's greedy output."""
    prompt = prompt_of(5)
    expected = target.generate(prompt, 30, temperature=0)
    got, _ = speculative_generate(target, draft, prompt, 30, gamma=gamma, temperature=0)
    assert torch.equal(got, expected)
 def test_a_perfect_draft_is_always_accepted(target):
    out, stats = speculative_generate(
        target, target, prompt_of(4), 40, gamma=4, temperature=1.0,
        generator=torch.Generator().manual_seed(0),
    )  # fmt: skip
    assert out.shape == (1, 44)
    assert stats.acceptance_rate > 0.97
    assert stats.tokens_per_target_call > 3.5  # ~gamma + 1 per target forward
 def test_stats_are_consistent(target, draft):
    _, stats = speculative_generate(target, draft, prompt_of(3), 25, gamma=4, temperature=1.0)
    assert stats.new_tokens == 25
    assert stats.target_calls == stats.rounds
    assert 0 <= stats.accepted <= stats.proposed
    assert stats.target_calls < 25  # fewer target passes than tokens, even with a poor draft
 def test_seeded_runs_are_reproducible(target, draft):
    runs = [
        speculative_generate(
            target, draft, prompt_of(3), 20, temperature=0.9, top_p=0.95,
            generator=torch.Generator().manual_seed(11),
        )[0]
        for _ in range(2)
    ]  # fmt: skip
    assert torch.equal(runs[0], runs[1])
 def test_output_distribution_is_exactly_the_targets():
    """The defining guarantee: speculation changes speed, never the distribution.
    Tiny vocabulary so the exact target distribution over the first two tokens (25 outcomes)
    and the marginal of the third can be computed by brute force and compared with samples.
    """
    v = 5
    t = make_model(GPTConfig(vocab_size=v, block_size=16, n_layer=2, n_head=2, n_embd=16), 4, 0.4)
    d = make_model(GPTConfig(vocab_size=v, block_size=16, n_layer=1, n_head=2, n_embd=16), 9, 0.4)
    prompt = torch.tensor([[1, 3]])
    def dist(model, tokens):
        with torch.no_grad():
            return logits_to_probs(model(torch.tensor([tokens]))[0][:, -1], 1.0)[0]
    def exact(model):
        joint = torch.zeros(v, v)
        third = torch.zeros(v)
        for a, b in itertools.product(range(v), repeat=2):
            w = dist(model, [1, 3])[a] * dist(model, [1, 3, a])[b]
            joint[a, b] = w
            third += w * dist(model, [1, 3, a, b])
        return joint, third
    joint_t, third_t = exact(t)
    joint_d, _ = exact(d)
    # Power check: the draft really is a different distribution, so "just sample the draft"
    # (the bug this test exists to catch) would be detected below.
    assert (joint_t - joint_d).abs().max() > 0.04
    n = 2500
    g = torch.Generator().manual_seed(123)
    joint = torch.zeros(v, v)
    third = torch.zeros(v)
    for _ in range(n):
        out, _ = speculative_generate(t, d, prompt, 3, gamma=4, temperature=1.0, generator=g)
        a, b, c = out[0, 2:].tolist()
        joint[a, b] += 1 / n
        third[c] += 1 / n
    assert (joint - joint_t).abs().max() < 0.03
    assert (third - third_t).abs().max() < 0.03
 def test_top_k_support_is_respected(target, draft):
    """Every emitted token must lie in the target's own top-k set given its prefix."""
    g = torch.Generator().manual_seed(5)
    for trial in range(40):
        out, _ = speculative_generate(
            target, draft, prompt_of(4, trial), 6, gamma=3, temperature=1.0, top_k=3, generator=g
        )
        for pos in range(4, 10):
            with torch.no_grad():
                top = target(out[:, :pos])[0][0, -1].topk(3).indices.tolist()
            assert int(out[0, pos]) in top
 def test_validation(target, draft):
    with pytest.raises(ValueError):
        speculative_generate(target, draft, prompt_of(3).repeat(2, 1), 5)  # batch > 1
    with pytest.raises(ValueError):
        speculative_generate(target, draft, prompt_of(60), 10)  # exceeds block_size
    with pytest.raises(ValueError):
        speculative_generate(target, draft, prompt_of(3), 5, gamma=0)
    other = make_model(GPTConfig(vocab_size=V + 1, block_size=64, n_layer=1, n_head=2, n_embd=16))
    with pytest.raises(ValueError):
        speculative_generate(target, other, prompt_of(3), 5)
 def test_training_mode_is_restored(target, draft):
    target.train()
    speculative_generate(target, draft, prompt_of(3), 5, temperature=0)
    assert target.training and not draft.training
    target.eval()
 def test_temperature_zero_probs_are_one_hot():
    probs = logits_to_probs(torch.tensor([[0.1, 3.0, -2.0]]), temperature=0)
    assert probs.tolist() == [[0.0, 1.0, 0.0]]