跳到主要内容
推理优化

第5章:Speculative Decoding

理解投机解码的核心原理(Draft + Verify)、Self-Draft 方案(Medusa/EAGLE-2)及其收益边界与限制

Speculative Decoding 投机解码 Medusa EAGLE Rejection Sampling

LLM Decode 是严格串行的——一次只能生成一个 token,GPU 算力大量空闲。Speculative Decoding 用”实习生先草拟、主编一次审阅”的思路打破这个串行限制:用便宜的 Draft 模型批量猜 K 个 token,用昂贵的 Target 模型一次性验证。本章深入它的数学正确性、各种 Draft 方案(独立小模型 / Medusa / EAGLE-2)、收益边界和工程坑点。

📑 目录


1. 核心思想:猜 + 验证

1.1 朴素 Decode 的瓶颈

T1: model(x) → token_1
T2: model(x, token_1) → token_2
T3: model(x, token_1, token_2) → token_3
...

每次 Decode 都是一次 forward,N 个 token 要 N 次 forward。GPU 算力实际只用了一小部分(memory bound)。

1.2 Speculative Decoding 流程

1. Draft 模型(小,快)在 1 次 forward 里猜 K 个 token: d_1, d_2, ..., d_K
2. Target 模型(大,准)用 1 次 forward 同时算所有候选位置:
   p_1 = P(_|x)        ← 用于验证 d_1
   p_2 = P(_|x, d_1)
   ...
   p_K+1 = P(_|x, d_1, ..., d_K)
3. 从前往后逐个验证:
   - d_i 被 p_i 接受?接受
   - 第一个被拒绝的位置:用 (p_i - q_i) 重新采样一个,后续全部丢弃
4. 合计接受 m 个 token(0 ≤ m ≤ K),进入下一轮

关键点:Target 一次 forward 处理 K+1 个位置(因为是已知序列上的并行算)——比串行 K 次 forward 快得多。

1.3 加速比

设接受率 α\alpha(每个猜的 token 被接受的概率),Draft 单次开销与 Target 比 cc:

Speedup=1αK+1(1α)(1+Kc)\text{Speedup} = \frac{1 - \alpha^{K+1}}{(1 - \alpha)(1 + Kc)}
  • α=0.7\alpha = 0.7, K=4K = 4, c=0.1c = 0.1:Speedup ≈ 2.5×
  • α=0.9\alpha = 0.9, K=5K = 5, c=0.05c = 0.05:Speedup ≈ 3.5×

2. Rejection Sampling:正确性的数学保证

2.1 验证规则

设 Target 概率 pp,Draft 概率 qq。对于每个猜的 token dd:

  • 如果 p(d)q(d)p(d) \ge q(d),接受
  • 否则,以概率 p(d)/q(d)p(d) / q(d) 接受

如果拒绝,则从修正分布 p(x)=max(0,p(x)q(x))/Zp'(x) = \max(0, p(x) - q(x)) / Z 中重新采样一个 token(ZZ 为归一化)。

2.2 正确性证明

可以证明,这种 rejection sampling 接受的 token 严格服从 Target 模型 pp 的分布——数学上等价于直接从 Target 采样。

意味着:Speculative Decoding 不改变模型的输出分布、不改变模型的能力——它只是用并行换串行,纯加速,无质量损失。

2.3 直觉

把 Target 分布画成”大山”,Draft 分布画成”小山”:

  • 大部分位置 Draft 和 Target 重叠 → 直接接受
  • Draft 估高的部分 → 概率性接受
  • Draft 漏掉的部分 → 拒绝时从”差额”重采样

整个流程是精确的统计等价,不是近似。


3. Draft 模型方案

3.1 独立小模型(经典方案)

用一个相同 vocab 的小模型作 Draft:

Target推荐 Draft
LLaMA-70BLLaMA-7B / TinyLLaMA
Qwen-72BQwen-1.8B
GPT-3自训 small GPT

要求:相同 tokenizer,相似分布(同家族微调过的)。

from vllm import LLM
llm = LLM(
    model="meta-llama/Llama-3-70B",
    speculative_model="meta-llama/Llama-3-8B",
    num_speculative_tokens=5,
    use_v2_block_manager=True,
)

3.2 优劣

维度独立 Draft
接受率中-高(70-85%)
内存多占一份 Draft 模型
训练已有现成模型
维护简单

3.3 N-gram Draft(零成本)

更简陋:用 prompt 中已出现的 N-gram 当 Draft(假设输出会重复 prompt 中的某些短语):

Prompt:  "...The quick brown fox jumps over the lazy dog..."
Decode 中遇到 "quick brown" → 猜下一个是 "fox"

零参数,但只在重复性强的场景(代码、文档生成)有效。vLLM 也支持。


4. Medusa:多头解码

4.1 思路

不需要外部 Draft 模型——给 Target 模型加几个并行的 Decoding Head,一次性预测未来 K 个位置:

input → Target backbone → hidden state h

        ┌──────┬──────┬──────┐
        Head 0 Head 1 Head 2 Head 3   ← 各自预测下一个、下两个、...
        ↓      ↓      ↓      ↓
       t+1    t+2    t+3    t+4

4.2 训练

Medusa 头作为 Target 模型的”扩展”,在 Target 输出上训练几层 MLP,损失 = 多个位置的 cross-entropy。

