跳转到主要内容

投机解码：突破 decode 一次只出一个 token 的限制

#LLM #Inference #Speculative Decoding #EAGLE #Medusa #Draft Model

1. 核心问题：Decode 为什么慢？

你已经知道：

  • Decode 每步只生成 1 个 token
  • 每步都要读一遍所有模型权重（memory-bound）
  • 生成 100 个 token = 读 100 次权重 = 100 次显存访问
传统 decode (Llama-3-70B, FP16):
  每 token: 读 140GB 权重 → 做一次前向 → 输出 1 token
  100 tokens: 读 140GB × 100 = 14TB 数据搬运
  
  能不能一次读权重，输出多个 token？

答案就是投机解码 (Speculative Decoding)。

2. 投机解码的核心思想

2.1 类比：先猜后验

传统方式（逐字翻译）:
  翻译官看一个词 → 翻译一个词 → 看下一个词 → 翻译 → ...
  每个词都要完整思考一次

投机解码（助手先猜，翻译官批量审核）:
  助手快速猜 5 个词 → 翻译官一次性审核 5 个词
  → 前 3 个对了，第 4 个错了
  → 接受前 3 个 + 翻译官修正第 4 个 = 一次得到 4 个词

2.2 正式流程

角色:
  Draft Model (小模型): 快速但不太准，如 Llama-3-1B
  Target Model (大模型): 慢但准确，如 Llama-3-70B

步骤:
  1. Draft Model 自回归生成 K 个候选 token: [t₁, t₂, t₃, t₄, t₅]
     → 很快（小模型，每步 ~1ms）

  2. Target Model 一次性验证这 K 个 token
     → 把 [context, t₁, t₂, t₃, t₄, t₅] 作为输入做一次 forward
     → 得到每个位置的概率分布 [p₁, p₂, p₃, p₄, p₅]
     → 这是一次 prefill 操作（compute-bound，GPU 利用率高）

  3. 从左到右逐个验证:
     t₁ 被接受？ ✓ (draft 概率和 target 概率匹配)
     t₂ 被接受？ ✓
     t₃ 被接受？ ✓
     t₄ 被接受？ ✗ (拒绝，用 target 的分布重新采样得到 t₄')
     t₅ 丢弃（t₄ 已经被拒绝，后面的都无效）

  4. 最终输出: [t₁, t₂, t₃, t₄']  ← 一次得到 4 个 token！

  5. 回到步骤 1，从 t₄' 继续

2.3 为什么能加速？

传统 decode 生成 4 个 token:
  Target forward × 4 = 4 次大模型推理

投机解码生成 4 个 token:
  Draft forward × 5 = 5 次小模型推理（很快，可忽略）
  Target forward × 1 = 1 次大模型推理（验证）
  
  总计 ≈ 1 次大模型推理 → 加速 ~4 倍！

实际加速取决于 acceptance rate（接受率）：

  • 接受率高（draft 和 target 很像）→ 加速明显，2-3x
  • 接受率低（draft 和 target 差异大）→ 加速有限，1.2-1.5x

3. 关键保证：输出质量完全不变

这是投机解码最精妙的地方：验证过程保证输出分布与直接用大模型生成完全一致。

3.1 验证算法（Speculative Sampling）

对于 draft token t_i，draft 模型给出概率 q(t_i)，target 模型给出概率 p(t_i):

接受概率 = min(1, p(t_i) / q(t_i))

如果 p(t_i) ≥ q(t_i):  100% 接受（target 比 draft 更喜欢这个 token）
如果 p(t_i) < q(t_i):   以 p/q 的概率接受

被拒绝时，从修正分布中重新采样:
  p'(t) = normalize(max(0, p(t) - q(t)))

数学上可以证明：这个过程产生的 token 分布与直接从 target 模型采样完全相同。

这意味着：投机解码是无损加速，不牺牲任何质量。

4. 主流投机解码变体

4.1 Draft Model 方式（经典）

用一个同系列的小模型作为 draft:
  Target: Llama-3-70B
  Draft:  Llama-3-8B 或 Llama-3-1B

优点: 简单直接，不需要额外训练
缺点: 需要加载两个模型，占额外显存
      小模型和大模型的分布可能差异大 → 接受率低

4.2 EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency)

核心思想: 不用独立的 draft 模型，而是在 target 模型上加一个轻量级的
         预测头（1-2 层 Transformer），用 target 的隐藏状态来预测下一个 token

优点:
  - 不需要额外的完整模型
  - 预测头很小（<1% 模型参数）
  - 接受率高（直接用 target 的特征）
  - 加速 2-3x

缺点:
  - 需要训练预测头（但很快，几小时）

EAGLE-2 进一步引入了动态 draft 长度，根据置信度决定猜多少个 token。

4.3 Medusa

核心思想: 在 target 模型的最后一层加多个并行的预测头
         每个头预测不同位置的 token

         Head 1 → 预测 position +1
         Head 2 → 预测 position +2
         Head 3 → 预测 position +3

优点:
  - 单次 forward 就能得到多个位置的候选
  - 不需要额外模型
  - 可以用 tree attention 验证多条候选路径

缺点:
  - 需要训练 Medusa heads
  - 远距离预测准确率下降

4.4 对比总结

方法额外参数需要训练典型加速接受率
Draft Model完整小模型否1.5-2.5x中等
EAGLE/EAGLE-2<1%是（几小时）2-3x高
Medusa~5%是1.5-2.5x中高
Self-Speculative0否1.3-1.8x中等

5. 投机解码的工程考量

5.1 什么时候用投机解码？

适合:
  ✓ 延迟敏感的场景（实时对话、代码补全）
  ✓ 单请求或低并发场景
  ✓ Target 模型很大（70B+），decode 很慢

不太适合:
  ✗ 高并发场景（batch 已经很大，GPU 利用率已经高了）
  ✗ 小模型（decode 本身就快，加速空间小）
  ✗ 显存紧张（draft model 需要额外显存）

5.2 Draft 长度 K 的选择

K 太小 (如 K=2):
  每次最多接受 2 个 token，加速有限

K 太大 (如 K=10):
  后面的 token 接受率越来越低
  draft 生成时间增加
  验证的 prefill 计算量增加

最优 K 通常在 3-7 之间，取决于:
  - Draft 和 Target 的相似度
  - 生成内容的可预测性（代码 > 创意写作）
  - 硬件配置

5.3 vLLM 中使用投机解码

# 使用独立 draft model
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3-70B \
  --speculative-model meta-llama/Llama-3-8B \
  --num-speculative-tokens 5

# 使用 EAGLE
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3-70B \
  --speculative-model eagle-llama3-70b \
  --spec-decoding-method eagle

6. 关键要点总结

┌──────────────────────────────────────────────────────────┐
│  投机解码核心认知                                          │
├──────────────────────────────────────────────────────────┤
│  1. 小模型猜 K 个 token → 大模型一次验证 → 一次出多个       │
│  2. 验证算法保证输出分布完全一致 = 无损加速                  │
│  3. 加速倍数 ≈ 平均接受长度，通常 2-3x                     │
│  4. EAGLE 是当前最优方案（高接受率 + 低额外开销）            │
│  5. 适合低并发延迟敏感场景，高并发时收益递减                  │
│  6. 本质是把 memory-bound 的 decode 转化为 compute-bound    │
│     的 verification（一次 prefill 式验证）                  │
└──────────────────────────────────────────────────────────┘

7. 延伸阅读

修改历史1 次提交