Skip to content

README.md

Latest commit

 

History

History
104 lines (76 loc) · 3.2 KB

README.md

File metadata and controls

104 lines (76 loc) · 3.2 KB
Raw

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)