摘要
- 推測解碼是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×
接受率與加速比關係
| 接受率alpha |
γ=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工具和文本差異對比,快速處理推測解碼的除錯資料。
延伸閱讀