SpecInfer: Accelerating Large Language Model Serving with Tree-based Speculative Inference and Verification

URL

• paper: https://arxiv.org/pdf/2305.09781
• code: https://github.com/flexflow/FlexFlow/

TL;DR

• 本文利用了 Speculative decoding 思想实现了一个 SpecInfer 大语言模型推理加速系统

• 和标准 Speculative decoding 使用 seq 组织预测的 token 不同，SpecInfer 使用 tree 数据结构组织预测的 token

• SpecInfer 中的 target 模型可以并行验证整个 token tree 的正确性，而不是 token seq，且可以保证和原模型输出完全一致

Algorithm

增量式解码、序列式投机解码、树式投机解码三者对比

左图表示传统增量式推理 LLM 的解码方式
右图从上到下依次表示：增量式解码、序列式投机解码、树式投机解码

候选树生成和验证流程

SSM (small speculative model) 给出候选树
target model 对候选树进行验证

如何对树式结构并行解码

注意：这的 “并行” 是指可以一次解码一棵树而不需要将树深度优先搜索得到多组序列然后一一解码

1. 使用 kv cache 对解码过程进行加速是一个常见的方法，但树式结构无法直接使用 kv cache（因为要动态剪枝，如图左侧所示）
2. 因此本文提出一种新颖的方法使得树式结构可以并行解码，这个方法包括两个改进：
   1. 对树进行 深度优先搜索，使得树变成一个序列
   2. 将传统的上三角 attention mask 变成一个 topology-aware mask，即通过 mask 来确定 token 两两之间的可见性（如图右侧所示），这种可见性包含：
      1. 每个 token 看不到之后的 token（矩阵下三角置 $-\infty$）
      2. 每个 token 看不到不在其拓扑路径上的 token

总体算法流程

1. 先用 small speculate model 推理 sequence 得到候选树 $N$
2. 在使用 LLM 树式并行解码 sequence 得到大模型预测结果 $O$
3. 用贪心算法或随机匹配算法对候选树 $N$ 逐个节点验证（接受或拒绝）
4. 重复直到 <EOS> 出现

贪心验证算法和随机验证算法

• 这里的两种算法和 https://zhangzhe.space/2024/08/14/Fast-Inference-from-Transformers-via-Speculative-Decoding/ 里的流程基本一致，不一样的地方仅仅在于 tree / sequence 两种候选 token 的组织形式

贪心验证

1. $u$ 表示当前节点，从根节点开始广度优先遍历
2. $p_v=u$ 表示 $v$ 是 $u$ 的子节点
3. $t_v=O(u)$ 表示大模型和投机模型预测结果一致

随机验证

1. $u$ 表示当前节点，从根节点开始广度优先遍历
2. 当在 [0, 1] 之间随机均匀采样得到的随机数 $r \lt \frac{P_{LLM}}{P_{SSM}}$ 时，接受这个节点，否则拒绝

Thought

• 这篇论文站在 speculate decoding 的肩膀上，加入了树结构的候选列表，算是 speculate decoding 的超集
• topology-aware mask 和 decoder attention mask 结合的方法还是挺巧妙的
• 由于是树结构存储候选词，那么多个 SSM 是天然适配的，算是个工程角度很好用的 trick