自投机解码：层跳过与拒绝采样实现大模型推理加速

推理侧的成本最后都会落在同一个数字上：每生成一个 token，你得把整份权重从显存里重新读一遍。做推理优化这些年，我见过太多团队第一反应是换卡、上更激进的量化，但真正把端到端延迟压下去的，往往是改解码策略。Self-Speculative Decoding（自投机解码）是我最愿意推荐给团队先试的一类方案：它不改模型结构、不动权重精度、不需要额外部署一个小模型，靠“先猜几个、再一次性验证”的方式，把原本一步一停的自回归解码变成批量并行校验，在大语言模型推理加速这件事上，通常能拿到 1.3 到 2 倍的实际收益。

它解决的核心问题很具体：低并发场景下 GPU 算力大量闲置，瓶颈卡在显存带宽上，而自投机解码恰好绕开了这个瓶颈。这篇文章适合三类人看——正在做推理服务部署、被延迟指标追着跑的工程师；想搞懂解码层优化原理、准备做技术选型的算法同学；以及已经试过独立草稿模型方案、被额外显存和双模型运维劝退的人。下面我按“为什么慢、怎么破、怎么算、怎么写、怎么排错”的顺序,把这件事从头拆一遍。

1. 自回归解码为什么慢：先把瓶颈算清楚

1.1 每生成一个 token，账到底花在哪

Decoder-only 架构生成第 t 个 token 的过程，本质上是一次完整前向：把当前 sequence 喂进去，跑完所有层，取最后一个位置的 logits，采样出一个 token，拼回序列，然后再来一遍。麻烦的地方在于这 t 次前向无法并行——第 t+1 次的输入依赖第 t 次的输出，形成一条严格的串行链。串行本身不是问题，问题是每一次串行步骤里，你都重复做了一件极其昂贵的事：把全部权重从显存搬到计算单元。

拿一个 7B 级别的模型举例，单步浮点运算量大约是 2 倍参数量，也就是 14 GFLOPs 量级；而单步需要读取的权重是 7B 乘 2 字节（fp16），大约 14 GB。两者一除，算术强度约等于 1 FLOP/Byte。这个数字意味着什么？一张主流数据中心卡的算力带宽比通常在几百这个量级，而 1 远远落在拐点左侧，属于典型的 memory-bound 区间。换句话说，做 batch=1 的矩阵向量乘法时，计算单元基本在空转，时间全花在等数据搬运上。

按 2 TB/s 的有效带宽粗算，14 GB 权重读一遍就是 7 毫秒左右，这基本框定了单请求 7B 模型 fp16 的解码速度上限，再叠加 kernel 启动、归一化、注意力等零碎开销，实际每秒几十个 token 是很正常的。序列变长之后 KV cache 的读取量会涨上来，但短到中等长度下，权重读取仍然是绝对主项。理解了这一点，优化方向就很清楚了：不能指望“少算”，得想办法“少读”——一次读权重，多产出几个 token。

1.2 用算术强度解释 GPU 为什么在空转

我习惯用 roofline 模型来跟团队解释这件事。设峰值算力 F、显存带宽 B，一个算子的算术强度 I 如果低于 F/B，它就是带宽受限的，加算力没用，只有加带宽或者提高数据复用才有用。解码阶段的矩阵向量乘法，I 约等于 1，而入门的推理卡 F/B 都有几十上百，差距不是一点半点。

打个生活化的比方：你有一个手速极快的厨师，但食材存放在很远的仓库，每次只能搬一件。让他一次只做一道菜，他绝大部分时间都在往返路上。解法只有两个——要么一次做一百份（等价于加大 batch），要么让他在等食材的时候顺手把配菜也处理了，主厨只负责最后把关。自投机解码走的是第二条路：用便宜的“助手”先猜几步，主厨一次性批量验收。

这里有个反直觉但很关键的推论：投机解码并不会减少总计算量，反而增加了——草稿要算一遍，验证还要再算一遍。它减少的是权重读取的串行次数。所以它只在 memory-bound 区间有效，一旦上了大 batch 变成 compute-bound，收益会迅速消失甚至变成负的。这个判断标准决定了你该不该上这套方案，后面讲适用边界时会再展开。

2. 自投机解码的核心思路拆解

2.1 先猜后验：投机解码的基本框架

