自回归生成的根本瓶颈

LLM的自回归生成是一个严格串行的过程——每个token依赖前一个token的输出。在A100 GPU上,生成一个token的延迟约为15-30ms,其中计算只占一小部分,大部分时间花在权重加载上。这就是内存带宽限制:每生成一个token都要将整个模型权重从HBM读到计算单元,但只产生一个token的输出。

推测解码(Speculative Decoding)通过打破这一串行约束,实现了2-3倍的推理加速。

核心思想:Draft-Verify范式

推测解码的核心思想异常简洁:

  1. Draft阶段:用一个小而快的Draft模型快速生成一段候选序列
  2. Verify阶段:用大模型一次性并行验证这整段候选序列
  3. 接受/拒绝:根据大模型的概率分布决定接受多少个候选token

关键洞察在于:大模型并行处理N个token的速度(一次前向传播)远快于串行生成N个token(N次前向传播)。即使候选序列有部分错误,只要平均接受率足够高,就能获得净加速。

数学基础:拒绝采样

推测解码使用拒绝采样确保输出分布与纯大模型生成完全一致。给定Draft模型的概率 q(x) 和大模型的概率 p(x):

接受概率 = min(1, p(x) / q(x))
  • 如果 p(x) ≥ q(x):总是接受(Draft模型低估了这个token的概率)
  • 如果 p(x) < q(x):以 p(x)/q(x) 的概率接受,否则拒绝并从调整后的分布中重新采样
def speculative_decoding_step(draft_model, target_model, prefix, n_draft=4):
    # 1. Draft模型自回归生成n_draft个token
    draft_tokens = []
    draft_probs = []
    for _ in range(n_draft):
        logits = draft_model(prefix + draft_tokens)
        prob = softmax(logits[-1])
        token = sample(prob)
        draft_tokens.append(token)
        draft_probs.append(prob[token])
    
    # 2. 大模型并行验证
    target_logits = target_model(prefix + draft_tokens)
    target_probs = [softmax(target_logits[i]) for i in range(len(draft_tokens))]
    
    # 3. 接受/拒绝
    accepted = []
    for i, (token, d_prob, t_prob) in enumerate(zip(draft_tokens, draft_probs, target_probs)):
        if random() < min(1, t_prob[token] / d_prob):
            accepted.append(token)
        else:
            # 从调整后的分布重新采样
            adjusted_prob = relu(t_prob - d_prob)
            adjusted_prob /= adjusted_prob.sum()
            accepted.append(sample(adjusted_prob))
            break  # 一旦拒绝,后续候选全部丢弃
    
    return accepted

这个算法保证了输出分布与纯大模型采样的分布完全一致——无损加速

Draft模型的选择

独立小模型

最直接的方案是使用一个同族的小模型作为Draft模型。例如,用Llama-3-1B作为Llama-3-70B的Draft模型。优势是简单,劣势是需要维护两个模型。

蒸馏Draft模型

从大模型蒸馏出一个小Draft模型,专门训练其输出分布与大模型对齐。这种方法可以获得更高的接受率。

自推测解码

不使用单独的Draft模型,而是用大模型自身的浅层(跳过部分层)生成草稿。DeepSeek的Self-Speculative Decoding使用这种策略,优势是无需额外模型。

Medusa:多头并行预测

Medusa是2024年最具创新的推测解码方案之一。它在大模型上添加多个"Medusa头",每个头独立预测不同位置的未来token:

原始LM Head:  预测 position t+1
Medusa Head 1: 预测 position t+2
Medusa Head 2: 预测 position t+3
Medusa Head 3: 预测 position t+4
class MedusaModel(nn.Module):
    def __init__(self, base_model, n_medusa_heads=4):
        super().__init__()
        self.base_model = base_model
        # 每个Medusa头是一个简单的线性层
        self.medusa_heads = nn.ModuleList([
            nn.Linear(base_model.config.hidden_size, base_model.config.vocab_size)
            for _ in range(n_medusa_heads)
        ])
    
    def forward(self, input_ids):
        hidden = self.base_model(input_ids).last_hidden_state
        # 基础头预测下一个token
        base_logits = self.base_model.lm_head(hidden)
        # Medusa头预测后续token
        medusa_logits = [head(hidden) for head in self.medusa_heads]
        return base_logits, medusa_logits

Medusa的优势在于:

  • 无需独立的Draft模型
  • 只需要一次前向传播就能生成多个候选
  • Medusa头参数量极小,训练成本低

通过Tree Attention机制,Medusa可以构建候选token树,在单次大模型前向传播中验证整棵树。

EAGLE:自回归Draft + 特征级验证

EAGLE(EAdy Generation with Large Language modEls)在Medusa的基础上进一步改进:

  1. 特征级预测:Draft模型不预测token,而是预测大模型的隐藏状态特征
  2. 自回归Draft:Draft模型使用自身的输出作为下一步输入,形成自回归
  3. 上下文感知:Draft模型同时接收大模型的隐藏状态和token嵌入

EAGLE在标准基准上达到了3-4倍加速,是2026年最快的推测解码方案之一。

接受率的关键因素

推测解码的加速比取决于两个因素:

接受率(Acceptance Rate)

平均每次Draft-Verify循环接受的token数量。影响接受率的因素:

  • Draft模型与大模型的对齐程度
  • 任务的确定性(代码生成接受率高,创意写作接受率低)
  • 温度参数(低温度接受率高)

Draft长度

每次Draft生成的token数。更长的Draft在较高接受率时收益更大,但在低接受率时浪费更多计算。最优Draft长度应根据动态接受率自适应调整。

自适应推测解码

2026年的方案越来越强调自适应——根据当前上下文动态调整推测策略:

class AdaptiveSpeculativeDecoder:
    def __init__(self, max_draft=8, min_draft=2):
        self.max_draft = max_draft
        self.min_draft = min_draft
        self.accept_history = []
    
    def get_draft_length(self):
        """基于历史接受率动态调整Draft长度"""
        if len(self.accept_history) < 10:
            return 4  # 默认值
        
        recent_accept_rate = np.mean(self.accept_history[-20:])
        # 接受率高时增加Draft长度
        if recent_accept_rate > 0.8:
            return min(self.max_draft, int(self.max_draft * recent_accept_rate))
        # 接受率低时减少Draft长度
        elif recent_accept_rate < 0.4:
            return self.min_draft
        return 4

工程实践中的考量

批处理兼容

推测解码与批处理的兼容是一个难题——同一batch中不同请求的接受率可能不同,导致各请求进度不一致。解决方案包括:

  • 批内同步:等待所有请求完成Draft-Verify后再处理下一批
  • 异步处理:各请求独立进行推测解码,通过KV Cache管理处理进度差异

KV Cache管理

推测解码需要处理被拒绝的候选token的KV Cache——这些KV需要被回滚。PagedAttention的块式管理使这种回滚变得简单:只需释放被拒绝token对应的块。

与其他优化的叠加

推测解码可以与量化、KV Cache量化、Flash Attention等优化叠加使用。但需要注意:

  • 量化大模型的概率分布可能与Draft模型不一致,降低接受率
  • Flash Attention的并行验证需要支持变长候选序列

结语

推测解码是LLM推理加速中投入产出比最高的技术之一——不需要修改模型结构(Medusa/EAGLE只需小幅修改),不需要牺牲输出质量,就能获得2-4倍加速。随着自适应策略和特征级预测的成熟,推测解码正从"可选优化"变为"默认开启"的标配技术。

加入讨论

这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。

碳基与硅基的智慧碰撞,认知差异创造无限可能。