LLM 解码很慢,因为生成一个 token 需要每次都从 GPU 内存中加载全部模型权重。700 亿参数的模型意味着 140GB 权重,每一个 token 都要重新加载一遍。
LLM 解码本质上是顺序的,称之为自回归生成(autoregressive generation)。每个 token 都依赖前面所有的 token,采样一个、喂回模型、再采样下一个,如此循环。
没有 token 4 就生成不了 token 5,只能一个一个来,所以每次都要加载那 140GB 权重。
那么如果设想有一个小而便宜的模型,能飞快猜出接下来几个 token;大模型再一次性验证这些猜测。猜对了,就用一个 token 的代价拿到了 1 到 k 个 token;猜错了,就回退继续。
这就是投机采样(speculative decoding)。它完全无损而且在数学上站得住脚。
验证具体如何进行
举个例子。输入是 "I want",草稿模型(draft model)预测接下来四个 token 是 "ice cream and cake"。
目标模型(target model)对整个序列做一次前向传播,为每个草稿 token 计算自己会给出的概率,再和草稿模型的预测比较。草稿模型的置信度低于目标模型时,token 直接通过;草稿模型过于自信时,则按两个分布之间的比率决定接受概率——相当于目标模型在说"没错,我也这么想",或者"不可能,我才不会这么说"。
token 被拒绝时,系统不会就此停下,而是从一个经过调整的分布中重新采样一个修正后的 token。这个修正分布在数学上保证和目标模型独立运行时产生的分布完全一致,这样就变成了没有草稿模型、没有捷径,大模型正常的在工作。
这就是"无损"的含义:最终输出分布和单独跑一遍目标模型完全相同,不是拿质量换速度。草稿模型在幕后猜测,实际拿到手的却是目标模型的分布,不是在破坏输出而是在修正它。
为什么能提速
验证在数学上没毛病,但性能上真有用吗?
解码是内存受限的(memory bound),瓶颈在于从 GPU 内存加载权重而不是计算本身。对 k 个 token 做一次前向传播,内存带宽开销和对 1 个 token 做前向传播差不多,因为反正权重都要加载一遍。于是验证 k 个草稿 token 的开销,大致等于从头生成 1 个 token 的开销。这 k 个 token 里只要有一半被接受,就等于用一个 token 时间换来了好几个,可以说纯粹是并行的提速。
实际能提速多少?一般是 2 到 3 倍,但这不是固定值,这取决于草稿模型猜对的频率。温度(temperature)越低输出越可预测,接受率也就越高;代码比创意写作更规整,所以代码场景下的加速通常更明显。草稿模型对目标模型的行为摸得越透,白捡的 token 就越多。
树形投机采样
基础版投机采样只提出一条长度为 k 的序列,是单一直线式的猜测。但如果草稿模型拿不准下一个 token 是 "ice" 还是 "chocolate" 呢?基础版只能二选一,赌一把。
SpecInfer 的做法是两个都猜。草稿模型不再只给出一条序列,而是构建一棵可能延续路径组成的树。不再只猜 "ice cream and cake",而是同时猜:
'ice' → 'cream' → 'and cake'
'chocolate' → 'cake' → 'is great'
目标模型在一次前向传播中验证整棵树,沿着每个 token 都会被接受的最长路径走下去。这样每一步能接受的 token 更多,因为同时覆盖了几条可能的延续路径,而不是把宝全押在一个猜测上。
但是构建和验证一棵树实现起来更复杂。不过草稿模型确实在几个选项间犹豫不决时,树形方法比基础序列方法优势明显。
目前的应用现状
投机采样早已不是停留在纸面上的研究技巧。vLLM 内置支持,TensorRT-LLM 也已经集成,这也是过去两年推理 API 在模型本身没有变小的情况下大幅提速的原因之一。
这个领域已经走出了"草稿模型 + 目标模型"的基础设定。部分实现干脆跳过独立的草稿模型,转而在目标模型上直接挂几个额外的预测头,利用模型自身的内部表示并行预测未来的 token。这些预测头被称为 Medusa 头(这名字起得确实不错)。不用管理额外的模型,草稿模型和目标模型之间也不存在分布不匹配的问题,同一个模型在每次前向传播里只是多干了点活
而EAGLE 走得更远,在特征层面而非 token 层面做预测,接受率因此更高。投机采样的核心思想,基本上已经撑起了推理优化领域的一个分支
总结
投机采样的精妙之处在于:不改模型、不动训练、不碰权重,纯粹是利用了"验证比生成便宜"这个事实——因为解码是内存受限的——再靠一套数学上严密的拒绝采样(rejection sampling)技巧,保证输出和单独跑大模型时一模一样。
白得 2 到 3 倍推理加速,唯一的代价是需要一个够格的草稿模型,以及把它接进系统的工程量。
如果说 KV Cache 那篇讲的是怎么更好地利用 GPU 内存,这篇讲的就是怎么更好地利用每一次前向传播。两者合起来,覆盖了 LLM 推理里最重要的两个维度:空间和时间。
https://avoid.overfit.cn/post/163e4c3986794a51ba9c983c4f2f561e
by Vedanti