老字号后厨有位手艺顶尖的主厨，却有个怪规矩：一道菜做完，才肯开下一道火。客人点十道菜，他就得在灶台前忙十轮。AI 大模型生成文字时也是这副脾气——每生成一个 token 都要把整套模型从头跑一遍，号称“并行之王”，实际上是一句一句“挤”出来的。 后来店里来了几个跟了主厨多年的帮厨。主厨炒菜的锅气一冒出来，帮厨们就同时报出三道菜名：下一道、下下道、下下下道。主厨尝一口确认没错，三道菜一起端上桌——原本要开三次火的功夫，一趟就完成了。这就是 Medusa 的朴素直觉：让模型自己长出几个“预测头”同时猜后面的 token，再用树状验证一次性验收，把串行几步压成一步。 这篇故事会带你从“为什么慢”讲到“怎么变快”，从预测头怎么训练、树状验证怎么工作，一直聊到工程上怎么落地。看完你会发现，加速大模型推理这件事，原来可以像给主厨配帮厨一样简单。

岗位关联

Medusa 是少数不用换模型、不用加显存就能白捡 2-3 倍速度的推理加速技术。掌握它，你就能在生产环境里让大模型“跑得更快、等得更短”。

目录
  1. 可视化演示
  2. 解码瓶颈
  3. Medusa原理
  4. 预测头训练
  5. 树状验证
  6. 对比投机解码
  7. 工程落地
  8. AI 应用场景
  9. Python 演示
  10. 常见误区
  11. 面试题

可视化演示

演示快速排序算法的核心分区操作：选择基准元素，通过左右指针比较交换，将数组划分为小于和大于基准的两部分。

本节配有交互动效（进入视口后自动加载，约 3KB）。

选择基准 选取数组末尾元素作为基准，左右指针从两端向中间扫描 pivot = arr[high] 当前展示基准元素的选择过程

解码瓶颈

想象一座老字号后厨：主厨手艺了得，但规矩是“一道菜做完才能开下一道火”。客人点了十道菜，他就得在灶台前干等十轮。每道菜从洗锅到出锅，主厨都要把整套流程走一遍——这正是大模型自回归解码的写照：给定已经写出来的文字，模型一次前向计算只能产生下一个 token（可以理解成一个字或一个词）。要生成一百个字的回答，就得让模型“开火”一百次。

真正拖慢速度的，往往不是计算本身，而是“搬食材”的功夫。大模型的参数（权重）动辄几十上百 GB，每次生成一个 token 都要把这些权重从显存里完整读一遍。读权重的时间远大于算那一小步的时间——这就是内存带宽瓶颈。GPU 的计算单元常常闲着，等数据从显存“搬”过来。

所以总耗时可以写成一次前向的时间乘以串行步数：

既然每道菜都要花同样长的“开工时间”，那最直接的提速办法，就是让主厨一锅同时做出好几道菜。Medusa 要做的，正是这件事。

Ttotal=N×tstep\displaystyle T_{\text{total}} = N \times t_{\text{step}}
总耗时 = 生成 token 数 × 单步延迟。N 越大越慢，t_step 由显存带宽主导。Medusa 的思路就是从 N 下手——让一步能干出好几步的活。
  1. T_total总耗时，客人从点完菜到全部上齐所花的时间。
    TtotalT_{\text{total}}
  2. N要生成的 token 总数，也就是菜的盘数。
    NN
  3. t_step单次前向延迟，主厨做一道菜的完整时间。
    tstept_{\text{step}}

Medusa原理

给主厨配三个跟了十几年的帮厨。主厨炒菜的锅气（模型最后一步的隐藏状态）一冒出来，帮厨们立刻各报一道菜名：下一道、下下道、下下下道。没让主厨多开一次火，三道菜的方案已经摆在案板上。

Medusa 的做法也是如此：在主干模型最后一个隐藏层后面，挂上 K 个轻量“预测头”。每个预测头是一个小小的 MLP（多层感知机），输入同一个隐藏状态 h，各自输出未来第 k 个 token 的概率分布。主模型只需前向一次拿到 h，K 个头就能并行吐出 K 个草稿 token。

