Speculative Decoding with Open-Weights Models: Architecting Ultra-Low Latency LLM Pipelines

The Autoregressive Bottleneck in Modern LLM Inference

In the era of massive transformer models, deploying large language models (LLMs) with tens or hundreds of billions of parameters poses a fundamental engineering challenge: memory bandwidth saturation. Autoregressive decoding generates tokens sequentially, one by one. For each generated token, the entire model's weights must be loaded from High Bandwidth Memory (HBM) into the GPU SRAM.

Because of this, LLM generation at batch size 1 is strictly memory bandwidth bound, not compute bound. A GPU capable of hundreds of teraflops of compute sits mostly idle, waiting for weights to transfer. To break this bottleneck, machine learning engineers are increasingly turning to speculative decoding (also known as assisted generation). This architectural pattern allows us to utilize smaller, open-weights "draft" models to accelerate the execution of larger "target" models without sacrificing any mathematical accuracy in the output distribution.

This article dives deep into the system architecture, mathematical underpinnings, and implementation mechanics of speculative decoding using modern open-weights model families.

The Core Mechanics of Speculative Decoding

Speculative decoding relies on a simple observation: verifying the validity of $K$ tokens in parallel takes virtually the same amount of time as generating a single token. This is because running a forward pass on a sequence of length $K$ allows the GPU to utilize its tensor cores efficiently, shifting the operation from memory-bound to compute-bound.

The execution pipeline consists of two models:

The Draft Model ($M_D$): A small, fast, and lightweight model (e.g., a 1B to 3B parameter model) that can quickly generate speculative tokens.
The Target Model ($M_T$): The larger, high-capacity model (e.g., a 70B parameter model) whose outputs we want to preserve.

The Speculative Step-by-Step Loop

Speculation Phase: The draft model $M_D$ autoregressively generates $K$ candidate tokens (e.g., $K = 4$ or $5$) from the current context. This is fast because $M_D$ has a small memory footprint.
Verification Phase: The target model $M_T$ runs a single parallel forward pass on the original prompt plus the $K$ drafted tokens. It calculates the probability distributions for all $K+1$ token positions simultaneously.
Acceptance Phase: An acceptance algorithm determines how many of the $K$ drafted tokens are mathematically valid according to the target model's distribution. If a token is rejected, the loop is truncated, the first corrected token is appended, and the remaining speculative tokens are discarded.
KV Cache Synchronization: The Key-Value (KV) caches of both models are rolled back to align with the accepted sequence length.

By leveraging this architecture, we can achieve speedups of $1.5\times$ to $3\times$ depending on the alignment of the draft and target models.

Mathematical Integrity: Rejection Sampling for LLMs

To ensure that speculative decoding does not degrade the quality of the generated text, the output must follow the exact probability distribution of the target model $M_T$. This is achieved through a modified version of rejection sampling.

Let $p(x)$ be the probability distribution of token $x$ output by the target model $M_T$, and let $q(x)$ be the probability distribution output by the draft model $M_D$. For each speculative token $x_i$ generated by the draft model:

We accept $x_i$ with probability: $$\alpha = \min\left(1, \frac{p(x_i)}{q(x_i)}\right)$$
If $x_i$ is accepted, we proceed to evaluate $x_{i+1}$.
If $x_i$ is rejected, we discard all subsequent tokens ($x_{i+1}$ to $x_K$) and sample the next token from the modified distribution: $$p'(x) = \max\left(0, p(x) - q(x)\right)$$ normalized over all possible tokens. This guarantees that the final sampled token is statistically identical to a token generated directly by the target model.

Engineering the Architecture: KV Cache Synchronization

One of the most complex aspects of implementing speculative decoding in production-grade inference engines (such as vLLM or TensorRT-LLM) is managing the Key-Value (KV) cache.

During standard generation, the KV cache grows by 1 slot per step. In speculative decoding, the draft model generates $K$ tokens, appending $K$ entries to its KV cache. The target model then processes these $K$ tokens in a single step, appending $K$ entries to its own KV cache.