标准投机解码的流程可以拆成三步。第一步，一个轻量的草稿模型自回归生成 γ 个候选 token，这 γ 步很便宜，因为草稿模型小。第二步，目标模型把这 γ 个 token 连同已有前缀一次性喂进去，做一次前向，因为注意力机制天然支持并行处理，这一趟就能拿到每个位置上的完整分布。第三步，从左到右逐个比对草稿 token 和目标模型给出的分布，在第一个不匹配的位置截断，接受前面所有匹配的，再在该位置补一个修正 token。

为什么这样做结果不变？因为验证阶段的接受和拒绝规则是严格按概率设计的，最终输出的分布和目标模型单独解码时完全一致。这是投机解码最有价值的一点——它叫“无损加速”，不是用质量换速度。很多工程团队听到“加速”第一反应是担心输出变差，这个顾虑在投机解码上不成立，前提是你把拒绝采样那套逻辑实现对了。

代价也很明确：你需要第二个模型。额外的显存占用、额外的部署单元、额外的版本对齐工作，还有最麻烦的一点——草稿模型和目标模型的 tokenizer 必须完全一致，词表对不上这套机制直接崩。我在实际项目里见过好几次，方案在实验环境跑得挺好，一上生产就因为双模型运维成本被砍掉。这个落地障碍，恰恰是自投机解码要解决的问题。

2.2 为什么叫“自”投机：草稿从哪来

自投机解码的核心思路是：不用第二个模型，从目标模型自身“抠”出一个廉价的近似版本当草稿。草稿和验证共享同一份权重、同一个 embedding、同一个输出头，tokenizer 天然一致，显存增量接近于零，部署形态和普通推理完全一样。这就是“自”这个字的来源——草稿来自模型自己。

具体怎么从自己身上抠，主流有三条路线。层跳过是最常见的一条：Decoder 是逐层堆叠的，前 k 层的计算量只占总量的 k/L，那就用前 k 层加上 embedding 和输出头组成一个草稿器，生成 γ 个 token，再用完整 L 层去验证。草稿成本 c 近似等于 k/L，控制起来很直观。早期退出是在中间层挂分类头，判断“这个位置已经足够自信”就提前结束，属于动态深度。近似注意力是草稿阶段用更粗糙的注意力计算，比如只看局部窗口、或者用更低的数值精度，验证时恢复完整精度。

还有一条偏门但有意思的路线，不依赖结构差异，而是靠并行迭代猜测来逼近结果，理论上单次迭代就能出多个 token，但收敛性高度依赖输入，接受率波动大，我试过几次，稳定性不如层跳过。综合来看，层跳过在实现复杂度和收益之间平衡得最好，也是我建议优先尝试的方向。

2.3 三条技术路线的工程对比

把这几条路线放在工程视角下对齐看，差异会非常明显。下面这张表是我根据实际项目经验整理的，里面的接受率是经验区间，具体数值和模型、任务、是否做过对齐训练关系很大，不要当成硬指标。

这张表最值得看的一列是“草稿成本”。很多人一开始会本能地选接受率最高的方案，但加速比是接受率和成本共同决定的，接受率 0.9 配成本 0.8 的组合，实际表现往往不如接受率 0.7 配成本 0.2。我在一个项目里就吃过这个亏——早期选了近似注意力的方案，接受率很漂亮，但草稿本身把大部分时间吃掉了，端到端只快了 8%，基本等于白干。判断标准在第四节会给出公式。

3. 正确性怎么保证：拒绝采样的数学

3.1 贪心解码：argmax 比对就够

温度设为 0 的贪心解码是最简单的情形。草稿模型提出一个 token x，目标模型在这个位置算出的分布里，如果 x 恰好是概率最大的那个，就接受；否则丢弃这个 token，直接采用目标模型的 argmax 结果，并且截断后续所有草稿。这个规则很直观，实现也就几行代码。

但这里有个细节容易被忽略：贪心解码下，接受判定的依据是“草稿 token 是否等于目标分布的众数”，而不是“目标模型给这个 token 的概率是否超过某个阈值”。我看过一些实现拿概率阈值做判断，逻辑上说得通，但会破坏无损性——它可能接受一个非众数的 token，输出就和基线不一致了。这类偏差在短文本上看不出来，长文本上会表现为偶发的用词怪异或者逻辑跳步，非常难定位。