这些预测头非常轻，结构参考了 Llama 的前馈层：先做一层归一化，再接两层线性变换，中间用 SiLU 激活函数。增加的参数相比主干模型几乎可以忽略不计。数学上，第 k 个头的输出是：

帮厨是主厨亲手带出来的，报的菜名自然贴合主厨口味。下一节就讲这些帮厨是怎么“带出来”的。

pk=softmax(Wk h+bk)\displaystyle p_k = \text{softmax}(W_k\, h + b_k)
第 k 个预测头对隐藏状态 h 做一次线性变换加 softmax，得到“未来第 k+1 个 token 是谁”的概率分布。这就是帮厨们同时报菜名的过程。
  1. h最后一个 token 的隐藏状态，相当于主厨出锅时的那股“锅气”，所有预测头都看它。
    hh
  2. W_k 和 b_k第 k 个预测头自己的权重和偏置，帮厨们各自的小本本。
    Wk, bkW_k,\ b_k
  3. softmax把一堆得分变成概率，让“报菜名”变成有把握的“打分排名”。
    softmax(⋅)\text{softmax}(\cdot)

预测头训练

帮厨怎么练出来？不是换掉主厨，而是主厨照常炒菜，帮厨在旁边观摩记录：主厨在每种情境下会做哪道菜。主厨的心法（参数）完全不动，帮厨只更新自己的小本本。

训练 Medusa 预测头时，主干模型被完全冻结。训练数据有两种常见来源：一是通用文本语料，让预测头学会“正常人接着会说什么”；二是自蒸馏——用主模型自己生成一批文本，拿这些样本来教预测头，让它们模仿主厨自己的口味偏好，连采样时的随机小习惯也一并学去。每个预测头 k 的损失函数是：

训练时还有个实用技巧：随机丢弃部分预测头位置（stochastic dropping），逼着每个头独立学会在不同上下文下工作，避免头与头之间过度依赖。因为只训练几个小 MLP，通常只需要一小部分数据和几小时 GPU 时间，训练成本比微调主模型低两三个数量级。

Lk=−log⁡P(xt+k∣ht;θk)\displaystyle L_k = -\log P(x_{t+k} \mid h_t; \theta_k)
让第 k 个预测头在看到当前位置隐藏状态 h_t 时，尽量给“真实的未来第 k+1 个 token”打高分。这就是帮厨对着标准答案改错题。
  1. L_k第 k 个头的训练损失，帮厨做错题时被扣的分。
    LkL_k
  2. x_{t+k}当前位置往后数第 k+1 个真实 token，也就是标准答案。
    xt+kx_{t+k}
  3. h_t当前位置的隐藏状态，帮厨观察到的“当前锅气”。
    hth_t
  4. θ_k第 k 个预测头的参数，帮厨自己小本本上的内容。
    θk\theta_k

树状验证

帮厨报了菜名，主厨总不能全信。传统的投机解码，主厨得把每道菜都重新做一遍来确认；Medusa 的主厨有绝活：一勺伸进几盘菜，一次尝完——这就是树状验证。

预测头不止猜一条路。比如第一个头猜“今天天气不错”，第二个头可能猜“今天天气不好”，两个头各带分支继续往下猜，就组成一棵“候选树”。困难在于：怎么一次前向验证整棵树？Medusa 构造了一个树状注意力掩码（tree attention mask）：每个候选节点只能看到自己路径上的祖先 token，看不到兄弟分支的内容。这样一来，所有候选节点可以在同一次前向里并行算出各自的概率，和草稿一一比对。

验证时从树根出发，找“连续匹配”的最长前缀。匹配到哪就接受哪一段，一旦遇到第一个错误节点，就从那个位置重新开始下一轮推测。用公式表达，一次验证的接受长度是：

主厨一勺尝三盘，省掉了三趟来回。但留心对比就会发现：这和“请个二厨先做一版”的投机解码不太一样。下一节专门掰扯这个区别。