However, if the target model rejects the draft token at index $j$ (where $j < K$), we must:

Roll back the target model's KV cache by discarding all entries from index $j$ onward.
Roll back the draft model's KV cache to match the exact same state.
Synchronize the state of both caches before starting the next cycle.

Without efficient, in-place memory manipulation of cache tensors, the overhead of copying and resizing KV cache matrices can easily negate the latency gains of speculative decoding.

Implementing Speculative Decoding: A Technical Walkthrough

Below is a simplified, high-level PyTorch implementation illustrating the core speculative loop and the rejection sampling logic.

import torch
import torch.nn.functional as F

def sample_rejection(target_probs, draft_probs, draft_token):
    """
    Evaluates whether to accept or reject a drafted token.
    """
    p = target_probs[draft_token]
    q = draft_probs[draft_token]
    
    # Acceptance probability
    alpha = torch.minimum(torch.tensor(1.0), p / (q + 1e-9))
    u = torch.rand(1).item()
    
    if u < alpha:
        return True, None
    else:
        # Adjust target distribution for rejection sampling
        adjusted_distribution = torch.clamp(target_probs - draft_probs, min=0.0)
        adjusted_distribution /= (adjusted_distribution.sum() + 1e-9)
        new_token = torch.multinomial(adjusted_distribution, 1)
        return False, new_token

def speculative_decoding_step(draft_model, target_model, input_ids, K=4):
    # 1. Draft model generates K tokens autoregressively
    draft_input = input_ids.clone()
    draft_tokens = []
    draft_distributions = []
    
    for _ in range(K):
        outputs = draft_model(draft_input)
        next_token_logits = outputs.logits[:, -1, :]
        probs = F.softmax(next_token_logits, dim=-1)
        next_token = torch.multinomial(probs, 1)
        
        draft_tokens.append(next_token.item())
        draft_distributions.append(probs.squeeze(0))
        draft_input = torch.cat([draft_input, next_token], dim=-1)
        
    # 2. Target model parallel validation run
    target_outputs = target_model(draft_input)
    target_logits = target_outputs.logits[:, -(K+1):, :]
    target_distributions = F.softmax(target_logits, dim=-1).squeeze(0)
    
    accepted_tokens = []
    is_rejected = False
    
    # 3. Validation loop
    for i in range(K):
        token_i = draft_tokens[i]
        target_probs_i = target_distributions[i]
        draft_probs_i = draft_distributions[i]
        
        accepted, fallback_token = sample_rejection(target_probs_i, draft_probs_i, token_i)
        
        if accepted:
            accepted_tokens.append(token_i)
        else:
            # Append the corrected token from adjusted distribution
            accepted_tokens.append(fallback_token.item())
            is_rejected = True
            break
            
    # 4. If all tokens accepted, sample the (K+1)th token from the target model's last distribution
    if not is_rejected:
        last_target_probs = target_distributions[-1]
        final_token = torch.multinomial(last_target_probs, 1).item()
        accepted_tokens.append(final_token)
        
    return accepted_tokens

Optimizing Draft-Target Selection and Tradeoffs

To achieve optimal inference speedups, system architects must carefully balance model selection and pipeline parameters:

Draft/Target Size Ratio: Empirically, the draft model should be roughly $10\times$ to $20\times$ smaller than the target model. If the draft model is too large, draft generation latency dominates. If it is too small, the acceptance rate drops, resulting in frequent rejections and high validation overhead.
Vocabulary Alignment: Speculative decoding requires both models to share the exact same tokenizer and vocabulary. If they do not, expensive and complex mapping layers must be inserted to align token IDs, which degrades performance and accuracy.
Domain Drift: If your application is highly domain-specific (e.g., writing Rust code or medical diagnostics), the draft model must be fine-tuned on the same domain as the target model. A generic draft model will have a low acceptance rate on specialized tasks, causing the pipeline to fall back to the slow, target-only generation speeds.

By leveraging open-weights models like the Llama-3 or Mistral families, engineering teams can pair a highly capable model with its smaller distilled sibling, transforming raw compute power into highly responsive, real-time AI interfaces.