【Deepseek 学投机推理】Sequoia

随着大型语言模型（LLMs）的使用日益广泛，对这些模型进行高效推理变得愈发重要。尽管推测解码（speculative decoding）最近成为加速推理的一个有前景的方向，但现有方法在扩展到更大的推测预算、适应不同超参数和硬件方面仍存在局限性。本文介绍了Sequoia，一种可扩展、鲁棒且硬件感知的推测解码算法。为了实现更好的可扩展性，Sequoia引入了一种动态规划算法，以找到推测标记的最优树结构。为了实现稳健的推测性能，Sequoia采用了一种新颖的采样和验证方法，在不同解码温度下均优于以往的工作。最后，Sequoia引入了一种硬件感知的树优化器，通过自动选择给定硬件平台上的标记树大小和深度来最大化推测性能。Sequoia将Llama2-7B、Llama2-13B和Vicuna-33B在A100 GPU上的解码速度分别提升了高达4.04倍、3.73倍和2.27倍。在L40的卸载设置中，Sequoia实现了低至0.56秒/标记的精确Llama2-70B推理延迟，这比我们优化的卸载系统（5.6秒/标记）快9.96倍，比DeepSpeed-Zero-Inference [2]快9.7倍，比Huggingface Accelerate [16, 45]快19.5倍。代码可在https://github.com/Infini-AI-Lab/Sequoia获取。

随着大型语言模型（LLMs）的广泛采用 [3, 7, 43]，高效地部署这些 LLM 变得愈发重要。然而，加速 LLM 推理具有挑战性，因为生成单个新标记需要访问 LLM 的所有参数 [34]。由于这一 I/O 瓶颈，硬件在生成过程中未能得到充分利用。这一问题在小批量推理和基于卸载的推理场景中更为严重，在这些情况下，生成一个标记所花费的时间与在现代 GPU 上处理包含数百或数千个标记的提示所需时间相当。

为了解决这一挑战，最近的研究引入了推测解码（speculative decoding），以在保留 LLM 输出分布的同时加速 LLM 推理 [5, 25, 28, 40]。这些方法利用一个或多个草稿模型（draft models）来预测 LLM 的输出；预测结果被组织成一个标记树，其节点表示不同的推测标记序列。随后，通过 LLM 的一次前向传播并行验证这些推测标记的正确性。使用标记树（而非单一序列）可以通过为每个标记位置提供多种选项来增加 LLM 接受的标记数量。

图 1：Sequoia 是一种可扩展的推测解码方法。左图：Sequoia 的树构建算法能够生成其平均标记数量（在验证后）随树规模持续增长的树结构，而现有的树结构则趋于饱和（asymptote）。这使得 Sequoia 在如卸载等高度受限于内存的场景中表现远优于现有方法。右图：一个对比可视化图，展示了 Sequoia 树结构与其他常见手工设计树结构的区别。

尽管已有大量关于基于树的推测解码方法的研究 [28, 40]，但我们的实验表明它们存在重要局限性。首先，我们观察到现有的标记树构建算法在小规模树上表现良好，但在大规模树上并非最优。例如，SpecInfer 通过 k 个独立序列构建标记树，其拓扑结构导致可接受的预期标记数量受限，且与树规模无关（图 1）。其次，现有标记树的采样与验证算法无法在不同推理超参数配置下保持良好性能；例如，SpecInfer [28] 和 SpecTr [40] 在低温条件下（图 3）表现较差，这是由于它们可能以高草稿模型概率重复采样错误标记。最后，现有系统无法针对任意硬件配置有效优化推测树的规模和形状（表 3）。这是因为现有模型对推测解码加速的描述 [25, 40] 假设验证时间为常数，而这一假设对大规模推测树不成立，导致这些模型无法有效选择最佳树维度。

本文旨在回答以下研究问题：如何设计一种最优的基于树的推测解码方法，以在现代硬件上实现最大加速？实现这一目标需要解决若干技术挑战。首先，对于任意树规模和深度，我们必须能够高效搜索指数级增长的树拓扑空间，以找到最大化预期生成标记数量的最优结构。其次，必须设计一种在不同推理超参数下表现优异的树采样与验证流程，避免重复采样错误标记，同时保持正确的输出分布。最后，针对任意硬件配置，我们必须能够选择与其最佳树维度匹配的树规模和深度，以实现最大加速。

