Medusa — 让大模型自己同时猜好几个 token

待复核

是什么

Medusa 是一种让大语言模型自己同时猜下一个、下下一个、下下下个 token，再一次性验证的推理加速方法。日常类比：原本你写字是”一笔一画”，Medusa 让你左手右手一起动，同时写出后面三个字，再回头检查哪些写对了——对的留下，错的擦掉重写。

技术形式：在原模型最后一层 hidden state 上外接 K 个并行的小 head（默认 K=5，每个就是单层 MLP + 残差）。head k 负责预测位置 t+k+1 的 token。一次前向就拿到 5 个候选位置的 top-k token，拼成树形候选集，一次再前向验证哪些续写合理。

代码改动量极小（几百行 Python 集成进 vLLM / TGI），训练只要冻结 base 训 head（Vicuna-7B 单卡 A100 几小时）。

为什么重要

LLM 推理慢的根因是自回归解码——一次前向只产 1 个 token，KV cache 大半时间在等。各种加速方案里 Medusa 是工程性价比最高的一档：

比经典投机解码（Leviathan 2023）省一个模型：经典方案需要训一个对齐分布的 draft 小模型，部署两个模型对齐两套权重；Medusa 只在大模型本身长几个头，省掉 draft 模型的所有对齐工程
比 SpecInfer 简单：specinfer-2023 用一棵 token 树 + 一群小 draft 模型，Medusa 把”draft 模型”换成”多头并行预测”，一个模型搞定
2.2-3.6x 端到端加速：Vicuna 7B / 13B / 33B 上实测；不依赖 draft 模型质量
被主流引擎接住：vLLM / TensorRT-LLM 等已支持 Medusa 类多头投机，作为可选加速路径

不理解 Medusa，下面这些事都没法解释：

为什么 vLLM 配置里有 num_speculative_tokens 这一项
为什么 LLM 推理引擎评测都开始报”接受率”（acceptance rate）这个指标
为什么 2024 年之后小厂自训模型也敢上 spec decoding——门槛降到了”训几个 head”

核心要点

Medusa 的工程设计可以拆成 三件事：

多头并行预测：在最后一层 hidden state 之后，并联 K 个 Medusa head。head k 是一个单层 MLP（带残差到原 hidden state），输出 t+k+1 位置的 logits。一次前向同时拿到 K 个未来位置的概率分布。
树形注意力（tree attention）：把每个 head 的 top-k 候选拼成一棵 prefix tree。比如 head1 给出 {the, a}，head2 给出 {cat, dog}，组合出 4 条候选续写。用一个精心设计的 attention mask 让这 4 条续写在一次前向里被并行验证——不用跑 4 次。
typical acceptance：经典投机解码用严格的 rejection sampling 保证分布完全一致；Medusa 改成”按典型度阈值”接受概率合理的 token，工程更简单、效果几乎不变。

两种训练方式：

Medusa-1：冻结 base 模型，只训 head（Vicuna-7B 几小时单卡 A100），加速 ~2.18x
Medusa-2：联合训练 head + base（损失加权融合），加速 ~2.83x，但要更多算力

为什么”多个头”能省时间？关键在于自回归不是物理必然，是建模惯例——LLM 在 t 位置的 hidden state 已经”包含了未来几个 token 的概率信息”，只是原本只有一个 lm_head 去读它。Medusa 的洞察是：让多个 head 各读一份，每个 head 学习把 hidden state 解码成”未来第 k 个 token”——理论上当然不如老老实实每步重新前向准，但准到 60-80% 就够用。

实践案例

案例 1：默认配置下一次能”押”几个 token

Vicuna-7B 默认：5 个 head + 每 head top-5 + 树拼成约 64 条路径。逐步看一轮解码：

读当前 hidden state：大模型对已生成前缀做一次前向，得到位置 t 的表示
多头并行猜：5 个 head 各自给出未来位置的 top-k token，拼成候选树
一次验证前向：用树形 attention mask 并行算整棵树在大模型下的真实概率
按前缀留下：从根往叶找最长可接受前缀；平均约接受 2.5 个 token，错的截断后下一轮重猜

吞吐近似翻倍，关键是步骤 2–3 用一次（加一次验证）前向换多个 token。

案例 2：跟 vLLM 集成是什么样

from vllm import LLM, SamplingParams

