首页 / 资讯中心 / 文章详情

投机解码原理与工程实践:从自回归瓶颈到无损加速

投机解码原理与工程实践:从自回归瓶颈到无损加速 ★ FEATURED ARTICLE
1. 从自回归的一个接一个说起为什么大模型推理慢得让人抓狂如果你部署过任何基于Transformer架构的大语言模型一定对那种明明GPU利用率不高但就是快不起来的体验印象深刻。输入一段提示词模型开始一个字一个字往外蹦每个token的生成都要等上一个token算完才能开始。这种串行模式就是自回归解码的本质也是大模型推理延迟居高不下的根源。我最早接触这个问题是在一个对话系统项目里。当时用了一张推理卡跑一个中等规模的模型单次对话响应时间稳定在1.5秒以上用户体感就是卡。我盯着监控面板看了很久发现计算单元的利用率只有30%左右显存带宽倒是跑满了。这说明瓶颈不在算力而在访存——每生成一个token模型都要把全部权重从显存里读一遍计算量却小得可怜。这就是典型的memory-bound场景。自回归解码的数学形式很简洁给定前缀序列模型预测下一个token的概率分布采样得到结果后拼接到序列末尾再重复这个过程。每一步的计算量是固定的但步与步之间无法并行因为第t步的输入依赖于第t-1步的输出。这种依赖关系像一条铁链把整个生成过程锁死了。有人可能会想那能不能一次预测多个token可以但直接让模型一次输出多个位置的结果会破坏概率分布的准确性。因为语言模型的条件概率是逐位定义的你没法在不改变分布的前提下跳步。这就是投机解码要解决的核心矛盾既要保持输出分布严格不变又要打破串行瓶颈。投机解码的思路其实很直觉用一个小的、快的模型先猜出接下来几个token然后用大的、慢的目标模型去验证这些猜测。如果猜对了就一次性接受多个token如果猜错了就回退到第一个错误位置从那里重新采样。关键在于这个验证过程要保证最终输出的分布和直接用目标模型逐token采样完全一致。听起来像先斩后奏但数学上可以做到严丝合缝。这个方案之所以被称为开山奠基之作是因为它第一次系统性地给出了投机解码的数学框架和算法流程后续的Medusa、Lookahead Decoding、EAGLE等工作都是在它的基础上做改进。理解了这个原始方案再看后面的变体会有一种万变不离其宗的通透感。2. 投机解码的数学骨架接受-拒绝采样如何保证分布无损2.1 从拒绝采样说起要理解投机解码为什么能保分布得先搞清楚拒绝采样Rejection Sampling的基本逻辑。假设你想从某个复杂分布p(x)中采样但直接采样很困难。你可以找一个容易采样的提议分布q(x)先从q(x)中采一个样本x然后以一定的概率决定接受还是拒绝这个样本。只要接受概率设计得当最终接受的样本就服从p(x)。投机解码把这个思想用在了token级别。设目标模型为M_p草稿模型为M_q。对于当前位置M_q先给出一个猜测token记为x。然后M_p计算自己对该token的概率p(x)M_q的概率q(x)是已知的。接受这个猜测的概率设为min(1, p(x)/q(x))。如果接受就输出x如果拒绝就从修正后的分布中重新采样一个token。这个修正分布是norm(max(0, p - q))也就是把p和q的差值截断到非负后归一化。为什么是这个形式因为要保证最终输出的边缘分布恰好是p。推导过程不复杂对于任意token x最终输出x的概率等于从q中采到x且被接受的概率加上从q中采到任意token被拒绝后从修正分布中采到x的概率。前者是q(x) * min(1, p(x)/q(x))后者是(1 - 接受率) * 修正分布(x)。两项相加经过化简恰好等于p(x)。这就是数学严密性的来源。2.2 多token猜测的链式验证单token的接受-拒绝只是热身投机解码真正的威力在于一次猜测多个token。草稿模型连续生成γ个token形成一个候选序列。目标模型对这γ个位置并行计算概率然后从左到右逐个验证。验证过程是这样的对于第i个猜测token计算接受概率min(1, p_i/q_i)。如果接受继续验证下一个如果拒绝则停止从修正分布中重新采样一个token作为该位置的输出并且丢弃后面所有猜测。如果γ个token全部被接受那就额外从目标模型采样一个token作为奖励这样一次前向就能产出γ1个token。这里有个细节值得注意目标模型对这γ个位置的并行计算和逐token计算的结果是完全等价的。因为Transformer的注意力机制在训练时就是并行的推理时之所以串行是因为没有已知的后续token。现在草稿模型提供了候选token目标模型就可以像训练时那样一次性算出所有位置的概率。这也是投机解码能加速的前提——验证的并行化不会引入额外的计算误差。2.3 加速比的理论上界投机解码的加速比取决于两个因素草稿模型和目标模型的速度比以及猜测的接受率。设草稿模型生成γ个token的时间为T_q目标模型验证这γ个token的时间为T_p。由于验证是并行的T_p大约等于目标模型生成1个token的时间。那么一轮的总时间是T_q T_p产出的token数是接受长度加1。如果接受率为α平均接受长度约为(1 - α^(γ1))/(1 - α)。加速比就是期望产出token数除以单token生成时间。当α接近1时加速比可以接近γ1当α很低时加速比可能小于1也就是反而变慢。所以草稿模型的选择很关键它要足够快同时和目标任务足够相似才能保证高接受率。我实测下来草稿模型参数量在目标模型的1/10到1/5之间时接受率通常能维持在0.7以上。如果草稿模型太小虽然速度快但猜测质量差接受率会掉到0.5以下加速效果大打折扣。如果草稿模型太大验证开销上去了整体收益也不明显。这个比例不是绝对的还要看具体任务和训练数据分布。3. 草稿模型怎么选不是越小越好也不是越像越好3.1 独立小模型方案最直接的做法是单独训练一个同架构的小模型作为草稿模型。比如目标模型是70亿参数草稿模型用10亿参数两者用相同的tokenizer和训练数据。这种方案的好处是工程上简单草稿模型可以独立部署和更新不影响目标模型。但这里有个坑小模型和大模型的输出分布往往有系统性偏差。小模型倾向于生成更安全的token概率分布更平坦导致在某些位置上的q(x)远大于p(x)接受概率被压低。我遇到过一种情况草稿模型在句首总是倾向于生成高频词而目标模型更愿意生成多样化的开头结果就是每轮第一个token的接受率特别低整体加速比上不去。解决办法之一是在训练草稿模型时做蒸馏让它的输出分布尽量逼近目标模型。具体做法是用目标模型对大量文本打分然后让草稿模型去拟合这些软标签。蒸馏损失通常用KL散度温度系数设高一些比如2到4可以让分布更平滑减少极端值的影响。实测下来蒸馏后的草稿模型接受率能提升10到15个百分点。3.2 自蒸馏与多头草稿如果不想维护两个模型可以考虑自蒸馏方案在目标模型内部加几个轻量的预测头让它们负责猜测后续token。这些预测头共享目标模型的底层表示只增加很少的参数。因为底层表示已经包含了丰富的语义信息预测头的猜测质量通常比独立小模型更好。这种方案的代表是Medusa它在目标模型的最后一层隐藏状态上接了多个解码头每个头预测不同偏移位置的token。训练时只更新这些头目标模型本身冻结。推理时多个头并行给出候选序列然后交给目标模型验证。因为共享了底层计算草稿部分的额外开销很小整体加速比很可观。不过多头方案有个限制每个头只能预测一个固定偏移的位置灵活性不如自回归草稿模型。如果某个位置的猜测被拒绝后面的头就浪费了。所以实际部署时通常会把多头输出和树状注意力结合起来让验证过程能处理分支结构提高候选利用率。3.3 草稿模型与目标模型的协同训练更激进的做法是让草稿模型和目标模型联合训练在训练阶段就考虑推理时的接受率。具体来说损失函数里除了常规的语言建模损失还加一项鼓励草稿模型生成容易被目标模型接受的token。这一项可以设计成可微的近似接受率比如用p(x)/q(x)的期望作为奖励信号。这种协同训练能让两个模型的分布更对齐但工程复杂度也更高。需要同时加载两个模型显存占用翻倍训练时间也拉长。我的经验是如果目标模型本身已经训练好了重新做协同训练的成本太高不如直接做蒸馏来得划算。只有在从零开始训练一套模型时协同训练才有明显优势。4. 工程落地中的真实瓶颈显存、批处理与调度4.1 显存占用的隐性成本投机解码在理论上能加速但实际部署时第一个撞上的墙往往是显存。草稿模型虽然小但它的权重、KV Cache、中间激活值都要占显存。如果目标模型已经快把显存吃满了再加一个草稿模型可能直接OOM。我试过一个方案是把草稿模型量化到INT8甚至INT4牺牲一点猜测质量来换显存空间。实测下来INT8量化的草稿模型接受率下降约3到5个百分点但显存占用减少一半整体吞吐反而提升了。INT4量化更激进接受率掉10个百分点以上需要谨慎评估。另一个思路是让草稿模型和目标模型共享KV Cache。因为两个模型处理的是同一个前缀序列KV Cache的内容高度相似。如果架构允许可以把草稿模型的KV投影到目标模型的维度直接复用。这样能省下不少显存但需要修改模型结构不是所有场景都能做。4.2 批处理下的接受率波动单条请求下投机解码的表现和批处理场景差别很大。批处理时不同请求的接受长度不一样有的请求一轮接受5个token有的只接受1个。如果按最长接受长度来同步短的那些请求就会空等浪费计算资源。解决办法是动态批处理每轮验证后把已经完成的请求移出批次把新请求加进来。但这样会导致批次大小频繁变化GPU利用率不稳定。更精细的做法是给每个请求维护一个剩余预算当某个请求的接受长度超过预算时就把它挂起等下一轮再处理。这样能在吞吐和延迟之间找一个平衡点。我实测过一个动态批处理方案在混合负载下比静态批处理的吞吐高30%左右。关键是要把调度逻辑做轻量不能引入太多CPU开销。如果调度本身成了瓶颈那就得不偿失了。4.3 与连续批处理的兼容性现在主流的推理框架都支持连续批处理Continuous Batching也就是在生成过程中动态加入新请求。投机解码和连续批处理的结合需要小心处理。因为投机解码的验证过程是批量的新加入的请求可能还没有草稿token需要先单独跑一轮草稿生成。我的做法是给每个请求维护一个状态机刚加入时处于草稿生成状态生成完候选token后转入验证状态验证通过后回到草稿生成如此循环。调度器根据状态来决定哪些请求可以一起做验证。这样虽然增加了一些状态管理开销但整体吞吐比串行处理高得多。还有一个细节是KV Cache的管理。投机解码在验证失败时需要回退也就是丢弃部分KV Cache。如果KV Cache是预分配的连续内存回退只需要移动指针如果是动态分配的回退可能涉及内存碎片整理。我倾向于预分配固定大小的KV Cache每个请求一个槽位回退时只调整有效长度不释放内存。这样虽然浪费一些显存但避免了碎片化带来的性能抖动。5. 实测数据与调参经验从理论加速比到真实收益5.1 不同温度下的接受率变化温度参数对投机解码的影响很大。温度越高采样越随机草稿模型和目标模型的分布差异被放大接受率下降。我做过一组对比实验温度设为0.1时接受率能到0.85以上温度升到0.7接受率掉到0.65左右温度到1.0时接受率只有0.5出头。这意味着在需要高多样性的场景下投机解码的加速效果会打折扣。如果业务允许可以适当降低温度来换取更高的加速比。但要注意温度降低会影响生成质量可能让输出变得重复和呆板。我的建议是根据任务类型来定事实问答、代码生成这类任务可以低温创意写作、对话这类任务还是保持正常温度接受加速比低一些。5.2 猜测长度γ的选取γ是每轮草稿模型生成的token数。γ太小加速比上不去γ太大草稿生成时间线性增长而且后面的token被接受的概率越来越低。理论上存在一个最优γ让期望加速比最大。我实测下来γ在4到8之间比较合适。具体取值要看草稿模型和目标模型的速度比如果草稿模型很快γ可以取大一些如果草稿模型本身就不快γ取4左右就够了。另外接受率高的任务可以取大γ接受率低的任务取小γ。有些实现会根据历史接受率动态调整γ这个思路很好但调整策略要平滑避免频繁震荡。5.3 与量化、蒸馏等技术的叠加效果投机解码可以和量化、蒸馏等技术叠加使用但叠加效果不是简单相加。我试过目标模型用INT8量化、草稿模型用INT4量化的组合加速比确实比单独用量化高但接受率下降明显。原因是量化误差让两个模型的分布偏差更大验证通过率降低。蒸馏和投机解码的叠加效果更好。如果草稿模型已经用目标模型蒸馏过接受率本身就高再叠加投机解码加速比能到2.5倍以上。我的经验是先把草稿模型蒸馏好再考虑量化最后上投机解码。顺序反过来调优难度会大很多。还有一个容易被忽略的点是tokenizer的一致性。草稿模型和目标模型必须用完全相同的tokenizer否则token的切分方式不同概率没法直接比较。我见过有人用不同tokenizer的模型做投机解码结果接受率极低排查了半天才发现是tokenizer的问题。6. 从开山之作到后续演进投机解码的边界与延伸投机解码的原始方案给出了一个干净的数学框架但它也有自己的边界。比如它假设草稿模型和目标模型处理的是同一个前缀如果前缀本身很长草稿模型的KV Cache也会占用不少显存。再比如它没有考虑草稿模型本身的生成延迟如果草稿模型太慢整体加速比会被拖累。后续的改进工作大致沿着几个方向展开。一是改进草稿模型的结构让它更轻量、更准确比如用EAGLE方案在特征层做自回归预测而不是在token层。二是改进验证策略比如用树状注意力一次验证多个候选分支提高候选利用率。三是把投机解码和调度系统深度结合让不同请求的草稿生成和验证过程能交错执行进一步提升硬件利用率。这些改进各有侧重但核心思想没有变用低成本的方式生成候选用高精度的方式验证候选在保证分布无损的前提下打破串行依赖。理解了原始方案的数学骨架再看这些变体就能很快抓住它们的创新点和适用场景。我在实际项目中的体会是投机解码不是银弹它的收益高度依赖于任务特性、模型组合和工程实现。在动手之前先花时间分析一下你的瓶颈到底在哪里如果是显存带宽受限投机解码大概率有效如果是算力本身就不够那可能要先考虑模型压缩或硬件升级。另外草稿模型的选择和调优需要反复实验不要指望一套参数打天下。最后别忘了监控接受率这个核心指标它直接决定了加速比的上限一旦发现接受率持续偏低就要回头检查草稿模型和目标模型的分布对齐情况。
阅读完成 · 觉得有帮助?
咨询建站