How to Implement Speculative Decoding for Faster LLM Responses

Large language models have transformed how we build applications, but inference latency remains a persistent bottleneck. Each token generated by an autoregressive LLM requires a full forward pass through the model, and for long responses this sequential cost adds up quickly. Speculative decoding is a technique that can dramatically reduce latency without changing the output distribution of the model. In this tutorial, you will learn what speculative decoding is, why it matters, how to implement it from scratch, and how to apply it using popular libraries.

What Is Speculative Decoding?

Speculative decoding is an inference acceleration method that uses a small, fast "draft" model to propose candidate tokens, which are then verified in parallel by a larger "target" model. Because modern GPUs are underutilized when generating a single token at a time, batching several candidate tokens into one forward pass lets us extract far more value from each computation.

The key insight is that verification is cheaper than generation. When the target model verifies a sequence of draft tokens, it produces logits for every position in a single pass. If the draft model guessed correctly, we effectively generated multiple tokens for the cost of one forward pass. If it guessed wrong, we simply discard the tokens after the first mismatch and continue — and crucially, the final output distribution is mathematically identical to what the target model would have produced on its own.

Why It Matters

Latency reduction: Typical speedups range from 2x to 3x for chat workloads, with higher gains on highly predictable text.
No quality loss: The target model's distribution is preserved exactly through rejection sampling.
Better GPU utilization: Transforms memory-bandwidth-bound single-token generation into compute-bound batched verification.
Drop-in compatibility: Works with most transformer-based LLMs without retraining.

The Algorithm in Detail

At each step, the draft model autoregressively generates k candidate tokens. These tokens, along with their probabilities, are passed to the target model. The target model performs a single forward pass over the entire candidate sequence and produces its own probability distribution at each position. For each draft token, we apply rejection sampling:

If the draft token's probability under the draft model is less than or equal to its probability under the target model, accept it.
Otherwise, accept it with probability p_target / p_draft, and reject with the remaining probability.
On the first rejection, sample a replacement token from a rescaled distribution max(0, p_target - p_draft), normalized.
If all draft tokens are accepted, sample one bonus token from the target model's distribution at the final position.

This procedure guarantees that the accepted tokens follow the target model's distribution exactly.

A Minimal Implementation from Scratch

Below is a simplified but complete implementation in PyTorch. It assumes both models expose a logits method that returns logits for a sequence of token IDs.

import torch
import torch.nn.functional as F

def sample_from_logits(logits, temperature=1.0):
    if temperature <= 0:
        return torch.argmax(logits, dim=-1)
    probs = F.softmax(logits / temperature, dim=-1)
    return torch.multinomial(probs, num_samples=1).squeeze(-1)

def speculative_decode(target_model, draft_model, input_ids, max_new_tokens=128, k=4):
    """
    Generate tokens using speculative decoding.
    
    Args:
        target_model: large model with .logits(ids) -> logits [seq, vocab]
        draft_model: small model with .logits(ids) -> logits [seq, vocab]
        input_ids: starting token IDs as a 1D LongTensor
        max_new_tokens: total tokens to generate
        k: number of draft tokens per speculation round
    
    Returns:
        Tensor of generated token IDs
    """
    generated = input_ids.clone()
    tokens_generated = 0
    
    while tokens_generated < max_new_tokens:
        # 1. Draft model proposes k tokens autoregressively
        draft_tokens = []
        draft_probs = []
        current = generated.clone()
        
        for _ in range(k):
            d_logits = draft_model.logits(current)[-1]  # last position
            d_probs = F.softmax(d_logits, dim=-1)
            next_token = torch.multinomial(d_probs, num_samples=1)
            draft_tokens.append(next_token)
            draft_probs.append(d_probs)
            current = torch.cat([current, next_token])
        
        draft_tokens = torch.stack(draft_tokens)  # [k]
        
        # 2. Target model verifies all k tokens in one forward pass
        verify_input = torch.cat([generated, draft_tokens])
        t_logits = target_model.logits(verify_input)
        # Logits at positions corresponding to each draft token
        t_logits_for_draft = t_logits[generated.shape[0] - 1 : generated.shape[0] - 1 + k]
        t_probs = F.softmax(t_logits_for_draft, dim=-1)  # [k, vocab]
        
        # 3. Rejection sampling
        accepted = 0
        for i in range(k):
            d_tok = draft_tokens[i]
            p_d = draft_probs[i][d_tok]
            p_t = t_probs[i][d_tok]
            
            if p_t >= p_d:
                # Always accept
                generated = torch.cat([generated, d_tok.unsqueeze(0)])
                accepted += 1
                tokens_generated += 1
                if tokens_generated >= max_new_tokens:
                    break
            else:
                # Accept with probability p_t / p_d
                r = torch.rand(1).item()
                if r < (p_t / p_d).item():
                    generated = torch.cat([generated, d_tok.unsqueeze(0)])
                    accepted += 1
                    tokens_generated += 1
                else:
                    # Reject: sample from rescaled distribution
                    rescaled = torch.clamp(t_probs[i] - draft_probs[i], min=0)
                    if rescaled.sum() > 0:
                        rescaled = rescaled / rescaled.sum()
                    else:
                        rescaled = t_probs[i]
                    new_token = torch.multinomial(rescaled, num_samples=1)
                    generated = torch.cat([generated, new_token])
                    tokens_generated += 1
                    break
        
        # 4. If all accepted, sample a bonus token from target
        if accepted == k and tokens_generated < max_new_tokens:
            bonus_logits = t_logits[-1]
            bonus_token = sample_from_logits(bonus_logits)
            generated = torch.cat([generated, bonus_token.unsqueeze(0)])
            tokens_generated += 1
    
    return generated

