Transformer(2) - Speculative Decoding

transformer本身是一个自回归模型，意味着必须算出来第N个token后再接着算第N+1个token。为了更高效地利用gpu算力，业界就有人提出了一个好办法，可不可以用一个小很多的模型，专门用来预测未来的K个token，这个模型不需要很强的能力，只需要可以高准确率地根据上下文预测未来的K个token即可。预测完了以后再由主模型进行并行的验证。如果运气好，假设猜的K个token都对，那么相当于主模型并行的生成了K个token，而不是像原来那样迭代K次。 假如运气不好，第M个token猜错了，M之后的都要舍弃，但是M之前的所有token依然是没问题的。这个优化手段就叫“推测解码”，具体由很多的实践，eagle3，dflash，deepseek的dspark等

那么这个提效从gpu层面看具体是怎么工作的呢。

我们可以认为gpu本质上是一个矩阵计算器，他内部众多的计算单元在工作时并不是直接从内存(HBM)取数据，而是从内部自带的SRAM中读取，真实的计算数据要从内存中加载到SRAM中，这个加载速度相对于gpu的运算速度是非常慢的，SRAM是在制造芯片时就和gpu的运算单元一起光刻到晶元上的，物理距离就很近，带宽也大。而HBM是封装阶段和gpu通过硅中介连接到一起的。通信带宽和延迟都远不如SRAM。假如给定上下文ABCD我们要算后面10个token，就得先把ABCD序列得到的矩阵，发送到SRAM，然后分批从内存里搬运模型参数到SRAM中进行计算得到E，然后ABCDE序列的矩阵发送到SRAM（不考虑KV cache），再重复刚才的过程。这个过程中把数据从内存搬到SRAM的耗时在总耗时中的占比能达到90%。所以很多情况下gpu都在闲着等数据

在有推测解码的情况下，对于给定ABCD这个序列，我们先快速预测出后面10个token“EFGHIJKLMN”，假设预测的都对。那么实际的计算过程就是我们把“ABCDEFGHIJKLMN”这个序列发到SRAM里，然后分批从内存里搬运模型参数到SRAM中进行一次前向传播，transformer模型天然会产出每个位置上下一个token的概率，我们只需要判断主模型输出的每个位置上预测token的概率，和草稿模型给的每个预测token的概率就可以决定是否接受草稿模型预测的token。假如运气好，猜的token全被接受，就直接输出结果。这么一来同样是产生10个token，有推测解码的情况下从内存往gpu里搬模型参数的次数就变成一次了，结果就是gpu运算的耗时在总耗时中的占比大大提高了。

那么推测解码的性能提升效果怎么样呢，以deepseek为例，搭载了dspark的v4 flash模型速度最快能提升85%，v4 pro则更低。那么为什么连速度翻倍都没做到呢？哪怕只猜未来两个token，也应该速度翻倍啊？那是因为推测的准确率并不是100%，而且有错误推测的时候会浪费一定的算力，还得考虑草稿模型本身的耗时。综合下来之后就是我们看到的提效结果。

效率是提高了，但是如何保证外挂的推测解码不会影响模型的输出呢。

假如主模型采用的是贪婪解码，也就是说主模型每次预测下一个token时都取概率最大的那个token。那么草稿模型给的token如果不是概率最大的那个，直接就判错了。 如果主模型用的是随机采样，情况则稍微复杂一点，但是也会保证不影响主模型的概率分布，也就是说具体到某一次前向推理，草稿有可能会改变主模型的输出，但是总体来看不影响主模型的概率分布