Skip to main content
UltraInstinct
Back to latest articles
Artificial Intelligence14 min read

How Speculative Decoding in LLM Serving Actually Works

Explore how speculative decoding in LLM serving accelerates token generation via draft models, tree attention verification, Triton kernels, and math bounds.

Featured visual representing How Speculative Decoding in LLM Serving Actually Works
Advertisement
Google AdSense In-Article Contextual Slot

Introduction

Modern deep learning systems deployed for autoregressive text generation are fundamentally bounded by high-bandwidth memory (HBM) latency rather than raw matrix-compute saturation. In standard token-by-token generation, every forward pass reads tens or hundreds of billions of parameters from DRAM across memory channels to process an effective token length of one. Implementing Speculative Decoding in LLM Serving has emerged as the premier architectural paradigm for breaking this memory-wall bottleneck. By utilizing a smaller, computationally inexpensive draft mechanism to propose candidate token sequences and subsequently executing parallel verification across the target model in a single forward pass, inference engines convert memory-bound vector-matrix operations into compute-dense matrix-matrix operations without degrading the mathematical distribution of the output text.

Standard Autoregressive Decoding (Memory-Bound):
Step 1: [Read Weights O(W)] -> Generate Token 1
Step 2: [Read Weights O(W)] -> Generate Token 2
Step 3: [Read Weights O(W)] -> Generate Token 3
Total DRAM Traffic: 3 * W bytes

Speculative Decoding (Compute-Bound Verification):
Draft:  [Fast Draft Model]  -> Propose [t1, t2, t3] (Tiny DRAM Traffic)
Target: [Read Weights O(W)] -> Verify  [t1, t2, t3] Simultaneously
Total DRAM Traffic: ~1 * W bytes (Yields up to 3 tokens per target memory read)

At its core, the speedup of speculative decoding leverages arithmetic intensity. In an autoregressive step with batch size $B=1$, the operational intensity is approximately 1 FLOP per byte of weight transferred, leaving tensor processing units and streaming multiprocessors severely underutilized. Speculative execution bundles $K$ candidate tokens into a single evaluation pass of the target model. This raises operational intensity by a factor proportional to $K$, allowing the hardware to perform meaningful parallel work while amortizing weight transfer overhead.

However, moving from theory to an industrial deployment introduces concrete systems challenges: non-deterministic output risks, KV tensor synchronization overheads, complex tree-mask generation, and latency penalties when acceptance rates degrade. Understanding how these components synchronize within modern runtimes is essential for maximizing serving efficiency.

Mathematical Foundations of Speculative Decoding in LLM Serving

