【LLM推理与Serving④】Speculative Decoding投机解码：打破自回归的一次只出一个token

大家好，我是你的AI Infra学习搭子。今天我们来聊一个在LLM推理优化领域非常有意思的技术——Speculative Decoding（投机解码）。如果你觉得大模型生成速度慢、一个字一个字往外蹦太煎熬，那这篇文章一定要看完，因为它讲的就是怎么让大模型"一次蹦好几个字"。

---

一、开篇引入：为什么Decode阶段这么慢？

1.1 自回归的"原罪"

用过ChatGPT的同学都有体会：大模型生成回答的时候，是一个字一个字往外蹦的。这不是厂商故意吊你胃口，而是由LLM的自回归（Autoregressive） 本质决定的。

简单来说，LLM每生成一个token（大概是半个到一个汉字），都要：

1. 把前面所有token喂进模型
2. 模型跑一遍完整的forward（几十层Transformer）
3. 从输出概率分布里采样出下一个token
4. 把这个新token加进去，重复上述过程

生成N个token，模型就要跑N次完整的forward。这就像你写文章，每写一个字都要把全文重读一遍再想下一个字——听起来很蠢，但LLM就是这么工作的。

1.2 Decode阶段的瓶颈：内存带宽，而非计算

更让人"难受"的是，Decode阶段GPU的利用率其实很低。为什么？

• Prefill阶段：一次性处理整个prompt，是大矩阵乘法，计算密集（compute-bound），GPU利用率高。
• Decode阶段：每次只输入1个新token，是小批量GEMV操作，需要频繁从显存读KV Cache和模型权重，变成了内存带宽瓶颈（memory-bandwidth-bound）。

打个比方：Prefill像是工厂流水线全开，工人忙不过来；Decode像是流水线只加工一个零件，大部分时间都在搬物料。

所以问题来了：能不能让大模型一次验证多个token，减少forward的次数？

答案是：能！这就是Speculative Decoding要做的事。

---

二、核心思想：用小模型快速猜测，用大模型并行验证

2.1 一个生活中的类比

先讲个故事帮你理解。假设你是一个大公司的CEO（大模型Target），每天要审批很多文件。你水平很高但处理速度慢。你有个助理（小模型Draft），水平不如你但动作很快。

传统方式：每一份文件都由你亲自审批，一份一份来，很慢。

投机方式：

1. 助理先快速审批一堆文件，给出他的意见
2. 你把这一堆文件拿过来，一次性快速复核
3. 助理批对的，直接通过（省了你的时间）
4. 助理批错的，你纠正过来，然后继续下一批

虽然你还是要亲自看，但你一次看一堆的时间，跟你一次看一份的时间差不了太多（因为文件内容已经摆在你面前了）。只要助理正确率不算太低，整体效率就能提升。

2.2 Speculative Decoding的核心思路

把上面的故事翻译成LLM术语：

• Draft Model（草稿模型）：一个较小的模型（比如7B），推理速度快，但生成质量相对低。
• Target Model（目标模型）：我们真正想用的大模型（比如70B），质量高但速度慢。

流程是这样的：

1. Draft Model快速自回归生成K个候选token（比如5个）
2. 把这K个token一次性输入Target Model
3. Target Model并行计算每个位置的概率分布（因为KV已经在缓存里了，一次forward就能算出所有位置的logits）
4. 逐个检查：Draft猜的token，Target是否"认可"
5. 认可的token直接输出，不认可的地方由Target重新采样，然后开始下一轮

关键洞察：Target Model一次验证K个token的成本，远小于它自回归生成K个token的成本。 因为前者只需要1次forward，后者需要K次。

---

三、Speculative Sampling经典框架

3.1 Draft + Target 两模型架构

经典的Speculative Sampling由DeepMind在2023年的论文《Fast Inference from Transformers via Speculative Decoding》中提出。

整个系统包含两个模型：

