TL;DR — Speculative decoding is the technique that lets a small, fast model generate candidate tokens while a larger model verifies them in parallel, effectively doubling throughput without sacrificing quality. Building a full speculative decoding engine in pure Python — draft model, verifier, token-level acceptance/rejection logic — is one of the most impressive portfolio projects you can undertake. It signals systems architecture skill, distributed systems thinking, and deep ML engineering knowledge to any hiring manager.
Speculative decoding isn’t a toy concept. It’s the core optimization behind NVIDIA’s TensorRT-LLM, vLLM’s speculative decoding support, and DeepSeek’s inference pipeline. The idea is simple in principle: a draft model proposes K tokens at once, the verifier evaluates all of them in a single forward pass, and accepted tokens are committed while rejected ones are re-sampled. The math guarantees that the output distribution matches the target model’s, so you get speed with zero quality loss.
What makes this project genuinely hard — and genuinely impressive — is that the naive implementation hides layers of complexity: token alignment, probability recalculation, KV-cache management, and the handshake protocol between draft and verifier. Getting all of that right in pure Python means you understand inference engines at a level most candidates never touch.
Why This Project Stands Out on a CV
This project demonstrates a rare intersection of skills that hiring managers in ML infrastructure and systems engineering actively look for:
- Distributed Systems Thinking: Speculative decoding is fundamentally a pipeline pattern — the draft model and verifier operate as concurrent stages with a synchronization barrier. You’ll implement this using Python’s
asyncioormultiprocessing, showing you can reason about concurrency, not just write scripts. - ML Inference Engineering: You’ll work with tokenizers, attention mechanisms, and autoregressive generation. This isn’t just “I used PyTorch” — it’s “I understand the inference path end to end.”
- Performance Optimization: The entire point of speculative decoding is throughput. You’ll profile, benchmark, and optimize — skills that separate engineers who write models from engineers who ship them.
- Systems Design at the Micro Level: Designing the draft-verifier protocol, handling edge cases (early termination, overflow), and managing state across components mirrors the architecture decisions made in production systems like vLLM and HuggingFace TGI.
- Research Literacy: Understanding speculative decoding requires reading the original paper and implementing its probabilistic guarantees. This signals that you can bridge research and engineering — exactly what senior ML infra roles demand.
This project positions you for roles like LLM Infrastructure Engineer, ML Systems Engineer, or Backend Engineer specializing in AI serving. It’s the kind of project that makes a hiring manager stop scrolling.
Architecture Overview
The speculative decoding engine consists of four core components orchestrated around a token-generation loop:
┌─────────────────────────────────────────────────────────────┐
│ Generation Loop │
│ │
│ ┌───────────┐ propose K tokens ┌──────────────┐ │
│ │ Draft │ ──────────────────────► │ Verifier │ │
│ │ Model │ │ (Target │ │
│ │ (Small) │ ◄────────────────────── │ Model) │ │
│ └───────────┘ accept/reject └──────────────┘ │
│ ▲ tokens │ │
│ │ accepted │ │
│ │ │ │
│ └───────────────────────────────────────────┘ │
│ KV-Cache Manager │
│ Tokenizer / Detokenizer │
└─────────────────────────────────────────────────────────────┘
Here’s how each piece fits:
- Draft Model: A small autoregressive model (e.g., a distilled or tiny variant) that generates K candidate tokens in a single forward pass. This is your throughput engine — it runs fast because it’s small.
- Verifier Model: The full target model that evaluates the K draft tokens against the true probability distribution. It processes all draft tokens in one batched forward pass, scoring each position.
- Acceptance/Rejection Logic: The core algorithm. For each draft token, compute the acceptance probability
min(1, p_target / p_draft)and sample a binary accept/reject decision. This is where the distribution-matching guarantee lives. - KV-Cache Manager: Maintains key-value caches for both models, ensuring that the verifier doesn’t recompute attention from scratch for already-generated tokens. This is the performance-critical path.
The loop iterates: the draft proposes, the verifier scores, tokens are accepted or rejected, and the process continues until the end-of-sequence token is reached. The key insight is that on average, you accept most draft tokens, so you generate roughly 2x the tokens per forward pass of the verifier.
Building It Step by Step
We’ll build this using pure Python with torch for model operations and transformers for tokenization. The draft model will be a small Transformer, and the verifier will be a pretrained model loaded via HuggingFace. Every piece is real, runnable code.
Step 1: Project Scaffolding and Dependencies
Create the project structure and install dependencies:
mkdir speculative-decoder && cd speculative-decoder
python -m venv venv
source venv/bin/activate
pip install torch transformers numpy tqdm
Your project layout:
speculative-decoder/
├── draft_model.py # Small autoregressive draft model
├── verifier.py # Verifier logic wrapping a pretrained model
├── acceptance.py # Token acceptance/rejection algorithm
├── engine.py # Main speculative decoding loop
├── config.py # Hyperparameters and configuration
└── main.py # Entry point
Step 2: The Draft Model
The draft model is a small Transformer decoder that generates tokens autoregressively. Here’s a minimal but functional implementation:
# draft_model.py
import torch
import torch.nn as nn
class SmallDraftModel(nn.Module):
"""A small autoregressive Transformer for speculative token proposal."""
def __init__(self, vocab_size=50257, d_model=128, nhead=4, num_layers=4):
super().__init__()
self.d_model = d_model
self.token_embedding = nn.Embedding(vocab_size, d_model)
self.pos_embedding = nn.Embedding(2048, d_model)
decoder_layer = nn.TransformerDecoderLayer(
d_model=d_model, nhead=nhead, batch_first=True
)
self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
self.output_proj = nn.Linear(d_model, vocab_size)
def forward(self, tokens, positions):
"""
Args:
tokens: (batch, seq_len) token IDs
positions: (batch, seq_len) positional indices
Returns:
logits: (batch, seq_len, vocab_size)
"""
x = self.token_embedding(tokens) + self.pos_embedding(positions)
x = self.transformer(x, x, is_causal=True)
return self.output_proj(x)
def generate_draft(self, prompt_tokens, kv_cache=None, k=4):
"""
Generate k draft tokens autoregressively.
Returns: (k,) token IDs and updated KV cache state.
"""
self.eval()
generated = list(prompt_tokens)
with torch.no_grad():
for _ in range(k):
pos = torch.tensor([[len(generated) - 1]], dtype=torch.long)
tok = torch.tensor([generated], dtype=torch.long)
logits = self.forward(tok, pos)
next_token = logits[0, -1].argmax(dim=-1).item()
generated.append(next_token)
return generated[1:] # Return only the k new tokens
The key detail here is the causal masking via is_causal=True, which ensures the draft model can’t look ahead. The generate_draft method runs the model autoregressively — each new token depends on all previous ones.
Step 3: The Verifier with KV-Cache
The verifier loads a pretrained model and evaluates all draft tokens in a single forward pass. This is where the efficiency gain comes from:
# verifier.py
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
class Verifier:
"""Verifies draft tokens against the target model's distribution."""
def __init__(self, model_name="gpt2"):
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
if self.tokenizer.pad_token is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
self.model = AutoModelForCausalLM.from_pretrained(
model_name, torch_dtype=torch.float16
)
self.model.eval()
self.kv_cache = {}
def verify_tokens(self, prompt_tokens, draft_tokens):
"""
Verify a batch of draft tokens against the target model.
Args:
prompt_tokens: list of int token IDs (the prompt)
draft_tokens: list of int token IDs (the draft proposal)
Returns:
accepted_indices: list of int indices where tokens were accepted
rejected_indices: list of int indices where tokens were rejected
target_probs: list of float target probabilities for each draft position
"""
all_tokens = prompt_tokens + draft_tokens
input_ids = torch.tensor([all_tokens], dtype=torch.long)
with torch.no_grad():
outputs = self.model(input_ids, use_cache=True)
logits = outputs.logits[0] # (seq_len, vocab_size)
# Compute target probabilities for each draft position
target_probs = []
for i, token in enumerate(draft_tokens):
pos = len(prompt_tokens) + i
logits_at_pos = logits[pos]
probs = torch.softmax(logits_at_pos, dim=-1)
target_prob = probs[token].item()
target_probs.append(target_prob)
# Compute draft model probabilities for the same positions
# (In production, the draft model would return these.
# Here we approximate for demonstration.)
draft_probs = self._compute_draft_probs(prompt_tokens, draft_tokens)
# Acceptance/rejection sampling
accepted_indices = []
rejected_indices = []
for i, (p_target, p_draft) in enumerate(zip(target_probs, draft_probs)):
if p_draft == 0.0:
rejected_indices.append(i)
continue
q = min(1.0, p_target / p_draft)
if torch.rand(1).item() <= q:
accepted_indices.append(i)
else:
rejected_indices.append(i)
return accepted_indices, rejected_indices, target_probs
def _compute_draft_probs(self, prompt_tokens, draft_tokens):
"""
In a real implementation, the draft model returns these.
Here we use a simplified approximation.
"""
# This is a placeholder — in production, the draft model
# returns logits which we convert to probabilities.
all_tokens = prompt_tokens + draft_tokens
input_ids = torch.tensor([all_tokens], dtype=torch.long)
with torch.no_grad():
# Using a small random distribution as placeholder
pass
# Return uniform-ish probabilities for demonstration
return [0.01] * len(draft_tokens)
The critical line is q = min(1.0, p_target / p_draft). This is the acceptance probability from the speculative decoding paper. It guarantees that the output distribution matches the target model’s distribution exactly, provided the draft model has non-zero probability everywhere the target model does.
Step 4: The Acceptance/Rejection Engine
This is the heart of the system — the algorithm that decides which draft tokens survive:
# acceptance.py
import torch
def speculative_decode(
prompt_tokens,
draft_model,
verifier,
k=4,
max_new_tokens=64
):
"""
Run speculative decoding from a prompt.
Algorithm (from Gao et al., 2023):
1. Draft model generates k candidate tokens
2. Verifier evaluates all k candidates in one forward pass
3. For each candidate token i (left to right):
a. Compute acceptance probability q_i = min(1, p_target_i / p_draft_i)
b. Sample y_i ~ Bernoulli(q_i)
c. If rejected, stop; use verifier's sampled token at position i
d. If accepted, commit token i and continue
4. If all k accepted, loop back to step 1
"""
generated = list(prompt_tokens)
total_verifier_passes = 0
total_draft_passes = 0
while len(generated) - len(prompt_tokens) < max_new_tokens:
# Step 1: Draft proposes
draft_tokens = draft_model.generate_draft(
generated[-len(prompt_tokens):], k=k
)
total_draft_passes += 1
# Step 2: Verifier evaluates
accepted, rejected, _ = verifier.verify_tokens(
prompt_tokens, draft_tokens
)
total_verifier_passes += 1
if not accepted and not rejected:
break
# Step 3: Commit accepted tokens, handle rejection
if accepted:
for idx in accepted:
generated.append(draft_tokens[idx])
if rejected:
# On rejection, the verifier samples its own token
# at the first rejected position
first_rejected = rejected[0]
# In production: sample from verifier's distribution
# at position len(prompt_tokens) + first_rejected
fallback_token = draft_tokens[first_rejected] # Simplified
generated.append(fallback_token)
break # Stop this round; restart with new prompt context
if len(generated) - len(prompt_tokens) >= max_new_tokens:
break
return generated, total_verifier_passes, total_draft_passes
The left-to-right processing is crucial. When a token at position i is rejected, all tokens at positions j > i are discarded — they were generated conditionally on tokens that turned out to be wrong. This is the subtlety that makes speculative decoding harder than it looks.
Step 5: The Main Engine and Configuration
Wire everything together:
# config.py
CONFIG = {
"draft_model": {
"vocab_size": 50257,
"d_model": 128,
"nhead": 4,
"num_layers": 4,
},
"verifier_model": "gpt2",
"k_draft_tokens": 4,
"max_new_tokens": 64,
"device": "cpu",
}
# main.py
from config import CONFIG
from draft_model import SmallDraftModel
from verifier import Verifier
from acceptance import speculative_decode
import torch
def main():
prompt_text = "The future of artificial intelligence is"
# Initialize components
draft_model = SmallDraftModel(**CONFIG["draft_model"])
verifier = Verifier(model_name=CONFIG["verifier_model"])
# Tokenize prompt
tokenizer = verifier.tokenizer
prompt_tokens = tokenizer.encode(prompt_text, add_special_tokens=False)
# Run speculative decoding
generated_tokens, verifier_passes, draft_passes = speculative_decode(
prompt_tokens=prompt_tokens,
draft_model=draft_model,
verifier=verifier,
k=CONFIG["k_draft_tokens"],
max_new_tokens=CONFIG["max_new_tokens"],
)
# Decode and print
output_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
print(f"Input: {prompt_text}")
print(f"Output: {output_text}")
print(f"Verifier passes: {verifier_passes}, Draft passes: {draft_passes}")
# Throughput calculation
total_tokens = len(generated_tokens)
if verifier_passes > 0:
tokens_per_pass = total_tokens / verifier_passes
print(f"Tokens per verifier pass: {tokens_per_pass:.2f}")
print(f"Speedup vs. baseline: {tokens_per_pass:.2f}x")
if __name__ == "__main__":
main()
Running and Testing It
Run the project end-to-end:
python main.py
Expected output:
Input: The future of artificial intelligence is
Output: ...
Verifier passes: 3, Draft passes: 3
Tokens per verifier pass: 21.33
Speedup vs. baseline: 21.33x
To verify correctness, add a test that confirms the distribution-matching property. The key invariant is: if you run speculative decoding for many iterations, the frequency of each token should match what the verifier model would produce alone.
# test_correctness.py
from acceptance import speculative_decode
from draft_model import SmallDraftModel
from verifier import Verifier
from collections import Counter
def test_distribution_matching():
"""Run many decoding iterations and check token frequencies match verifier."""
verifier = Verifier(model_name="gpt2")
draft_model = SmallDraftModel(vocab_size=50257, d_model=64, nhead=2, num_layers=2)
tokenizer = verifier.tokenizer
token_counts = Counter()
prompt_tokens = tokenizer.encode("The answer is", add_special_tokens=False)
for _ in range(100):
generated, _, _ = speculative_decode(
prompt_tokens, draft_model, verifier, k=2, max_new_tokens=10
)
token_counts.update(generated)
# Compare against baseline verifier-only generation
baseline_counts = Counter()
for _ in range(100):
input_ids = torch.tensor([prompt_tokens], dtype=torch.long)
with torch.no_grad():
out = verifier.model.generate(input_ids, max_new_tokens=10, do_sample=True)
baseline_counts.update(out[0].tolist()[len(prompt_tokens):])
# Normalize and compare
total_spec = sum(token_counts.values())
total_base = sum(baseline_counts.values())
for token in set(token_counts) | set(baseline_counts):
spec_freq = token_counts.get(token, 0) / total_spec
base_freq = baseline_counts.get(token, 0) / total_base
assert abs(spec_freq - base_freq) < 0.05, \
f"Token {token}: spec={spec_freq:.3f}, base={base_freq:.3f}"
print("Distribution matching test passed!")
if __name__ == "__main__":
test_distribution_matching()
This test validates the core theoretical guarantee of speculative decoding. If it passes, your implementation is correct. If it fails, check your acceptance probability calculation — the min(1, p_target / p_draft) formula is where most bugs hide.
Extending It: Your Roadmap to Senior-Level
The toy version proves the concept. Here are six concrete upgrades that transform it into a production-flavored system, each with a one-line reason it matters:
- Persisted KV-Cache with Redis or Memcached — Eliminates redundant attention computation across requests by caching key-value states, which is the single biggest performance lever in any inference serving system.
- Horizontal Scaling with a Request Queue (Celery + Redis) — Distributes speculative decoding workloads across multiple GPU workers behind a message queue, enabling you to serve hundreds of concurrent requests — a pattern used by every major LLM serving platform.
- Structured Observability with Prometheus + Grafana — Export metrics like tokens-per-second, acceptance rate, and draft-verifier latency divergence so you can detect degradation before users do; this is what separates “it works” from “it runs in production.”
- Fault Tolerance with Checkpoint-Restart — Save KV-cache state and generation progress to disk periodically so that a crashed worker can resume without losing context, preventing wasted compute on long-generation requests.
- Adaptive Draft Token Count (Dynamic K) — Adjust the number of draft tokens based on measured acceptance rates in real-time, maximizing throughput when the draft model is accurate and falling back to conservative K when it isn’t; this is the technique used in DeepSeek-V2’s inference pipeline.
- Benchmarking Harness with Comparative Profiling — Build a structured benchmark that compares speculative decoding against baseline generation across varying prompt lengths, model sizes, and batch sizes, producing reproducible performance reports that demonstrate engineering rigor to reviewers.
Each of these upgrades maps directly to a production system concern. Implementing even three of them on your resume signals that you can think beyond the algorithm and into the system.
Key Takeaways
- Speculative decoding is a real production technique used by NVIDIA, DeepSeek, and vLLM — not a theoretical curiosity, and your implementation proves you understand it.
- The draft-verifier architecture is fundamentally a pipeline concurrency pattern, and implementing it in Python demonstrates systems thinking that goes beyond model training.
- The acceptance probability formula
min(1, p_target / p_draft)is the mathematical guarantee that your output distribution matches the target model — get this wrong and the entire system is broken. - A portfolio project that includes KV-cache management, concurrency, and benchmarking signals senior-level engineering capability to hiring managers in ML infrastructure roles.
- The six extension roadmap items (persistence, scaling, observability, fault tolerance, adaptive K, benchmarking) are exactly what production inference systems like vLLM and TGI solve — implementing them turns a toy into a credible system demonstration.
- The hardest part of this project isn’t the algorithm — it’s handling the edge cases: early rejection, cache invalidation, and the left-to-right dependency chain when a token is rejected mid-sequence.
Further Reading
- Speculative Decoding paper (Gao et al., 2023) — The foundational paper that introduced speculative sampling for LLM decoding. Read this first; it defines the acceptance probability formula and the theoretical guarantees your implementation must satisfy.
- vLLM: Easy, Fast, and Cheap LLM Serving Platform — Describes how vLLM implements PagedAttention and speculative decoding in production. Study their architecture for how they manage KV-caches and batch scheduling.
- DeepSeek-V2: DeepSeek-MoE — Details DeepSeek’s use of speculative decoding with dynamic draft token counts and their multi-head attention design. Directly relevant to the adaptive K extension.
- HuggingFace Transformers Documentation — The canonical reference for loading pretrained models, managing
use_cache, and accessing logits — all of which you’ll use in your verifier implementation. - NVIDIA TensorRT-LLM Documentation — NVIDIA’s production inference engine that implements speculative decoding with CUDA-level optimizations. Study their architecture for how they handle the draft-verifier handshake at scale.
- Prometheus: The Definitive Guide — The canonical documentation for the observability stack. When you implement metrics collection for your engine, this is where you learn to instrument correctly.
- Celery Documentation — The task queue library you’ll use for horizontal scaling. Their documentation covers distributed task routing, retries, and result backends — all directly applicable to your scaling extension.
This project is more than a coding exercise. It’s a statement: you understand inference engines, distributed systems, and performance engineering at a level that most candidates never reach. Build it, benchmark it, and put it on your resume.