LLM Inference & Speculative Decoding: From Roofline Physics to Lossless Acceleration
Autoregressive generation is inherently sequential: to sample token $x_{t+1} \sim p_{\text{target}}(\cdot \mid x_{\le t})$, a 70B parameter transformer must stream all 140 GB of BF16 weights from High-Bandwidth Memory (HBM) into on-chip SRAM just to perform a single vector-matrix multiplication per layer. At small batch sizes, over 98% of GPU compute capacity sits completely idle waiting for weight bytes to cross the memory bus.
Speculative Decoding (Leviathan et al., 2023; Chen et al., 2023) breaks this sequential bottleneck without changing a single weight or sacrificing a single bit of output quality. By using a fast draft model $q_{\text{draft}}$ to propose $\gamma$ candidate tokens and verifying all $\gamma + 1$ positions in a single parallel forward pass of the target model $p_{\text{target}}$, we ride the Roofline curve from the memory-bound regime toward the compute-bound plateau – while mathematically guaranteeing that the output distribution is 100% identical to standard sampling from $p_{\text{target}}$.
0. The Roofline Physics of Prefill vs. Decode
Consider a linear layer $Y = X W$ with weight matrix $W \in \mathbb{R}^{d_{\text{in}} \times d_{\text{out}}}$ in BF16 ($2\text{ bytes/param}$) processing $N$ tokens ($N = B \cdot S_{\text{step}}$). The matrix multiplication executes $2 N d_{\text{in}} d_{\text{out}}$ FLOPs and transfers approximately $2 d_{\text{in}} d_{\text{out}}$ bytes of weights (for $N \ll d_{\text{in}}, d_{\text{out}}$):
On an H100 SXM5 ($\text{Peak}_{\text{FLOP/s}} = 989\text{ TFLOP/s}$, $\text{BW}_{\text{HBM}} = 3.35\text{ TB/s}$), the hardware ridge point is $N_{\text{ridge}} = \frac{989}{3.35} \approx 295\text{ tokens}$. For $B=1$ decoding ($N=1$), verifying $\gamma+1 = 8$ tokens in parallel costs practically the exact same wall-clock latency as decoding $1$ token because weight loading dominates!
1. KV-Cache Engineering & Paged Block Allocation
Before accelerating token verification, an inference engine must eliminate dynamic tensor concatenation (torch.cat([k_past, k_new], dim=2)), which triggers $\mathcal{O}(S^2)$ HBM memory copies across an $S$-step generation and fragments up to 60% of GPU memory when sequence lengths are unknown in advance.
Inspired by operating system virtual memory, PagedAttention (Kwon et al., 2023) partitions the KV cache into fixed-size physical blocks of $B_{\text{block}} = 16$ tokens. Each request maintains a logical-to-physical Block Table $\mathcal{T}[r, b_{\text{logical}}] \to b_{\text{physical}}$, eliminating external fragmentation and enabling zero-copy prefix sharing (and instant KV rollback when speculative tokens are rejected!):
💻 PyTorch Implementation: Paged KV-Cache Manager with Instant Speculative RollbackPure PyTorch
import torch
class PagedKVCacheManager:
"""
Pre-allocated Paged KV-Cache supporting O(1) append, prefix sharing, and
O(1) speculative token rollback (simply decrementing seq_len and freeing trailing blocks).
"""
def __init__(self, num_blocks: int, block_size: int, n_kv_heads: int, head_dim: int, device="cpu"):
self.block_size = block_size
self.k_pool = torch.zeros((num_blocks, n_kv_heads, block_size, head_dim), device=device)
self.v_pool = torch.zeros((num_blocks, n_kv_heads, block_size, head_dim), device=device)
self.free_blocks = list(range(num_blocks - 1, -1, -1))
self.block_tables: dict[int, list[int]] = {}
self.seq_lens: dict[int, int] = {}
def allocate_request(self, req_id: int) -> None:
self.block_tables[req_id] = []
self.seq_lens[req_id] = 0
def append_kv(self, req_id: int, k_tokens: torch.Tensor, v_tokens: torch.Tensor) -> None:
"""k_tokens, v_tokens: [n_kv_heads, num_new_tokens, head_dim]"""
num_new = k_tokens.size(1)
cur_len = self.seq_lens[req_id]
for i in range(num_new):
pos = cur_len + i
logical_idx = pos // self.block_size
offset = pos % self.block_size
if logical_idx >= len(self.block_tables[req_id]):
phys_block = self.free_blocks.pop()
self.block_tables[req_id].append(phys_block)
else:
phys_block = self.block_tables[req_id][logical_idx]
self.k_pool[phys_block, :, offset, :] = k_tokens[:, i, :]
self.v_pool[phys_block, :, offset, :] = v_tokens[:, i, :]
self.seq_lens[req_id] = cur_len + num_new
def rollback_to(self, req_id: int, target_len: int) -> None:
"""O(1) speculative rollback when draft tokens are rejected."""
assert target_len <= self.seq_lens[req_id]
needed_blocks = (target_len + self.block_size - 1) // self.block_size if target_len > 0 else 0
while len(self.block_tables[req_id]) > needed_blocks:
freed = self.block_tables[req_id].pop()
self.free_blocks.append(freed)
self.seq_lens[req_id] = target_len
def gather_contiguous_kv(self, req_id: int) -> tuple[torch.Tensor, torch.Tensor]:
"""Reconstructs [n_kv_heads, seq_len, head_dim] view for verification."""
total_len = self.seq_lens[req_id]
if total_len == 0:
return (torch.empty(self.k_pool.size(1), 0, self.k_pool.size(3), device=self.k_pool.device),
torch.empty(self.v_pool.size(1), 0, self.v_pool.size(3), device=self.v_pool.device))
k_list, v_list = [], []
for idx, phys in enumerate(self.block_tables[req_id]):
valid = min(self.block_size, total_len - idx * self.block_size)
k_list.append(self.k_pool[phys, :, :valid, :])
v_list.append(self.v_pool[phys, :, :valid, :])
return torch.cat(k_list, dim=1), torch.cat(v_list, dim=1)
2. Speculative Decoding: Exact Modified Rejection Sampling Proof
Suppose at prefix $x_{\le t}$, a fast draft model $q(x) = q_{\text{draft}}(x \mid x_{\le t})$ proposes a token $\tilde{x} \sim q(x)$, and our target model computes $p(x) = p_{\text{target}}(x \mid x_{\le t})$ in parallel across all $\gamma$ drafted positions. How do we accept or reject $\tilde{x}$ so that the resulting sample $x'$ is distributed identically to $p(x)$?
First-Principles Proof of Losslessness
Let us prove that the marginal probability $\mathbb{P}(X = x)$ of emitting any token $x \in V$ under this two-step procedure equals $p(x)$:
- Probability of accepting $x$ directly from the draft proposal: $$\mathbb{P}(\text{draft proposes } x \text{ and accepted}) = q(x) \cdot \min\!\left(1,\; \frac{p(x)}{q(x)}\right) = \min\!\bigl(q(x),\; p(x)\bigr)$$
- Overall acceptance probability $\alpha$: Summing over all $v \in V$, the probability that some drafted token is accepted is $\alpha = \sum_{v \in V} \min(q(v), p(v)) = 1 - \frac{1}{2}\|p - q\|_1 = 1 - D_{\text{TV}}(p, q)$, where $D_{\text{TV}}$ is the Total Variation distance. Thus, the rejection probability is $1 - \alpha = \sum_{v \in V} (p(v) - q(v))_+$.
- Total probability of emitting $x$: Either $x$ was drafted and accepted, OR the draft was rejected (with probability $1 - \alpha$) and $x$ was resampled from $p_{\text{res}}(x)$: $$\mathbb{P}(X = x) = \min\!\bigl(q(x), p(x)\bigr) + (1 - \alpha) \cdot \frac{\bigl(p(x) - q(x)\bigr)_+}{1 - \alpha} = \min\!\bigl(q(x), p(x)\bigr) + \max\!\bigl(0, p(x) - q(x)\bigr) \equiv p(x) \quad \blacksquare$$
where $\alpha = 1 - D_{\text{TV}}(p, q)$ is the mean acceptance rate and $c = T_{\text{draft}} / T_{\text{target}}$ is the draft-to-target latency cost ratio. For $\alpha = 0.8, \gamma = 5, c = 0.05$, $\mathbb{E}[\text{tokens/step}] = 3.69$ and wall-clock speedup is $\mathbf{2.95\times}$.
💻 PyTorch Implementation: Lossless Speculative Decoding Step with Residual ResamplingPure PyTorch
import torch
import torch.nn.functional as F
def sample_from_probs(probs: torch.Tensor) -> torch.Tensor:
"""Samples 1 token ID per batch row from normalized probability distribution [B, V]."""
return torch.multinomial(probs, num_samples=1).squeeze(-1)
def speculative_decoding_step(
draft_tokens: torch.Tensor, # [B, gamma] candidate token IDs proposed by draft model
draft_probs: torch.Tensor, # [B, gamma, V] draft probabilities q_i(v)
target_probs: torch.Tensor, # [B, gamma + 1, V] target probabilities p_i(v) from single forward pass
) -> tuple[list[torch.Tensor], torch.Tensor]:
"""
Executes exact Leviathan et al. (2023) speculative verification across batch B.
Returns (accepted_sequences_per_batch, num_emitted_tokens_per_batch).
"""
B, gamma, V = draft_probs.shape
emitted_batches = []
num_emitted = torch.zeros(B, dtype=torch.long, device=draft_tokens.device)
for b in range(B):
accepted_list = []
all_accepted = True
for i in range(gamma):
x_tilde = draft_tokens[b, i]
q_val = draft_probs[b, i, x_tilde].clamp(min=1e-12)
p_val = target_probs[b, i, x_tilde]
# Accept with probability min(1, p(x_tilde) / q(x_tilde))
u = torch.rand((), device=draft_tokens.device)
if u < torch.minimum(torch.ones_like(p_val), p_val / q_val):
accepted_list.append(x_tilde)
else:
# Reject: sample correction token from residual (p - q)_+ / sum((p - q)_+)
residual = torch.clamp(target_probs[b, i, :] - draft_probs[b, i, :], min=0.0)
residual_probs = residual / residual.sum().clamp(min=1e-12)
x_corr = torch.multinomial(residual_probs, num_samples=1).squeeze(0)
accepted_list.append(x_corr)
all_accepted = False
break
# Bonus token if all gamma draft tokens were accepted!
if all_accepted:
x_bonus = torch.multinomial(target_probs[b, gamma, :], num_samples=1).squeeze(0)
accepted_list.append(x_bonus)
seq_out = torch.stack(accepted_list)
emitted_batches.append(seq_out)
num_emitted[b] = seq_out.numel()
return emitted_batches, num_emitted
3. Multi-Draft Tree Speculation (Medusa, EAGLE & DeepSeek MTP)
When a single chain of $\gamma$ draft tokens hits a low-confidence branch at step 2, steps $3 \dots \gamma$ are wasted. Tree-based speculative decoding (Medusa, Cai et al., 2024; EAGLE-2, Li et al., 2024) drafts a tree of candidate continuations flattened into a single sequence of length $M$ (e.g., $M = 16$ nodes) and verifies the entire tree in one target forward pass using a custom 2D Tree Causal Mask:
Every tree node $i$ uses its tree depth for RoPE positional encoding ($t_0 + \text{depth}(i)$) and attends strictly to the shared prompt plus its own ancestor path from the root.
💻 PyTorch Implementation: Tree-Attention Mask & Depth-Aligned RoPE BuilderPure PyTorch
import torch
def build_tree_attention_mask(parent_indices: list[int], device="cpu") -> tuple[torch.Tensor, torch.Tensor]:
"""
Builds the exact 2D causal tree attention mask [M, M] and RoPE depth offsets [M]
from a parent pointer array where parent_indices[0] = -1 (root).
"""
M = len(parent_indices)
mask = torch.full((M, M), float("-inf"), device=device)
depths = torch.zeros(M, dtype=torch.long, device=device)
for i in range(M):
cur = i
depth = 0
while cur != -1:
mask[i, cur] = 0.0
cur = parent_indices[cur]
if cur != -1:
depth += 1
depths[i] = depth
return mask, depths
4. Failure Modes & Production Inference Diagnostics
softmax(logits / T) underflows in BF16/FP16 so that $\sum (p - q)_+$ evaluates to $0.0$, causing NaN in torch.multinomial.Fix: For $T = 0$ (greedy decoding), bypass division by $q(x)$ entirely: accept $\tilde{x}_i$ if and only if $\tilde{x}_i = \arg\max_v p_i(v)$, and emit $\arg\max_v p_i(v)$ on the first mismatch. For $T \gt 0$, always compute $p, q$, and $(p-q)_+$ in FP32.
Fix: Dynamically throttle $\gamma$ as a function of active batch size $B$: use $\gamma = 5$ for $B \le 8$, $\gamma = 2$ for $B \in [9, 32]$, and disable speculation ($\gamma = 0$) when $B \gt 32$.
- Acceptance Rate $\alpha$
- Typically $\alpha \in [0.75, 0.90]$ for code/math (high determinism) and $\alpha \in [0.60, 0.75]$ for creative text.
- Optimal Draft Depth $\gamma$
- $\gamma \in [4, 6]$ for chain speculation; up to $M = 32$ tree nodes for EAGLE-2 / Medusa tree verification.
- DeepSeek-V3 MTP
- Uses native Multi-Token Prediction (MTP) heads trained jointly with the base model, reaching $\alpha \approx 0.85\text{–}0.90$ for $1.8\times$ end-to-end TPS acceleration.
5. Runnable End-to-End Verification: Proving Zero Distribution Drift
Below is a self-contained statistical verification script that runs $10{,}000$ speculative decoding steps with a deliberately mismatched draft distribution $q(x)$ against target distribution $p(x)$ and proves that the empirical distribution of accepted/resampled tokens converges to $p(x)$ with Total Variation distance $\lt 0.01$, while also verifying PagedKVCacheManager rollback and build_tree_attention_mask.
🧪 Runnable Verification Script: Lossless Distribution Equivalence & Paged RollbackVerified PyTorch
import torch
def speculative_decoding_step(
draft_tokens: torch.Tensor,
draft_probs: torch.Tensor,
target_probs: torch.Tensor,
) -> tuple[list[torch.Tensor], torch.Tensor]:
"""Executes exact Leviathan et al. (2023) speculative verification across batch B."""
B, gamma, V = draft_probs.shape
emitted_batches = []
num_emitted = torch.zeros(B, dtype=torch.long, device=draft_tokens.device)
for b in range(B):
accepted_list = []
all_accepted = True
for i in range(gamma):
x_tilde = draft_tokens[b, i]
q_val = draft_probs[b, i, x_tilde].clamp(min=1e-12)
p_val = target_probs[b, i, x_tilde]
u = torch.rand((), device=draft_tokens.device)
if u < torch.minimum(torch.ones_like(p_val), p_val / q_val):
accepted_list.append(x_tilde)
else:
residual = torch.clamp(target_probs[b, i, :] - draft_probs[b, i, :], min=0.0)
residual_probs = residual / residual.sum().clamp(min=1e-12)
x_corr = torch.multinomial(residual_probs, num_samples=1).squeeze(0)
accepted_list.append(x_corr)
all_accepted = False
break
if all_accepted:
x_bonus = torch.multinomial(target_probs[b, gamma, :], num_samples=1).squeeze(0)
accepted_list.append(x_bonus)
seq_out = torch.stack(accepted_list)
emitted_batches.append(seq_out)
num_emitted[b] = seq_out.numel()
return emitted_batches, num_emitted
class PagedKVCacheManager:
def __init__(self, num_blocks: int, block_size: int, n_kv_heads: int, head_dim: int, device="cpu"):
self.block_size = block_size
self.k_pool = torch.zeros((num_blocks, n_kv_heads, block_size, head_dim), device=device)
self.v_pool = torch.zeros((num_blocks, n_kv_heads, block_size, head_dim), device=device)
self.free_blocks = list(range(num_blocks - 1, -1, -1))
self.block_tables: dict[int, list[int]] = {}
self.seq_lens: dict[int, int] = {}
def allocate_request(self, req_id: int) -> None:
self.block_tables[req_id] = []
self.seq_lens[req_id] = 0
def append_kv(self, req_id: int, k_tokens: torch.Tensor, v_tokens: torch.Tensor) -> None:
num_new = k_tokens.size(1)
cur_len = self.seq_lens[req_id]
for i in range(num_new):
pos = cur_len + i
logical_idx = pos // self.block_size
offset = pos % self.block_size
if logical_idx >= len(self.block_tables[req_id]):
phys_block = self.free_blocks.pop()
self.block_tables[req_id].append(phys_block)
else:
phys_block = self.block_tables[req_id][logical_idx]
self.k_pool[phys_block, :, offset, :] = k_tokens[:, i, :]
self.v_pool[phys_block, :, offset, :] = v_tokens[:, i, :]
self.seq_lens[req_id] = cur_len + num_new
def rollback_to(self, req_id: int, target_len: int) -> None:
assert target_len <= self.seq_lens[req_id]
needed_blocks = (target_len + self.block_size - 1) // self.block_size if target_len > 0 else 0
while len(self.block_tables[req_id]) > needed_blocks:
freed = self.block_tables[req_id].pop()
self.free_blocks.append(freed)
self.seq_lens[req_id] = target_len
def gather_contiguous_kv(self, req_id: int) -> tuple[torch.Tensor, torch.Tensor]:
total_len = self.seq_lens[req_id]
if total_len == 0:
return (torch.empty(self.k_pool.size(1), 0, self.k_pool.size(3), device=self.k_pool.device),
torch.empty(self.v_pool.size(1), 0, self.v_pool.size(3), device=self.v_pool.device))
k_list, v_list = [], []
for idx, phys in enumerate(self.block_tables[req_id]):
valid = min(self.block_size, total_len - idx * self.block_size)
k_list.append(self.k_pool[phys, :, :valid, :])
v_list.append(self.v_pool[phys, :, :valid, :])
return torch.cat(k_list, dim=1), torch.cat(v_list, dim=1)
def build_tree_attention_mask(parent_indices: list[int], device="cpu") -> tuple[torch.Tensor, torch.Tensor]:
M = len(parent_indices)
mask = torch.full((M, M), float("-inf"), device=device)
depths = torch.zeros(M, dtype=torch.long, device=device)
for i in range(M):
cur = i
depth = 0
while cur != -1:
mask[i, cur] = 0.0
cur = parent_indices[cur]
if cur != -1:
depth += 1
depths[i] = depth
return mask, depths
def verify_speculative_inference():
torch.manual_seed(7)
V, gamma, N_trials = 8, 4, 10000
# Target distribution p and deliberately skewed draft distribution q
p_true = torch.tensor([0.32, 0.24, 0.16, 0.10, 0.08, 0.05, 0.03, 0.02])
q_skew = torch.tensor([0.15, 0.30, 0.20, 0.15, 0.05, 0.05, 0.05, 0.05])
# Sample N_trials draft chains of length gamma from q_skew
draft_probs = q_skew.view(1, 1, V).expand(N_trials, gamma, V)
target_probs = p_true.view(1, 1, V).expand(N_trials, gamma + 1, V)
draft_tokens = torch.multinomial(q_skew, num_samples=N_trials * gamma, replacement=True).view(N_trials, gamma)
emitted_seqs, lengths = speculative_decoding_step(draft_tokens, draft_probs, target_probs)
# Check empirical distribution of first emitted token vs exact p_true
first_tokens = torch.stack([seq[0] for seq in emitted_seqs])
counts = torch.bincount(first_tokens, minlength=V).float()
p_empirical = counts / N_trials
tv_distance = 0.5 * (p_empirical - p_true).abs().sum().item()
# Check theoretical expected tokens/step = (1 - alpha^(gamma+1)) / (1 - alpha)
alpha_theo = torch.minimum(p_true, q_skew).sum().item()
expected_len_theo = (1.0 - alpha_theo ** (gamma + 1)) / (1.0 - alpha_theo)
empirical_len = lengths.float().mean().item()
# Check PagedKVCacheManager append & rollback
cache = PagedKVCacheManager(num_blocks=16, block_size=4, n_kv_heads=2, head_dim=8)
cache.allocate_request(req_id=1)
k_dummy = torch.randn(2, 10, 8)
v_dummy = torch.randn(2, 10, 8)
cache.append_kv(1, k_dummy, v_dummy)
cache.rollback_to(1, target_len=6)
k_gat, _ = cache.gather_contiguous_kv(1)
assert torch.allclose(k_gat, k_dummy[:, :6, :]), "Paged KV rollback mismatch!"
# Check Tree Attention mask
tree_mask, depths = build_tree_attention_mask([-1, 0, 0, 1, 1, 2])
assert depths.tolist() == [0, 1, 1, 2, 2, 2], "Tree depth error!"
assert tv_distance < 0.015, f"Distribution drift detected: TV={tv_distance}"
print(f"[PASS] Empirical vs Target TV Distance : {tv_distance:.5f} (Lossless!)")
print(f"[PASS] Mean Tokens/Step (Emp vs Theo) : {empirical_len:.3f} vs {expected_len_theo:.3f} (alpha={alpha_theo:.2f})")
print(f"[PASS] Paged KV Rollback & Tree Mask : Verified (6 retained tokens, tree depths {depths.tolist()})")
if __name__ == "__main__":
verify_speculative_inference()
6. Speculative Decoding Paradigms: Head-to-Head
| Paradigm | Draft Mechanism | Extra VRAM Overhead | Mean Acceptance $\alpha$ | Wall-Clock Speedup |
|---|---|---|---|---|
| Standalone Draft Model (Leviathan et al.) | Separate small LLM (e.g. 1B drafting for 70B) | Draft weights + separate KV cache (~2–4 GB) | $0.70 – 0.82$ | $1.8\times – 2.4\times$ |
| Medusa (Cai et al.) | $K$ independent MLP heads on top of target hidden state $h_t$ | Tiny (< 2% of target weights, single KV cache) | $0.65 – 0.78$ (Tree) | $2.1\times – 2.8\times$ |
| EAGLE-2 (Li et al.) | 1-layer autoregressive feature-level draft head conditioned on $(f_t, e_{t+1})$ | Minimal (~1 transformer layer) | $0.82 – 0.92$ (Dynamic Tree) | $2.8\times – 3.6\times$ |
| DeepSeek-V3 MTP | Sequential Multi-Token Prediction modules trained during pre-training | 1 MTP module (~1 layer) | $0.85 – 0.90$ | $1.8\times$ (even at production batch sizes) |