┌─────────────────────────────────────────────────────────────┐
│                        输入 prompt                           │
└──────────────────────────┬──────────────────────────────────┘
                           │
              ┌────────────▼────────────┐
              │     Draft Model        │
              │   (小模型，速度快)       │
              │  自回归生成 K 个候选     │
              └────────────┬────────────┘
                           │ K 个候选 token
              ┌────────────▼────────────┐
              │    Target Model        │
              │   (大模型，质量高)       │
              │  并行验证所有候选 token  │
              │  + Rejection Sampling   │
              └────────────┬────────────┘
                           │
              ┌────────────▼────────────┐
              │   接受的 token + 重采样  │
              │   (分布与原模型完全一致)  │
              └─────────────────────────┘

3.2 算法流程详解

设：

• $p(x)$ 是Target Model的分布（我们想要的真实分布）
• $q(x)$ 是Draft Model的分布（近似分布）
• $\gamma$ 是每轮Draft生成的候选token数（也叫投机窗口大小）

算法步骤：

1. Draft Model自回归生成$\gamma$个候选token：$x_1,x_2,\dots ,x_{\gamma }$
2. 将这$\gamma$个token输入Target Model，得到每个位置$i$的分布$p_i(x)$和$q_i(x)$
3. 从第一个token开始，逐个执行Rejection Sampling：
   • 以概率$\min \left(1,\frac{p_i(x_i)}{q_i(x_i)}\right)$接受$x_i$
   • 如果接受，继续检查下一个
   • 如果拒绝，从修正分布$p_i^{\prime }(x)\propto \max (0,p_i(x)-q_i(x))$中重新采样，停止本轮
4. 如果全部$\gamma$个都被接受，再额外采样一个token（本轮共生成$\gamma +1$个）
5. 用已接受的token作为新前缀，重复以上过程

3.3 代码示例：一个简化版投机解码

import torch
import torch.nn.functional as F

def speculative_decode_step(draft_model, target_model, input_ids, gamma=5):
    """
    单轮 Speculative Decoding
    
    Args:
        draft_model: 小模型（Draft Model）
        target_model: 大模型（Target Model）
        input_ids: 当前已生成的 token 序列
        gamma: 每轮投机生成的候选 token 数
    
    Returns:
        accepted_tokens: 本轮接受的 token 列表
        n_accepted: 接受的 token 数量
    """
    # ===== Step 1: Draft Model 生成 gamma 个候选 token =====
    draft_tokens = []
    draft_logits_list = []
    
    current_ids = input_ids.clone()
    for _ in range(gamma):
        # Draft Model 自回归生成
        with torch.no_grad():
            outputs = draft_model(current_ids)
            next_logits = outputs.logits[:, -1, :]  # [batch, vocab]
        
        # 采样下一个 token
        next_token = torch.multinomial(F.softmax(next_logits, dim=-1), num_samples=1)
        draft_tokens.append(next_token.item())
        draft_logits_list.append(next_logits)
        
        # 追加到输入，用于下一个位置的生成
        current_ids = torch.cat([current_ids, next_token], dim=-1)
    
    # ===== Step 2: Target Model 并行验证所有候选 =====
    with torch.no_grad():
        # 一次性把 draft token 都喂给 target model
        # 因为 KV Cache 的存在，一次 forward 就能算出所有位置的 logits
        target_outputs = target_model(current_ids)
        target_logits = target_outputs.logits  # [batch, seq_len, vocab]
    
    # 提取候选位置对应的 target logits
    # input_ids 长度为 L，候选位置是 L, L+1, ..., L+gamma-1
    seq_len = input_ids.shape[1]
    target_logits_candidates = target_logits[:, seq_len-1:seq_len-1+gamma, :]
    
    # ===== Step 3: Rejection Sampling 逐个验证 =====
    accepted_tokens = []
    
    for i in range(gamma):
        draft_token = draft_tokens[i]
        p_logits = target_logits_candidates[:, i, :]  # target 分布
        q_logits = draft_logits_list[i]               # draft 分布
        
        p_probs = F.softmax(p_logits, dim=-1)
        q_probs = F.softmax(q_logits, dim=-1)
        
        # 接受概率: min(1, p(x) / q(x))
        p_x = p_probs[0, draft_token].item()
        q_x = q_probs[0, draft_token].item()
        accept_prob = min(1.0, p_x / max(q_x, 1e-10))
        
        # 随机决定是否接受
        if torch.rand(1).item() < accept_prob:
            accepted_tokens.append(draft_token)
        else:
            # 被拒绝，从修正分布重新采样
            residual = torch.clamp(p_probs - q_probs, min=0)
            residual = residual / residual.sum(dim=-1, keepdim=True)
            resampled = torch.multinomial(residual, num_samples=1)
            accepted_tokens.append(resampled.item())
            break  # 停止本轮
    
    # ===== Step 4: 如果全部接受，额外采样一个 token =====
    if len(accepted_tokens) == gamma:
        final_logits = target_logits[:, -1, :]
        final_probs = F.softmax(final_logits, dim=-1)
        final_token = torch.multinomial(final_probs, num_samples=1)
        accepted_tokens.append(final_token.item())
    
    return accepted_tokens, len(accepted_tokens)

