diff --git a/README.md b/README.md
index e183ca0..980100f 100644
--- a/README.md
+++ b/README.md
@@ -1,55 +1,78 @@
 # microgpt
 
-A GPT-style language model written from scratch in PyTorch, with the pieces that matter in
-production LLM systems: rotary position embeddings, RMSNorm, SwiGLU, an incremental-decoding
-**KV cache**, nucleus sampling, and a test suite that proves each piece is correct.
+A GPT-style language model written from scratch in PyTorch, plus the inference techniques that
+matter in production LLM serving: a **KV cache**, **grouped-query attention**, **continuous
+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
-written out explicitly (~220 lines of model code) so every line can be reasoned about.
+No `nn.Transformer`, no Hugging Face. Attention, masking, caching, batching and sampling are all
+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 |
-| RoPE | `model.py` | Relative position without a learned table; cached keys keep their absolute rotation |
-| RMSNorm + SwiGLU, tied embeddings | `model.py` | The Llama-style recipe at toy scale |
+| Causal attention, RoPE, RMSNorm, SwiGLU, tied embeddings | `model.py` | The Llama recipe at toy scale; mask offset by cache length |
+| Grouped-query / multi-query attention | `model.py` | Cache stores only the KV heads; q is regrouped so matmul broadcasts instead of copying K/V |
 | 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 |
-| Trainer | `train.py` | AdamW with decoupled weight decay (no decay on norms), warmup + cosine LR, grad clipping |
-| Benchmark | `scripts/benchmark.py` | Measures the cache speedup instead of asserting it |
+| Continuous batching | `batching.py` | Left-padded shared cache, per-row RoPE positions, key mask, mid-flight admission, column trimming |
+| Speculative decoding | `speculative.py` | Exact target distribution (rejection sampling + residual), cache rollback |
+| 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 == full forward** — prefill then one-token-at-a-time decoding reproduces the full
-  forward pass to `1e-5`; so does multi-token continuation (which exercises the offset causal mask).
-* **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`.
-* **Sampling** — top-k and top-p keep exactly the intended tokens; empirical frequencies match
-  the softmax distribution; seeded runs are reproducible.
-* **Learning** — the model can overfit a tiny batch (loss drops by more than 70%).
+**KV cache.** Prefill-then-decode reproduces the full forward pass to `1e-5`, including
+multi-token continuation (which exercises the offset causal mask). Perturbing future tokens never
+changes earlier logits. RoPE preserves norms and its attention scores depend only on `m - n`.
+
+**Grouped-query attention.** A GQA model must equal a multi-head model whose K/V heads are
+*copies* of the GQA ones (weights are copied across and logits compared, both full-forward and
+token-by-token). Cache bytes are asserted to scale exactly with the number of KV heads.
+
+**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,
-1200 steps (~5.6 min on CPU):
+### Training (character-level Tiny Shakespeare)
+
+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,
-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.
+Cached throughput stays roughly flat as the sequence grows while uncached throughput decays.
 
-## Design notes and known limitations
+### Grouped-query attention
 
-* **Sliding window past `block_size`.** The uncached path always sees the last `block_size`
-  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
-  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.
-  Both are natural next steps and would also shrink the KV cache.
-* **Character-level, small scale.** The point is correctness and clarity, not sample quality.
-  Samples are Shakespeare-flavoured but not coherent.
+8 query heads, 6 layers, 256-dim, cache holding 496 tokens:
+
+| KV heads | variant | params | KV cache | vs MHA |
+|---:|:--|---:|---:|---:|
+| 8 | MHA | 4.84M | 5.8 MiB | 1.00x |
+| 4 | GQA | 4.45M | 2.9 MiB | 0.50x |
+| 2 | GQA | 4.25M | 1.5 MiB | 0.25x |
+| 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
-- [ ] Batched generation with per-sequence cache lengths
-- [ ] Speculative decoding with a draft model, verified to preserve the target distribution
-- [ ] BPE tokenizer
+- [ ] BPE tokenizer trained from scratch, with round-trip tests
+- [ ] Paged KV cache (fixed-size blocks, block table) to remove the `torch.cat` reallocation
+- [ ] Speculative decoding inside the batch engine
+- [ ] Load a public checkpoint into this architecture and match Hugging Face logits
diff --git a/scripts/bench_speculative.py b/scripts/bench_speculative.py
new file mode 100644
index 0000000..a95d265
--- /dev/null
+++ b/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()
diff --git a/src/microgpt/generate.py b/src/microgpt/generate.py
index c3a1441..75421b7 100644
--- a/src/microgpt/generate.py
+++ b/src/microgpt/generate.py
@@ -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,
diff --git a/src/microgpt/sampling.py b/src/microgpt/sampling.py
index 6d9d159..214e7fd 100644
--- a/src/microgpt/sampling.py
+++ b/src/microgpt/sampling.py
@@ -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
-    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)
+    probs = logits_to_probs(logits, temperature, top_k, top_p)
     return torch.multinomial(probs, num_samples=1, generator=generator)
diff --git a/src/microgpt/speculative.py b/src/microgpt/speculative.py
new file mode 100644
index 0000000..23335a5
--- /dev/null
+++ b/src/microgpt/speculative.py
@@ -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
diff --git a/tests/test_speculative.py b/tests/test_speculative.py
new file mode 100644
index 0000000..e724aae
--- /dev/null
+++ b/tests/test_speculative.py
@@ -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]]