本文提出 Sequoia ，一种可扩展、鲁棒且硬件感知的推测解码算法。如图 1 所示，Sequoia 相比增量解码可实现高达 10 倍加速，并通过以下关键技术解决上述挑战。

在第3.1节中，为解决第一个挑战，我们将树构建问题形式化为约束优化问题，并采用动态规划算法来发现最优的推测标记树。通过理论分析和实验验证，我们证明该树结构生成的标记数量是无界的，且其规模大致随树的尺寸呈对数增长。

• 在第3.2节中，针对第二个挑战，我们基于SpecInfer [28] 算法进行改进，通过从草稿模型中进行无放回采样（sampling without replacement） ，避免草稿模型重复犯错，同时保持目标模型的输出分布。我们从理论上证明，这种新的采样与验证方法能够在高、低温条件下均实现高接受率，并通过实验证实这一结论。

• 在第3.3节中，为解决最后一个挑战，我们提出一种硬件感知的树优化器，将验证时间建模为待验证标记数量的硬件相关函数，并基于此函数求解最优树形状和深度。我们证明，相较于与硬件无关的方法，该方法能够显著提升加速效果。

在第4节中，我们通过大量端到端实验和消融实验验证了Sequoia的有效性。我们在Hugging Face（及Accelerate框架）[16, 45] 的基础上，结合CUDA Graphs [31, 32] 实现了Sequoia。实验表明：单A100 GPU上 ，Sequoia对Llama2-7B模型的推理速度提升高达4.04倍 ；

L40 GPU的卸载场景中 ，Llama2-70B模型的推理延迟可降低至0.56秒/标记 ，而当前最先进的卸载系统（如DeepSpeed-Zero-Inference [2]）的延迟为5.5秒/标记，Huggingface Accelerate [16, 45] 的CPU卸载API则高达11秒/标记。

此外，消融实验表明：

1. 树结构的可扩展性 ：与k个独立序列相比，Sequoia的树结构在每个解码步骤中可多生成33%的标记

   （树规模≤512），展现了更优的扩展性；

2. 采样与验证的鲁棒性：Sequoia的算法对超参数（如温度、top-p）选择具有鲁棒性，相比SpecInfer和top-k采样方法，速度分别提升65%和27% ；

3. 硬件感知优化 ：Sequoia的硬件感知树优化器能自动为不同硬件选择最佳树规模和深度，从而最大化加速效果。

2 背景

在此，我们回顾基于树的推测解码方法。具体而言，我们将讨论现有方法如何选择推测树的结构（第2.1节）、其用于采样与验证标记树的算法（第2.2节），以及这些方法如何自动选择标记树形状（第2.3节）。

2.1 树构建

现有方法主要采用的树结构由从根节点分支出的k个独立序列 （每个序列长度为L）组成（根节点对应当前前缀[x₁,x₂,...,xₙ₋₁]）。

SpecTr 论文额外探讨了任意分支模式（k₁,k₂,...,kₜ），但实验表明其性能未优于独立序列。

Medusa 构建了一个完整的k叉树，这提高了每层的成功率，但在中等标记预算下无法形成深度树 [4]。

2.2 树采样与验证

我们现回顾 SpecInfer [28] 、SpecTr [40] 、朴素采样 [28] 及 top-k采样¹ 的标记树采样与验证方法：

采样方式 ：

SpecInfer、SpecTr 和朴素采样均从草稿模型中进行独立同分布采样（有放回） 。

top-k 采样则从草稿模型中选择概率最高的前k个标记 。

验证逻辑 ：

SpecInfer与SpecTr ：比较草稿模型与目标模型对采样标记的概率，以决定是否接受该标记。

朴素采样与top-k采样 ：从目标模型分布中采样一个标记，若其与推测树中的标记匹配则接受。

验证流程 ：

所有方法均以递归方式 从根节点开始验证树结构，仅在每个节点应用的验证算法不同。

SpecInfer方法 ：

SpecInfer通过迭代验证从一个或多个草稿模型中采样的标记。与原始推测解码方法 [25] 类似，它通过比较草稿模型与目标模型的概率来决定是否接受标记。

值得注意的是，尽管SpecInfer允许从k个不同草稿模型生成k个子节点，但本文聚焦于仅有一个草稿模型的常见场景 。因此，我们对比的是SpecInfer的单草稿模型版本（即对单模型进行k次采样）。

