1. 自回归解码的瓶颈到底卡在哪
大模型推理慢,很多人第一反应是"模型太大、算力不够"。这个判断只对了一半。真正让推理速度上不去的,是Transformer解码阶段那个绕不开的串行结构——每生成一个token,都要等前一个token算完才能开始。GPU里成千上万个计算核心,在这个环节大部分时间是闲着的。
我拿一个具体的场景来说明。假设你在部署一个70亿参数级别的对话模型,输出一段200个token的回复。整个解码过程要跑200轮前向传播,每一轮的计算量其实很小——就是一次矩阵乘向量。问题在于,这200轮之间是严格串行的,第N轮必须等第N-1轮的结果。GPU的并行能力在这种"一次只算一个token"的模式下完全发挥不出来,显存带宽成了真正的瓶颈,而不是浮点算力。
这就是所谓的内存墙问题。解码阶段每生成一个token,需要把整个模型的权重从显存里读一遍。70亿参数的模型,FP16精度下大约14GB权重,每生成一个token就要完整读一次。而实际参与计算的时间极短,大部分时间花在等数据从显存搬到计算单元上。算力利用率可能只有个位数百分比。
业界解决这个问题的思路大致分几个方向。一是批处理,把多个请求攒在一起同时解码,摊薄权重读取成本,但单请求的延迟并没有降低。二是模型压缩,量化、剪枝、蒸馏,减少权重大小,但会损失精度。三是改变解码范式,让一次前向能产出多个token——这正是投机解码切入的角度。
投机解码的核心洞察很朴素:与其让大模型一个token一个token地挤,不如找个小模型先快速"猜"出一串候选token,然后让大模型一次性并行验证这些候选。如果猜对了,就一次性接受多个token,相当于用一次大模型前向换来了多个token的输出。这个思路之所以成立,是因为验证比生成便宜得多——验证N个候选token只需要一次前向传播,而逐个生成N个token需要N次前向。
但这里有个关键问题:猜错了怎么办?如果小模型猜的token大模型不认可,那这次验证就白费了。所以投机解码的数学严密性就体现在这里——它必须保证最终输出的概率分布和原始大模型完全一致,不能因为引入了小模型就改变了输出质量。这是它区别于普通"草稿+修正"方案的根本所在。
2. 投机解码的数学骨架:为什么分布能严格保持
2.1 拒绝采样给出的理论保证
投机解码最精妙的地方,在于它用一套拒绝采样机制,把"小模型猜、大模型验"这个流程包装成了一个数学上严格等价的采样过程。最终产出的token序列,其概率分布与直接用大模型自回归采样完全一致。这不是近似,是精确等价。
具体来说,设大模型为 $q(x)$,小模型为 $p(x)$。小模型先采样出一个候选token $x \sim p$。然后计算一个接受概率:
$$\min\left(1, \frac{q(x)}{p(x)}\right)$$
以这个概率决定是否接受该候选。如果接受,就输出这个token;如果拒绝,就从修正分布中重新采样一个token,修正分布为:
$$\text{norm}\left(\max(0, q(x) - p(x))\right)$$
这套机制的理论依据来自拒绝采样。可以证明,经过这个过程产出的token,其边际分布恰好等于 $q(x)$。我第一次推导这个结论的时候觉得挺反直觉的——明明小模型参与了生成过程,怎么最后分布还能和大模型一模一样?关键在于那个修正分布的设计,它恰好补偿了小模型"猜偏"的部分。
2.2 接受率的直觉理解
接受概率 $\min(1, q/p)$ 这个式子值得多琢磨一下。当 $q(x) \geq p(x)$ 时,接受概率为1,意思是只要大模型觉得这个token的概率不低于小模型给的,就照单全收。当 $q(x) < p(x)$ 时,接受概率为 $q/p$,意思是小模型高估了这个token,那就按比例打折接受。
这背后的直觉是:小模型如果和大模型"意见一致",那它的猜测就可信;如果小模型过度自信地推了一个大模型不太认可的token,就要以一定概率拒绝掉。接受率越高,加速效果越好。而接受率取决于小模型和大模型的分布有多接近——这也是为什么草稿模型的选择如此重要。
2.3 多token验证的并行化
单token的接受拒绝只是基础。投机解码真正的加速来自一次验证多个token。小模型连续生成 $\gamma$ 个候选token,大模型对这 $\gamma$ 个token做一次并行前向,同时得到每个位置的概率分布。然后从左到右逐个应用接受拒绝规则,一旦某个位置被拒绝,后面所有候选全部丢弃,并从该位置的修正分布重新采样一个token。
这个过程有个重要性质:被接受的前缀越长,加速比越高。最理想情况是小模型猜的 $\gamma$ 个token全被接受,那一次大模型前向就产出了 $\gamma+1$ 个token(最后一个来自修正分布或额外采样)。最差情况是第一个就被拒绝,那这次大模型前向只产出1个token,和普通解码一样,还多花了小模型的计算。
所以投机解码的收益是个期望值问题。设平均接受长度为 $\alpha$,则期望加速比大致为 $\frac{\alpha+1}{\text{小模型开销占比} + 1}$。当 $\alpha$ 在2到3之间,小模型开销又足够小的时候,2到3倍的加速是很现实的。
3. 草稿模型怎么选:不是越小越好
3.1 草稿模型与大模型的能力匹配度
很多人第一反应是草稿模型越小越好,毕竟它只是用来"猜"的。这个想法有道理但不完整。草稿模型太小,和大模型的分布差距就大,接受率会很低,猜了等于白猜。草稿模型太大,虽然接受率高,但它自己的推理开销也上去了,摊薄了加速收益。
我实测下来的经验是,草稿模型参数量控制在大模型的十分之一到二十分之一之间比较合适。比如70亿的主模型,配一个3亿到7亿的草稿模型。这个比例下,草稿模型单次前向的开销大约是大模型的5%到10%,而接受率通常能维持在60%到80%之间,综合加速比比较理想。
但参数量只是参考,更关键的是两个模型的训练数据分布要接近。如果草稿模型是在通用语料上训的,主模型是在代码语料上微调的,那在代码生成场景下草稿模型的猜测质量会很差。理想情况下,草稿模型应该和主模型同源,或者至少在相似数据上训练过。
3.2 同源蒸馏 vs 独立小模型
获取草稿模型有两条主流路径。一是同源蒸馏,用主模型的输出分布去训练一个小模型,让它在token级别上尽量逼近主模型。这种方式的接受率通常最高,因为蒸馏的目标就是让两个分布接近。但代价是需要额外的训练流程和算力。
二是直接用现成的小模型,比如拿一个同系列的小尺寸版本当草稿。这种方式零训练成本,开箱即用,但接受率取决于两个模型的"默契程度"。同系列的小模型通常表现不错,因为训练数据和方法一致。
还有一种更激进的思路是自投机,不引入额外的草稿模型,而是用主模型自身的浅层或部分层来做草稿。这种方式省显存,但实现复杂度高,加速比也不如独立草稿模型稳定。
| 草稿模型方案 | 接受率 | 额外显存 | 训练成本 | 适用场景 |
|---|---|---|---|---|
| 同源蒸馏小模型 | 高(75%-90%) | 中 | 高 | 对加速比要求高的生产环境 |
| 同系列现成小模型 | 中高(60%-80%) | 中 | 无 | 快速验证、资源有限 |
| 自投机(浅层草稿) | 中(50%-70%) | 低 | 无 | 显存紧张的单卡部署 |
| 多草稿模型集成 | 高 | 高 | 中 | 追求极致加速比 |
3.3 草稿长度的调优
草稿长度 $\gamma$ 是另一个关键超参。$\gamma$ 太小,加速效果不明显;$\gamma$ 太大,后面的token被拒绝的概率累积上升,浪费草稿模型的计算。经验值通常在4到8之间。
调优方法很直接:在目标场景的数据上跑一批样本,统计不同 $\gamma$ 下的平均接受长度和端到端延迟,选延迟最低的那个。我一般会从 $\gamma=4$ 开始,以2为步长往上试,找到拐点。不同任务的最优值差异挺大——翻译任务因为输出确定性高,$\gamma$ 可以设大一些;开放式对话因为多样性高,$\gamma$ 要小一些。
4. 工程落地的几个硬骨头
4.1 KV Cache的同步问题
投机解码在工程实现上最麻烦的地方是KV Cache的管理。普通自回归解码时,每生成一个token就往KV Cache里追加一组键值对,逻辑很清晰。但投机解码里,草稿模型生成的候选token在被验证之前,它们的KV Cache状态是"待定"的。
如果候选被接受了,这些KV Cache要正式并入主模型的缓存;如果被拒绝了,这些缓存要回滚。实现上通常的做法是:主模型验证时,对候选token重新计算KV,而不是复用草稿模型的缓存。因为草稿模型和主模型的维度、层数可能不同,缓存没法直接共享。这意味着验证阶段的计算量比看起来要大一些,但相比逐个生成还是划算的。
另一个细节是位置编码的处理。候选token的位置是连续的,验证时要保证位置编码正确对应。如果实现时位置编码算错了,接受率会莫名其妙地下降,而且很难排查。我踩过一次这个坑,查了大半天才发现是位置偏移了一位。
4.2 批处理下的投机解码
生产环境不可能一次只服务一个请求,批处理是必须的。但投机解码和批处理的结合有个天然矛盾:不同请求的接受长度不一样,有的请求一次接受了5个token,有的只接受了1个。这导致批内各序列的进度不一致,需要动态调整。
常见的处理方式是按序列独立维护草稿和验证状态,在批维度上做mask。被拒绝的序列在后续验证步骤中屏蔽掉,避免无效计算。这增加了实现的复杂度,但能保证批处理的吞吐优势不被破坏。
还有一种简化策略是固定草稿长度,所有序列都猜 $\gamma$ 个token,验证后统一截断到各自的实际接受长度。这种方式实现简单,但会有一些计算浪费。在延迟敏感的场景下,动态方案更优;在吞吐优先的场景下,固定方案够用。
4.3 与量化、连续批处理的叠加
投机解码不是孤立的优化,它需要和现有的推理优化技术叠加使用。和量化叠加时要注意:草稿模型和主模型的量化精度最好一致,否则分布差异会变大,接受率下降。我试过主模型用INT8、草稿模型用FP16的组合,接受率比两者都用FP16时低了将近10个百分点。
和连续批处理叠加时,投机解码的验证步骤要嵌入到调度循环里。由于验证是一次前向产出多个token,调度器需要相应调整每次迭代的token预算分配。这块的实现因框架而异,需要仔细阅读所用推理框架的调度逻辑。
提示:在正式部署前,务必在目标硬件和目标数据集上做端到端的延迟和吞吐测试。投机解码的理论加速比和实际加速比之间可能有较大差距,尤其是当草稿模型的开销被低估时。
5. 实测数据与调优记录
5.1 测试环境与基线设定
我在一台单卡环境下做了组对比测试,主模型是一个70亿参数的对话模型,FP16精度。草稿模型选了同系列的3亿参数版本,草稿长度从2试到10。测试数据是200条多轮对话,覆盖问答、摘要、代码生成三类任务。基线是不开投机解码的普通自回归采样,batch size设为1,模拟低并发低延迟场景。
先看基线数据:普通解码下,平均每token延迟约45毫秒,输出200个token的回复需要9秒左右。GPU利用率在解码阶段只有12%上下,显存带宽占用率倒是接近饱和,印证了前面说的内存墙问题。
5.2 不同草稿长度下的加速比
固定草稿模型为3亿参数版本,只调草稿长度 $\gamma$,结果如下:
| 草稿长度 | 平均接受长度 | 端到端延迟 | 加速比 | 接受率 |
|---|---|---|---|---|
| 2 | 1.6 | 6.2s | 1.45x | 80% |
| 4 | 2.8 | 4.1s | 2.20x | 70% |
| 6 | 3.5 | 3.6s | 2.50x | 58% |
| 8 | 3.8 | 3.9s | 2.31x | 48% |
| 10 | 3.9 | 4.4s | 2.05x | 39% |
可以看到 $\gamma=6$ 时端到端延迟最低,加速比2.5倍。$\gamma$ 继续增大,虽然平均接受长度还在涨,但草稿模型自身的开销增长更快,反而拖累了整体延迟。这个拐点位置和任务类型强相关,代码生成任务因为确定性高,拐点会往右移;开放式对话则往左移。
5.3 草稿模型规模的影响
固定 $\gamma=6$,换不同规模的草稿模型:
| 草稿模型规模 | 接受率 | 草稿开销占比 | 加速比 |
|---|---|---|---|
| 1.5亿 | 52% | 3% | 1.9x |
| 3亿 | 58% | 6% | 2.5x |
| 7亿 | 72% | 14% | 2.3x |
| 13亿 | 81% | 25% | 1.8x |
3亿参数版本在这个场景下综合表现最好。7亿版本接受率更高,但草稿开销上去了,净加速反而略低。13亿版本虽然接受率超过80%,但草稿模型本身的开销已经吃掉了大部分收益。这个结果再次说明草稿模型不是越大越好,也不是越小越好,存在一个甜点区。
5.4 几个影响接受率的隐藏因素
测试过程中发现几个容易被忽略的因素。温度参数对接受率影响很大,温度越高,采样随机性越强,草稿模型和主模型的分布差异被放大,接受率明显下降。温度从0.7提到1.0,接受率掉了大约15个百分点。所以高温采样场景下,投机解码的收益会打折扣。
提示词长度也有影响。长提示词下,主模型和草稿模型对上下文的理解差异更容易暴露,接受率会略降。但这个影响比温度小得多,通常在5个百分点以内。
还有一个反直觉的发现:输出长度对加速比的影响不是线性的。短输出(50token以内)因为验证的固定开销占比高,加速比偏低;中等长度(100到300token)加速比最高;超长输出(500token以上)因为后期上下文变长,草稿模型质量下降,加速比又会回落。
6. 这套方案适合谁,不适合谁
投机解码不是万能药,它有明确的适用边界。适合的场景是:单请求延迟敏感、有可用的草稿模型、主模型推理时显存带宽是瓶颈而非算力。典型如在线对话服务、代码补全、实时翻译。
不太适合的场景:一是已经做了大规模批处理、GPU算力已经打满的情况,此时瓶颈在算力不在带宽,投机解码的收益有限;二是主模型本身很小(比如1B以下),草稿模型的空间被压缩得很小,加速比上不去;三是对输出确定性要求极高的场景,虽然投机解码分布严格等价,但实现bug可能导致细微偏差,需要额外验证。
从投入产出比看,投机解码的工程改造量中等,主要成本在草稿模型的获取和KV Cache逻辑的改造。如果已经有同系列的小模型可用,落地周期可以很短。如果要从头蒸馏草稿模型,那训练成本需要纳入考量。
我在实际项目里的体会是,投机解码最适合作为推理优化的"第二层"——第一层先把量化、算子融合这些基础优化做扎实,第二层再上投机解码榨取额外加速。跳过第一层直接上投机解码,往往事倍功半,因为基础瓶颈没解决,投机解码的收益会被掩盖。
最后分享一个排查技巧:如果发现加速比远低于预期,先别急着调草稿长度,先检查草稿模型和主模型的tokenizer是否完全一致。tokenizer不一致会导致token id对不上,接受率会低得离谱,而且日志里不一定有明显报错。这个坑我见过不止一次,排查起来很费时间,但原因其实很简单。