tinydraft

Draft-model speculative decoding implemented from first principles — Qwen2.5-0.5B drafting for Qwen2.5-1.5B, on a laptop. Sibling project of tinybatch (the serving engine; this is the decoding accelerator).

The idea

LLM decoding is memory-bandwidth-bound: producing one token requires streaming all the weights through the chip, so verifying k+1 tokens in one forward pass costs barely more than producing one. Speculative decoding exploits that asymmetry (Leviathan et al. 2023; Chen et al. 2023):

1. A small draft model proposes k tokens autoregressively (cheap — 3× fewer params here).
2. The large target model scores all proposals + one bonus position in a single forward pass.
3. Rejection sampling keeps the output distribution exactly the target's: accept proposal x with probability min(1, p(x)/q(x)); on rejection, resample from norm(max(0, p − q)); both KV caches roll back to the accepted prefix.

It is an exact acceleration, not an approximation — with temperature 0 the rule collapses to exact-match-vs-target-argmax, so greedy speculative output is token-identical to plain greedy decoding.

Correctness (tested, not asserted)

• tests/test_rejection_math.py — the rejection-sampling identity, verified empirically: 200k draws through accept/resample match the target distribution to 5e-3 even with a badly mismatched draft. Runs in CI, no models needed.
• tests/test_lossless.py — end-to-end: greedy speculative generation equals greedy baseline generation token-for-token across prompts and k, with both real models (CPU/fp32 for bit-stable numerics).

The KV-cache engineering is the subtle part: per-model processed-length tracking, DynamicCache cropping on rejection, and the fully-accepted-window case where the draft has never seen the last proposal or the bonus token (see _CachedModel in core.py).

Benchmarks

python benchmarks/bench.py [--target ...] [--ks 2 3 4] — Apple M5, MPS/fp16, greedy, 6 prompts × 96 tokens (full data: results_1.5b.json, results_3b.json):

target (draft = Qwen2.5-0.5B)	baseline	best speculative	acceptance
Qwen2.5-1.5B (3× params)	29.3 tok/s	33.0 tok/s (1.13×, k=2)	69.4%
Qwen2.5-3B (6× params)	17.0 tok/s	25.5 tok/s (1.50×, k=3)	61.3%

Two lessons the sweep makes concrete:

• Speedup scales with the draft/target cost ratio. The same draft and near-identical acceptance yields 1.13× against a 3×-larger target but 1.50× against 6× — production deployments use 10–100× ratios (e.g. 1B drafting for 70B) precisely because of this curve.
• Bigger k is not better. On the 1.5B target, k=6 degrades to 0.99×: acceptance falls with depth (47.5%), so later proposals are usually wasted draft compute. The optimal k balances acceptance decay against verification amortization.

Run it

python3.12 -m venv .venv && source .venv/bin/activate
pip install -e ".[dev]"

pytest tests/test_rejection_math.py     # fast, no downloads
pytest                                  # full suite (downloads both models)
python benchmarks/bench.py              # speedup + acceptance-rate table

from tinydraft import generate, load_models

target, draft, tok = load_models()
tokens, stats = generate(target, draft, tok, "Why is LLM decoding memory-bound?", k=4)
print(tok.decode(tokens), stats.summary())

Honest limitations

Single-sequence decoding only (no batching — that's tinybatch's job; production systems combine both), greedy/temperature sampling without top-p on the speculative path, and no tree/EAGLE-style multi-candidate drafting. Absolute speedups are laptop-scale; the acceptance-rate math and the losslessness proof are hardware-independent.

License

MIT