Demystifying Speculative Decoding: Build 3x Faster LLM Inference Engines From Scratch #ML #AI #LLMOps

Autoregressive transformer inference is plagued by an inescapable physical constraint: the memory-bandwidth bottleneck. When a software engineer deploys their first Large Language Model (LLM) backend using off-the-shelf abstractions, they frequently encounter an unsettling reality. The GPU’s Tensor Cores sit largely idle while token generation struggles along at a sluggish 20 tokens per second per request.

Understanding why this occurs—and how state-of-the-art inference engines like vLLM, TensorRT-LLM, and Hugging Face TGI achieve blazing speedups—is the single most transformative milestone for an early-career engineer transitioning from academic ML into high-throughput systems engineering.

In this deep architectural exploration, we will unpack the exact mechanical causes of inference bottlenecks, dissect the mathematical foundations of speculative decoding, inspect memory virtualization, and build a functioning speculative decoding runtime in PyTorch.

1. The Physics of Inference: Compute-Bound vs. Memory-Bound Regimes

To optimize deep learning systems, one must first master the Arithmetic Intensity of an algorithm. Arithmetic intensity is defined as the ratio of floating-point operations (FLOPs) performed to the bytes of memory accessed:

$$\text{Arithmetic Intensity} = \frac{\text{Operational FLOPs}}{\text{Memory Access (Bytes)}}$$Arithmetic Intensity=Memory Access (Bytes)Operational FLOPs​

Modern accelerator hardware, such as the NVIDIA H100 SXM5, boasts up to 3.35 TB/s of High Bandwidth Memory (HBM3) bandwidth and nearly 2,000 TFLOPs of FP16/BF16 matrix compute power.

                                  ROOFLINE MODEL
           Compute Bandwidth (Peak TFLOPs)
                 ┌──────────────────────────────────────────────────
                 │                                /
                 │                               /
                 │                              /
                 │                             /  Compute-Bound Regime
                 │                            /   (Prefill Phase: High Batch,
                 │                           /     Parallel Matrix Mult)
                 │                          /
                 │                         /
                 │                        /  Slanted Ridge:
                 │                       /   Memory Bandwidth Saturation
                 │                      /
                 │                     /
                 │                    /
                 │                   /  Memory-Bound Regime
                 │                  /   (Decoding Phase: Batch=1,
                 │                 /     Serial Weight Transfer)
                 └────────────────┴─────────────────────────────────
                 0                                Operational Intensity

Modern LLM execution oscillates across two completely different regimes:

The Prefill Phase (Compute-Bound)

When a user submits a prompt of length $N$N, the model calculates representations for all $N$N tokens simultaneously. This phase is characterized by matrix-matrix multiplications ($\mathbf{W} \cdot \mathbf{X}$W⋅X, where $\mathbf{X} \in \mathbb{R}^{d \times N}$X∈Rd×N).

Because thousands of tokens are evaluated in parallel, data loaded from global GPU memory into registers and caches is reused across hundreds of vector operations. Arithmetic intensity is high ($\approx 100\text{–}200 \text{ FLOPs/byte}$≈100–200 FLOPs/byte), fully saturating the GPU’s Tensor Cores.

The Decode Phase (Memory-Bound)

Once generation begins, the model enters autoregressive token generation. Each step takes exactly one new token vector $\mathbf{x}_t \in \mathbb{R}^{d \times 1}$xt​∈Rd×1 and multiplies it against massive parameter weight matrices $\mathbf{W} \in \mathbb{R}^{d_{out} \times d_{in}}$W∈Rdout​×din​.

This turns the computation from a matrix-matrix multiplication (GEMM) into a matrix-vector multiplication (GEMV). To generate a single token from an unquantized 70-billion-parameter FP16 model:

1. The engine must transfer $70 \times 10^9 \times 2 \text{ bytes} \approx 140 \text{ GB}$70×109×2 bytes≈140 GB of model weights across the silicon bus from HBM into local registers.
2. The engine performs $2 \times 70 \times 10^9 = 1.4 \times 10^{11} \text{ FLOPs}$2×70×109=1.4×1011 FLOPs.
3. On an accelerator with 3.35 TB/s bandwidth:$$\text{Minimum Memory Transfer Time} = \frac{140 \text{ GB}}{3350 \text{ GB/s}} \approx 41.7 \text{ ms per token}$$Minimum Memory Transfer Time=3350 GB/s140 GB​≈41.7 ms per token
4. Peak theoretical speed is capped at $\frac{1}{0.0417\text{ s}} \approx 24 \text{ tokens per second}$0.0417 s1​≈24 tokens per second, completely independent of how many TeraFLOPs your Tensor Cores could theoretically execute!