A=max⁡{m:∀ i≤m, x^t+i=xt+i}\displaystyle A = \max\{m : \forall\, i \le m,\ \hat{x}_{t+i} = x_{t+i}\}
A 是能连续匹配上的最大长度。前 m 个草稿 token 都和真实解码一致，就一次接受 m 个；遇到第一个不一致就回退。树越宽，一次吃下的 token 越多，但验证的计算量也越大，所以 K 一般取 3 到 5。
  1. A本次验证接受的长度，主厨一次确认了几道菜。
    AA
  2. m连续匹配的深度，从树根往下数第几层还都对得上。
    mm
  3. x̂_{t+i} 与 x_{t+i}草稿 token 与真实 token，帮厨报的菜名和主厨实际做的菜。
    x^t+i=xt+i\hat{x}_{t+i} = x_{t+i}

对比投机解码

另一家店用的方法是投机解码：请隔壁二厨按自己的口味先做一版菜，主厨再复刻确认。问题有二：二厨手艺和主厨不一样，做十道常被驳回六七道——传菜确认的过程白忙活；二厨还占一个灶台，即额外加载一个小模型，白吃显存。

传统投机解码里，草稿由独立的小模型生成，接受概率有一个经典公式：

Medusa 本质上可以理解为投机解码的一个“自我猜测”变体：它把草稿的职责从独立小模型换成了轻量预测头。预测头和主干共享 embedding 和隐藏状态，不占额外模型显存；额外开销只有几个小 MLP 的前向，几乎可以忽略。自家帮厨报菜名，主厨点头的概率自然高。

α=min⁡(1, ptarget(x)pdraft(x))\displaystyle \alpha = \min\left(1,\ \frac{p_{\text{target}}(x)}{p_{\text{draft}}(x)}\right)
α 是草稿被主模型接受的概率。当草稿模型的判断（p_draft）和目标模型（p_target）差距越大，拒绝率越高，加速越少。Medusa 因为草稿来自目标模型自己的预测头，p_draft 和 p_target 几乎一致，所以 α 接近 1。
  1. α接受概率，主厨点头放行的比例。
    α\alpha
  2. p_target(x)目标模型给 token x 的概率，主厨的真实口味。
    ptarget(x)p_{\text{target}}(x)
  3. p_draft(x)草稿模型给 token x 的概率，二厨或帮厨的口味。
    pdraft(x)p_{\text{draft}}(x)

工程落地

帮厨练成之后，正式编入后厨排班。不用新灶台，不用换主厨，老厨房直接改造——这正是 Medusa 工程落地最吸引人的地方：单卡即可部署，不需要额外加载草稿模型，因此特别适合显存紧张的生产环境。

实践中，Medusa 已经集成进 Hugging Face Transformers 生态，也能配合 vLLM、TensorRT-LLM 等主流推理框架使用。论文中的实测数据（以 Llama-2-7B 等模型为例）显示，加速比普遍在 2.2 到 3.6 倍之间，常见水平是 2 到 3 倍；后续的 Medusa-2 加入少量微调后还能更高。

加速比可以粗略表达为：

落地时还有几个细节要留意：预测头数量 K 需要调优，K 太大候选树会指数膨胀；推理时的采样参数（temperature、top-p）最好和训练时保持一致，否则帮厨的“口味”会对不上主厨；上线前要在目标 GPU 上预热做 benchmark，用真实的请求长度测加速比。后厨改造完毕，同样的客人，等菜时间缩短一半以上。

Speedup=NN/Kˉ+c\displaystyle \text{Speedup} = \frac{N}{N / \bar{K} + c}
加速比 = 串行步数 ÷ 实际需要的验证次数。平均每次验证能接受 K̄ 个 token，验证次数就从 N 降到 N/K̄，c 是候选树构造等固定开销。K̄ 越大、c 越小，加速越明显。
  1. Speedup加速倍数，同样一桌菜快了多少倍。
    Speedup\text{Speedup}
  2. K̄平均每次验证接受的 token 个数，帮厨平均一次报对几道菜。
    Kˉ\bar{K}
  3. c候选树构建和分支管理的固定开销，传菜喊话的功夫。
    cc

AI 应用场景

聊天机器人与 Agent 低延迟响应

阶段：推理优化 / 服务部署

解决的问题：传统自回归解码下，用户问一句复杂问题，模型要串行生成几百个 token，首字延迟和整段返回时间都长。Medusa 把串行步数减少一半以上，用户等待时间从“盯着转圈”变成“话刚到嘴边”。

