# Implementing Speculative Decoding for Faster LLM Inference: From Theory to Production

> Accelerate LLM inference up to 3x with speculative decoding. Learn to implement this technique from theory to production using a draft model and rejection sampling.

- Repository: [Rohit Ghumare/ai-engineering-from-scratch](https://github.com/rohitg00/ai-engineering-from-scratch)
- Tags: deep-dive
- Published: 2026-07-26

---

**Speculative decoding accelerates large language model inference by 3× or more using a small draft model to propose candidate tokens that a target model verifies in parallel, preserving exact output distribution via the Leviathan-Kalai-Matias rejection sampling rule.**

Implementing speculative decoding for faster LLM inference requires decoupling token generation into drafting and verification phases. The `rohitg00/ai-engineering-from-scratch` repository provides a complete implementation spectrum—from mathematical harnesses demonstrating statistical correctness to production-grade KV cache management—showing how to achieve sub-linear latency for autoregressive generation.

## How Speculative Decoding Works

The architecture follows a **draft-verify-resample** loop that minimizes expensive forward passes through the target model.

### Draft Phase

A lightweight **draft model** proposes *k* tokens using its own probability distribution. In the curriculum implementation, this is a synthetic model that mirrors the target distribution with tunable alignment parameters (see `DraftModel` in [`phases/19-capstone-projects/14-speculative-decoding-server/code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/19-capstone-projects/14-speculative-decoding-server/code/main.py)). The draft model executes multiple inexpensive autoregressive steps to build a candidate token sequence.

### Verify Phase

The **target model** evaluates the entire draft sequence in a single batched forward pass (implemented as `TargetModel.verify`). Each token is accepted only if the target model’s probability for that token exceeds a scaled threshold based on the draft probability; otherwise, verification halts at the first rejection. This batched verification replaces *k* serial target model calls with one parallel computation.

### Resample Phase

When the target model rejects a draft token at position *i*, it samples a correction token from the residual distribution (the normalized difference between target and draft probabilities). The accepted prefix up to position *i-1* is emitted immediately, and the next draft begins from the new context including the correction token.

## Mathematical Foundation and Throughput Analysis

The expected throughput of speculative decoding follows a geometric series based on the acceptance rate α and draft length *K*:

```

E[t] = (1 - α^(K+1)) / (1 - α)

```

This calculation is implemented in the `expected_tokens` function in [`phases/10-llms-from-scratch/25-speculative-decoding/code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/10-llms-from-scratch/25-speculative-decoding/code/main.py). When α ≈ 0.75 and *K* = 4, the function returns approximately **3.2 tokens per verification**, yielding a theoretical **3× speed-up** compared with standard greedy decoding.

The **Leviathan-Kalai-Matias rejection rule** guarantees that the first emitted token of each speculative step follows the exact target distribution, preserving statistical fidelity. This mathematical guarantee is empirically verified in the harness through the `verify_distribution` function, which measures total-variation distance between speculative and target outputs.

## Core Architectural Components

The implementation relies on four primitives that manage the draft-verify lifecycle.

### Draft Model

The `DraftModel` class generates candidate token sequences using a lightweight architecture. In the server scaffold at [`phases/19-capstone-projects/14-speculative-decoding-server/code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/19-capstone-projects/14-speculative-decoding-server/code/main.py), the draft model accepts an `alpha_hint` parameter to simulate varying degrees of alignment with the target distribution, enabling performance testing across different accuracy scenarios.

### Target Model Verification

The `TargetModel` exposes a `verify` method that accepts a draft sequence and returns both the accepted prefix length and a correction token. This method implements the batched forward pass that evaluates draft quality against the target distribution's true probabilities.

### KV Cache Management

Production implementations require sophisticated cache handling. The `KVBuffer` class in [`phases/10-llms-from-scratch/15-speculative-decoding-eagle3/code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/10-llms-from-scratch/15-speculative-decoding-eagle3/code/main.py) tracks logical cache length for the verifier, enabling efficient rollback when rejection occurs mid-sequence. This prevents the computational cost of recomputing key-value tensors for accepted tokens while maintaining correct attention state.

### Scheduler Loop

The `speculative_decode` function (referenced in the server implementation) orchestrates the complete flow: initializing the random generator, invoking the draft model, calling target verification, emitting accepted tokens, and managing cache state between iterations.

## Implementing the Verification Loop

To verify the statistical correctness of your implementation, use the toy harness to measure distribution alignment:

```python
import numpy as np
from phases.10_llms_from_scratch.25_speculative_decoding.code.main import (
    make_target, make_draft, speculative_step, verify_distribution
)

rng = np.random.default_rng(0)
target = make_target(vocab=32, rng=rng)
draft = make_draft(target, alpha_hint=0.75, rng=rng)

tv_plain, tv_spec = verify_distribution(target, draft, K=4, n_samples=20000, rng=rng)
print(f"TV plain vs target: {tv_plain:.4f}")
print(f"TV spec vs target: {tv_spec:.4f}")

```

This reproduces the empirical total-variation distances that validate the rejection sampling rule.

Calculate expected throughput for your specific acceptance rate:

```python
from phases.10_llms_from_scratch.25_speculative_decoding.code.main import expected_tokens

alpha = 0.75
K = 4
print(f"Expected tokens per verification: {expected_tokens(alpha, K):.2f}")

```

Outputs `3.20`, confirming the theoretical speed-up for these parameters.

## Production-Ready Implementation with Eagle-3

For full-scale deployment, the Eagle-3 implementation provides complete KV cache bookkeeping:

```bash
python phases/10-llms-from-scratch/15-speculative-decoding-eagle3/code/main.py

```

This script showcases the production loop where `KVBuffer` manages cache rollback during speculative steps. The implementation handles the mathematical justification for the rejection rule while maintaining efficient memory usage during batched verification.

## End-to-End Server Scaffold

To run the complete speculative decoding server that ties components together and reports speed-up metrics:

```bash
python phases/19-capstone-projects/14-speculative-decoding-server/code/main.py

```

The script prints comparisons between baseline greedy decoding and speculative decoding across various draft-alignment values and draft lengths (e.g., α=0.60/k=2, α=0.75/k=4), measuring the ratio of target model calls required for equivalent output generation.

## Summary

- **Speculative decoding** achieves faster LLM inference by drafting tokens with a cheap model and verifying them in parallel with the target model.
- The **Leviathan-Kalai-Matias rejection rule** guarantees exact preservation of the target distribution while enabling 3× speed-ups with α=0.75 and K=4.
- **KVBuffer** in the Eagle-3 implementation manages cache rollback efficiently, making speculative decoding practical for production use.
- The `expected_tokens` function calculates theoretical throughput using the formula `E[t] = (1 - α^(K+1)) / (1 - α)`.
- Complete implementations range from the mathematical harness in [`phases/10-llms-from-scratch/25-speculative-decoding/code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/10-llms-from-scratch/25-speculative-decoding/code/main.py) to the production server in [`phases/19-capstone-projects/14-speculative-decoding-server/code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/19-capstone-projects/14-speculative-decoding-server/code/main.py).

## Frequently Asked Questions

### How does speculative decoding preserve the exact target distribution?

The **Leviathan-Kalai-Matias rejection rule** guarantees statistical fidelity by accepting draft tokens only when the target model's probability exceeds a scaled threshold, and resampling from the residual distribution upon rejection. This ensures the first emitted token of each speculative step follows the target distribution exactly, as proved in the `speculative_step` implementation of the toy harness.

### What speed-up can I expect from speculative decoding?

With a draft acceptance rate α ≈ 0.75 and draft length K = 4, the `expected_tokens` function in [`phases/10-llms-from-scratch/25-speculative-decoding/code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/10-llms-from-scratch/25-speculative-decoding/code/main.py) calculates approximately 3.2 tokens per verification. This yields a **3× speed-up** over standard greedy decoding where each token requires a separate forward pass.

### How is KV cache rollback handled when drafts are rejected?

The production implementation uses the `KVBuffer` class to track logical cache length for the verifier, enabling efficient rollback when the target model rejects a draft token partway through the sequence. This mechanism, found in [`phases/10-llms-from-scratch/15-speculative-decoding-eagle3/code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/10-llms-from-scratch/15-speculative-decoding-eagle3/code/main.py), prevents costly cache recomputation while maintaining correct attention state.

### What makes a good draft model for speculative decoding?

According to the server scaffold in [`phases/19-capstone-projects/14-speculative-decoding-server/code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/19-capstone-projects/14-speculative-decoding-server/code/main.py), an effective draft model mirrors the target distribution with high alignment (α), typically implemented as a smaller transformer or synthetic model with tunable alignment. The `DraftModel` class demonstrates this with configurable `alpha_hint` values to simulate different alignment scenarios and measure resulting speed-ups.