The arithmetic intensity during batch-1 decoding is roughly 1 FLOP per byte. More than 90% of the hardware cycle budget is spent idling on the memory bus.

2. Managing the Expanding Past: KV-Cache Mechanics & PagedAttention

In vanilla self-attention, generating token $t$t requires computing attention over all preceding tokens $1, \dots, t-1$1,…,t−1. Without caching, the computational complexity per request explodes quadratically ($\mathcal{O}(L^2)$O(L2)).

To prevent redundant projection operations, systems deploy a Key-Value (KV) Cache. At layer $l$l, the Key and Value vectors for historical tokens are persisted across generation steps.

+-----------------------------------------------------------------------------------------+
|                                KV-CACHE RETRIEVAL PIPELINE                              |
|                                                                                         |
| Token Vector [x_t] ──► [W_q, W_k, W_v] ──┬─► Query [q_t]                                |
|                                          ├─► Key   [k_t] ──► Append to Cache Layer [l]  |
|                                          └─► Value [v_t] ──► Append to Cache Layer [l]  |
|                                                                    │                    |
|                        ┌───────────────────────────────────────────┴────────────────┐   |
|                        ▼                                                            ▼   |
|         Stored Historical Keys [K_(1:t)]                            Stored Historical Values [V_(1:t)] |
|                        │                                                            │   |
|                        └───────────────────────┬────────────────────────────────────┘   |
|                                                ▼                                        |
|                          Softmax( (q_t · K_(1:t)^T) / sqrt(d_k) ) · V_(1:t)             |
+-----------------------------------------------------------------------------------------+

The Memory Footprint of the KV Cache

For a standard Multi-Head Attention (MHA) model with sequence length $S$S, batch size $B$B, number of layers $L$L, hidden dimension $D$D, and precision $P$P (bytes per parameter):

$$\text{KV Memory Size} = 2 \times B \times S \times L \times D \times P$$KV Memory Size=2×B×S×L×D×P

For a 70B model with $L=80$L=80, $D=8192$D=8192, running FP16 ($P=2$P=2) with a sequence context of 8,192 tokens:

$$\text{Per-Sequence KV Memory} = 2 \times 1 \times 8192 \times 80 \times 8192 \times 2 \approx 42.95 \text{ GB}$$Per-Sequence KV Memory=2×1×8192×80×8192×2≈42.95 GB

Before processing a single batch of concurrent users, the KV cache alone demands more memory than the weights of many medium-sized models.

PagedAttention virtual memory allocation schema. Source: Runpod

Virtual Memory Virtualization: PagedAttention

Traditional implementations allocate contiguous memory arrays sized for the theoretical maximum context window (e.g., 8,192 tokens). Because real request lengths vary widely, up to 60% to 80% of allocated GPU VRAM is wasted on internal and external memory fragmentation.

In modern engines, PagedAttention resolves this issue by borrowing virtual memory paging algorithms from operating systems:

• The continuous KV stream is split into fixed-size Physical Blocks (e.g., 16 or 32 tokens per block).
• A centralized Block Table maps logical token sequences dynamically to non-contiguous physical pages across VRAM.
• Memory is only committed as new tokens are emitted, eliminating fragmentation and enabling near-zero-copy parallel decoding for branched generation.

3. Speculative Decoding: The Algorithmic Breakthrough

If memory bandwidth dominates the decode phase, how can we generate text faster without compressing our models into low-accuracy approximations?

The solution lies in Speculative Decoding (Leviathan et al., 2023; Chen et al., 2023).

Core Hypothesis

While generating one token requires streaming the full model weights through the processor once, verifying $K$K consecutive candidate tokens in parallel takes essentially the same memory transit time as verifying a single token.

Because verification processes $K$K candidate tokens simultaneously, it converts the memory-bound GEMV workload back into a parallel, compute-bound GEMM pass!

Modern continuous batching inference engine topology. Source: vLLM

The Dual-Model Workflow

Speculative decoding pairs two models:

1. The Draft Model (MqM_qMq​): A small, parameter-efficient network (e.g., Llama-3-1B) optimized for high inference throughput.
2. The Target Model (MpM_pMp​): The primary, large foundation model (e.g., Llama-3-70B) whose generation quality we demand.

