投机推理算法的常见优化手段

在前文我介绍了投机推理这门主流大模型推理加速技术的算法原理，本文更进一步记录在这个基础算法框架上衍生出来的优化手段。

树状草稿

树状草稿最早是在SpecInfer这篇论文里提出的，已经提出就被后续的Medusa和Eagle采纳。其核心思想就是在起草阶段不要每次只生成一条草稿序列，而是生成多条候选草稿序列，每条草稿序列都有相同的前缀，组织成一棵树。如此一来，在验证过程中就能最大化接收草稿token，从而提高草稿质量。

草稿树的生成

怎样生成一颗草稿书就有不少讲究。最早的SpecInfer就提出了两种不同的方法：基于合并的方法和基于扩展的方法：

• 基于合并的方法使用k个小模型并行推理，这k个小模型可以是同一个小模型的不同蒸馏版本、不同微调版本，用以生成具有一定多样性的草稿。每个小模型独立做自回归推理生成多条候选序列后，将所有候选序列按公共前缀合并，这就形成了一颗草稿树。
• 基于扩展的方法则只使用一个小模型，但每次小模型的前向推理并不是只取概率值最高的一个token，而是取topK个token，下一步基于这K个token再做topK扩展。一个显而易见的问题是指数爆炸问题——如果设定每次起草阶段让小模型推理m次，每次取topK个token，那就会产生\(K^{(m+1)}\)个草稿token，只要K和m稍微一大，比如\(K=3, m=5\)，那就会产生\(3^6 = 729\)个token，那么大模型验证阶段就要一次计算729个token对应的togits，时延开销太大了。对此，SpecInfer提出的补丁是让K逐步递减，再推理开始就设置好每一步K的值，比如起草的第一步可以设置\(K=3\)，第二、三步让\(K=2\)，后面的\(K\)则为1，那就只会生成\(3 + 3*2 + 3*2*2 + 3*2*2 + 3*2*2 = 45\)个token。这个补丁的弊端是草稿树的结构从一开始就固定好了，无法动态适配。

后续基于草稿树的投机方案，包括Medusa和Eagle，基本延续了基于扩展的静态草稿树生成思路。而Eagle2的一个核心贡献就是提出了一套让草稿树从静态变成了动态构建的方：每一轮迭代并不是对草稿树最底层么欸一个叶节点做扩展，而是支队底层最可能的topK个叶节点做扩展。假设\(K=5\)，那么即便第二轮这颗草稿树的的叶节点有25个，Eagle2叶只会对其中最可能的5个叶节点做扩展。至于怎么选这topK个叶节点，则是把每条路径上每个节点的概率累乘起来，选择累乘概率最高的叶节点——这也等价于每次旨在累乘概率最高的几条路径上做扩展。

草稿树的验证——树形掩码

在起草阶段获得一颗草稿树以后，怎么对其做验证，也是一个难点。这里面的关键是包括mask、kvcache、postion id在内的所有相关组件都要与草稿树做适配。假设我们获得了以下这么一颗草稿树：

        root
       /    \
      A      B
     / \      \
    C   D      E

首先我们要按照深度优先或者广度优先原则将这颗二维的树展平成一维的序列：[root, A, B, C, D, E]，然后哦通过构造特殊的树状掩码来让后面的token只能看到它的祖先。比如token B只能看到root而看不到A，token C只能看到A和root但看不到B。数学上就是要定义这么一个矩阵\(M\)，其中\(M_{ij}=1\) 当且仅当\(i\)是\(j\)的祖先token，否则\(M_{ij}=0\)：

   root	A	B	C	D	E
root 1	0	0	0	0	0
A	 1	1	0	0	0	0
B	 1	0	1	0	0	0
C	 1	1	0	1	0	0
D	 1	1	0	0	1	0
E	 1	0	1	0	0	1

准备好所有输入以后，一同送入大模型做一次前向推理：

logits = llm.forward(draft_tree, tree_mask, position_ids, kvcache)

得到的会是每个树节点位置 \(i\) 的概率分布\(P(i| \texttt{ancester}(i))\)， 表示"在节点\(i\)对应的上下文里，大模型认为下一个 token 应该是什么"。拿到logits以后，就可以根据前文所述的验证方法来验证哪些草稿token是正确的。草稿序列的校验和草稿树的家教严区别仅在于前者只需要便利序列中的每一个token，而草稿树的校验则需要便利每一层的每一个token。

起草-验证并行

大小模型投机推理的范式是小模型生成草稿token然后大模型验证token，这过程里小模型在生成的时候大模型不干事，大模型验证的时候小模型不干事。一个自然的想法是让小模型和大模型并行起来，不要浪费。理想很丰满，现实很骨感。传统投机推理中起草和验证之所以串xing，是因为他们之间有数据依赖：小模型的起草需要完整的已验证的token序列，而大模型的验证需要完整的草稿token序列。想要让起草和验证并行起来，就必须先克服起草和验证两个过程之间的数据依赖。可惜的是他们之间的数据依赖是内生的，因此目前的研究多是利用一些启发式或者trick。

腾讯AI团队的Pearl这篇论文提供了一个思路：

• Pre-Verify：在 草稿模型起草的同时，主模型提前验证第一个草稿token。如果第一个 token 就被拒，直接跳过本次验证阶段，避免无效计算。
• Post-Verify：如果上一轮的所有草稿 tokens 被全部接受，草稿模型可以假设这一轮所有草稿token也会被接受，那么在主模型验证这一轮的同时，继续生成下一批 tokens，实现流水线并行。

动态草稿长度

经典投机推理的另一个问题时固定的草稿生成长度和验证长度，比如说Eagle3和DFlash都会设置一个固定的每次推理生成多少个草稿token。这的的问题是再某些领域草稿质量高（比如数学、代码）而某些领域草稿质量低（比如纯聊天），如我们的一个模型服务要兼顾所有场景所有类型的画，那不可避免的有些时候会生成很多低质量的草稿token，它们既浪费了草稿生成时间，也浪费了验证时间，引发了传统投机推理的一大痛点：高并发的时候收益迅速降低，甚至开启投机推理特性以后时延和吞吐反而下降了。一个自然而然的想法就是，动态调整草稿生成的长度，在投机推理不容易生成草稿token的时候少产生，少验证，反之。

对于这个问题，阿里的千问应用团队近期提出了一套系统化的解决方案ECHO。通过对投机推理系统做profiling分析，团队发现在高并发条件下“起草-验证”中验证花的时间是最多的，是整个投机推理过程的性能瓶颈。其核心思想是先定好每一次验证作多验证多少个token，这是一个全局预算\(K_{max}\)，然后一个batch内所有token的草稿长度共用这一个预算。理所当然的就是，有些请求生成的草稿token多，剩下的请求可以生成的token比较少。具体哪些请求能生产多一些，哪些要生成少一些，就有两个核心组件：置信度门控和全局预算调度器来决定。

• 在最开始的时候，ECHO会根据业务数据集离线得出一些检查点——草稿树的一些特定深度层级——一般来说过了这些深度后token的采纳率会大幅降低
• 实际推理时，当某个请求的一条分支在门控检查点的置信度高时，调度器会继续加深该请求的草稿树，生成更多层级的 token，否则则截断该分支，停止继续向下生成。
• 如果某个请求因低置信度被截断，但全局预算仍有剩余，就可以在被截断的深度上增加该层的草稿token数量（加宽），以提高猜中目标 token 的概率。
• 在高并发场景下，低置信度请求释放的预算会被重新分配给其他高置信度请求，让它们生成更深的草稿树