Speculative Decoding for LLM Inference

An implementation of speculative sampling to accelerate large language model inference using a draft-target model architecture.

Overview

Speculative decoding speeds up LLM text generation by using a small "draft" model to propose multiple tokens at once, which are then verified in parallel by a larger "target" model. This approach can achieve 2-3x speedup while maintaining identical output distribution to standard autoregressive sampling.

Key Concepts

The algorithm works in three steps per iteration:

1. Draft Generation: A small, fast model generates K candidate tokens autoregressively
2. Parallel Verification: The large target model scores all K tokens in a single forward pass
3. Modified Rejection Sampling: Tokens are accepted/rejected based on probability ratios between models

The mathematical guarantee: the final output distribution is identical to sampling directly from the target model.

Requirements

pip install torch transformers

Usage

Basic Example

from decoder import SpeculativeDecoder

# Initialize with target and draft models
speculative_decoder = SpeculativeDecoder(
    target_model_name="gpt2-medium",
    draft_model_name="gpt2",
    device="cuda",
    temperature=0.8,
    top_p=0.9,
)

# Generate text
text = speculative_decoder.generate(
    prompt="Once upon a time",
    max_length=128,
    k=4,  # Number of draft tokens per iteration
)

Comparing with Baseline

# Baseline: target model only
baseline_text, baseline_time = speculative_decoder.generate_baseline(
    prompt="Once upon a time",
    max_length=128,
)

# Speculative decoding
spec_text = speculative_decoder.generate(
    prompt="Once upon a time",
    max_length=128,
    k=4,
)

# Compare speedup
print(f"Speedup: {baseline_time / spec_time:.2f}x")

Parameters

• target_model_name: HuggingFace model ID for the large target model
• draft_model_name: HuggingFace model ID for the small draft model
• temperature: Sampling temperature (default: 1.0)
• top_p: Nucleus sampling parameter (default: 0.9)
• k: Number of draft tokens per iteration (default: 4)

Algorithm Details

Nucleus Sampling

Implements top-p sampling where tokens with cumulative probability > p are filtered out before sampling.

Modified Rejection Sampling

For each draft token with probability p(x) from draft and q(x) from target:

• Accept with probability min(1, q(x)/p(x))
• If rejected, resample from adjusted distribution (q - p)+
• If all K tokens accepted, sample one bonus token from target

Performance Metrics

The implementation tracks:

• Acceptance rate: percentage of draft tokens accepted
• Tokens per iteration: average tokens generated per step
• Time breakdown: draft model vs target model time
• Overall speedup compared to baseline

Implementation Notes

• Both models must share the same tokenizer vocabulary
• Draft model should be from the same architecture family as target
• Larger parameter gaps generally mean higher speedup potential
• Acceptance rate depends on how well draft model approximates target

Reference

Based on "Accelerating Large Language Model Decoding with Speculative Sampling" (Chen et al., 2023)