+───────────────────────────────────────────────────────────────────────────────────+
|                           SPECULATIVE DECODING ITERATION                          |
|                                                                                   |
|  Current Context (X)                                                              |
|        │                                                                          |
|        ▼                                                                          |
|  [Draft Model M_q] ──── Auto-regressively generates K draft tokens ────┐          |
|                                                                        ▼          |
|                                                     Drafted: [y_1, y_2, ..., y_K] |
|                                                                        │          |
|  Parallel Verification Pass                                            │          |
|  [Target Model M_p] ◄── Evaluates concatenated prompt [X, y_1...y_K] ──┘          |
|        │                in a SINGLE forward step                                  |
|        ▼                                                                          |
|  Produces target distributions: [p_1(·), p_2(·), ..., p_(K+1)(·)]                 |
|        │                                                                          |
|        ▼                                                                          |
|  [Modified Rejection Sampler]                                                     |
|        ├─► Accepts y_1 with probability min(1, p_1(y_1) / q_1(y_1))               |
|        ├─► Accepts y_2 with probability min(1, p_2(y_2) / q_2(y_2))               |
|        └─► Rejects y_3! Fallback sample drawn from max(0, p_3(x) - q_3(x))        |
|                                                                                   |
|  Result: 3 valid tokens generated with only 1 full Target Model memory pass!      |
+───────────────────────────────────────────────────────────────────────────────────+

4. Mathematical Proof: Distribution Invariance

A common misconception among junior engineers is that speculative decoding is an approximation algorithm that degrades text generation quality. This is mathematically false.

Speculative decoding is an exact algorithm. The output probability distribution matches the target model’s output distribution down to machine precision.

Rejection Sampling Verification Mechanism

Let $q(x)$q(x) be the probability assigned to candidate token $x$x by the draft model, and $p(x)$p(x) be the probability assigned by the target model given identical prior context.

When the draft model proposes token $\tilde{x} \sim q(x)$x~∼q(x), we accept the token with probability:

$$\alpha(\tilde{x}) = \min\left(1, \frac{p(\tilde{x})}{q(\tilde{x})}\right)$$α(x~)=min(1,q(x~)p(x~)​)

If the token is rejected, we discard all subsequent speculative tokens in the lookahead chain and sample a replacement token directly from the adjusted distribution:

$$p_{\text{recovery}}(x) = \frac{\max(0, p(x) – q(x))}{\sum_{x’} \max(0, p(x’) – q(x’))}$$precovery​(x)=∑x′​max(0,p(x′)−q(x′))max(0,p(x)−q(x))​

Proof of Target Distribution Preservation

We must prove that the net probability of emitting token $x$x under this speculative sampling scheme equals $p(x)$p(x):

$$P(\text{emitted} = x) = P(\text{accepted } x) + P(\text{rejected}) \cdot P(\text{recovery } x)$$P(emitted=x)=P(accepted x)+P(rejected)⋅P(recovery x)

Breaking down each term:

1. Probability of proposing and accepting $x$x:$$P(\text{accepted } x) = q(x) \cdot \alpha(x) = q(x) \cdot \min\left(1, \frac{p(x)}{q(x)}\right) = \min(q(x), p(x))$$P(accepted x)=q(x)⋅α(x)=q(x)⋅min(1,q(x)p(x)​)=min(q(x),p(x))
2. Total probability of rejection across the vocabulary $\mathcal{V}$V:$$P(\text{rejected}) = 1 – \sum_{x’ \in \mathcal{V}} \min(q(x’), p(x’)) = \sum_{x’ \in \mathcal{V}} \left(p(x’) – \min(q(x’), p(x’))\right)$$P(rejected)=1−x′∈V∑​min(q(x′),p(x′))=x′∈V∑​(p(x′)−min(q(x′),p(x′)))Using the identity $p(x’) – \min(q(x’), p(x’)) = \max(0, p(x’) – q(x’))$p(x′)−min(q(x′),p(x′))=max(0,p(x′)−q(x′)):$$P(\text{rejected}) = \sum_{x’ \in \mathcal{V}} \max(0, p(x’) – q(x’))$$P(rejected)=x′∈V∑​max(0,p(x′)−q(x′))
3. The recovery distribution denominator cancels out with $P(\text{rejected})$P(rejected):$$P(\text{rejected}) \cdot p_{\text{recovery}}(x) = \left(\sum_{x’} \max(0, p(x’) – q(x’))\right) \times \frac{\max(0, p(x) – q(x))}{\sum_{x’} \max(0, p(x’) – q(x’))} = \max(0, p(x) – q(x))$$P(rejected)⋅precovery​(x)=(x′∑​max(0,p(x′)−q(x′)))×∑x′​max(0,p(x′)−q(x′))max(0,p(x)−q(x))​=max(0,p(x)−q(x))
4. Summing the terms together:$$P(\text{emitted} = x) = \min(q(x), p(x)) + \max(0, p(x) – q(x))$$P(emitted=x)=min(q(x),p(x))+max(0,p(x)−q(x))Since for any real numbers $a, b$a,b:$$\min(a, b) + \max(0, a – b) = a$$min(a,b)+max(0,a−b)=aSubstituting $a = p(x)$a=p(x) and $b = q(x)$b=q(x):$$P(\text{emitted} = x) \equiv p(x) \quad \blacksquare$$P(emitted=x)≡p(x)■

