先说一个可能很多做推理部署的人都踩过的坑为了让一个 7B 模型在单卡上跑得更快我去接了外挂式的投机采样结果模型性能没提多少反而多养了一个草稿模型。那个草稿模型要单独蒸馏、单独调对齐、单独占显存还得在服务端多维护一套调度逻辑。后来我看到 MEDUSA 这套多头树状注意力投机解码方案才意识到另一种玩法不需要独立草稿模型让主干模型自己长出几个“解码头”用树状方式并行验证多条候选路径一次前向就能无损拿到 2.83 倍左右推理加速。这篇文章把我从原理、训练、推理到调参踩坑的全过程整理出来适合已经在做 LLM 推理加速、或者正准备改造服务端推理管线的朋友参考。1. 草稿模型的坑我踩了三个月后为什么转向多头方案1.1 独立草稿模型的维护成本比想象中高传统投机解码的思路是这样的拿一个很小的模型当“草稿员”让它先快速生成接下来若干 token 的候选然后大模型一次性审核这些候选能对上多少就接受多少。理论上如果草稿员和大模型行为足够接近大模型每生成一个 token 的开销可以摊薄到几个 token 上。但问题也藏在这个“足够接近”里。小模型不是天生就会模仿大模型的你得拿大量大模型生成的数据去蒸馏它蒸馏完还得评估它和大模型的 token 级重合率。更麻烦的是主干模型一旦换了量化方式或者你改了采样参数草稿员和主干的“默契”可能就变了。我在公司里接过一个内部推理服务草稿模型单独占了差不多 2GB 显存虽然看起来不多但在线上并发一高这 2GB 直接挤占了 KV cache 的位置批量推理的吞吐掉了一大截。另一个隐性成本是部署复杂度。草稿模型和主干模型如果放在同一个服务进程里要同时维护两份 tokenizer、两份模型实例、两份 batch scheduler如果分开部署又多了草稿模型的请求延迟。我当时排查过一个线上问题草稿模型服务本身偶尔抖动导致投机解码整体变慢最终回滚回普通自回归解码了。折腾一圈性能收益被运维成本吃掉了大半。1.2 MEDUSA 的核心思路不请外援让主干自己出多份答卷MEDUSA 换了个角度。它不引入第二个模型而是在主干模型最后一层隐藏状态后面接若干个小头部每个头部负责预测“当前位置之后第 k 个 token”的候选。说得直观一点主干解码到第 i 步时头部 1 预测第 i1 步的 top 候选头部 2 预测第 i2 步的 top 候选以此类推。等到真正验证时仍然由主干模型自己对这些候选做一次前向打分。这个思路的关键在于“无独立草稿模型”。训练这些头的时候主干权重完全冻结只更新轻量 head推理的时候同一个模型既当“出题人”又当“阅卷人”。所谓无损指的是最终接受哪个 token 的判断标准始终来自主干自身输出不再经过一个近似分布的中间人。只要头部预测的路径和主干真实生成结果不一致就会被拒绝并回退到主干重新选择所以不会引入额外的分布偏移。我第一次看到这个方案时脑子里蹦出来的类比是这样的传统投机解码像是让实习生先写一份草稿领导再审MEDUSA 则是让领导自己同时列几版提纲然后挑一版往下写。前者要培训实习生后者只需要提高领导列提纲的速度。对推理服务来说省掉的不仅是显存还有训练、部署、监控的整条链。2. 多头树状注意力一次前向验证一整片假设2.1 多个解码头是怎么预测“未来 token”的每个 Medusa head 的构造不复杂我在实现里用的是两层 MLP 加 RMSNorm输入是主干最后一层输出的 hidden state输出维度是词表大小。假设主干当前已经解码到位置 i此时最后一层输出的 hidden state 是 h_i那么第 k 个 head 会计算head_k_logits head_k(h_i) # shape: [batch, vocab_size] top_candidates topk(head_k_logits, ktop_k_per_head)这里有一个容易误解的地方head k 并不是拿“自己预测的 token”作为输入再去预测下一步因为那会变成串行解码性能没有优势。它只是在复用同一个 h_i一次算出多个未来位置的候选分布。正因为多个 head 可以并行计算候选路径的数量才可能撑起一棵足够宽的树。训练时head k 的目标是在真实 token 序列上做 teacher forcing。给定序列的第 i 个 token 位置第 k 个 head 的监督标签是第 ik 个 token。损失就是模型在目标 token 上的交叉熵。我实验时会对不同 k 的 head 按 k 值加权比如保留更多损失给更近的头因为预测近未来的难度低可信度也高树的前几层如果选错后面整棵树的接受率都会崩。候选树的生成方式也很朴素。第一个 head 先给出 top-N 候选作为树的第一层节点每个第一层节点继续用第二个 head 的 top-M 候选扩展出第二层节点到第三层、第四层同理。若不做裁剪节点数会指数爆炸所以我通常限制树的层数为 4节点总数上限设为 60 左右并用路径的累积 log 概率做 beam 剪枝。剪枝逻辑简单有效如果一条路径前三步的累积概率已经很低那它几乎不可能被主干接受没必要再往下展开。2.2 树掩码的一次前向所有候选同时验证树搭好之后下一步是把这些候选节点一次性喂给主干做前向。常规做法是把所有候选节点当作一个 batch 里的若干条序列和原始上下文拼在一起同时做注意力计算。但这里有一个关键细节树里的每个节点只能看到自己的祖先节点不能看到同层或别的路径上的兄弟节点。如果直接按普通 batch 并行处理注意力会把所有节点等同看待候选路径之间就会“互相剧透”验证结果完全失真。所以必须构建一棵树掩码。实现上的做法是为每个候选节点记录它的祖先位置然后生成一个全局 attention mask。假设原始上下文长度为 C候选节点一共 N 个那么前向时的 token 序列长度是 CN注意力矩阵是一个 (CN) x (CN) 的 mask。原始上下文部分保持因果 mask新加入的候选节点部分只有祖先关系和自身位置是 True其余全是 False。在我自己的工程实现里构建树掩码的核心函数大概长这样def build_tree_attention_mask(context_len, node_ancestors): n_nodes len(node_ancestors) seq_len context_len n_nodes mask torch.zeros(seq_len, seq_len, dtypetorch.bool) mask[:, :context_len] True # 上下文部分互相可见 mask[:context_len, :context_len] torch.tril(torch.ones(context_len, context_len, dtypetorch.bool)) for node_id, ancestors in enumerate(node_ancestors): abs_id context_len node_id mask[abs_id, :context_len] True for anc in ancestors: mask[abs_id, context_len anc] True mask[abs_id, abs_id] True return mask生成这个 mask 之后把整棵树的节点输入主干一次前向就能得到每个节点位置的原始 logits。再拿这些 logits 和候选 token 逐个比对从根路径开始找最长的一段连续匹配前缀这段前缀就是可以接受的新 token 数。如果有节点匹配不上就从最后一个不匹配位置回退重新采样。这里最反直觉的一点是验证速度并不和候选节点数成线性关系。一个 60 个节点的候选树前向耗时可能只相当于普通解码 2 到 3 个 token 的耗时因为它本质上只是一次大的 batch 前向。而传统投机解码里草稿模型的串行生成仍然是一个一个来的。这也是 MEDUSA 加速比能做到比较夸张的主要原因。3. 2.83 倍加速是怎么测出来的3.1 加速比计算公式和实验配置很多人看到“2.83 倍”第一反应是怀疑第二反应是问怎么复现。先说公式。设 baseline 自回归解码平均生成一个 token 的耗时是 T_baseMEDUSA 一个完整循环的耗时是 T_tree每个循环平均接受的 token 数是 A那么有效生成速度大概是 A / T_tree。加速比可以写成speedup T_base / (T_tree / A) A * T_base / T_tree所以真正决定加速比的是两个变量循环耗时 T_tree以及平均接受长度 A。接受长度又取决于候选择优质量、树宽、采样温度和训练质量。我在测 2.83 倍时的配置是这样的使用一个约 7B 规模的开源模型量化方式为 FP16单卡环境下跑。输入上下文长度固定为 128解码总长度 256采样温度为 0贪心解码。候选树规模设置为 4 层、每层 top-8 左右经剪枝后节点总数约 64。推理侧开了 CUDA Graph并把树掩码预计算后缓存住。T_base 在相同输入长度下大概是 42 毫秒每 tokenT_tree 一次循环在 70 毫秒上下平均接受长度约 2.9算出来加速比大约就是 2.8 到 3.0 之间。这个结果让我印象很深树验证的耗时只比单步生成多百分之六十多却平均换回了接近 3 个 token。核心原因是单步自回归生成时显存带宽和计算单元都没有吃满而树验证一次喂进去几十个候选反而把 GPU 利用率拉高了。3.2 影响加速比的关键参数与调优顺序第一个调参是树的层数和宽度。层数太多远距离预测准确率会断崖式下跌宽度太宽低概率候选浪费显存和计算。我自己的经验是 3 到 4 层比较均衡每层候选数量前几层可以多一些后几层少一些。树节点总数建议控制在 40 到 80 之间少于 40 接受长度不够高多于 80 显存和耗时增长反而抵消收益。第二个调参是训练头的数据分布。Medusa head 是在主干模型自身的输出分布上训练的所以训练数据最好由这个主干模型生成而不是直接用开源文本。我一开始偷懒直接拿通用语料训练结果头预测的 token 和主干实际偏好差异很大平均接受长度只有 1.4加速比被拖到 1.5 倍以下。重新用目标模型生成训练语料之后接受长度一下子就到了 2.5 以上。第三个调参是采样参数。温度越低模型生成的 token 越确定head 也越容易猜中。如果线上服务必须用 temperature0.8 甚至更高的随机采样接受长度会下降加速比可能只有 1.5 倍左右。我通常会在高温度场景下把每个 head 的候选数加大给接受路径更多余量。3.3 什么场景下加速比会明显缩水2.83 倍不是万能数字。我后来把输入长度从 128 拉到 1024同样配置下加速比直接掉到 2.1 倍左右。原因是 prefill 阶段的耗时占比变高投机解码只优化 decode 阶段对 prefill 没有帮助。如果业务里每次请求都会带很长的上下文必须把 prefill 优化也考虑进来否则总延迟提升有限。并发升高时又是另一种情况。我最初在 batch1 的离线脚本里测出 2.8 倍高高兴兴接进服务端结果在全并发压测时只有 1.6 倍。原因很简单基线的 GPU 已经在处理多个并发请求了硬件资源本身被利用得比较充分树验证想要再“偷偷并行”就没那么多空余算力可以占。所以在“GPU 已经跑满”的高吞吐场景投机解码的收益会被稀释而在“低并发、高延迟敏感”的服务场景加速比反而更接近单测结果。这个经验也是我后来跟业务方对齐预期时反复强调的不要把一个 batch1 的加速比直接拿去承诺线上 P99 提升。先看你的在线请求实际并发度再看输入和输出长度分布最后再决定要不要上这类方案。4. 从训练到推理的落地过程4.1 训练解码头冻结主干只学小头我训练 Medusa head 时没有改动主干权重只用目标模型生成的文本做监督。流程大概是这样的先准备一批提示词让模型用目标采样参数生成完整回复形成一个“模型自生成的语料”。然后再把语料分成长度为 512 的片段每个片段内取前 256 个 token 作为上下文后 256 个 token 作为 head 的监督目标。训练时主干只负责输出 hidden state不参与反向传播的梯度更新。代码里可以这样设置for param in backbone.parameters(): param.requires_grad False for head in medusa_heads: for param in head.parameters(): param.requires_grad True损失函数是多个 head 的交叉熵之和。注意这里不是让 head 直接预测“原句的下一个 token”而是让第 k 个 head 预测“从当前 token 往后第 k 个 token”。所以在构造标签时我会把训练样本错位 k 步for k, head in enumerate(medusa_heads, start1): shift_logits head(hidden_states[:, :-k]) shift_labels labels[:, k:]训练步数不用太多我在单个 7B 模型上用一个批次大小 8、几千条样本训了大约两到三个小时就收敛了。关键监控指标是每个 head 在验证集上的 top-5 准确率。如果第一个 head 的 top-5 准确率低于 80%后面几层的准确率会更差整棵树基本白搭。我当时把训练序列长度设成 512发现第二个 head 的 top-5 准确率比 256 长度版本高了 5 个百分点因为更长的上下文让模型预测更稳定。4.2 推理主循环候选树构建、验证、接受推理阶段的主循环比训练稍微绕一点但拆开看也就四步构建候选树、生成树掩码、一次性前向、找最长接受前缀。第一步在当前位置把 hidden state 取出来过一遍所有 Medusa head得到每组 top 候选。第二步把候选按路径展开成树同时记录每个节点的祖先列表。第三步把原始上下文和整棵树的候选节点拼在一起构造 batch 输入再套上树掩码做一次前向。第四步拿到每个候选节点位置的输出 logits和目标候选 token 对比从根节点开始按顺序检查匹配情况。伪代码如下while generated_length max_len: hidden backbone.encode(context) head_logits [head(hidden) for head in medusa_heads] tree build_candidate_tree(head_logits, top_k_list) mask build_tree_attention_mask(len(context), tree.node_ancestors) batch_tokens context tree.node_tokens logits backbone.forward(batch_tokens, attention_maskmask) accept_len find_longest_common_prefix(tree.paths, logits) if accept_len 0: extend context by accepted tokens else: reject first candidate, sample replacement from logits, extend by 1 update kv_cache accordingly最耗时的树构建模块是 beam 剪枝因为它要在 Python 层面对候选路径做排序。如果节点数只有几十个Python 层面的开销还好但节点一旦上百构建树的时间就会占掉一个循环的五分之一。我的做法是把 beam 剪枝的逻辑用 numpy 向量化候选 logits 直接按矩阵形式处理好避免逐条 for 循环。4.3 工程集成KV 缓存、动态 batch 和 CUDA Graph一个很容易出错的点是 KV cache 的位置映射。普通自回归解码里每次前向只输入一个新 tokenKV cache 会自然往后追加。但树验证时候选节点并不是按顺序追加的它们的 KV cache 位置需要按照“原始上下文长度 节点在树中的前序遍历序号”来分配。我的做法是在构建树的同时生成一份位置索引前向之后按这份索引更新 KV cache并只保留被接受那条路径对应的缓存。没被接受的路径缓存直接丢弃。这个操作不能用简单的切片必须写成 gather 操作否则缓存会错乱。第一次接手这个项目时我在这个位置 debug 了一整天最后是逐个打印候选节点的位置索引才看出问题。CUDA Graph 对这类任务的收益很大。树验证的前向结构是固定的只要候选树最大节点数不变就能把整个前向流程捕获成静态图。我启用 CUDA Graph 之后T_tree 从 110 毫秒降到了 70 毫秒接受长度不变加速比直接提升四成。代价是显存开销增加需要额外预留一块显存给 graph 池。如果显存吃紧可以只用 graph 包装模型主体前向把候选树构建和匹配逻辑留在 CPU 侧。5. 常见问题与排查技巧实录5.1 显存炸了树宽剪枝和激活优化树节点一多显存压力立刻上来。一个 7B 模型在 FP16 下权重占 14GB再加几十个候选节点做前向激活值可能需要额外 2 到 4GB。如果卡是 24GB再塞 KV cache 就很容易爆。我踩坑后的第一反应是缩小树宽把 top-k 从 10 减到 6节点总数从 64 减到 40显存立刻降了 1.5GB加速比只损失 0.2 倍左右。如果还希望保持树的覆盖度可以动态调整前两个 head 保留较大候选数后两个 head 各自只保留 top-4让树的形状从“等宽”变成“前宽后窄”。另外可以把树节点的 hidden state 在 CPU 上构建GPU 只负责前向部分减少 CPU-GPU 之间的传输。5.2 加速比低于预期先从“接受长度”查起如果线上加速比明显低于单测第一件事是打点看平均接受长度。如果接受长度已经很长但加速比还是不行瓶颈大概率在树验证的 T_tree 上如果接受长度本身就低则要回头检查 head 训练质量和采样参数。我在一次排查中发现某个模型的 embedding 层和 head 层的 hidden size 不匹配导致加载权重时随机初始化了一部分。表面上看训练 loss 在下降但头部权重几乎没学到有效信息平均接受长度只有 1.1。检查方式很简单把“输入相同 context 时 head 的预测结果”和“主干下一 token 的真实分布”做一次 top-5 对比重合率太低就能定位问题。还有一次是温度参数设置成了 0.0但解码时启用了 top-p0.9 的截断。head 训练时也用 top-p 截断但截断后候选里出现了大量不具备足够累积概率的尾巴 token导致接受率大幅下降。最后统一了训练和解码时的采样配置接受长度才恢复正常。5.3 输出分布偏移检查树掩码而不是怀疑模型“无损”的前提是验证逻辑严格等价于自回归解码。如果验证后输出的文本和不用投机解码时不一致优先检查树掩码是否正确。常见错误是某个节点能看到同一个 batch 里的其他候选节点导致验证 logits 偏高本应拒绝的路径被错误接受。我建议在调试阶段写一个对照实验固定一段输入分别用普通自回归解码和 MEDUSA 解码跑同一个随机种子逐 token 比较输出。如果出现不一致可以逐步缩小到某个候选路径上打印该节点的祖先列表和 attention mask 对应行。大多数时候都是掩码构造时少加了一行祖先关系属于低级但隐蔽的错误。5.4 与并发批处理系统的兼容性处理MEDUSA 默认假设每次只处理一条序列的树验证。如果你要把多个请求拼成一个 batch树和树之间不能互相干扰每条请求的树掩码都要独立构造。一种比较实用的做法是保留一个全局最大的树节点数然后用 padding token 把短请求补齐到统一长度。显存代价会高一些但调度逻辑会简单很多。更复杂的情况是动态 batch一个请求正在构建候选树另一个请求可能已经验证完准备接收 token。这种情况下我建议引入请求级状态机把“构建树”和“树前向”拆成两个队列而不是把一个请求的整个循环做成一个原子操作。这部分改造量较大如果你的服务已经用了成熟的 continuous batching 框架最好先评估框架是否允许自定义 attention mask 和 KV cache 位置映射再决定是否直接嵌入。5.5 我最后保留的几个务实经验依赖关系上训练数据里建议混合一部分更“保守”的通用语料防止 head 对线上特定 prompt 风格过拟合。如果你同时服务多种业务场景最好为每个场景存一份专用 head 权重推理时按请求路由。我曾经为了省事只训了一份头结果某个代码生成场景接受长度很高另一个客服问答场景只有 1.3整体加速比被严重拉平。另外一个细节是候选树的第一层一定要吃满 top-k 的配额不要把预算分给后面几层。因为第一层如果覆盖不到主干真正会选的 token后面层再准也没机会被验证。我习惯把第一层的候选数设置成后两层的两倍这样同样数量的节点加速比能多出来零点几倍。动手做这个项目之后我最大的体会是投机解码这类方案真正的难点从来不是“树注意力怎么写”而是你愿不愿意花时间把训练数据分布、候选树结构、采样参数、KV cache 映射这几个环节串起来调。任何一环松懈都会把 2.83 倍的理想数字打回 1.5 倍的原形。希望这篇记录能帮你少踩几个坑尤其是树掩码和 KV cache 位置映射这两个隐蔽雷区。
