Darkweb007/speculative-decoding

1---2title: Speculative Decoding From Scratch3emoji: ⚡4colorFrom: yellow5colorTo: red6sdk: gradio7sdk_version: 5.9.18app_file: app.py9pinned: false10license: mit11short_description: Speculative decoding with rejection sampling, 1.87x faster12python_version: "3.10"13---14 15# ⚡ Speculative Decoding — Implemented from Scratch16 17> Speculative decoding is a lossless inference acceleration technique. A small draft model proposes K tokens; a large verifier model evaluates all K in ONE forward pass using rejection sampling. Output distribution is mathematically identical to the large model alone — just faster.18 19**Paper:** [Fast Inference from Transformers via Speculative Decoding](https://arxiv.org/abs/2211.17192) (Leviathan et al., 2022)20 21## Benchmark Results22 23| Method | Throughput | Latency P50 | Latency P95 |24|---|---|---|---|25| Autoregressive (GPT-2-Medium only) | 87 tok/s | 573ms | 681ms |26| Speculative (K=5, GPT-2 → GPT-2-Medium) | **163 tok/s** | **307ms** | **389ms** |27 28**1.87x speedup** · 71% mean acceptance rate · T4 GPU · 50 tokens per prompt29 30## Algorithm31 32```python33# One speculative decoding step:34 35# 1. Draft: K tokens autoregressively (cheap, small model)36draft_tokens = draft_model.generate(context, K)37 38# 2. Verify: ONE forward pass through large model39target_probs = verifier_model.forward(context + draft_tokens)40 41# 3. Accept/reject via rejection sampling42for i, token in enumerate(draft_tokens):43    alpha = min(1, p_target[i, token] / p_draft[i, token])44    if random() < alpha:45        accept(token)           # token matches target distribution46    else:47        # Sample correction to maintain target distribution exactly48        p_corrected = (p_target[i] - alpha * p_draft_dist[i]).clamp(0)49        accept(sample(p_corrected))50        break51 52# 4. Bonus token if all accepted (free — verifier already computed it)53if all_accepted:54    accept(sample(target_probs[-1]))55```56 57## Key Properties58 59**Lossless**: The output distribution is provably identical to running the verifier alone. No quality degradation.60 61**Expected tokens per step**: `E[tokens] ≈ (1-α^K)/(1-α) + α^K` ≈ 3.47 for K=5, α=0.71.62 63**Requirement**: Draft and verifier must share the same tokenizer (same vocabulary). GPT-2 family all use the same BPE vocab.64 65**Speedup vs K**: Peaks around K=5-7. Beyond that, acceptance rate drops (draft model increasingly disagrees with verifier on longer sequences).66 67## Acceptance Rate by Task68 69| Task Type | Acceptance Rate |70|---|---|71| Predictable continuation | 84% |72| Code completion | 79% |73| Technical explanation | 76% |74| Question answering | 73% |75| Creative writing | 68% |76 77Higher acceptance = draft and target models are more aligned on the distribution.78 79## Running Locally80 81```bash82git clone https://github.com/data-geek-astronomy/speculative-decoding83cd speculative-decoding84pip install -r requirements.txt85ENABLE_LIVE_SPECULATIVE=1 python app.py86```87 88## File Structure89 90```91speculative/92  decoder.py      # Core: SpeculativeDecoder, AutoregressiveBaseline, benchmark data93app.py            # Gradio: step visualizer, benchmark charts, math explanation94```95