这段代码展示了投机解码的核心逻辑。注意几个关键点：

• Draft Model是自回归一个一个生成的（因为它不知道自己猜的对不对）
• Target Model只需要一次forward就能验证所有候选（KV Cache的功劳）
• Rejection Sampling保证了输出分布和纯用Target Model生成完全一样

---

四、正确性保证：Rejection Sampling数学原理

4.1 为什么要Rejection Sampling？

你可能会问：直接让Target Model选概率最高的token不就行了？为什么要搞这么复杂的采样？

答案是：为了保证输出分布和纯用Target Model生成完全一致（无偏）。

如果直接用argmax或者普通采样，Draft Model的偏见会被引入，导致输出分布偏离Target Model。而Rejection Sampling是一种数学上严格的方法，能保证最终分布等于$p(x)$。

4.2 Rejection Sampling原理

Rejection Sampling的基本思想是：用一个容易采样的提议分布$q(x)$，加上接受/拒绝机制，来采样目标分布$p(x)$。

数学上可以证明，按以下流程采样：

1. 从$q(x)$中采样一个候选$x$
2. 以概率$\min \left(1,\frac{p(x)}{M\cdot q(x)}\right)$接受$x$（M是一个常数，保证$p(x)\le M\cdot q(x)$）
3. 如果拒绝，重复步骤1

最终接受的样本服从分布$p(x)$。

在Speculative Decoding中，因为$p(x)$和$q(x)$都是概率分布（和为1），而且我们只需要验证Draft猜的那个token是否被接受，所以情况更简单。

4.3 接受率与加速比

设单个token的平均接受率为$\alpha$，那么每轮期望接受的token数是多少呢？

对于$\gamma$个候选token，全部被接受的概率是${\alpha }^{\gamma }$，此时额外生成1个token，总共$\gamma +1$个。

期望接受数的公式为：

$$\mathbb{E}[\text{accepted}]=\frac{1-{\alpha }^{\gamma +1}}{1-\alpha }$$

推导过程：这是一个几何分布的变形，第一个拒绝的位置服从几何分布，期望接受数为${\sum }_{k=0}^{\gamma }{\alpha }^k=\frac{1-{\alpha }^{\gamma +1}}{1-\alpha }$。

加速比（Speedup）的简化估算：

$$\text{Speedup}\approx \frac{\mathbb{E}[\text{accepted}]}{1+\frac{t_{draft}}{t_{target}}\cdot \gamma }$$

其中$t_{draft}$是Draft Model跑一步的时间，$t_{target}$是Target Model跑一步的时间。

直观理解：