还有一个工程上的常见妥协：有些实现为了省事，验证阶段也用 argmax 近似，接受率会略微偏高，因为草稿的贪心结果和目标的贪心结果更容易撞上。代价是极端情况下分布有微小偏移，如果业务对输出一致性要求严格，不建议这么做。

3.2 采样解码：为什么要做拒绝采样

温度大于 0 的时候，事情就复杂了。设草稿分布为 q，目标分布为 p，草稿提出 token x，正确的做法是以 min(1, p(x)/q(x)) 的概率接受它；如果拒绝，则从一个修正分布里重新采样，这个修正分布正比于 max(0, p - q) 归一化之后的结果。

为什么需要这个修正？直观理解：草稿可能在某几个 token 上过度自信（q 大于 p），这些位置要按比例打折；同时它在另一些 token 上估计不足（p 大于 q），这些位置要补偿。修正分布正好承担补偿的角色。数学上可以证明，接受分支和拒绝分支合起来的输出分布精确等于 p，这就是无损性的来源。推导过程不复杂，核心是把“接受 x 的概率”写成 min(q(x), p(x))，再把它和拒绝后补采样的部分相加，正好还原出 p(x)。

这里有个坑我必须强调：直接接受草稿 token、不做任何修正，代码能跑、输出看起来也通顺，但分布已经偏了。表现出来是重复片段变多、罕见词变少、整体风格向草稿靠拢。这个问题极难靠人工检查发现，因为单条输出看不出异常。工程上必须做分布验证，方法在第六节给。

3.3 一个容易踩的坑：随机数与 batch 内对齐

采样解码涉及两处随机性：接受判定用的均匀随机数，以及拒绝后重采样用的分布采样。如果同一个 batch 里所有序列共用同一个随机数，接受行为会高度相关，分布立刻偏掉。正确的做法是每个位置独立生成随机数，形状和候选 token 的 batch 维度对齐。

batch 内对齐是另一个麻烦点。一轮验证里，有的序列接受了 4 个 token，有的一个都没接受，有的在中途遇到终止符。实现上要做三件事：验证前向把长度补齐到最大候选数，用 mask 标记哪些位置有效；验证结束后按每个序列的实际接受长度裁剪 KV cache，多出来的部分必须丢掉；被拒绝位置之后的草稿全部作废，因为那些 token 是基于错误前缀猜出来的，没有意义。

KV cache 裁剪这一步最容易出问题。如果裁剪不干净，后续 token 的注意力会看到那些已经作废的 token，输出会变得莫名其妙，而且错误会随着生成长度累积。我调试过一个案例，现象是生成到 200 字左右突然开始胡言乱语，最后定位就是 KV cache 长度和实际序列长度差了 2 个位置，非常隐蔽。建议在实现里加一个断言，每轮验证后校验两者的长度一致。

4. 参数怎么定：加速比模型的推导

4.1 期望接受长度的闭式解

先做一个简化假设：每个草稿 token 被接受的概率相互独立且都等于 α。现实中这个假设不成立——靠后的位置接受率会下降，因为前面的误差会传导——但这个模型足以指导参数选择，方向是对的。

草稿长度设为 γ。至少接受 k 个 token 的概率是 α 的 k 次方。那么期望接受的草稿 token 数是各阶概率之和，从 1 到 γ 累加，得到 α(1-α^γ)/(1-α)。别忘了验证阶段即使第一个草稿就被拒绝，目标模型也会在修正分布上采样出一个 token，所以每轮迭代必定至少前进 1 个位置。期望前进长度就是：

E[len] = 1 + Σ(k=1..γ) α^k = (1 - α^(γ+1)) / (1 - α)

这个式子很漂亮，代入几个数感受一下。α=0.8、γ=4 时，E 约等于 3.36，也就是说一趟验证平均能推进 3.36 个 token，而基线要跑 4 趟才能达到同样效果。

4.2 加速比公式与参数扫描

单轮迭代的成本要分开算。草稿阶段跑了 γ 次浅层前向，每次成本是 c，合计 γc。验证阶段跑一次完整前向，处理 γ+1 个位置，权重读取和单 token 前向相同，所以近似记为 1。这样得到朴素版的加速比：

S = (1 - α^(γ+1)) / ((1-α) × (1 + γc))

