AI/ML Inference Speculative Decoding Systems ML PyTorch

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}}$):

Linear Layer Arithmetic Intensity & Hardware Ridge Point
$$\mathcal{I}(N) = \frac{2 N d_{\text{in}} d_{\text{out}}}{2 d_{\text{in}} d_{\text{out}} + 2 N (d_{\text{in}} + d_{\text{out}})} \approx N \;\text{FLOPs/byte} \quad (\text{for } N \ll d_{\text{model}})$$

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!

Arithmetic Intensity I ≈ Active Tokens N (Log Scale, FLOPs / Byte) Attainable TFLOP/s (Log) I_ridge ≈ 295 (H100) HBM Bandwidth Ceiling (3.35 TB/s) Tensor Core Peak (989 TFLOP/s BF16) Standard Decode (N=1) ~3.3 TFLOP/s (0.3% Util!) Speculative Verify (N=γ+1) Same HBM weight load, γ× FLOPs! Prefill / Continuous Batch N > 295: Compute-Bound
Figure 1: Roofline model of LLM inference. Because $B=1$ autoregressive decoding sits at $\mathcal{I} \approx 1\text{ FLOP/byte}$, verifying $\gamma+1$ draft tokens in parallel climbs the bandwidth slope for "free" wall-clock time.

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!):

Logical Position $t$ to Physical KV-Block Offset Mapping
$$b_{\text{logical}}(t) = \left\lfloor \frac{t}{B_{\text{block}}} \right\rfloor, \qquad \text{offset}(t) = t \bmod B_{\text{block}}, \qquad b_{\text{phys}}(r, t) = \mathcal{T}\!\left[r,\; b_{\text{logical}}(t)\right]$$
💻 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)$?

Speculative Acceptance Rule & Residual Distribution
$$\text{Accept } \tilde{x} \sim q(x) \text{ with probability } \beta(\tilde{x}) = \min\!\left(1,\; \frac{p(\tilde{x})}{q(\tilde{x})}\right)$$ $$p_{\text{res}}(x) = \frac{\max\!\bigl(0,\; p(x) - q(x)\bigr)}{\sum_{v \in V} \max\!\bigl(0,\; p(v) - q(v)\bigr)} = \frac{\bigl(p(x) - q(x)\bigr)_+}{1 - \alpha}$$

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)$:

  1. 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)$$
  2. 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))_+$.
  3. 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$$
The Bonus Token Guarantee: If all $\gamma$ drafted tokens $\tilde{x}_1, \dots, \tilde{x}_\gamma$ are accepted, the target model's parallel forward pass on $[\;x_{\le t}, \tilde{x}_1, \dots, \tilde{x}_\gamma\;]$ has already computed the next-token logits $p(\cdot \mid x_{\le t}, \tilde{x}_{1:\gamma})$ at position $\gamma+1$! Thus every verification step yields at least $1$ valid token and up to $\gamma + 1$ valid tokens.
Expected Tokens Per Target Forward Pass & Wall-Clock Speedup
$$\mathbb{E}[\text{tokens/step}] = \frac{1 - \alpha^{\gamma + 1}}{1 - \alpha}, \qquad \text{Speedup}(\alpha, \gamma, c) = \frac{1 - \alpha^{\gamma + 1}}{(1 - \alpha)(1 + \gamma c)}$$

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}$.

Step 1: Fast Draft Autoregressive Rollout (γ = 4 steps, q_draft) Prefix: "The" x̃₁ = "gradient" x̃₂ = "vanishes" x̃₃ = "quickly" x̃₄ = "during" Step 2: Single Parallel Target Forward Pass p_target([Prefix, x̃₁, x̃₂, x̃₃, x̃₄]) ✓ Accept x̃₁ u₁ < p(x̃₁)/q(x̃₁) ✓ Accept x̃₂ u₂ < p(x̃₂)/q(x̃₂) ✗ Reject x̃₃ Resample (p₃−q₃)_+ ✂ Discard x̃₄ Rollback KV Cache Step 3: Committed Lossless Output (3 tokens generated in 1 target forward pass!) "The" → "gradient" (✓) → "vanishes" (✓) → "exponentially" (Resampled from p_res at pos 3)
Figure 2: Speculative Decoding verification cycle. Positions 1 and 2 are accepted; position 3 is rejected and immediately replaced by an exact sample from $(p_3 - q_3)_+$ without requiring another target model pass.
💻 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:

Tree-Attention Ancestor Mask Condition
$$M_{\text{tree}}[i, j] = \begin{cases} 0 & \text{if node } j \in \text{Ancestors}(i) \cup \{i\} \\ -\infty & \text{otherwise} \end{cases}, \qquad \text{PosID}[i] = t_0 + \text{depth}(i)$$

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

1. Numeric Underflow in $(p - q)_+$ at Low Temperature ($T \to 0$): When temperature $T \lt 0.05$, 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.
2. Speculative Negative Speedup at High Batch Sizes ($B \ge 64$): When continuous batching already saturates $N = B \ge N_{\text{ridge}} \approx 295$, multiplying the active tokens per forward pass by $(\gamma + 1)$ pushes the target model deep into the compute-bound regime, increasing step latency linearly with $\gamma$.
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)