具体案例：一个客服 Agent 接入 Llama-2-7B，原生成 200 token 需要约 2 秒。部署 Medusa 后同一段回复约 0.7 秒返回，用户流失率明显下降，且回答内容与之前完全一致（贪心解码下不变）。

离线批处理与大规模评测提速

阶段：离线推理 / 成本优化

解决的问题：批量跑摘要、代码生成、评测集打分时，GPU 资源被串行解码白白占用。同样的单卡，用 Medusa 能在相同时间内产出 2 到 3 倍的 token，直接折算成 GPU 成本下降。

具体案例：某团队每天用 LLaMA-2-13B 跑一万条代码生成任务，单卡耗时 14 小时。接入 Medusa 后降到 5 小时左右，省下的时间用来跑更多测试用例。

单卡私有化部署

阶段：模型部署 / 边缘推理

解决的问题：传统投机解码需要额外加载一个小模型，显存小的机器根本塞不下。Medusa 不引入任何额外模型，只多了几个轻量预测头，单卡甚至边缘设备都能跑，且天然适配私有化交付场景。

具体案例：一家政企客户要求模型完全内网部署，显卡只有一张 24G 的 A10。用 Medusa 改造后，模型照样装得下，回答速度还比原来快了一倍多。

Python 演示

python92 行
import torchimport torch.nn as nn V, D, K, LEN = 16, 32, 3, 20  # vocab, hidden, heads, seq lentorch.manual_seed(0) class TinyLM(nn.Module):    def __init__(self):        super().__init__()        self.embed = nn.Embedding(V, D)        self.rnn = nn.LSTM(D, D, batch_first=True)        self.out = nn.Linear(D, V)    def logits(self, x):        h, _ = self.rnn(self.embed(x))        return self.out(h)    def hidden(self, x):        h, _ = self.rnn(self.embed(x))        return h class MedusaHeads(nn.Module):    def __init__(self):        super().__init__()        self.heads = nn.ModuleList([nn.Sequential(            nn.Linear(D, D), nn.SiLU(), nn.Linear(D, V)) for _ in range(K)])    def forward(self, h):        return [head(h[:, -1]) for head in self.heads] data = torch.stack([torch.tensor([(s + i) % V for i in range(LEN)])                    for s in torch.randint(0, V, (800,))])model = TinyLM()heads = MedusaHeads() def fit(m, opt, mode):    opt.zero_grad()    if mode == 'lm':        lg = m.logits(data)[:, :-1]        loss = nn.functional.cross_entropy(lg.reshape(-1, V), data[:, 1:].reshape(-1))    else:        with torch.no_grad():            h = m.hidden(data)        loss = None        for k, head in enumerate(heads.heads):            cur = nn.functional.cross_entropy(                head(h[:, :-k-1]).reshape(-1, V), data[:, k+1:].reshape(-1))            loss = cur if loss is None else loss + cur    loss.backward()    opt.step() opt1 = torch.optim.Adam(model.parameters(), lr=1e-3)for _ in range(60):    fit(model, opt1, 'lm') opt2 = torch.optim.Adam(heads.parameters(), lr=1e-3)for _ in range(100):    fit(model, opt2, 'head') def decode(prompt, n=12):    with torch.no_grad():        for _ in range(n):            nxt = model.logits(prompt)[:, -1].argmax(-1, keepdim=True)            prompt = torch.cat([prompt, nxt], 1)    return prompt def medusa_decode(prompt, n=12):    steps = 0    with torch.no_grad():        while prompt.shape[1] < 3 + n:            steps += 1            h = model.hidden(prompt)[:, -1]            drafts = [head(h).argmax(-1).item() for head in heads.heads]            truth = model.logits(prompt)[:, -1].argmax(-1).item()            if drafts[0] != truth:                prompt = torch.cat([prompt, torch.tensor([[truth]])], 1)                continue            accept = [truth]            for d in drafts[1:]:                cand = torch.cat([prompt, torch.tensor([accept])], 1)                nxt = model.logits(cand)[:, -1].argmax(-1).item()                if d == nxt:                    accept.append(d)                else:                    break            prompt = torch.cat([prompt, torch.tensor([accept])], 1)    return prompt, steps if __name__ == "__main__":    p = torch.randint(0, V, (1, 3))    a = decode(p)    b, steps = medusa_decode(p)    b = b[:, :a.shape[1]]    print("samples equal:", torch.equal(a, b))    print("autoregressive fwds: 12, medusa fwds:", steps)
目的：用一个小型自回归模型模拟“大模型”，训练 3 个 Medusa 预测头，对比普通解码与 Medusa 多头解码的步数与输出一致性。
  1. self.heads = nn.ModuleList([...])创建三个帮厨：3 个轻量 MLP 预测头。
  2. head(h[:, :-k-1]).reshape(-1, V), data[:, k+1:]第 k 个头学习“看当前锅气、猜后面第 k+1 道菜”。
  3. h = model.hidden(prompt)[:, -1]主厨炒完当前这道菜，把锅气递给帮厨们。
  4. drafts = [head(h).argmax(-1).item() for head in heads.heads]三个头同时报菜名，各猜一个 token。
  5. truth = model.logits(prompt)[:, -1].argmax(-1).item()主厨亲自确认第一道菜到底是什么。
  6. if drafts[0] != truth:第一道就猜错，当场回退重来。
  7. for d in drafts[1:]:第一道对了，继续往后续的菜逐一验证。
  8. prompt = torch.cat([prompt, torch.tensor([accept])], 1)连续正确的多个 token 一次端上桌，这就是省下的串行步骤。
  9. torch.equal(a, b)验证 Medusa 的输出和普通解码完全一致，加速不牺牲质量。