• 接受率$\alpha$越高，加速比越大
• Draft Model越快（$t_{draft}/t_{target}$越小），加速比越大
• $\gamma$不是越大越好，太大的话Draft生成的时间会抵消收益

4.4 举个例子

假设：

• Target是70B模型，Decode一步需50ms
• Draft是7B模型，Decode一步需10ms
• 接受率$\alpha =0.7$
• 投机窗口$\gamma =5$

计算：

• 期望接受数 = $\frac{1-0.7^6}{1-0.7}\approx 2.95$ 个token/轮
• 每轮时间 = Draft生成5个 + Target验证1次 = $5\times 10+50=100$ms
• 加速比 = $2.95/(100/50)=2.95/2=1.475$ 倍

也就是说，生成速度提升了约47.5%。这还只是保守估计，实际中因为Target验证多个token的增量成本很小，加速比可能更高。

---

五、Medusa：多头解码，免Draft模型

5.1 痛点：需要额外维护一个Draft模型

经典Speculative Decoding有个工程上的麻烦：你需要同时部署两个模型。这意味着：

• 显存里要放两份模型权重
• 要选合适的Draft模型（不是随便一个小模型都行）
• 两套推理链路，维护成本高

有没有办法不用额外的Draft模型？

有！这就是Medusa（美杜莎）方案，由卡内基梅隆大学和Together AI在2024年提出。

5.2 Medusa核心思想

Medusa的想法很巧妙：在原有模型上加几个"多头"，让模型自己预测未来多个token。

传统模型:  隐状态 → LM Head → 下一个 token 的概率分布

Medusa:   隐状态 → LM Head → 第 1 个未来 token 的概率分布
                 ↘ Medusa Head 1 → 第 2 个未来 token 的概率分布
                 ↘ Medusa Head 2 → 第 3 个未来 token 的概率分布
                 ↘ ...
                 ↘ Medusa Head k → 第 k+1 个未来 token 的概率分布

每个Medusa Head负责预测"往前数第几个token"。比如Head 1预测下下个token，Head 2预测下下下个token，以此类推。

5.3 训练与验证

训练阶段：用原模型生成的语料，冻结主干，只训练Medusa Heads。因为是预测未来token，本质上是个有监督的分类任务。

推理阶段：

1. 当前step跑完后，用所有Medusa Heads同时预测未来多个token
2. 构建一棵候选树（不同位置有不同候选）
3. 下一个step时，并行验证树中的所有候选路径
4. 选择最长的可接受路径

5.4 优缺点

优点：

• 不需要额外的Draft模型，显存开销小（只多几个Head）
• 实现简洁，一套模型搞定
• 因为共享了前面层的计算，Draft生成几乎是"免费"的

缺点：

• 需要额外的训练步骤（虽然只训几个Head）
• 预测更远的token，准确率会下降
• 树形验证比线性验证更复杂

---

六、EAGLE-2：动态Draft Tree

6.1 从线性到树形

经典的Speculative Decoding是"线性"的——Draft Model一条路走到黑，生成一个token序列。但如果中间某个token猜错了，后面全白费。

EAGLE-2（Extrapolation Algorithm for Greater Language-model Efficiency）提出了一个更激进的想法：既然要猜，为什么不多猜几条路？

它构建一棵Draft Tree，同时探索多个可能的token序列：

          当前位置
             │
         ┌───┴───┐
         ▼       ▼
        候选A   候选B    (第1层：2个候选)
       /   \
      ▼     ▼
    候选C  候选D         (第2层：部分分支继续扩展)

这样，即使某条路径在第一层就被拒绝了，另一条路径可能走得更远。

6.2 树形投机的优势与代价

优势：

• 同样的验证成本下，期望接受的token数更多
• 尤其适合"容易预测"的文本段（比如固定格式、常见短语）
• 树的结构可以动态调整，自适应难度

代价：

• 实现复杂度大幅提升（树的构建、验证、结果回退）
• Draft阶段的计算量增加（要生成更多候选）
• 内存占用增加（树中每个节点都需要KV状态）