Loss = Σ_k λ_k * CE(head_k(h_t), t+k+1)

只训新增的 Head,Target 主体冻结,训练成本极小(单卡几小时)。

4.3 Medusa Tree

不是顺序猜 K 个,而是按 Top-K 维护一棵候选树,Target 一次 forward 验证所有路径:

root
 ├── token_a (头0 top1)
 │    ├── token_x (头1 top1)
 │    └── token_y (头1 top2)
 └── token_b (头0 top2)
      ├── ...

接受率比简单串行猜更高。

4.4 收益

  • ~2-3× 加速
  • 不需要额外 Draft 模型,工程简单
  • vLLM、SGLang 都支持

5. EAGLE-2:动态 Draft Tree

5.1 EAGLE 的核心创新

不只用 Target 的 hidden state,还用 Target 已生成的 features 当 Draft 的输入——更接近 Target 的分布,接受率高。

Standard:  draft 用 token id 当输入
EAGLE:     draft 用 Target 的 hidden state 当输入,精度高

5.2 EAGLE-2 的动态 Tree

EAGLE-1 用静态 Tree(每个深度 fixed top-k)。 EAGLE-2 根据 Draft 的置信度动态决定每个分支扩展多少——置信度高的多扩,低的少扩。

接受率从 EAGLE-1 的 ~75% 提升到 ~85%,加速 3-4×。

5.3 用法

# vLLM 集成 EAGLE
llm = LLM(
    model="meta-llama/Llama-3-70B",
    speculative_model="yuhuili/EAGLE-LLaMA3-Instruct-70B",
    num_speculative_tokens=8,
)

6. 收益边界与不赚的场景

6.1 接受率取决于 Draft / Target 一致性

场景接受率加速
代码生成(可预测)85-95%3-4×
翻译(确定性)80-90%2.5-3×
通用对话65-75%1.5-2×
创造性写作(高 temperature)50-65%1.2-1.5×
高 temp(temperature > 1)40-55%< 1.2×

6.2 不赚的场景

❌ 高 batch size 已经吃满 GPU

如果 batch 已经把 GPU 算力打满(compute bound),Target 的”一次验证 K 个位置”不再省时间——反而浪费 Draft 的开销。

低 batch (memory bound):  Speculative 收益高
高 batch (compute bound): Speculative 可能负收益

6.3 不赚的场景

❌ Draft 模型本身太大

Draft 模型每次 forward 也是开销。如果 Draft 占 Target 的 30%+,加速比反而变差。

经验:Draft 应该是 Target 的 5-15%(70B 的 Target 用 4-10B 的 Draft)。

6.4 不赚的场景

❌ 量化叠加导致 Draft 精度暴跌

Target 量化没问题,但 Draft 量化后接受率可能暴跌——量化对小模型更敏感。

6.5 不赚的场景

❌ Continuous Batching 中调度复杂

不同请求的 K 不一样,padding 损失算力;接受数不同,Continuous Batching 调度变复杂——需要专门优化。


7. 工程实战:vLLM Speculative

7.1 配置

from vllm import LLM, SamplingParams

# 方案 1:独立 Draft
llm = LLM(
    model="meta-llama/Llama-3-70B",
    speculative_model="meta-llama/Llama-3-8B",
    num_speculative_tokens=5,
)

# 方案 2:N-gram(零成本)
llm = LLM(
    model="meta-llama/Llama-3-70B",
    speculative_model="[ngram]",
    ngram_prompt_lookup_max=5,
    num_speculative_tokens=5,
)

# 方案 3:Medusa / EAGLE
llm = LLM(
    model="meta-llama/Llama-3-70B",
    speculative_model="path/to/eagle-checkpoint",
    num_speculative_tokens=8,
)

7.2 性能监控

vLLM 提供 acceptance rate 监控:

outputs = llm.generate(...)
for out in outputs:
    metrics = out.metrics
    print(f"Acceptance rate: {metrics.spec_decode_acceptance_rate}")

7.3 调优经验

现象调整
接受率 < 50%Draft 模型不够好,换更大的 Draft 或 EAGLE
K 太大 latency 高减小 num_speculative_tokens(典型 4-8)
K 太小没收益增大,但要看接受率是否撑住
高 temperature 收益消失接受 spec 只在低温采样有用
高并发负收益降低 batch 触发 spec 的阈值,或关掉 spec

✅ 自我检验清单

  • 核心思想:能向同事解释”小模型猜、大模型验”的并行化思路
  • Rejection Sampling 正确性:能解释为什么 Spec Decoding 数学上不改变输出分布
  • 加速公式:能用 α,K,c\alpha, K, c 估算理论加速比
  • Draft 选型:能为 LLaMA-70B 选合适的 Draft 模型大小
  • Medusa:能解释 Medusa 头的训练目标和优势
  • EAGLE-2:能解释 EAGLE-2 用 hidden state 当 draft 输入的好处
  • 接受率场景:能列出代码 / 翻译 / 对话 / 创意写作的典型接受率
  • 不赚的场景:能列出至少 3 种 Spec Decoding 收益不佳或负收益的情况
  • 工程调优:能用 vLLM 启用 spec,监控 acceptance rate 并调优 K

📚 参考资料

论文

代码