常见误区

误区“Medusa 和投机解码一样，都是先跑一个小模型生成草稿，再让大模型验证。所以我也得准备一个小模型来用 Medusa。”

正解Medusa 虽然属于推测解码家族，但它不需要任何独立小模型。草稿由大模型自己的预测头生成，预测头只是几个极轻的 MLP 层（几 MB 级别），不像投机解码要加载一个完整的草稿模型（几 GB 级别）。它更像是“自产自销”：主厨自己带帮厨，而不是请隔壁二厨。

为何常见：因为两者共享“草稿 + 验证”的框架，很多人想当然认为草稿必须来自另一个模型。

误区“既然一个头能多猜一个 token，那我装 8 个、10 个头，速度不就翻好几倍？”

正解预测头越多，候选树的分支组合指数增长，验证一次前向的计算量也水涨船高；更关键的是，越往后的预测头准确率越低，接受率快速衰减。实践中最优 K 通常在 3 到 5，继续增加往往收益很小甚至变慢。就像帮厨再多，主厨一次也尝不了十几盘菜。

为何常见：直觉上觉得“并行数量 = 加速倍数”，忽略了验证成本和接受率这两个约束。

误区“在模型上加了东西，输出肯定和原来不一样了，质量会打折扣。”

正解Medusa 的两个设计保证了输出不变：一是主干模型完全冻结，预测头只是“外部挂件”；二是所有草稿都必须经过主模型的树状验证，验证时以主模型自己的 logits 为准，贪心解码下输出与原始模型严格一致。真正改变输出的，是采样参数变了，而不是 Medusa 本身。

为何常见：人们习惯把“加模块”和“改模型”划等号，忽略了“只做草稿、不做决策”的验证机制。

误区“推理加速肯定等于省显存，我用了 Medusa 就能塞进更小的显卡。”

正解Medusa 的主要收益是降低延迟——减少串行前向次数，而不是减少显存。它相比投机解码确实省下了草稿模型的显存，但相对“裸模型”基本不省显存。如果目标是省显存，应该考虑量化、剪枝或 KV Cache 优化，而不是 Medusa。

为何常见：加速和优化经常被混为一谈，大家容易把“提速”带来的间接效果当成直接收益。

面试题

Q1. Medusa 的预测头加在主干模型的哪里？

单选题 · easy

  • A.输入词嵌入之后
  • B.中间某一层之后
  • C.最后隐藏层之后← 正确答案
  • D.输出层之前

答案：C。Medusa 预测头加在最后一个隐藏层之后，输入是最后一个 token 的隐藏状态 h。加在别的位置拿不到全局语义，输出层之前已经算 logits 了。就像帮厨要看主厨出锅时的“锅气”，而不是看食材刚下锅的样子。

Q2. Medusa 与传统投机解码最核心的区别是什么？

