概念-Medusa

概述

Medusa 是一种高效的推测解码（Speculative Decoding）替代方案，旨在加速大语言模型（LLM）的自回归推理。与传统的推测解码需要一个小型草稿模型（Draft Model）不同，Medusa 通过在原始模型上添加多个”Medusa 头”（Medusa Heads），直接从同一个基础模型并行预测多个未来 token，从而消除了对独立草稿模型的依赖。

推测解码背景

自回归推理的瓶颈

LLM 的标准自回归推理是顺序的：逐个生成 token，每个 token 依赖前一个 token。这导致推理延迟与输出长度成正比，且 GPU 计算单元在大部分时间处于空闲状态（内存带宽瓶颈）。

推测解码（Speculative Decoding）原理

推测解码的核心思想是”用草稿模型猜，用目标模型验”：

1. 草稿生成（Draft）：使用小模型快速生成 $k$ 个候选 token。
2. 并行验证（Verify）：目标模型一次前向传播验证所有 $k$ 个候选 token。
3. 接受与拒绝（Accept/Reject）：基于 token 匹配概率接受或拒绝候选 token。

问题：需要维护两个模型，草稿模型的选择和部署增加了复杂性。

Medusa 架构

核心思想

Medusa 不要额外的草稿模型，而是在原始 LLM 的最后一层隐藏状态之上，添加多个预测头（Medusa Heads），每个头预测一个未来位置的 token。

架构细节

假设原始 LLM 在位置 $t$ 的隐藏状态为 $h_t$：

• Medusa Head 1：预测位置 $t+1$ 的 token（相当于原始 LM Head）
• Medusa Head 2：预测位置 $t+2$ 的 token
• Medusa Head 3：预测位置 $t+3$ 的 token
• …
• Medusa Head k：预测位置 $t+k$ 的 token

每个 Medusa Head 是一个简单的线性层 + Softmax，输出维度为词汇表大小 $V$：

$P_i(\cdot |h_t)=\text{Softmax}(W_i\cdot h_t+b_i),i=1,2,...,k$

训练 Medusa Heads

Medusa Heads 的训练非常高效：

1. 冻结原始 LLM 的所有参数。
2. 使用标准文本语料，对每个位置 $t$，以位置 $t+i$ 的实际 token 作为 Head $i$ 的目标。
3. 仅训练 Medusa Heads 的权重（通常只需几小时）。

推理流程

1. 原始 LLM 前向传播一次，得到隐藏状态 $h_t$。
2. 所有 Medusa Heads 同时预测 $t+1,t+2,...,t+k$ 的候选 token。
3. 使用树状注意力（Tree Attention）并行验证所有候选序列。
4. 接受匹配的 token，拒绝不匹配的，继续下一轮。

Lookahead Decoding（前瞻解码）

Lookahead Decoding 是另一种无草稿模型的推测解码方案，由 lm-sys（FastChat 团队）提出。

核心思想

Lookahead Decoding 利用雅可比迭代（Jacobi Iteration）和 n-gram 匹配来生成候选 token：

1. 使用随机 token 初始化一个长度为 $k$ 的序列。
2. 并行计算每个位置的 token 分布。
3. 通过雅可比迭代逐步收敛。
4. 使用 n-gram 匹配决定接受哪些 token。

与 Medusa 的关键区别

2026年采用现状

2026 年，Medusa 已成为生产环境中最受欢迎的推测解码方案之一，因为它在加速比和部署复杂度之间取得了最佳平衡。vLLM 和 SGLang 的原生支持使得 Medusa 几乎成为默认配置。

技术深度对比

加速比影响因素

• 模型大小：大模型（70B+）的加速比通常更高（2.0x - 3.0x），因为自回归瓶颈更严重。
• 任务类型：代码生成、翻译等结构化任务加速比高，创意写作加速比低。
• Medusa Head 数量：通常 3-5 个 Head 最优，过多会增加验证开销。
• 批次大小：大 batch 场景下加速比下降，因为 GPU 利用率已经较高。

Medusa 的局限性

1. 需要训练：虽然训练成本低，但仍需额外步骤。
2. 输出分布偏移：Medusa Heads 的预测可能与原始模型略有偏差（可通过阈值控制）。
3. 长序列退化：随着序列变长，Medusa Heads 的预测精度下降。
4. 模型特定：Medusa Heads 与特定模型绑定，切换模型需重新训练。

相关笔记

• 推测解码（深挖）：推测解码的深度技术分析
• 第20章 · 模型推理部署：模型推理优化与部署
• 概念-vLLM：vLLM 推理引擎，PagedAttention 与连续批处理
• 概念-模型服务：模型服务基础设施（Ray Serve、Triton 等）
• 第10章 · LLM 原理与架构演进：LLM 自回归推理原理

来源

• Medusa 论文: “Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads” (2024): https://arxiv.org/abs/2401.10774
• Lookahead Decoding 论文: “Break the Sequential Dependency of LLM Inference Using Lookahead Decoding” (2024): https://arxiv.org/abs/2402.02057
• Medusa GitHub 仓库: https://github.com/FasterDecoding/Medusa
• vLLM Speculative Decoding 文档: https://docs.vllm.ai/en/latest/features/spec_decode.html
• SGLang Medusa 集成: https://github.com/sgl-project/sglang
• 推测解码综述: https://arxiv.org/abs/2302.01318