摘要
- 推测解码是LLM推理加速的杀手锏:2-3×加速比,精度零损失,2026年已成为推理服务标配
- 3大推测策略:草稿模型(Draft Model)、Medusa多头、树状推测(Tree Speculation),各有最佳场景
- 草稿模型选型关键:草稿模型速度需≥5×目标模型,接受率≥70%才能实现2×+加速
- Medusa多头无需额外模型:在目标模型上添加预测头,训练成本低、部署简单
- 本文提供vLLM+推测解码部署方案与自研草稿模型训练实战
目录
推测解码:LLM推理加速的杀手锏
自回归解码的瓶颈
LLM推理的核心瓶颈是自回归解码的串行性:每生成1个token都需要完整的前向传播,GPU利用率极低。
| 模型规模 |
单token延迟 |
GPU利用率 |
吞吐量 |
| 7B |
15ms |
25% |
40 tok/s |
| 14B |
28ms |
18% |
22 tok/s |
| 70B |
80ms |
12% |
8 tok/s |
| 405B |
250ms |
5% |
2.5 tok/s |
推测解码核心思想
┌──────────────────────────────────────────────────────────────┐
│ 推测解码 vs 自回归解码 │
│ │
│ 自回归解码 (逐token生成) │
│ ┌──────────────────────────────────────────────────────┐ │
│ │ [BOS] → f() → t1 → f() → t2 → f() → t3 → f() → t4 │ │
│ │ 4次前向传播,4个token │ │
│ └──────────────────────────────────────────────────────┘ │
│ │
│ 推测解码 (草稿+验证) │
│ ┌──────────────────────────────────────────────────────┐ │
│ │ 草稿模型快速生成: t1' t2' t3' t4' (1次前向) │ │
│ │ 目标模型并行验证: [t1' t2' t3' t4'] → f() (1次前向) │ │
│ │ 假设t1't2't3'正确,t4'错误: │ │
│ │ 接受t1't2't3',拒绝t4',修正为t4 │ │
│ │ 2次前向传播,4个token → 2×加速 │ │
│ └──────────────────────────────────────────────────────┘ │
└──────────────────────────────────────────────────────────────┘
2026年推测解码方案对比
| 方案 |
加速比 |
精度损失 |
额外显存 |
部署复杂度 |
| 草稿模型 |
2-3× |
零 |
30-50% |
中 |
| Medusa多头 |
1.5-2.5× |
零 |
5-10% |
低 |
| Eagle树状推测 |
2.5-3.5× |
零 |
10-15% |
中 |
| 自推测(Self-Spec) |
1.5-2× |
零 |
0% |
低 |
| 投机采样+量化 |
2-3.5× |
极小 |
15-25% |
中 |
推测解码原理与数学基础
接受率与加速比
推测解码的加速比取决于草稿token的接受率(acceptance rate):
理论加速比 = γ / (1 + (1-α) × γ)
其中:
- γ = 草稿长度(每次推测的token数)
- α = 接受率(草稿token被目标模型接受的概率)
示例:
- γ=5, α=0.8: 加速比 = 5 / (1 + 0.2×5) = 5/2 = 2.5×
- γ=5, α=0.9: 加速比 = 5 / (1 + 0.1×5) = 5/1.5 = 3.3×
- γ=5, α=0.6: 加速比 = 5 / (1 + 0.4×5) = 5/3 = 1.67×
- γ=3, α=0.8: 加速比 = 3 / (1 + 0.2×3) = 3/1.6 = 1.875×
接受率与加速比关系
| 接受率α |
γ=3 |
γ=5 |
γ=7 |
γ=10 |
| 0.5 |
1.2× |
1.25× |
1.27× |
1.25× |
| 0.6 |
1.36× |
1.47× |
1.52× |
1.47× |
| 0.7 |
1.58× |
1.82× |
1.94× |
1.96× |
| 0.8 |
1.88× |
2.5× |
2.92× |
3.33× |
| 0.9 |
2.31× |
3.33× |
4.38× |
5.88× |
验证算法实现
import torch
import torch.nn.functional as F
def speculative_decode(
target_model,
draft_model,
input_ids,
draft_length=5,
temperature=1.0,
):
draft_tokens = []
draft_probs = []
current_ids = input_ids.clone()
for _ in range(draft_length):
with torch.no_grad():
outputs = draft_model(current_ids)
next_token_logits = outputs.logits[:, -1, :]
next_token_prob = F.softmax(next_token_logits / temperature, dim=-1)
next_token = torch.multinomial(next_token_prob, num_samples=1)
draft_tokens.append(next_token)
draft_probs.append(next_token_prob)
current_ids = torch.cat([current_ids, next_token], dim=-1)
draft_token_ids = torch.cat(draft_tokens, dim=-1)
with torch.no_grad():
target_outputs = target_model(current_ids)
target_logits = target_outputs.logits[:, -(draft_length + 1):, :]
target_probs = F.softmax(target_logits / temperature, dim=-1)
accepted_tokens = []
n_accepted = 0
for i in range(draft_length):
draft_token = draft_token_ids[:, i]
draft_prob = draft_probs[i].gather(1, draft_token.unsqueeze(-1))
target_prob = target_probs[:, i].gather(1, draft_token.unsqueeze(-1))
accept_prob = torch.min(
torch.ones_like(target_prob),
target_prob / (draft_prob + 1e-10),
)
if torch.rand(1, device=accept_prob.device) < accept_prob:
accepted_tokens.append(draft_token)
n_accepted += 1
else:
residual_prob = F.relu(target_probs[:, i] - draft_probs[i])
residual_prob = residual_prob / residual_prob.sum(dim=-1, keepdim=True)
corrected_token = torch.multinomial(residual_prob, num_samples=1)
accepted_tokens.append(corrected_token)
break
if n_accepted == draft_length:
bonus_prob = target_probs[:, -1]
bonus_token = torch.multinomial(bonus_prob, num_samples=1)
accepted_tokens.append(bonus_token)
result = torch.cat(accepted_tokens, dim=-1)
return result, n_accepted
草稿模型选型与训练
草稿模型选型原则
| 原则 |
要求 |
原因 |
| 速度 |
≥5×目标模型 |
草稿生成不能成为瓶颈 |
| 接受率 |
≥70% |
低于70%加速比不足2× |
| 体积 |
≤1/5目标模型 |
显存开销可控 |
| 词表 |
与目标模型一致 |
避免词表映射损失 |
主流草稿模型搭配
| 目标模型 |
草稿模型 |
速度比 |
接受率 |
加速比 |
| LLaMA-3-70B |
LLaMA-3-8B |
6× |
75% |
2.3× |
| Qwen2.5-72B |
Qwen2.5-7B |
5.5× |
78% |
2.5× |
| DeepSeek-V3 |
DeepSeek-V2-Lite |
7× |
72% |
2.2× |
| Mistral-Large |
Mistral-7B |
6× |
76% |
2.4× |
| Gemma-2-27B |
Gemma-2-2B |
8× |
70% |
2.1× |
草稿模型训练
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer
from datasets import load_dataset
class DraftModelTrainer:
def __init__(self, target_model_path, draft_model_path, output_path):
self.target = AutoModelForCausalLM.from_pretrained(
target_model_path, torch_dtype=torch.float16, device_map="auto"
)
self.draft = AutoModelForCausalLM.from_pretrained(
draft_model_path, torch_dtype=torch.float16, device_map="auto"
)
self.tokenizer = AutoTokenizer.from_pretrained(target_model_path)
self.output_path = output_path
def generate_distillation_data(self, prompts, num_samples=50000):
distill_data = []
for prompt in prompts:
inputs = self.tokenizer(prompt, return_tensors="pt").to(self.target.device)
with torch.no_grad():
outputs = self.target.generate(
**inputs,
max_new_tokens=256,
temperature=0.7,
do_sample=True,
num_return_sequences=5,
)
for output in outputs:
text = self.tokenizer.decode(output, skip_special_tokens=True)
distill_data.append({"text": text})
if len(distill_data) >= num_samples:
break
return distill_data
def train(self, distill_data, epochs=3, lr=5e-5):
from trl import SFTTrainer, SFTConfig
sft_config = SFTConfig(
output_dir=self.output_path,
num_train_epochs=epochs,
per_device_train_batch_size=4,
gradient_accumulation_steps=8,
learning_rate=lr,
warmup_ratio=0.03,
bf16=True,
logging_steps=100,
save_strategy="epoch",
max_seq_length=2048,
)
trainer = SFTTrainer(
model=self.draft,
train_dataset=distill_data,
args=sft_config,
)
trainer.train()
self.draft.save_pretrained(self.output_path)
return self.draft
草稿模型接受率优化
| 优化手段 |
接受率提升 |
方法 |
| 蒸馏微调 |
+10-15% |
用目标模型输出训练草稿模型 |
| 词表对齐 |
+5-8% |
确保草稿与目标词表一致 |
| 温度对齐 |
+3-5% |
匹配草稿与目标的采样温度 |
| 上下文对齐 |
+5-10% |
草稿模型使用相同system prompt |
| 组合优化 |
+20-30% |
以上全部 |
Medusa多头推测实战
Medusa架构
┌──────────────────────────────────────────────────────────────┐
│ Medusa多头架构 │
│ │
│ 原始LLM Backbone │
│ ┌──────────────────────────────────────────────────────┐ │
│ │ Input → Transformer Layers → Hidden State │ │
│ └──────────────────────────────────────────────────────┘ │
│ ↓ │
│ 原始LM Head (预测t+1) │
│ ┌──────────┐ │
│ │ LM Head │ → P(t+1) │
│ └──────────┘ │
│ │
│ Medusa Head 0 (预测t+2) │
│ ┌──────────────┐ │
│ │ Linear+Softmax│ → P(t+2) │
│ └──────────────┘ │
│ │
│ Medusa Head 1 (预测t+3) │
│ ┌──────────────┐ │
│ │ Linear+Softmax│ → P(t+3) │
│ └──────────────┘ │
│ │
│ Medusa Head K-1 (预测t+K) │
│ ┌──────────────┐ │
│ │ Linear+Softmax│ → P(t+K) │
│ └──────────────┘ │
└──────────────────────────────────────────────────────────────┘
Medusa实现
import torch
import torch.nn as nn
from transformers import PreTrainedModel
class MedusaHead(nn.Module):
def __init__(self, hidden_size, vocab_size):
super().__init__()
self.linear = nn.Linear(hidden_size, vocab_size, bias=False)
def forward(self, hidden_states):
return self.linear(hidden_states)
class MedusaModel(nn.Module):
def __init__(self, base_model, num_heads=4):
super().__init__()
self.base_model = base_model
self.num_heads = num_heads
hidden_size = base_model.config.hidden_size
vocab_size = base_model.config.vocab_size
self.medusa_heads = nn.ModuleList([
MedusaHead(hidden_size, vocab_size)
for _ in range(num_heads)
])
for head in self.medusa_heads:
nn.init.xavier_uniform_(head.linear.weight)
def forward(self, input_ids, attention_mask=None):
outputs = self.base_model.model(
input_ids=input_ids,
attention_mask=attention_mask,
use_cache=True,
)
hidden_states = outputs.last_hidden_state
base_logits = self.base_model.lm_head(hidden_states)
medusa_logits = []
for head in self.medusa_heads:
medusa_logits.append(head(hidden_states))
return base_logits, medusa_logits, outputs.past_key_values
class MedusaDecoding:
def __init__(self, model, tokenizer, num_heads=4, top_k=10):
self.model = model
self.tokenizer = tokenizer
self.num_heads = num_heads
self.top_k = top_k
def generate(self, input_ids, max_new_tokens=256):
generated = input_ids.clone()
for _ in range(max_new_tokens):
base_logits, medusa_logits, _ = self.model(generated)
base_prob = torch.softmax(base_logits[:, -1, :], dim=-1)
next_token = torch.argmax(base_prob, dim=-1, keepdim=True)
candidates = [next_token]
for head_idx in range(self.num_heads):
head_prob = torch.softmax(
medusa_logits[head_idx][:, -1, :], dim=-1
)
top_tokens = torch.topk(head_prob, self.top_k)
candidates.append(top_tokens.indices[:, :1])
candidate_ids = torch.cat(candidates, dim=-1)
with torch.no_grad():
verify_outputs = self.model.base_model(candidate_ids)
verify_logits = self.model.base_model.lm_head(
verify_outputs.last_hidden_state
)
accepted = [next_token]
for i, cand in enumerate(candidates[1:]):
verify_prob = torch.softmax(verify_logits[:, i, :], dim=-1)
if verify_prob.gather(1, cand).item() > 0.1:
accepted.append(cand)
else:
corrected = torch.argmax(verify_prob, dim=-1, keepdim=True)
accepted.append(corrected)
break
new_tokens = torch.cat(accepted, dim=-1)
generated = torch.cat([generated, new_tokens], dim=-1)
if new_tokens.shape[-1] > 1:
break
return generated
Medusa训练
class MedusaTrainer:
def __init__(self, medusa_model, lr=1e-4):
self.model = medusa_model
self.optimizer = torch.optim.AdamW(
[p for p in self.model.medusa_heads.parameters()],
lr=lr,
weight_decay=0.01,
)
self.base_model = medusa_model.base_model
for param in self.base_model.parameters():
param.requires_grad = False
def train_step(self, input_ids, labels):
with torch.no_grad():
outputs = self.base_model.model(input_ids=input_ids)
hidden_states = outputs.last_hidden_state
total_loss = 0
for head_idx, head in enumerate(self.model.medusa_heads):
shift = head_idx + 1
head_logits = head(hidden_states[:, :-shift, :])
head_labels = labels[:, shift:]
loss = F.cross_entropy(
head_logits.reshape(-1, head_logits.size(-1)),
head_labels.reshape(-1),
)
total_loss += loss
total_loss /= len(self.model.medusa_heads)
total_loss.backward()
self.optimizer.step()
self.optimizer.zero_grad()
return total_loss.item()
| Medusa配置 |
头数 |
Top-K |
加速比 |
额外显存 |
| Medusa-2 |
2 |
5 |
1.5× |
3% |
| Medusa-4 |
4 |
10 |
2.0× |
5% |
| Medusa-4 |
4 |
20 |
2.2× |
5% |
| Medusa-8 |
8 |
10 |
2.5× |
10% |
树状推测与Eagle方案
Eagle树状推测原理
┌──────────────────────────────────────────────────────────────┐
│ Eagle树状推测 │
│ │
│ 传统线性推测: │
│ t1 → t2 → t3 → t4 → t5 (1条路径) │
│ │
│ Eagle树状推测: │
│ t1 │
│ / \ │
│ t2a t2b │
│ / \ / \ │
│ t3a t3b t3c t3d │
│ │
│ 多条候选路径并行验证,选择接受率最高的路径 │
│ 优势:即使某条路径被拒绝,其他路径仍可能被接受 │
└──────────────────────────────────────────────────────────────┘
Eagle实现
class EagleSpeculativeDecoder:
def __init__(self, target_model, draft_model, tree_width=4, tree_depth=3):
self.target = target_model
self.draft = draft_model
self.tree_width = tree_width
self.tree_depth = tree_depth
def build_speculation_tree(self, input_ids):
tree_nodes = [{"ids": input_ids.clone(), "parent": None}]
leaves = []
for depth in range(self.tree_depth):
new_nodes = []
for node in tree_nodes[-self.tree_width ** depth:]:
with torch.no_grad():
outputs = self.draft(node["ids"])
probs = torch.softmax(outputs.logits[:, -1, :], dim=-1)
top_k = torch.topk(probs, self.tree_width)
for i in range(self.tree_width):
child_ids = torch.cat([
node["ids"],
top_k.indices[:, i:i+1],
], dim=-1)
child = {
"ids": child_ids,
"parent": node,
"prob": top_k.values[:, i].item(),
"token": top_k.indices[:, i],
}
new_nodes.append(child)
if depth == self.tree_depth - 1:
leaves.append(child)
tree_nodes.extend(new_nodes)
return tree_nodes, leaves
def verify_tree(self, tree_nodes, leaves):
best_path = []
best_score = 0
for leaf in leaves:
path = []
node = leaf
while node["parent"] is not None:
path.append(node)
node = node["parent"]
path.reverse()
with torch.no_grad():
full_ids = leaf["ids"]
outputs = self.target(full_ids)
target_probs = torch.softmax(outputs.logits, dim=-1)
score = 1.0
accepted = 0
for i, step in enumerate(path):
target_prob = target_probs[:, -(len(path) - i), :].gather(
1, step["token"]
).item()
score *= target_prob
if target_prob > 0.3:
accepted += 1
else:
break
if score > best_score:
best_score = score
best_path = path[:accepted]
return best_path
树状推测 vs 线性推测
| 方案 |
候选路径 |
接受率 |
加速比 |
计算开销 |
| 线性γ=5 |
1 |
70% |
2.3× |
低 |
| 树w=2,d=3 |
8 |
82% |
2.8× |
中 |
| 树w=3,d=3 |
27 |
88% |
3.2× |
中高 |
| 树w=4,d=3 |
64 |
92% |
3.5× |
高 |
| Eagle(w=4,d=4) |
256 |
95% |
3.8× |
高 |
生产部署:vLLM+推测解码
vLLM推测解码配置
from vllm import LLM, SamplingParams
llm = LLM(
model="Qwen/Qwen2.5-72B-Instruct",
speculative_model="Qwen/Qwen2.5-7B-Instruct",
num_speculative_tokens=5,
speculative_max_model_len=4096,
gpu_memory_utilization=0.9,
tensor_parallel_size=4,
trust_remote_code=True,
)
sampling_params = SamplingParams(
temperature=0.7,
top_p=0.9,
max_tokens=512,
)
outputs = llm.generate(
["请解释量子计算的基本原理"],
sampling_params=sampling_params,
)
for output in outputs:
print(output.outputs[0].text)
vLLM推测解码性能
| 配置 |
目标模型 |
草稿模型 |
吞吐(tok/s) |
加速比 |
| 无推测 |
72B |
无 |
8 |
1× |
| 草稿推测 |
72B |
7B |
18 |
2.25× |
| 草稿推测 |
72B |
0.5B |
22 |
2.75× |
| Medusa-4 |
72B |
无(内置) |
16 |
2.0× |
| Eagle |
72B |
7B |
25 |
3.1× |
生产部署注意事项
| 注意事项 |
说明 |
解决方案 |
| 草稿模型显存 |
额外占用30-50% |
量化草稿模型到INT8 |
| 接受率波动 |
不同任务接受率差异大 |
动态调整草稿长度 |
| KV Cache管理 |
被拒绝token的KV需清理 |
vLLM自动管理 |
| 批处理兼容 |
推测解码与连续批处理冲突 |
vLLM已支持 |
| 首token延迟 |
推测解码增加首token延迟 |
对话场景关闭推测 |
动态推测策略
class DynamicSpeculativeConfig:
def __init__(self):
self.min_draft_length = 2
self.max_draft_length = 8
self.target_acceptance_rate = 0.75
self.current_draft_length = 5
self.history = []
def update(self, acceptance_rate):
self.history.append(acceptance_rate)
if len(self.history) > 100:
self.history = self.history[-100:]
avg_rate = sum(self.history) / len(self.history)
if avg_rate > 0.85:
self.current_draft_length = min(
self.current_draft_length + 1,
self.max_draft_length,
)
elif avg_rate < 0.65:
self.current_draft_length = max(
self.current_draft_length - 1,
self.min_draft_length,
)
return self.current_draft_length
总结与引流
关键要点回顾
- 推测解码是LLM推理加速的免费午餐:2-3×加速,精度零损失
- 草稿模型:选型关键是速度≥5×+接受率≥70%,蒸馏微调可提升20-30%接受率
- Medusa多头:无需额外模型,5%显存开销换2×加速,部署最简单
- Eagle树状推测:多路径并行验证,3.5×+加速,适合高吞吐场景
方案推荐
| 场景 |
推荐方案 |
加速比 |
| 快速上线 |
Medusa-4 |
2.0× |
| 通用推理 |
草稿模型(7B→72B) |
2.5× |
| 高吞吐 |
Eagle树状推测 |
3.5× |
| 显存紧张 |
Medusa+量化草稿 |
2.2× |
需要处理推理数据的格式转换?试试我们的JSON转YAML工具和文本差异对比,快速处理推测解码的调试数据。
延伸阅读