Speculative Decoding: 2.5x Speedup Without Quality Loss
Use a small, high-throughput draft model (e.g. Llama-3-8B) to speculate tokens, verified in parallel by the 70B target model in a single forward pass.
- •Derive the exact modified rejection sampling formula proving zero distribution degradation
- •Execute the speculative pipeline through our interactive latency & speedup simulator
- •Implement KV-cache speculative rollbacks and tree-based verification with PyTorch
- •Benchmark speedup multipliers across code generation and mathematical reasoning
Breaking the Autoregressive Memory Wall
Autoregressive generation in large models is severely memory-bandwidth bound: during decoding, generating a single token requires streaming all 140 GB of model weights from HBM to SRAM for arithmetic intensity under 1 FLOP/byte. Speculative decoding bypasses this bottleneck: an inexpensive draft model proposes K candidate tokens in rapid sequence, and the massive target model evaluates all K tokens concurrently in a single forward pass.
Speculative Decoding Verification & Rejection Pipeline
Draft model generates K tokens sequentially with low latency
Target model runs parallel forward pass on proposed sequence
Compare p_target(x) vs q_draft(x) with uniform random draw
Sample from max(0, p - q) on first rejected token to prevent drift
Retain KV pages for accepted tokens, prune rejected branch
Speculative Decoding Latency & Speedup Calculator
Production Speculative Engine Implementation Suite
Inspect the production-grade implementation of speculative decoding. The suite includes the core autoregressive speculative coordinator, vectorized rejection sampling with mathematical guarantees, and a multi-GPU latency benchmarking suite.
import torch
import torch.nn.functional as F
from typing import Tuple, List
class SpeculativeDecoder:
"""Production speculative decoding engine coordinating draft and target LLMs."""
def __init__(self, target_model, draft_model, tokenizer, k_lookahead: int = 4):
self.target = target_model
self.draft = draft_model
self.tokenizer = tokenizer
self.k = k_lookahead
@torch.inference_mode()
def generate_step(self, input_ids: torch.Tensor) -> Tuple[torch.Tensor, int]:
batch_size, prompt_len = input_ids.shape
draft_ids = input_ids.clone()
draft_logits_list = []
# 1. Draft phase: Auto-regressively sample K candidate tokens
for _ in range(self.k):
out = self.draft(draft_ids)
next_logits = out.logits[:, -1, :] # [batch, vocab]
draft_logits_list.append(next_logits)
next_token = torch.argmax(next_logits, dim=-1, keepdim=True)
draft_ids = torch.cat([draft_ids, next_token], dim=-1)
proposed_tokens = draft_ids[:, prompt_len:] # [batch, K]
q_draft = torch.stack([F.softmax(logits, dim=-1) for logits in draft_logits_list], dim=1)
# 2. Target phase: Parallel forward pass across all K candidates simultaneously
target_out = self.target(draft_ids)
target_logits = target_out.logits[:, prompt_len - 1 : prompt_len - 1 + self.k, :]
p_target = F.softmax(target_logits, dim=-1)
# 3. Vectorized Modified Rejection Sampling (Leviathan et al.)
accepted_tokens = []
rejected_index = None
for i in range(self.k):
token_i = proposed_tokens[0, i].item()
p_val = p_target[0, i, token_i].item()
q_val = q_draft[0, i, token_i].item()
u = torch.rand(1).item()
if u <= min(1.0, p_val / max(q_val, 1e-9)):
accepted_tokens.append(token_i)
else:
rejected_index = i
# Sample from adjusted residual distribution max(0, p - q)
residual = torch.clamp(p_target[0, i] - q_draft[0, i], min=0.0)
residual_prob = residual / residual.sum()
recovery_token = torch.multinomial(residual_prob, num_samples=1).item()
accepted_tokens.append(recovery_token)
break
# If all K tokens accepted, target model generates a bonus token for free!
if rejected_index is None:
bonus_token = torch.argmax(target_out.logits[:, -1, :], dim=-1).item()
accepted_tokens.append(bonus_token)
new_input_ids = torch.cat([input_ids, torch.tensor([accepted_tokens], device=input_ids.device)], dim=-1)
return new_input_ids, len(accepted_tokens)Because residual sampling draws from max(0, p - q), the output probability distribution is guaranteed to be 100% indistinguishable from sampling directly from the 70B target model.
Run this lesson's Triton kernels and multi-file code directly in an ephemeral pre-configured PyTorch 2.5 container.