The principal requirement of lossless speculative execution is that the target token distribution must remain strictly preserved. Let the large target model define a conditional probability distribution $P(x_t \mid x_{ \alpha_t$):** The token $\tilde{x}t$ is discarded along with all subsequent tokens $\tilde{x}{t+1}, \dots, \tilde{x}_K$. The engine terminates speculative evaluation for the step and samples a replacement token $x_t$ from the adjusted residual distribution:

$$P'(x) = \frac{\max\left(0, P(x) - Q(x)\right)}{\sum_{y \in \mathcal{V}} \max\left(0, P(y) - Q(y)\right)}$$

Where $\mathcal{V}$ denotes the full model vocabulary.

Distribution Alignment:

Case 1: P(x) >= Q(x)
       +-----------------------+ P(x)
       | Residual Area         |
+------+-----------------------+ Q(x)
| Always Accepted (alpha = 1)   |
+------------------------------+

Case 2: P(x) < Q(x)
+----------------+ Q(x)
| Truncated Area | (Rejected with probability 1 - P/Q)
+----------------+ P(x)
| Accepted Area  |
+----------------+

Exact Proof of Distributional Invariance

To prove that the marginal probability of emitting token $x$ under speculative sampling equals $P(x)$, we sum the joint probabilities across both acceptance and rejection pathways:

$$\mathbb{P}(X = x) = \mathbb{P}(\text{Accepted}) \cdot \mathbb{P}(X = x \mid \text{Accepted}) + \mathbb{P}(\text{Rejected}) \cdot \mathbb{P}(X = x \mid \text{Rejected})$$

Substituting the sampling definitions:

$$\mathbb{P}(X = x) = Q(x) \min\left(1, \frac{P(x)}{Q(x)}\right) + \left(1 - \sum_{y \in \mathcal{V}} Q(y) \min\left(1, \frac{P(y)}{Q(y)}\right)\right) \cdot P'(x)$$

Notice that:

$$Q(x) \min\left(1, \frac{P(x)}{Q(x)}\right) = \min(Q(x), P(x))$$

The probability of rejection across the vocabulary is:

$$1 - \sum_{y \in \mathcal{V}} \min(Q(y), P(y)) = \sum_{y \in \mathcal{V}} \left(Q(y) - \min(Q(y), P(y))\right) = \sum_{y \in \mathcal{V}} \max(0, Q(y) - P(y))$$

Because $\sum_y P(y) = \sum_y Q(y) = 1$, the total residual probability mass matches the total excess probability mass:

$$\sum_{y \in \mathcal{V}} \max(0, Q(y) - P(y)) = \sum_{y \in \mathcal{V}} \max(0, P(y) - Q(y))$$

Substituting these identities back into the residual formulation for $P'(x)$:

$$\mathbb{P}(X = x) = \min(Q(x), P(x)) + \max(0, P(x) - Q(x)) = P(x)$$

The recovered distribution is identically $P(x)$, validating that speculative rejection sampling provides zero degradation in generation quality.

The Speculation Data Path: Drafting, Tree Attention, and Verification

The execution path of a speculative serving system requires coordinating execution between the proposal generator and the target verifier. State modern runtimes leverage either an independent compact transformer, speculative heads attached to target layers (e.g., Medusa), or multi-token prediction layers (e.g., EAGLE).

+-------------------------------------------------------------------------+
| Speculative Execution Loop                                             |
|                                                                         |
|   +-----------------------------------------------------------------+   |
|   | 1. Draft Phase: Autoregressive expansion of K candidates        |   |
|   |    Draft Model produces candidate tree: T1 -> (T2a, T2b)        |   |
|   +-----------------------------------------------------------------+   |
|                                    |                                    |
|                                    v                                    |
|   +-----------------------------------------------------------------+   |
|   | 2. Speculation Topology Construction                            |   |
|   |    - Flatten candidate tree into packed sequence buffer         |   |
|   |    - Generate 2D Tree Attention Mask reflecting branch graph    |   |
|   +-----------------------------------------------------------------+   |
|                                    |                                    |
|                                    v                                    |
|   +-----------------------------------------------------------------+   |
|   | 3. Target Model Parallel Evaluation                             |   |
|   |    - Single batched prefill pass across proposed tokens         |   |
|   |    - Extract logits across all tree nodes simultaneously        |   |
|   +-----------------------------------------------------------------+   |
|                                    |                                    |
|                                    v                                    |
|   +-----------------------------------------------------------------+   |
|   | 4. Verification & State Rollback                                |   |
|   |    - Run vectorized rejection sampling on device                |   |
|   |    - Accept longest valid prefix trajectory                     |   |
|   |    - Truncate / Rollback invalid KV cache allocations           |   |
|   +-----------------------------------------------------------------+   |
|                                    |                                    |
|                                    +---> Loop repeats with accepted state
+-------------------------------------------------------------------------+

While distributed serving platforms leverage Disaggregated Inference Architecture: Paged KV-Cache Tiering, Chunked Prefill, and Asynchronous Overlap to decouple prompt processing from autoregression, speculative decoding addresses the fundamental memory-bandwidth bottleneck within the decode nodes themselves.

Tree-Based Speculation and Custom Causal Masks

Linear candidate drafting ($K$ sequential tokens) restricts the acceptance probability to a single path: if token 1 is rejected, the remaining $K-1$ tokens are completely discarded. Modern architectures replace linear drafting with tree-based speculative structures. By evaluating multiple diverging candidate branches simultaneously, the system significantly improves the likelihood that at least one branch achieves deep acceptance.

Verifying a candidate tree in a single target forward pass requires a modified causal attention mask. Unlike standard lower-triangular causal masks, a tree attention mask allows an arbitrary node $i$ to attend only to its strict ancestors in the tree.

Candidate Tree Graph:
        Root (T0)
        /      \
     T1          T2
    /  \          |
  T3    T4       T5

Adjacency/Mask Tensor (Row attends to Column):
      T0  T1  T2  T3  T4  T5
  T0   1   0   0   0   0   0
  T1   1   1   0   0   0   0
  T2   1   0   1   0   0   0
  T3   1   1   0   1   0   0  <-- Attends to Root and T1
  T4   1   1   0   0   1   0  <-- Attends to Root and T1
  T5   1   0   1   0   0   1  <-- Attends to Root and T2

In memory, nodes are flattened into a contiguous 1D token vector. The self-attention matrix applies the non-triangular binary mask $M \in {0, -\infty}^{N \times N}$ inside the scaled dot-product computation:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V$$

This guarantees that node $T_4$ does not leak attention activations to node $T_3$ or $T_2$, allowing independent branch paths to share a single execution context without mutual interference.

At scale, production engines integrate speculation loops with Disaggregated LLM Serving: Architecting Asynchronous Prefill-Decode Engines with Zero-Copy RDMA KV-Cache Migration to route speculative verification across clusters without encountering network synchronization stalls.

Custom Verification Kernel Implementation in Triton

To avoid host-device synchronization during the acceptance decision, the rejection sampling algorithm must be executed natively on the accelerator. Bringing logits back to the host CPU via PCIe stalls GPU execution pipelines.

The following Triton kernel performs parallel speculative rejection sampling across a batch of linear candidate chains. It computes the target probabilities, evaluates the acceptance criterion against draft probabilities, executes prefix scans to find the first rejection point, and samples the recovery token from the residual distribution using an in-place curand-style linear congruential generator (LCG).

import torch
import triton
import triton.language as tl

@triton.jit
def speculative_verification_kernel(
    Target_Logits_Ptr,   # [Batch, K, Vocab]
    Draft_Probs_Ptr,    # [Batch, K]
    Draft_Tokens_Ptr,   # [Batch, K]
    Rand_Uniform_Ptr,   # [Batch, K]
    Accepted_Tokens_Ptr,# [Batch, K + 1]
    Accepted_Count_Ptr, # [Batch]
    stride_tl_b, stride_tl_k, stride_tl_v,
    stride_dp_b, stride_dp_k,
    stride_dt_b, stride_dt_k,
    stride_ru_b, stride_ru_k,
    stride_at_b, stride_at_k,
    K: tl.constexpr,
    VOCAB_SIZE: tl.constexpr,
    BLOCK_VOCAB: tl.constexpr
):
    batch_id = tl.program_id(0)

    # Local storage for acceptance mask
    # A thread processes the verification of the candidate chain for one batch item
    accepted_prefix = True
    accepted_tokens_count = 0

    for k in range(K):
        draft_token = tl.load(Draft_Tokens_Ptr + batch_id * stride_dt_b + k * stride_dt_k)
        draft_prob = tl.load(Draft_Probs_Ptr + batch_id * stride_dp_b + k * stride_dp_k)
        rand_val = tl.load(Rand_Uniform_Ptr + batch_id * stride_ru_b + k * stride_ru_k)

        # 1. Compute target softmax over vocabulary for position k
        # Standard streaming max & sum for numerical stability
        max_logit = -float('inf')
        denom = 0.0

        for v_offset in range(0, VOCAB_SIZE, BLOCK_VOCAB):
            cols = v_offset + tl.arange(0, BLOCK_VOCAB)
            mask = cols < VOCAB_SIZE
            logits = tl.load(
                Target_Logits_Ptr + batch_id * stride_tl_b + k * stride_tl_k + cols * stride_tl_v,
                mask=mask,
                other=-float('inf')
            )
            block_max = tl.max(logits, axis=0)
            new_max = tl.maximum(max_logit, block_max)
            denom = denom * tl.exp(max_logit - new_max) + tl.sum(tl.exp(logits - new_max), axis=0)
            max_logit = new_max

        # 2. Extract specific target probability for the draft token
        target_token_logit = tl.load(
            Target_Logits_Ptr + batch_id * stride_tl_b + k * stride_tl_k + draft_token * stride_tl_v
        )
        target_prob = tl.exp(target_token_logit - max_logit) / denom

        # 3. Speculative Rejection Criterion
        alpha = tl.minimum(1.0, target_prob / draft_prob)
        is_accepted = (rand_val <= alpha) and accepted_prefix

        if is_accepted:
            tl.store(
                Accepted_Tokens_Ptr + batch_id * stride_at_b + accepted_tokens_count * stride_at_k,
                draft_token
            )
            accepted_tokens_count += 1
        else:
            if accepted_prefix:
                # First rejected token: sample from residual distribution
                # P'(x) = max(0, P(x) - Q(x)) normalized
                residual_sum = 0.0
                # Pass 1: Compute normalization constant
                for v_offset in range(0, VOCAB_SIZE, BLOCK_VOCAB):
                    cols = v_offset + tl.arange(0, BLOCK_VOCAB)
                    mask = cols < VOCAB_SIZE
                    logits = tl.load(
                        Target_Logits_Ptr + batch_id * stride_tl_b + k * stride_tl_k + cols * stride_tl_v,
                        mask=mask,
                        other=-float('inf')
                    )
                    p_vals = tl.exp(logits - max_logit) / denom
                    # Approximate Q(x) ~ 0 for non-draft tokens, draft_prob for draft_token
                    is_draft = (cols == draft_token)
                    q_vals = tl.where(is_draft, draft_prob, 0.0)
                    residual = tl.maximum(0.0, p_vals - q_vals)
                    residual_sum += tl.sum(residual, axis=0)

                # Pass 2: Sample using cumulative threshold
                rand_sample = rand_val * residual_sum
                acc_residual = 0.0
                selected_token = draft_token

                for v_offset in range(0, VOCAB_SIZE, BLOCK_VOCAB):
                    cols = v_offset + tl.arange(0, BLOCK_VOCAB)
                    mask = cols < VOCAB_SIZE
                    logits = tl.load(
                        Target_Logits_Ptr + batch_id * stride_tl_b + k * stride_tl_k + cols * stride_tl_v,
                        mask=mask,
                        other=-float('inf')
                    )
                    p_vals = tl.exp(logits - max_logit) / denom
                    is_draft = (cols == draft_token)
                    q_vals = tl.where(is_draft, draft_prob, 0.0)
                    residual = tl.maximum(0.0, p_vals - q_vals)
                    
                    # Accumulate within block
                    cumsum = tl.cumsum(residual, axis=0)
                    hit = (acc_residual + cumsum) >= rand_sample
                    if tl.sum(hit.to(tl.int32)) > 0:
                        first_idx = tl.argmax(hit.to(tl.int32), axis=0)
                        selected_token = v_offset + first_idx
                        break
                    acc_residual += tl.sum(residual, axis=0)

                tl.store(
                    Accepted_Tokens_Ptr + batch_id * stride_at_b + accepted_tokens_count * stride_at_k,
                    selected_token
                )
                accepted_tokens_count += 1
                accepted_prefix = False

    tl.store(Accepted_Count_Ptr + batch_id, accepted_tokens_count)

Engine orchestrators often pair speculative execution pipelines with High-Throughput Deterministic Actor Runtimes: Cache-Conscious Architectures, Lock-Free Rings, and Kernel-Bypass I/O to manage asynchronous token submission and KV state adjustments across multi-stream environments without incurring synchronization stalls.

def launch_speculative_verification(
    target_logits: torch.Tensor,
    draft_probs: torch.Tensor,
    draft_tokens: torch.Tensor,
    rand_uniform: torch.Tensor
):
    batch_size, k_len, vocab_size = target_logits.shape
    accepted_tokens = torch.empty((batch_size, k_len + 1), dtype=torch.long, device=target_logits.device)
    accepted_count = torch.empty((batch_size,), dtype=torch.int32, device=target_logits.device)

    grid = (batch_size,)
    speculative_verification_kernel[grid](
        target_logits, draft_probs, draft_tokens, rand_uniform,
        accepted_tokens, accepted_count,
        target_logits.stride(0), target_logits.stride(1), target_logits.stride(2),
        draft_probs.stride(0), draft_probs.stride(1),
        draft_tokens.stride(0), draft_tokens.stride(1),
        rand_uniform.stride(0), rand_uniform.stride(1),
        accepted_tokens.stride(0), accepted_tokens.stride(1),
        K=k_len,
        VOCAB_SIZE=vocab_size,
        BLOCK_VOCAB=1024,
        num_warps=8
    )
    return accepted_tokens, accepted_count

System Dynamics: Memory Bandwidth, Acceptance Rates, and Batch Overhead

Speculative decoding does not provide unconditional speedups across all serving regimes. Because the system trades additional arithmetic operations for reduced DRAM transfers, its utility is strictly governed by the arithmetic intensity threshold of the target hardware and the empirical acceptance rate $\alpha$.

Latency vs. Concurrency Trade-off

Latency (ms)
  ^
  |        / (Speculative Decoding Overhead Dominated)
  |       / 
  |      /       / (Standard Autoregressive Serving)
  |     /       /
  |    /       /
  |   /       /
  |  /-------/  <--- Break-Even Crossover Point (Compute-Bound Transition)
  | /       /
  |/       /
  +----------------------------------------------------> Batch Size
  Low Batch (Memory-Bound)          High Batch (Compute-Bound)
  [Speculation Accelerates]         [Speculation Degrades Throughput]

Analytical Speedup Model

Let $T_{\text{target}}$ denote the latency of a single forward pass of the target model with a batch size of 1. Because single-token decoding is memory-bandwidth bound:

$$T_{\text{target}} \approx \frac{\text{Weights}{\text{bytes}}}{\text{Bandwidth}{\text{HBM}}}$$

When verifying $K$ speculative tokens in a batch, the execution time shifts into a matrix-matrix regime ($[B, K] \times [K, D]$), modeled as:

$$T_{\text{verify}}(K) \approx \max\left(\frac{\text{Weights}{\text{bytes}}}{\text{Bandwidth}{\text{HBM}}}, \frac{2 \cdot K \cdot P_{\text{params}}}{\text{Peak FLOPS}}\right) + T_{\text{tree_overhead}}$$

Let $T_{\text{draft}}$ represent the cost of generating one draft token. Generating $K$ speculative tokens autoregressively requires $K \cdot T_{\text{draft}}$.

If the average acceptance rate per token is $\alpha$, the expected number of accepted tokens $\mathbb{E}[L]$ per verification cycle follows a truncated geometric sequence:

$$\mathbb{E}[L] = \sum_{j=1}^{K} \alpha^j = \frac{\alpha (1 - \alpha^K)}{1 - \alpha}$$

Including the final sample emitted upon rejection or completion, the average total tokens emitted per step is $\mathbb{E}[L] + 1$. The theoretical speedup $S$ is expressed by the ratio of standard autoregressive time to the amortized speculative cycle time:

$$S = \frac{(\mathbb{E}[L] + 1) \cdot T_{\text{target}}}{K \cdot T_{\text{draft}} + T_{\text{verify}}(K)}$$

Empirical Speedup Parametric Sweep (Target: 70B, Draft: 8B, HBM3e):

Acceptance Rate (alpha) | Speculation Depth (K) | Speedup Factor (S)
---------------------------------------------------------------------
0.90                    | 5                      | 2.85x
0.80                    | 5                      | 2.21x
0.70                    | 4                      | 1.68x
0.50                    | 3                      | 1.15x
0.35                    | 3                      | 0.88x (Slowdown!)

The High-Batch Contention Boundary

The common failure mode of speculative decoding in enterprise environments is deploying it under high concurrency. As batch size $B$ scales ($B \ge 64$ on modern accelerators like Nvidia Hopper or Blackwell architectures):

  1. Saturation of Memory Bandwidth: Standard autoregressive decoding naturally transitions from being memory-bound to compute-bound because the collective token activations sufficiently saturate the matrix tensor units:

$$\text{Intensity} = \frac{2 \cdot B \cdot P_{\text{params}}}{P_{\text{params}} \cdot \text{bytes_per_weight}} = \frac{2 B}{\text{bytes_per_weight}} \gg \frac{\text{Peak FLOPS}}{\text{Bandwidth}}$$

  1. Compute Multiplication: Verifying $K$ speculative tokens across $B$ streams increases the effective batch size evaluated by the target model to $B \times K$. If $B=64$ and $K=5$, the target model processes an effective prefill batch of 320 tokens per cycle. If the GPU is already compute-saturated at $B=64$, multiplying the operational load by 5 directly inflates $T_{\text{verify}}$ by a factor of ~4-5x, entirely eliminating the speedup.
  2. KV Cache Fragmentation: Reserving addressable block allocations for $K$ candidate tree branches per request multiplies instantaneous allocation pressure, triggering frequent cache evictions or context paging stalls.

Consequently, production inference engines employ dynamic speculative gating: speculative decoding is automatically engaged when request concurrency is low (single-stream low-latency SLA modes) and progressively throttled or disabled as serving queues expand and the cluster transitions to high-throughput compute-dense regimes.

Architectural Trade-offs in Production Serving

Engineers selecting a speculative serving framework face design decisions regarding draft proposal mechanics, memory topologies, and pipeline distribution.

Mechanism Structural Characteristics Latency Overhead ($T_{\text{draft}}$) Acceptance Rate ($\alpha$) Architectural Constraints
Independent Small Model (e.g., Llama-3.2-1B with Llama-3-70B) Distinct transformer running sequential autoregression. High (accumulates across $K$ serial forward steps). Moderate (distribution shifts between model families). Requires dedicated GPU SRAM/HBM space for draft weights.
Speculative Residual Heads (e.g., Medusa) Multiple lightweight heads attached to top hidden layers of the target model. Ultra-Low (single forward step generates tree heads). Moderate-to-High across structured tokens. Requires modifying or retraining model head weights.
Feature Recurrent Networks (e.g., EAGLE-2) Small secondary transformer predicting next hidden feature vectors directly. Low (processes embeddings rather than complete token paths). High (conditions proposals directly on top target activations). Requires custom tree mask compilation and state caching.
Prompt Lookup Decoding (N-gram heuristic) Zero-parameter string matching against prompt/context history. Near Zero (hash table lookups on CPU/host). Domain-specific (high in coding/summarization, poor in reasoning). Fails entirely on novel, low-entropy token generations.

Conclusion

Speculative decoding restructures the economics of LLM serving by directly addressing the memory-bandwidth limitations of modern hardware. By transforming sequential DRAM reads into parallel matrix compute via speculative proposal and rejection sampling, inference engines dramatically reduce per-token latencies for memory-bound workloads.

Achieving these performance gains in real-world infrastructure requires precise execution across the software stack:

  • Maintaining distributional invariance through GPU-native rejection sampling kernels.
  • Maximizing token candidate paths using tree-based causal masking.
  • Dynamically monitoring concurrency thresholds to prevent compute-bound degradation during high batching conditions.

As multi-token prediction heads become natively integrated into base model architectures, speculative decoding will evolve from an external serving optimization into a fundamental hardware-software execution standard across advanced AI systems.

References

  • Leviathan, Y., Kalman, M., & Matias, Y. (2023). Fast Inference from Transformers via Speculative Decoding. International Conference on Machine Learning (ICML). https://arxiv.org/abs/2211.17192
  • Chen, C., Borgeaud, S., Mensch, A., Sifre, L., Liang, Z., et al. (2023). Accelerating Large Language Model Decoding with Speculative Sampling. arXiv preprint. https://arxiv.org/abs/2302.01318
  • Li, Y., Wei, F., Zhang, C., & Zhang, H. (2024). EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty. International Conference on Machine Learning (ICML). https://arxiv.org/abs/2401.15077