This implementation prioritizes clarity over performance. In production, you would batch the draft model's autoregressive steps, use KV caching for both models, and avoid the Python-level loop over k by vectorizing the acceptance check.

Using Hugging Face Transformers

The Hugging Face transformers library supports speculative decoding natively through the AssistedGeneration API. You only need to provide a draft model via the assistant_model argument.

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_id = "meta-llama/Llama-2-13b-hf"
draft_id = "meta-llama/Llama-2-7b-hf"

tokenizer = AutoTokenizer.from_pretrained(model_id)
target_model = AutoModelForCausalLM.from_pretrained(
    model_id, torch_dtype=torch.float16, device_map="auto"
)
draft_model = AutoModelForCausalLM.from_pretrained(
    draft_id, torch_dtype=torch.float16, device_map="auto"
)

prompt = "Explain how transformers process sequential data."
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

outputs = target_model.generate(
    **inputs,
    assistant_model=draft_model,
    max_new_tokens=256,
    do_sample=True,
    temperature=0.7,
)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Under the hood, the library handles KV cache management, draft token proposal, parallel verification, and rejection sampling. The output is statistically identical to running the target model alone.

Using vLLM for Production Inference

For production deployments, vLLM offers a high-performance implementation. You launch the server with a draft model specified, and speculative decoding is applied automatically.

# Start the server with speculative decoding enabled
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-2-13b-hf \
    --speculative-model meta-llama/Llama-2-7b-hf \
    --num-speculative-tokens 5 \
    --port 8000

Then query it like any OpenAI-compatible endpoint:

import openai

client = openai.Client(base_url="http://localhost:8000/v1", api_key="dummy")

response = client.chat.completions.create(
    model="meta-llama/Llama-2-13b-hf",
    messages=[{"role": "user", "content": "Write a short poem about latency."}],
    max_tokens=200,
)
print(response.choices[0].message.content)

Best Practices

Choose a compatible draft model: The draft model should share the tokenizer and ideally the architecture family of the target model. A 7B draft for a 70B target is a common sweet spot.
Tune k carefully: Too few draft tokens underutilizes the verification pass; too many wastes compute when the draft model diverges. Values between 3 and 8 typically work well.
Match distributions: The closer the draft model's distribution is to the target's, the higher the acceptance rate. Fine-tuning the draft model on the target's outputs can help.
Use KV caching: Both models must cache their key-value states to avoid recomputation. Most libraries do this automatically, but custom implementations must handle it.
Profile before and after: Measure tokens-per-second and acceptance rate. If acceptance is below 50%, consider a better draft model or a smaller k.
Consider workload characteristics: Speculative decoding shines on predictable text like code completion or structured outputs. Highly creative sampling with high temperature may reduce acceptance rates.
Beware memory overhead: The draft model and its KV cache consume additional GPU memory. Ensure your deployment has headroom.

Common Pitfalls

One frequent mistake is using a draft model with a different tokenizer than the target. Even subtle differences in token boundaries break the algorithm because the rejection sampling assumes token-level alignment. Another pitfall is forgetting to handle the bonus token case when all draft tokens are accepted — omitting it leaves free performance on the table. Finally, mixing greedy decoding in the draft model with sampling in the target model can skew acceptance rates; keep the sampling strategy consistent or use the proper assisted-generation path that handles this correctly.

Conclusion

Speculative decoding is one of the most practical techniques for reducing LLM inference latency without sacrificing output quality. By pairing a fast draft model with a powerful target model and leveraging parallel verification, you can achieve 2-3x speedups on real workloads. Whether you implement it from scratch for learning, use the Hugging Face assistant_model API for prototyping, or deploy with vLLM for production scale, the principles remain the same: choose a compatible draft model, tune the number of speculative tokens, and measure acceptance rates to guide optimization. As LLMs continue to grow in size, techniques like speculative decoding will only become more essential for delivering responsive AI experiences.