但实际工程里有个优化能把验证成本压下来。草稿阶段算出的前 k 层 KV cache，在验证阶段可以完整复用——因为因果注意力保证前缀 token 在浅层的表示不会因为后面接了新 token 而改变。所以验证阶段只需要从第 k+1 层算到第 L 层，成本降到大约 1-c。修正后的公式变成：

S = (1 - α^(γ+1)) / ((1-α) × (1 - c + γc))

按修正公式，c 固定为 0.25 时，扫一遍 α 和 γ，能看出不少东西：

这张表最值得盯的是 α=0.5 那一行——γ 越大加速比越低，因为草稿成本已经超过了收益。这说明一个很实际的结论：接受率上不去的时候，把 γ 调大是纯亏。反过来，α 到 0.9 时，最优 γ 落在 4 到 6 之间，再大也没用了，因为期望长度的增长已经饱和。

由此可以提炼出两条调参规则。第一，α 越高，最优 γ 越大；α 低于 0.6 时建议 γ 取 2 到 3，先把成本压住。第二，c 越大，最优 γ 越小，草稿层数多到 50% 的时候，γ 取 2 到 3 就够了，别贪多。

4.3 实测区间与理论偏差

理论值算出来好看，实测普遍要打八折到八五折，主要有三个来源。一是接受概率并不独立，靠后位置的接受率会明显衰减，比如 γ=4 配置下前两个位置接受率 0.8，第四个位置可能只有 0.6。二是验证阶段处理 γ+1 个位置时，注意力部分的计算量确实比单 token 多，虽然不占大头但不是零。三是 batch 内长短不齐导致的 padding 浪费，在并发场景下这部分能吃掉 10% 以上。

我手上比较稳的一组数据：单卡、batch=1 到 4、fp16、7B 级模型、输出 128 token，草稿取 1/4 层数、γ=4、不做额外对齐训练，端到端加速 1.3 到 1.5 倍；如果做过轻量对齐训练把接受率从 0.65 提到 0.8，能到 1.7 到 1.9 倍。batch 提到 16 以上，加速比会掉到 1.1 倍以下，因为那时已经接近 compute-bound 区间了。

5. 动手实现：一个可跑的最小版本

5.1 环境与基线测量

测量这件事本身就有讲究。必须做预热，至少跑 3 轮丢弃，因为第一次调用会触发 kernel 编译和显存分配。计时前后都要做设备同步，否则测到的是 kernel 下发时间而不是执行时间，结果会虚低。重复次数不少于 5 次取中位数，避免偶发抖动。

先给基线代码，把时钟校准了再谈优化：

import time import torch from transformers import AutoModelForCausalLM, AutoTokenizer torch.set_grad_enabled(False) def load_model(path, dtype=torch.float16): tok = AutoTokenizer.from_pretrained(path) model = AutoModelForCausalLM.from_pretrained( path, torch_dtype=dtype, device_map="cuda" ) model.eval() return model, tok def bench_baseline(model, tok, prompt, max_new_tokens=128, rounds=5): ids = tok(prompt, return_tensors="pt").input_ids.cuda() # 预热 for _ in range(3): model.generate(ids, max_new_tokens=8, do_sample=False, pad_token_id=tok.eos_token_id) torch.cuda.synchronize() times = [] for _ in range(rounds): torch.cuda.synchronize() t0 = time.perf_counter() model.generate(ids, max_new_tokens=max_new_tokens, do_sample=False, pad_token_id=tok.eos_token_id) torch.cuda.synchronize() times.append(time.perf_counter() - t0) times.sort() return times[len(times) // 2]

跑完这一步，你手上应该有一个明确的基线数字，比如 128 token 花了 2.4 秒。后面所有优化都拿这个数字做分母，别拿实验室论文里的数字做参照。

5.2 草稿阶段的层跳过实现

核心思路是只跑前 k 层，然后直接接输出头。这里有个必须提前知道的事实：原版模型直接切前 k 层，接受率会低得让你怀疑人生。原因是中间层的隐藏状态从来没有被最终的归一化层和输出头“见过”，分布是错位的。我实测过几个模型，直接切层的接受率只有 0.4 到 0.5，加速比还不到 1.1。

改进办法是做一轮轻量的自蒸馏训练：让前 k 层加上归一化和输出头，去拟合完整模型的输出分布，训练量很小，几千条样本、一两个 epoch 就够。这一步做完，接受率能提到 0