单选题 · easy

  • A.使用更小的草稿模型
  • B.没有独立草稿模型← 正确答案
  • C.需要重新训练主干
  • D.只能用于贪心生成

答案：B。传统投机解码需要一个独立小模型当草稿，Medusa 直接用主模型自己的预测头做草稿，没有独立草稿模型。A 是投机解码的特点，C 说反了（主干被冻结），D 更是误解。

Q3. 训练 Medusa 预测头时，主干模型如何参与？

单选题 · mid

  • A.和预测头一起全量微调
  • B.完全冻结不更新参数← 正确答案
  • C.只更新词嵌入层参数
  • D.需要用强化学习训练

答案：B。训练时主干模型完全冻结，只更新预测头的参数。这也是 Medusa 训练成本低、且不改变原始模型行为的关键。全量微调或改词嵌入都会破坏主干分布，强化学习更是没必要。

Q4. Medusa 的树状验证一次前向能做什么？

单选题 · mid

  • A.同时验证所有候选分支← 正确答案
  • B.顺次验证每一条候选链
  • C.只验证第一个预测头
  • D.重新生成新的候选集合

答案：A。树状验证利用树注意力掩码，让所有候选节点共享前缀、互不干扰，在同一次前向里全部算出 logits，再批量比对。顺次验证就不叫加速了。

Q5. 在贪心解码场景下，Medusa 的输出与原始模型相比？

单选题 · hard

  • A.输出质量会有小幅下降
  • B.输出结果与原始模型一致← 正确答案
  • C.跳跃式生成内容更连贯
  • D.输出长度一定更短

答案：B。Medusa 的预测头只负责提供草稿，最终是否接受由主模型自己的 logits 决定，且只在链式匹配成功时才跳跃。因此贪心解码下输出与原始模型严格一致。其余选项都是对加速机理的误解。

Q6. Medusa 相对投机解码的优势有哪些？

多选题 · mid

  • A.无需加载额外的草稿模型← 正确答案
  • B.草稿分布与目标更一致← 正确答案
  • C.拒绝时回退代价更小← 正确答案
  • D.必须多卡并行才能运行

答案：ABC。Medusa 的草稿来自主模型自己的头，所以不需要额外草稿模型、分布一致、接受率高、回退代价小。D 恰恰相反，Medusa 单卡就能跑，这也是它的落地优势。

Q7. 多个预测头产生的候选组合靠什么一次验证？

单选题 · mid

  • A.动态规划的链式解码
  • B.树注意力掩码配合前向← 正确答案
  • C.贪心搜索逐条链路验证
  • D.额外的验证小模型打分

答案：B。所有候选组成一棵树，树注意力掩码保证每个候选节点只看到自己的祖先 token，一次前向验证整棵树。逐条验证会退回串行，额外验证模型则引入了投机解码的老问题。

Q8. Medusa 接受率高的根本原因是什么？

单选题 · hard

  • A.草稿由主模型自己的头生成← 正确答案
  • B.草稿模型训练数据更多
  • C.验证时使用了更强模型
  • D.预测头数越多接受越高

答案：A。草稿分布和目标分布几乎一致，所以接受率接近 1。这是 Medusa 相比投机解码最本质的优势。训练数据多少、验证模型强弱都不是原因，预测头数量更是与接受率无单调关系。

Q9. Medusa 工程落地时需要考虑哪些因素？

多选题 · hard

  • A.预测头数量 K 的取舍← 正确答案
  • B.采样参数与训练一致← 正确答案
  • C.树宽与显存的开销权衡← 正确答案
  • D.必须替换整个推理框架

答案：ABC。K 要调优，太大候选树膨胀；采样参数和训练不一致会导致草稿“口味”对不上；树宽增加会提升验证计算量。Medusa 已集成进主流框架，不需要替换整个推理框架，D 错误。

Q10. Medusa 加速的根源是什么？

单选题 · easy

  • A.提高单次矩阵乘法的速度
  • B.减少串行解码前向次数← 正确答案
  • C.剪除多余的注意力头
  • D.用整数运算替代浮点

答案：B。Medusa 通过多头预测加树状验证，让一次前向产生并验收多个 token，把 N 次串行前向压缩到 N/K 次左右。它不改变单次前向本身的速度，也不涉及剪枝和量化。