6.3 什么时候树形投机收益最大？

• 接受率高的场景：比如代码生成有固定模式，或者对话中有大量常见表达
• Target Model验证成本低：比如GPU显存带宽充足，多验证几个候选不怎么费时间
• Draft Model足够快：生成树的开销远小于Target验证的开销

---

七、收益边界：什么场景加速明显，什么场景反而变慢？

7.1 加速明显的场景

1. 小模型猜大模型，且领域匹配

   • 比如用同一系列的小模型做Draft（如Llama-2-7B猜Llama-2-70B）
   • 同架构、同训练数据的模型，接受率通常较高（60%-85%）

2. 生成内容可预测性高

   • 代码生成（语法固定）
   • 常见问答（FAQ类问题）
   • 固定格式输出（JSON、表格等）

3. 长输出场景

   • 写文章、生成代码、翻译长文
   • 输出越长，投机的累计收益越大

4. 低并发场景

   • 单用户或低并发时，GPU有空闲算力
   • Continuous Batching已经把GPU"喂饱"的高并发场景，投机的边际收益会下降

7.2 可能变慢的场景

1. Draft模型选得不好

   • 如果Draft和Target差距太大，接受率低于30%，可能得不偿失
   • Draft模型太慢，生成候选的时间超过了节省的时间

2. 输出很短

   • 比如只生成1-2个token的分类任务
   • 投机的启动成本（第一次验证）都赚不回来

3. 高并发、GPU已经打满

   • Continuous Batching已经让GPU利用率很高
   • 投机解码增加的Target验证开销可能和收益抵消

4. 内容随机性很强

   • 比如创意写作、诗歌生成
   • Draft猜不对几个token，加速比很低

7.3 关键参数选择

参数	作用	选择建议
$\gamma$（投机窗口）	每轮生成的候选数	太小收益低，太大Draft开销高。通常4-8比较常用
Draft模型大小	影响接受率和速度	通常是Target的1/10到1/3参数量。同系列优先
树的深度/宽度（树形投机）	树形探索的规模	深度2-4，宽度2-3比较常见

---

八、小结与思考

8.1 小结

Speculative Decoding是LLM推理优化中一个非常"聪明"的方向——它不是在压榨硬件，而是在算法层面改变了"一次只能生成一个token"的游戏规则。

核心要点回顾：

1. 核心思想：用快的小模型（Draft）猜，用慢的大模型（Target）并行验证，用更少的forward次数生成更多token。
2. 正确性保证：Rejection Sampling从数学上保证输出分布与Target Model完全一致。
3. 经典框架：Draft + Target 两模型架构，每轮生成$\gamma$个候选，期望接受$\frac{1-{\alpha }^{\gamma +1}}{1-\alpha }$个token。
4. 进阶方案：Medusa用额外预测头免去Draft模型；EAGLE-2用树形探索进一步提高接受数。
5. 收益边界：接受率越高、输出越长、Draft越快，加速越明显；高并发、短输出、高随机性场景收益有限。

8.2 思考问题

1. 为什么Speculative Decoding不会改变输出分布？ 如果把Rejection Sampling换成直接取Target Model的top-1，结果会有什么不同？
2. 假设Target Model是70B，Decode一步需50ms；Draft Model是7B，一步需10ms。如果接受率是0.6，$\gamma$取多少时加速比最大？（可以自己推导或写代码验证）
3. Medusa和经典的Draft+Target方案，你觉得哪个更适合在生产环境中落地？ 为什么？
4. Speculative Decoding和Continuous Batching是什么关系？ 它们能同时使用吗？会不会互相影响？
5. 如果让你设计一个自动调节$\gamma$的机制，你会根据哪些指标来动态调整投机窗口大小？

---

如果你觉得这篇文章有帮助，欢迎点赞收藏关注三连~ 下一篇我们来聊Prefix Cache前缀缓存，看看如何让重复的prompt不再重复计算。