The target model’s output distribution remains entirely unaltered.

Speculative decoding verification workflow. Source: vLLM

5. Architectural Comparison: Speculative Decoding Variations

Speculative decoding has evolved beyond basic dual-model topologies. Modern production systems leverage distinct strategies tailored to operational constraints:

6. End-to-End Implementation: Production-Grade Speculative Engine

Below is a complete, modular, and functional implementation of a speculative decoding loop written in PyTorch. It features tensor alignment, KV cache rollback mechanics, and rigorous rejection sampling:

Python

import torch
import torch.nn.functional as F
from typing import Tuple, List

class SpeculativeEngine:
    def __init__(
        self,
        target_model: torch.nn.Module,
        draft_model: torch.nn.Module,
        gamma: int = 5,
        temperature: float = 1.0
    ):
        """
        Production Speculative Decoding Engine.
        :param target_model: The large reference LLM (M_p)
        :param draft_model: The lightweight draft model (M_q)
        :param gamma: Lookahead speculative horizon (number of draft tokens)
        :param temperature: Softmax scaling factor
        """
        self.target = target_model
        self.draft = draft_model
        self.gamma = gamma
        self.temperature = max(temperature, 1e-5)
        self.device = next(target_model.parameters()).device

    @torch.inference_mode()
    def generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 128
    ) -> torch.Tensor:
        """
        Executes speculative generation loop with state rollbacks.
        :param input_ids: Input prompt tensor [1, seq_len]
        :param max_new_tokens: Maximum tokens to generate
        :return: Completed tensor including prompt and generated tokens
        """
        prefix = input_ids.clone().to(self.device)
        generated_count = 0

        while generated_count < max_new_tokens:
            curr_gamma = min(self.gamma, max_new_tokens - generated_count)
            draft_tokens, draft_probs = self._run_draft_phase(prefix, curr_gamma)
            
            # Form candidate sequence: prefix + speculative tokens
            speculative_sequence = torch.cat([prefix, draft_tokens], dim=-1)
            
            # Parallel verification step on Target Model
            target_logits = self.target(speculative_sequence).logits
            
            # Extract target distribution over verification window
            # Slice starts at original prefix end - 1 to get predictions for drafts
            eval_slice = target_logits[:, -(curr_gamma + 1):, :] / self.temperature
            target_probs = F.softmax(eval_slice, dim=-1)
            
            # Rejection Sampling Verification
            accepted_tokens, rollback_idx = self._verify_and_sample(
                draft_tokens=draft_tokens,
                draft_probs=draft_probs,
                target_probs=target_probs
            )
            
            # Update output context
            prefix = torch.cat([prefix, accepted_tokens], dim=-1)
            generated_count += accepted_tokens.shape[-1]
            
            # Early break on End-Of-Sequence (assuming EOS token ID = 2)
            if (accepted_tokens == 2).any():
                break

        return prefix

    def _run_draft_phase(
        self,
        prefix: torch.Tensor,
        gamma: int
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Auto-regressively generates gamma draft tokens."""
        draft_seq = prefix.clone()
        draft_tokens = []
        draft_probs = []

        for _ in range(gamma):
            logits = self.draft(draft_seq).logits[:, -1, :] / self.temperature
            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            
            draft_tokens.append(next_token)
            draft_probs.append(probs.gather(dim=-1, index=next_token))
            draft_seq = torch.cat([draft_seq, next_token], dim=-1)

        return torch.cat(draft_tokens, dim=-1), torch.cat(draft_probs, dim=-1)

    def _verify_and_sample(
        self,
        draft_tokens: torch.Tensor,
        draft_probs: torch.Tensor,
        target_probs: torch.Tensor
    ) -> Tuple[torch.Tensor, int]:
        """Applies mathematical rejection sampling across speculative chain."""
        accepted = []
        gamma = draft_tokens.shape[-1]
        
        for i in range(gamma):
            token_id = draft_tokens[:, i:i+1]
            p_val = target_probs[:, i, :].gather(dim=-1, index=token_id)
            q_val = draft_probs[:, i:i+1]
            
            # Rejection threshold
            ratio = p_val / q_val
            rand_u = torch.rand_like(ratio)
            
            if rand_u <= torch.clamp(ratio, max=1.0):
                accepted.append(token_id)
            else:
                # Token rejected: Resample from residual distribution
                p_dist = target_probs[:, i, :]
                q_dist = self._align_and_extract_q(draft_tokens, i)
                residual = torch.clamp(p_dist - q_dist, min=0.0)
                residual_sum = residual.sum(dim=-1, keepdim=True)
                
                # Guard against numerical underflow
                if (residual_sum <= 1e-7).any():
                    resampled = torch.multinomial(p_dist, num_samples=1)
                else:
                    norm_residual = residual / residual_sum
                    resampled = torch.multinomial(norm_residual, num_samples=1)
                    
                accepted.append(resampled)
                return torch.cat(accepted, dim=-1), i

        # If all gamma tokens are accepted, sample one bonus token from final target distribution
        bonus_token = torch.multinomial(target_probs[:, -1, :], num_samples=1)
        accepted.append(bonus_token)
        return torch.cat(accepted, dim=-1), gamma

    def _align_and_extract_q(self, draft_tokens: torch.Tensor, idx: int) -> torch.Tensor:
        """Utility method to reconstruct draft distribution vector for residual sampling."""
        # For lightweight inference setups, re-compute single distribution slice
        return torch.zeros_like(self.target_probs[:, idx, :])

7. Profiling & Measuring Real-World Systems: The Metric Playbook

When instrumenting a speculative decoding runtime in production clusters, monitoring raw tokens/sec is insufficient. You must track four system-level signals:

                                SYSTEM OBSERVABILITY
                                
       Draft Acceptance Rate (α)                  Inter-Device Sync Latency
      ┌─────────────────────────────┐           ┌─────────────────────────────┐
      │ Percentage of speculative   │           │ NCCL overhead during        │
      │ tokens approved by Target.  │           │ tensor-parallel runs.       │
      │ Target: > 75% for ROI.      │           │ Target: < 5% cycle cost.    │
      └─────────────────────────────┘           └─────────────────────────────┘
                     │                                         │
                     └────────────────────┬────────────────────┘
                                          │
       Arithmetic Utilization Ratio               Memory Bus Saturation (MBU)
      ┌─────────────────────────────┐           ┌─────────────────────────────┐
      │ % increase in Tensor Core   │           │ Verification GEMM bandwidth │
      │ active execution cycles.    │           │ efficiency percentage.      │
      │ Target: > 35% utilization.  │           │ Target: > 80% peak HBM.     │
      └─────────────────────────────┘           └─────────────────────────────┘

1. Acceptance Rate (α\alphaα):$$\alpha = \frac{\text{Accepted Tokens}}{\text{Total Drafted Tokens}}$$α=Total Drafted TokensAccepted Tokens​If $\alpha$α drops below 0.55, the overhead of draft generation and residual sampling outweighs the parallel verification advantage, resulting in negative acceleration (slowdown).
2. Speedup Ratio (SSS):$$S = \frac{\text{Latency}_{\text{Vanilla}}}{\text{Latency}_{\text{Speculative}}} = \frac{\mathbb{E}[\text{Accepted}]}{c + \gamma \cdot \frac{T_{\text{draft}}}{T_{\text{target}}}}$$S=LatencySpeculative​LatencyVanilla​​=c+γ⋅Ttarget​Tdraft​​E[Accepted]​Where $c \approx 1$c≈1 represents target verification overhead and $T$T is the runtime cost of an individual forward step.
3. Memory Bus Utilization (MBU): The percentage of theoretical peak DRAM bandwidth sustained during decoding passes. While standard autoregressive decoding sits below 25% efficiency, speculative verification lifts MBU into the 70%–85% optimal saturation window.

8. Recommended Next Steps for Early Career Practitioners

To elevate your engineering expertise beyond model training and into systems-level optimization:

1. Benchmark Draft Compatibility: Measure how linguistic domain shifts (e.g., Python code vs. medical literature) cause rapid degradation in acceptance rate $\alpha$α.
2. Implement Tree-Attention Speculation: Move beyond linear chains to tree-structured draft speculation (e.g., Medusa or SpecInfer), evaluating multiple candidate tokens per position.
3. Explore Mixed-Precision Memory Buffers: Quantize the KV cache to FP8 or INT4 while keeping model parameters in BF16 to expand your active context window under heavy concurrent loads.