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.

scripts/bench_speculative.py

@@ -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()