llm = LLM(
    model="lmsys/vicuna-7b-v1.3",
    speculative_model="FasterDecoding/medusa-vicuna-7b-v1.3",
    num_speculative_tokens=5,
)
out = llm.generate(["写一首关于秋天的诗"], SamplingParams(max_tokens=200))

逐部分解释：

model：底座 Vicuna；speculative_model：预训好的 Medusa head 权重
num_speculative_tokens=5：对应 K=5 个 head，每轮最多多猜 5 个位置
generate(...)：业务侧仍只调一次生成 API，不用改提示词或后处理

案例 3：树形 attention 的 mask 是什么样

假设 head1 候选 {A, B}，head2 候选 {X, Y}。prompt 末尾记作 P，一次前向输入：

P, A, B, X, Y, X, Y
   └head1┘ └head2─┘ └head2─┘
            (跟随 A)  (跟随 B)

mask 让「跟随 A 的 X」只看见 P+A，「跟随 B 的 X」只看见 P+B。一次前向拿到各候选在大模型下的真实概率，再和 head 候选对照裁决。相对 specinfer-2023：Medusa 用多头代替一群 draft 小模型，部署一份权重，接受策略用 typical acceptance 换工程简单。

踩过的坑

K 不是越大越好：K=10 比 K=5 慢，因为树形候选指数爆炸、attention mask 矩阵变大。论文实测 K=5 是 sweet spot。
head 训练数据必须跟 base 推理分布一致：用 ShareGPT 训的 head 在代码任务上接受率会跌。生产里要按业务场景采数据训 head。
接受率随生成长度衰减：开头几十 token 接受率高（2.5-3 个/次），到几百 token 后掉到 1.5-2 个/次——hidden state 跟”未来 token”的相关性随距离衰减。
跟 KV cache 量化不兼容：早期版本 Medusa head 假设 hidden state 是 fp16，跟 INT8/FP8 KV cache 配合需要额外的反量化开销。现在的实现已修。
小 batch 反而变慢：batch=1 + 短 prompt 时，多头预测的固定开销（树形 attention 的 mask 构建 + 多 head 矩阵乘）大于省下来的前向次数。生产里通常配 batch ≥ 4 才打开 Medusa。
temperature 不是 0 时 typical acceptance 调参重要：阈值太严接受率掉，太松生成质量降。论文给的默认 epsilon=0.09、delta=exp(-10) 是 Vicuna 上的经验值，换模型要重新扫。

适用 vs 不适用场景

适用：

服务器侧 LLM 推理加速（vLLM / TGI / TensorRT-LLM）
模型已固定、想压低 latency 的场景
中等规模（7B-70B）模型——更小没必要，更大投机收益边际下降

不适用：

端侧极小模型（1B 以下）——head 开销占比高
极低 batch size（≤1）+ 短 prompt 短 output——投机的固定开销反而拖慢
需要严格分布对齐的科研评测（要用经典 rejection sampling）

历史小故事（可跳过）

2023 年 1 月：Leviathan 等人提出经典投机解码（需 draft 模型）
2023 年 11 月：specinfer-2023 把 draft 升级成 token 树
2024 年 1 月：Medusa 出现，把 draft 模型彻底干掉，工程上一个模型搞定
2024 年中：vLLM / TensorRT-LLM 等接入，成为生产推理的常见选项之一

学到什么

加速的关键不是新算法，是把工程复杂度降下来——Medusa 数学上更弱，但工程上简单 10 倍，所以赢
多头预测 + 树形验证 是一个值得记住的模式，背后是”用一次前向的算力换多次前向的效果”
接受率是衡量投机解码的核心指标——加速比 ≈ 平均接受 token 数 + 1（验证那一步本身也产 1 个 token）
学术 → 工程的抽税：经典投机解码的严格 rejection sampling 是一个数学性质优美但工程麻烦的设计，Medusa 用 typical acceptance 换简单——这种”放弃 5% 数学严谨换 10x 工程便利”的取舍在系统设计里反复出现

关联

specinfer-2023 —— Medusa 的直接前辈，用多个小 draft 模型；Medusa 把它简化成多头
vllm —— 工程主战场：Medusa 的 head 在 vLLM 里被调度执行
flash-attention —— Medusa 的树形 attention 依赖 FlashAttention 内核高效跑变长 mask
tensorrt-llm-2023 —— 工业推理引擎的另一个宿主

反向链接

sglang-2024 —— SGLang — 把 LLM 程序当成共享前缀的树来跑