Speculative Decoding: глубокое погружение
opensourceaillmspeculative-decodingperformanceit
Введение: ускорение генерации
Проблема
Standard autoregressive generation:
Token 1: 50ms → Token 2: 50ms → Token 3: 50ms → ...
20 tokens = 1000ms
✗ Sequential: each token depends on previous
✗ GPU idle during token generation
✗ Limited by decode speed, not prefill
Goal: Generate multiple tokens per forward pass
Факт: Speculative decoding can achieve 2-3x speedup with no quality change (research, 2025).
Что такое Speculative Decoding?
Core Idea
Standard decoding:
Large model generates 1 token at a time
Each step: full forward pass of large model
Speculative decoding:
1. Small model (draft) generates N tokens quickly
2. Large model (target) verifies all N tokens in parallel
3. Accept or reject based on probability
Result: Multiple tokens verified per large model pass
Visual Comparison
Standard (GPT-4):
Step 1: [GPT-4 forward] → "Hello"
Step 2: [GPT-4 forward] → "Hello,"
Step 3: [GPT-4 forward] → "Hello, world"
Step 4: [GPT-4 forward] → "Hello, world!"
Total: 4 forward passes
Speculative (GPT-4 + small draft):
Draft: [small forward] → "Hello, world! I am"
Target: [GPT-4 forward] → verify all 4 tokens
Accept: "Hello, world!" ✓
Reject: "I am" ✗ (resample)
Total: 1 forward pass (vs 4)
Алгоритм
Naive Speculative Decoding
def speculative_decode(draft_model, target_model,
prompt, max_tokens=10):
tokens = [start_token]
output = []
for _ in range(max_tokens):
# Draft: generate N tokens quickly
draft_tokens = draft_model.generate(
tokens, n_tokens=4
)
# Target: verify all draft tokens in parallel
verified = []
for i, draft_token in enumerate(draft_tokens):
target_probs = target_model.forward(tokens)
if random.random() < target_probs[draft_token]:
verified.append(draft_token)
else:
# Reject: sample from target distribution
remaining = target_probs / (
target_probs.sum() -
target_probs[:draft_token].sum()
)
verified.append(sample(remaining))
break
tokens.extend(verified)
output.append(verified[-1])
return tokens
Acceptance Probability
Key insight: target distribution must cover draft distribution
Acceptance probability:
α = min(1, p_target(draft) / p_draft(draft))
Example:
Draft predicts token "world" with p=0.8
Target predicts token "world" with p=0.7
α = min(1, 0.7/0.8) = 0.875
→ Accept with 87.5% probability
If rejected: sample from target distribution
Correctness Proof
Why does this preserve the target distribution?
Case 1: Draft token accepted
P(accept) = p_draft * min(1, p_target/p_draft)
If p_target >= p_draft:
P(accept) = p_draft * 1 = p_draft
If p_target < p_draft:
P(accept) = p_draft * (p_target/p_draft) = p_target
Case 2: Draft token rejected, resample from target
P(resample) = (1 - accepted_prob) * (p_target / (1 - p_draft_cumulative))
Result: marginal distribution = p_target ✓
Draft Model Strategies
Strategy 1: Small LLM as Draft
Draft: 7B model (fast, runs on same GPU)
Target: 70B model
Pros:
- Same architecture, good alignment
- Can run on same hardware
Cons:
- Still requires GPU for draft
- Memory bandwidth limited
Speedup: 2-2.5x typical
Strategy 2: N-gram Speculation
Draft: Simple n-gram model (CPU)
Target: Any LLM
Pros:
- No GPU needed for draft
- Very fast on CPU
Cons:
- Limited to repetitive patterns
- Lower acceptance rate
Speedup: 1.5-2x typical
Strategy 3: Helper Head
Draft: Extra head on target model
Target: Main model
Architecture:
Input → [Transformer layers] → Main head
↓
[Helper head] → draft tokens
Pros:
- No extra model loading
- Perfect alignment
Cons:
- Extra training required
- Adds parameters
Strategy 4: Token Tree Drafting
Instead of linear draft, create a tree:
[draft root]
/ | \
A B C
/ \ / \
D E F G
Target verifies entire tree in parallel
Pros:
- More tokens per forward pass
- Better GPU utilization
Cons:
- Complex implementation
- Need resampling strategy
Implementation with vLLM
Using vLLM's Speculative Decoding
from vllm import LLM, SamplingParams
# Load target model
llm = LLM(
model="meta-llama/Llama-3-70B",
speculative_model="meta-llama/Llama-3-8B",
num_speculative_tokens=8,
tensor_parallel_size=4
)
# Generate with speculative decoding
sampling_params = SamplingParams(
temperature=0.7,
max_tokens=1024
)
output = llm.generate("Hello, world!", sampling_params)
print(output[0].outputs[0].text)
Performance Comparison
Without speculative:
Input: 10 tokens, Output: 100 tokens
Time: 2.5s
Throughput: 40 tok/s
With speculative (8 draft tokens):
Input: 10 tokens, Output: 100 tokens
Time: 1.1s
Throughput: 91 tok/s
Speedup: 2.27x
Implementation with Transformers
Custom Speculative Decoding
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
class SpeculativeDecoder:
def __init__(self, target_path, draft_path, device="cuda"):
self.target = AutoModelForCausalLM.from_pretrained(
target_path
).to(device)
self.draft = AutoModelForCausalLM.from_pretrained(
draft_path
).to(device)
self.tokenizer = AutoTokenizer.from_pretrained(
target_path
)
self.device = device
def generate(self, input_text, max_new_tokens=100,
n_draft=5):
inputs = self.tokenizer(
input_text, return_tensors="pt"
).to(self.device)
output = self.target.generate(
**inputs,
max_new_tokens=max_new_tokens,
do_sample=True,
temperature=0.7
)
return self.tokenizer.decode(
output[0], skip_special_tokens=True
)
Using LightLLM for Fast Drafting
from lightllm.server_api import APIClient
# Fast draft on CPU
draft_client = APIClient("http://localhost:8000")
# Generate draft tokens
draft_result = draft_client.generate(
prompt="Hello, world!",
max_new_tokens=8,
ignore_eos=True
)
# Verify with target model
target_probs = target_model.forward(draft_result)
Advanced Techniques
Medusa
Medusa: Simple framework that improves
inference efficiency
Key idea:
- Add parallel decoding heads to LLM
- Each head predicts next token
- All heads predict in parallel
Architecture:
[Base LLM] → hidden states
↓
[Head 0] → token t+1
[Head 1] → token t+2
[Head 2] → token t+3
Speedup: 2-3x with minimal changes
EAGLE
EAGLE: Speculative autoencoder for fast LLM decoding
Key idea:
- Extract features from shallow layers
- Use features to draft tokens
- Target model verifies in deep layers
Architecture:
Input → [Shallow layers] → Features
↓
[Feature extractor]
↓
Draft tokens
↓
[Deep layers of target] → Verify
Speedup: 2-3x, better than small LLM draft
Tree Attention
Tree Attention for parallel verification:
Draft tree:
[EOS]
/ \
[Hello] [World]
| |
[!] [I]
Forward pass on tree:
- Compute attention for all nodes
- Verify multiple paths simultaneously
- Accept/reject branches independently
Benefit:
- More efficient than linear verification
- Better GPU utilization
Performance Analysis
Acceptance Rate
Acceptance rate = accepted_draft_tokens / total_draft_tokens
Typical values:
Small LLM draft (8B → 70B): 50-70%
N-gram draft: 20-40%
Helper head: 60-80%
EAGLE: 55-75%
Higher acceptance = more speedup
Speedup Formula
Speedup = N_draft * acceptance_rate / (1 + rejection_penalty)
Where:
N_draft = number of draft tokens
rejection_penalty = overhead when rejection occurs
Optimal N_draft:
Too small: not enough parallelism
Too large: too many rejections
Optimal: 4-8 for most cases
Benchmark Results
Model: Llama-3-70B, GPU: A100-80GB
Method Speedup Quality
─────────────────────────────────────────────
Baseline 1.0x 1.00
Small LLM (7B) 2.1x 0.99
N-gram 1.5x 0.97
Medusa 2.4x 0.99
EAGLE 2.6x 0.99
Tree Attention 2.8x 0.98
Quality = similarity to baseline output
Trade-offs and Limitations
When Speculative Decoding Helps
✓ Long generation (100+ tokens)
✓ Repetitive patterns in text
✓ Same vocabulary (draft & target)
✓ GPU-bound (not memory-bound)
✓ Can fit both models on GPU
When It Doesn't Help
✗ Very short generation (< 20 tokens)
✗ Highly diverse vocabulary
✗ Memory bandwidth limited
✗ Can't fit both models
✗ Strict latency requirements (first token)
Memory Requirements
Standard:
70B model (FP16): ~140GB VRAM
Speculative:
70B model (FP16): ~140GB VRAM
8B model (FP16): ~16GB VRAM
Total: ~156GB VRAM
Solution:
- Quantize draft model (INT8/INT4)
- Use CPU for draft
- Use n-gram draft (no model needed)
Practical Tips
Choosing Draft Model
1. Architecture match:
Same architecture → better alignment → higher acceptance
2. Size ratio:
Draft should be 5-10x smaller than target
3. Quantization:
INT8 draft is fine (acceptance drops < 5%)
4. Warm-up:
Warm up draft model before benchmarking
Tuning Parameters
# Optimal draft length
def find_optimal_draft_length(acceptance_rate,
draft_time, target_time):
"""Find N that maximizes speedup"""
best_n = 1
best_speedup = 1
for n in range(1, 16):
# Expected accepted tokens
expected = n * acceptance_rate
# Speedup
speedup = expected / (
1 + (n * draft_time / target_time)
)
if speedup > best_speedup:
best_speedup = speedup
best_n = n
return best_n
Monitoring
class SpeculativeMonitor:
def __init__(self):
self.accepted = 0
self.total_draft = 0
self.forward_passes = 0
def record(self, n_draft, n_accepted):
self.total_draft += n_draft
self.accepted += n_accepted
self.forward_passes += 1
@property
def acceptance_rate(self):
return self.accepted / self.total_draft
@property
def speedup(self):
return self.total_draft / self.forward_passes
Заключение
Speculative decoding is one of the most practical ways to speed up LLM inference today. Key takeaways:
Best approaches:
- Small LLM draft (8B → 70B): 2-2.5x speedup
- EAGLE/Medusa: 2.5-3x speedup
- Tree attention: emerging, promising
Implementation:
- vLLM supports speculative decoding out of the box
- Transformers needs custom implementation
- Consider memory constraints
When to use:
- Long-form generation
- Batch inference
- Cost-sensitive deployments