SpecInfer的伪代码详见附录B.2。

1. 独立同分布采样（i.i.d. sampling） 独立 ：每次采样（生成标记）是相互独立的事件，前一次采样的结果不会直接影响下一次采样。 同分布 ：所有采样均基于草稿模型（draft
   model）的同一概率分布（例如，草稿模型输出的下一个标记的概率分布）。 例子 ： 假设草稿模型预测下一个标记的概率分布为：[0.5,
   0.3, 0.2]（对应标记A、B、C）。 每次采样时，均根据这一分布独立选择标记，可能出现多次连续选择A的情况。
2. 有放回（with replacement）
   放回机制 ：每次采样后，被选中的标记会被"放回"候选池，允许后续采样中再次选择同一个标记。
   影响 ：同一标记可能被多次选中，导致推测树中出现重复分支。
   对比 ：

无放回采样 ：一旦某个标记被选中，后续采样中该标记将被排除（如随机排列）。 有放回采样 ：标记可被重复选中，增加多样性但可能导致冗余。

3. 在推测解码中的具体表现 SpecInfer、SpecTr 和朴素采样 ： 这三种方法在生成推测树时，均通过草稿模型的输出分布进行独立、有放回的重复采样 。 优点 ：简单高效，允许快速生成多个候选标记序列。 缺点

：可能导致重复采样错误标记（例如，草稿模型对某个错误标记概率过高时，可能被多次选中）。

4. 与其他方法的对比 top-k 采样 ： 与上述方法不同，top-k 采样会无放回地选择概率最高的前k个标记 ，避免重复但可能限制多样性。 Sequoia 的改进 ： Sequoia 采用无放回采样

（见论文第3.2节），避免重复错误，同时保持输出分布的正确性。

朴素采样与top-k采样

对于标记树中的某个节点，朴素采样和top-k采样的验证算法首先从目标模型的分布 P(⋅∣x <n) 中在该节点进行采样，若采样结果与该节点的某个子节点一致，则接受该样本。

这一验证算法能够自然保持目标模型的输出分布 ，因为无论标记树如何生成，该算法始终从目标模型中直接采样（而非像SpecTr和SpecInfer那样从草稿模型中采样）。

这一观察促使我们在top-k采样方法中选择通过选取每个节点的前k个子节点来构建树 ，而非朴素采样中采用的k次独立同分布采样（有放回）。在第3.2节的实验中，我们使用top-k采样方法以更深入地理解该验证算法的局限性。

3. 与草稿模型无关的鲁棒性
   不依赖草稿模型的分布 ：
   SpecTr和SpecInfer等方法通过比较草稿模型与目标模型的概率来决定是否接受标记，这可能导致以下问题：
   草稿模型的分布可能与目标模型存在偏差（例如低估或高估某些标记的概率）。
   验证算法可能因草稿模型的错误推测而接受不符合目标模型分布的标记。
   朴素采样和top-k采样的优势 ：
   由于它们仅依赖目标模型的采样结果 ，完全规避了草稿模型的潜在偏差，从而保证输出分布的正确性。
4. 对比：为何SpecInfer等方法可能破坏分布 依赖草稿模型的比较逻辑 ： SpecInfer通过比较草稿模型和目标模型的概率阈值来接受标记。若草稿模型对某个错误标记的概率估计过高，可能导致该标记被错误接受，从而扭曲输出分布。
   朴素采样和top-k采样的"无偏性" ：
   它们直接从目标模型采样，即使树的结构由草稿模型生成，最终的接受决策仍基于目标模型的真实分布，因此不会引入额外偏
   减少重复标记：朴素采样由于是有放回的，可能导致同一标记被多次选中，生成重复的子节点。这在推测解码中可能效率低下，因为重复的标记无法提供更多有用的信息，反而浪费计算资源。而top-k采样选择前k个最高概率的标记，确保每个子节点都是独特的，且更可能符合目标模型的分布。

• 提高验证效率：top-k采样通过选择概率最高的标记，使得验证时更有可能找到与目标模型匹配的标记，从而提高接受率。而朴素采样可能生成低概率的标记，导致验证时被拒绝的概率增加，降低整体效率。

解释一下以上公式的计算推导过程