EasySpec: Layer-Parallel Speculative Decoding for Efficient Multi-GPU Utilization

作者/机构: Yize Wu (中国科学院软件研究所智能软件研究中心 & 中国科学院大学), Ke Gao (中国科学院软件研究所智能软件研究中心), Ling Li (中国科学院软件研究所智能软件研究中心 & 中国科学院大学), Yanjun Wu (中国科学院软件研究所智能软件研究中心)

A1 主要贡献

本文旨在解决在多GPU系统中使用推测解码（Speculative Decoding）时出现的资源利用率低下的问题。

核心问题：推测解码结合张量并行（Tensor Parallelism, TP）可以加速大型语言模型（LLM）的推理，但通常草稿模型（draft model）的最佳TP规模小于基础模型（base model），这导致在草稿生成阶段部分GPU处于空闲状态，造成计算资源浪费。这种低效源于草朵模型各层之间看似自然但并非必要的顺序执行依赖。

研究目标：提出一种层并行推测策略，旨在优化草稿生成阶段的多GPU利用效率，从而加速整个推理过程。

创新点 (EasySpec)：
本文提出了EasySpec，一种无需训练、即插即用的层并行推测策略，通过以下两个核心组件实现加速：

  1. 模糊推测 (Fuzzy Speculation)：EasySpec打破了草稿模型中连续注意力层之间的数据依赖关系。它允许多个层以最近的隐藏状态作为共同输入，在多个设备上同时并行执行，实现层级并行。这种“模糊”的推测方式虽然会引入近似误差，但能显著加快草稿生成速度，同时保持较高的推测准确率。

  2. 奖励校准 (Bonus Calibration)：为了防止模糊推测中产生的近似误差在推理过程中长期累积，EasySpec在每次“草稿-验证”迭代后增加了一个校准步骤。该步骤将已接受的令牌（tokens）和“奖励”令牌（bonus token）合并，通过一次标准的、层顺序的前向传播来重新计算并更新草-稿模型的键值缓存（Key-Value Cache），从而以极小的额外延迟代价消除了误差累积。

如下图1所示，EasySpec修改了标准的推测解码流程，通过模糊推测实现层并行加速，并在验证后增加奖励校准阶段来确保准确性。图2直观展示了与传统TP和标准推测解码（SD）相比，EasySpec如何通过层并行推测解决草稿阶段的GPU闲置问题，从而提升了多GPU系统的利用效率。

图1：EasySpec概览：标准推测解码与EasySpec的比较。主要区别在于：(1) 模糊推测使用层并行化来加速草稿生成。(2) 无论接受长度如何，模糊的KV项都会被丢弃，而推测解码会保留已接受的项。(3) 推测解码是两阶段的，而EasySpec是三阶段的，增加了一个用于奖励校准的阶段。
图1：EasySpec概览：标准推测解码与EasySpec的比较。主要区别在于：(1) 模糊推测使用层并行化来加速草稿生成。(2) 无论接受长度如何，模糊的KV项都会被丢弃，而推测解码会保留已接受的项。(3) 推测解码是两阶段的，而EasySpec是三阶段的，增加了一个用于奖励校准的阶段。
图2：分布式系统中不同解码算法的GPU运行时情况。推测解码（SD）实现了令牌级别的并行（基础模型块更少），但在草稿生成阶段导致了多GPU利用率不足。EasySpec通过层并行推测解决了这个问题。
图2：分布式系统中不同解码算法的GPU运行时情况。推测解码（SD）实现了令牌级别的并行（基础模型块更少），但在草稿生成阶段导致了多GPU利用率不足。EasySpec通过层并行推测解决了这个问题。

实验结果表明，EasySpec相较于原生解码（vanilla decoding）可实现最高4.17倍的峰值加速。具体而言，草稿生成阶段最多可加速1.62倍，而推测准确率的下降不超过7%，同时保持了基础LLM的原始输出分布。

A3 背景知识

2.1 推测解码

推测解码是一种两阶段的非自回归解码方法，用于加速推理。这两个阶段分别是草稿生成（drafting）和验证（verification），在推理过程中迭代执行。在时间步t，输入令牌序列为X。在草稿生成阶段，一个更小更快的模型M'作为草稿模型。M'自回归地运行n次，生成一个推测令牌序列$x'_{t+1}, ..., x'_{t+n}$以及概率分布$p'_{t+1}, ..., p'_{t+n}$。然后，原始的大模型M将整个序列作为输入，进行单次前向验证，其输出为$p_{t+1}, ..., p_{t+n+1}$。令牌$x'_i$的接受概率为$min(1, p_i(x_i)/p'(x'_i))$。验证采样后，m个被接受的令牌序列$x_{t+1}, ..., x_{t+m}$（与$x'_{t+1}, ..., x'_{t+m}$相同）将被附加到最终输出中，同时还有一个“奖励”令牌（bonus token）$x_{t+m+1}$（称为奖励是因为它是验证过程的副产品），可以从分布p中采样得到：

$$\begin{aligned} p = \begin{cases} norm(\max(0, p_{t+m+1} - p'_{t+m+1})) & (m < n) \\ p_{t+m+1} & (m = n) \end{cases} \end{aligned}$$

通过整个验证过程，最终的输出令牌序列与基础模型的原始分布保持一致。该算法的详细信息见附录A。

2.2 张量并行分布式推理

张量并行（Tensor parallelism, TP）是一种将大型模型部署在多个设备（GPU）上的技术。对于TP分布式推理，张量被分割成多个块，每个块都分派到一个唯一的GPU上。在计算操作期间，每个GPU首先独立地处理其各自的块，然后同步结果以确保最终结果与原始结果相同。在现代多GPU系统中，先进的硬件和通信协议使得TP可以用于推理加速。

最优TP规模主要受参数大小的影响。首先，TP规模必须能整除隐藏层维度和注意力头的数量。例如，在Qwen-2-0.5B-Instruct中，这些数字分别是4864和14，这将最大可行的TP规模限制为2。其次，计算工作负载随参数大小而变化，这进一步影响了TP的效率。关于这些设备级特性的详细讨论见2.3节。

2.3 设备特性

现代GPU具有强大的计算能力，因此计算延迟很短。因此，LLM推理中涉及的矩阵计算通常受限于内存带宽或控制流开销【索引14，SqueezeLLM: Dense-and-sparse quantization，2023，arXiv】。在多GPU系统中，计算也可能受到通信和同步开销的限制。

形式化地，我们将1个令牌前向传播的工作负载表示为W，在单个GPU上执行W的延迟表示为$T_{exe}(W)$，TP分布的额外开销表示为$T_{addi}$。s个令牌并行的工作负载是sW，而s规模TP的工作负载是W/s。如果W相对于计算能力来说相对较小，我们可以预期以下关系：

$$T_{e x e}(W) \approx T_{e x e}(s W)<s T_{e x e}(W)$$ $$T_{exe}(W/s) + T_{addi} > T_{exe}(W)$$

令牌级并行的有效性可以通过方程(2)来说明，推测解码正是利用这一点来加速。方程(3)突显了分布式推理中通信和控制流开销的影响，这限制了TP规模不能过大，并限制了草稿模型的多GPU利用率。

A2 方法细节

本节详细介绍EasySpec如何增强多GPU并行性以加速推测解码的草稿生成阶段。两个关键步骤——模糊推测和奖励校准——将在各小节中具体说明。

3.1 模糊推测

如第1节所述，传统的推测过程限制了各层必须顺序执行。算法1展示了N个连续层的执行过程。如果我们去掉算法中的$h'_i$，则有

$$h_{i+1}=h_i+Attnoutput_i+MLPoutput_i$$

，这说明了第i+1层的输入与第i层的输出之间存在数据依赖关系。换言之，在前一个注意力层（$Attnoutput_i$）和MLP层（$MLPoutput_i$）的输出计算完成之前，无法对上层i+1（以及i+2到N）进行任何操作。

输入：隐藏状态h，N个连续的注意力层Attn1,··· ,AttnN和MLP层MLP1,··· ,MLPN
h1 = h
for i = 1 to N do
  Attnoutputi = Attni(hi)  (顺序执行)
  h′i = hi + Attnoutputi
  MLPoutputi = MLPi(h′i)
  hi+1 = h′i + MLPoutputi
end for

算法1：标准推测

当草稿模型直接用于推理时，这种对执行顺序的限制是必要的。以不同方式执行模型会使输出概率分布p'变为p''，从而损害生成内容的质量。然而，当草稿模型仅用于推测时，情况就不同了。无论p''与原始分布p'的偏差有多大，只要我们将接受概率设置为$min(1, p/p'')$，理论上最终的输出分布将与基础模型保持一致。这为优化推测策略提供了一个潜在方向：如果输出分布p'可以被p''高效地近似，即$p'' \approx p'$，那么整个过程可能会加速，而最终输出保持无损。也就是说，一种稍微模糊但更快的推测方法可能优于精确的方法。

我们提出了一种有效的模糊推测近似方法。受【索引15，Infinigen: Efficient generative inference of large language models with dynamic kv cache management，2024，OSDI】的启发，我们观察到，通过将所有1 < i ≤ N的输入$h_i$替换为$h_1$，可以很好地近似注意力层的输出。在方程(4)中，隐藏状态的更新$h_{i+1} - h_i = Attnoutput_i + MLPoutput_i$相对不那么显著，因此$h_i$和$h_{i+1}$之间的余弦相似度可以接近1，从而所有$h_i$和$h_1$之间的相似度也很高（图3是这种现象的二维演示）。这种输入的高相似性使得其他隐藏状态也能够被精确地近似（理论解释见附录E.1，经验结果见表4）。

图3：hi和hi+1之间高余弦相似度的二维演示。
图3：hi和hi+1之间高余弦相似度的二维演示。

算法2展示了该方法的详细步骤。从模型架构的角度看，注意力层的输入不是最新的隐藏状态，而是来自前几层的隐藏状态（图1）。通过这种对层执行策略的修改，我们打破了N个连续注意力层之间的数据依赖关系：一旦计算出某个j的$h_j$，那么j到j+N之间的所有注意力层都可以执行，它们之间没有数据依赖，从而实现了层级并发。在多GPU系统中，这些注意力层可以在不同设备上同时运行（图2）。

输入：隐藏状态h，N个连续的注意力层Attn1,··· ,AttnN和MLP层MLP1,··· ,MLPN
h1 = h
for i = 1 to N do
  Attnoutputi = Attni(h1)  (并行执行)
end for
for i = 1 to N do
  h′i = hi + Attnoutputi
  MLPoutputi = MLPi(h′i)
  hi+1 = h′i + MLPoutputi
end for

算法2：模糊推测

记顺序执行和模糊推测在这N个注意力层上的执行时间分别为$T_{seq}$和$T_{fuzzy}$，我们有：

$$\begin{aligned} \begin{gathered} T_{seq}=NT_{exe}(A) \\ T_{fuzzy}=T_{exe}(A)+T_{addi} \end{gathered} \end{aligned}$$

其中A是一个注意力层的工作负载。与传统的TP方法不同，运行一个注意力层的延迟足够长，可以补偿额外的开销，即当N > 1时，$(N-1)T_{exe}(A) >> T_{addi}$，从而导致$T_{fuzzy} < T_{seq}$。MLP层仍按原始顺序执行，因为并行化它们会导致显著的近似误差和端到端性能下降（附录B）。

为了实现更高的接受率，EasySpec在草稿生成阶段还结合了树状注意力（tree attention）【索引16，Specinfer: Accelerating large language model serving with tree-based speculative inference and verification，2024，ASPLOS】，这是现有推测解码方法中常用的一种技术【索引7，Medusa: Simple llm inference acceleration framework with multiple decoding heads，2024，arXiv】【索引17，EAGLE: Speculative sampling requires rethinking feature uncertainty，2024，arXiv】【索引18，EAGLE-2: Faster inference of language models with dynamic draft trees，2024，arXiv】。树状注意力通过增加草稿令牌序列的数量来提高接受率。如图5所示，草稿令牌被构造成一棵树，每条路径对应一个序列。所有序列的推测和验证都可以使用一个二维的树状注意力掩码在一次前向传播中完成。

3.1.1 层并行策略

对于层并行大小N，我们排除了第一层和最后一层，并对1到N-1层、N到2N-1层等进行并行化，直到剩余的层（除了最后一层）不足以构成一个N-集。例如，对于Llama-3-8B-Instruct（32层）和N=4，我们将1-3层、4-7层、...、28-30层并行化，而第28、29和30层必须组成一个3-集，因为第31层是最后一层（见表1）。采用这种策略的原因是，在我们的实验中，并行化中间层（排除第一层和最后一层）能产生最佳性能。我们推测这是因为第一层和最后一层对近似误差更敏感。

需要注意的是，可能存在针对特定模型和/或任务的更好策略，但当前策略已经产生了足够令人满意的加速效果。此外，尝试其他策略简单且轻量，因为它不需要修改原始模型架构，只需更改配置即可。

表1：草稿模型的层并行策略示例。其他草稿模型的策略与3.1.1节中指定的相同。
表1：草稿模型的层并行策略示例。其他草稿模型的策略与3.1.1节中指定的相同。

3.2 奖励校准

形式上，令牌j的推测可以表示为$x'_j, p'_j, kv'_{j-1} = M'(x'_{i<j-1}, kv'_{i<j-1}, x'_{j-1})$</strong>。模糊推测会给$x'$, $p'$和$kv'$的值带来误差。$x'$和$p'$中的误差仅限于单次“草稿-验证”迭代内，因为被拒绝的令牌和所有输出分布在验证后都会被丢弃。然而，$kv'_{j-1}$中的误差会影响j之后的所有生成，导致长期误差累积，从而严重降低准确性。

为了防止这种误差累积，需要精确的$kv'$值。然而，获取这些值需要顺序运行模型的各层，而这正是我们试图通过模糊推测优化的部分。为了解决这个两难问题，我们提出了一种以最小额外延迟获取精确KV值的高效方法。在标准草稿生成阶段开始时，上一次迭代的奖励令牌$x_{t+m+1}$被输入到草稿模型中，用于推测$x'_{t+m+2}$。这是一个单令牌的前向传播，并且如2.1节所讨论的，也是内存受限的。因此，与验证阶段类似，这里也可以应用令牌并行。通过将已接受的令牌序列$x_{t+1}, \dots, x_{t+m}$与奖励令牌$x_{t+m+1}$连接起来，并将整个序列$x_{t+1}, \dots, x_{t+m}, x_{t+m+1}$重新输入到草稿模型中，可以通过一次顺序的前向传播产生精确的KV值$kv'_{t+1}, \dots, kv'_{t+m}$（以及$kv'_{t+m+1}$）。根据方程(2)，这样一次前向传播的延迟几乎等于一次单令牌的顺序前向传播。

我们基于上述方法提出了奖励校准（bonus calibration）。如图1所示，我们首先丢弃所有模糊的KV项，无论对应的令牌是被接受还是被拒绝（例如，被接受的项1和2也被丢弃）。然后，进行一次令牌并行的前向传播，用精确的值重新填充KV缓存。从执行的角度看，精确的KV值是奖励令牌的下一个令牌推测过程的副产品，类似于奖励令牌本身是如何由验证过程产生的一样。这就是我们称此过程为“奖励”的原因。奖励校准可以生成一个精确的候选令牌和分布，而后续轮次的推测可以恢复为模糊模式。总体延迟略有增加，从$nT_{fuzzy}$增加到$T_{seq} + (n-1)T_{fuzzy}$（n是推测长度），但误差的消除是至关重要且显著的。更详细的理论支持见附录E.2。

A4 实验环境

  • 模型与基准测试

    • 基础模型: Llama-3-70B-Instruct, Qwen2-72B-Instruct, Llama-3.3-70B-Instruct，以及任务特定的模型 Qwen2-Math-72B-Instruct 和 Qwen2.5-Coder-32B-Instruct。
    • 草稿模型: 使用相应基础模型同系列的0.5B/1.5B/3B/7B/8B版本。
    • 基准测试集: 涵盖多种任务，包括语言理解(MMLU)、代码生成(HumanEval)、数学推理(MATH)、指令遵循(IFEval)和多语言使用(MGSM)。所有实验均采用思维链(chain-of-thought)推理，最大生成令牌数为128。此外，还在Spec-Bench上进行了评估以与相关工作进行比较。

  • 硬件配置

    • GPU: 8 × A100 GPUs。

  • 软件与配置

    • 层并行大小(N): 默认为4，这是大多数测试模型的最佳值。
    • 张量并行(TP)大小: 草稿模型为1，基础模型为8，这是大多数情况下的最优配置。
    • 推测长度: 5。
    • 基线设置: 与张量并行分布式解码(TP)、分布式系统中的推测解码(+SD)、树状注意力(+tree)以及EAGLE-2进行比较。
    • 性能指标: 单批次推理的令牌吞吐量、每100个令牌的草稿和验证阶段运行时间、接受率。

A4 实验结果

4.1 主要结果

4.1.1 有效性

EasySpec在各类模型和数据集上均表现出稳定且显著的加速效果。如表2所示，在对Llama-3-70B-Instruct和Qwen2-72B-Instruct的测试中，EasySpec始终能够加速整个推测解码过程。实验发现，尽管基础模型的验证阶段可以通过TP加速，但草稿生成阶段往往成为性能瓶颈，占据了大量的总运行时间。EasySpec通过层并行化，将草稿阶段的运行时间最多减少了1.62倍，从而提升了整体吞吐量。与此同时，推测准确率的下降幅度不超过7%，这表明结合了奖励校准的模糊推测在所有评估场景中都足够精确。此外，在针对Qwen2-Math-72B-Instruct和Qwen2.5-Coder-32B-Instruct等任务特定模型，以及其他多种草稿-基础模型组合的测试中（见表7和表8），EasySpec同样展现了其有效性。

表2：两个主流模型在不同数据集上每100个token的运行时间、总加速比和接受率。'd'、'v'分别代表草稿和验证的运行时间（秒/100 token），'total'和α分别代表总加速比和接受率。更详细的数据见表9。
表2：两个主流模型在不同数据集上每100个token的运行时间、总加速比和接受率。'd'、'v'分别代表草稿和验证的运行时间（秒/100 token），'total'和α分别代表总加速比和接受率。更详细的数据见表9。

4.1.2 泛化性与稳定性

与使用微小专用草稿模型的方法相比，EasySpec表现出更强的泛化性和稳定性。为了验证这一点，我们将EasySpec与一个典型的微小草稿模型方法EAGLE-2进行了比较。EAGLE-2使用一个经过精心训练的单层草稿模型以追求极低的草稿延迟。如表3所示，在Spec-Bench上的评估结果显示：

  1. 性能更优：对于Llama-3-70B-Instruct，使用EasySpec加速的1B草稿模型在所有任务上的平均令牌吞吐量均高于EAGLE-2（约1B参数），表明如果存在同系列的小模型，结合EasySpec足以提供出色的性能。
  2. 稳定性更高：EasySpec的吞吐量方差显著低于EAGLE-2，这得益于其远高于后者的推测准确率（接近0.8 vs 不到0.4）。
  3. 泛化性更强：在Llama-3.3-70B-Instruct上，性能差距更加明显，进一步说明经过训练的草稿模型对其他基础模型的泛化能力较差，而同系列的小模型则保持了良好的兼容性。EasySpec作为一个无需训练或微调的即插即用方法，具有更高的适应性。
表3：EasySpec和EAGLE-2在Spec-Bench上对Llama-3-70B-Instruct的平均令牌吞吐量（方差）和接受率。吞吐量格式为均值（方差）。α代表接受率。
表3：EasySpec和EAGLE-2在Spec-Bench上对Llama-3-70B-Instruct的平均令牌吞吐量（方差）和接受率。吞吐量格式为均值（方差）。α代表接受率。

4.2 近似精度

模糊推测的可行性建立在对草稿模型隐藏状态的高精度近似之上。表4展示了精确隐藏状态与模糊近似隐藏状态之间的平均余弦相似度。实验使用Llama-3-8B-Instruct作为草稿模型，在MMLU基准上进行测试，评估了输入隐藏状态h、注意力层输出Attnoutput以及注意力头内部的查询(q)、键(k)和值(v)。结果显示，即使随着并行层数的增加（LP size从2到4），余弦相似度虽然有所下降，但仍保持在0.8以上，表明近似精度足够高。特别值得注意的是，查询（q）和键（k）的近似效果非常好，余弦相似度接近1，这对于维持其他隐藏状态和输出的精度至关重要，因为注意力权重的误差会呈指数级放大。

表4：精确与近似隐藏状态之间的平均余弦相似度。测试模型和任务为Llama-3-8B-Instruct和MMLU。
表4：精确与近似隐藏状态之间的平均余弦相似度。测试模型和任务为Llama-3-8B-Instruct和MMLU。

4.3 消融研究

通过消融研究验证了EasySpec各组件的有效性。我们对树状注意力的宽度、层并行（LP）大小以及是否使用奖励校准进行了实验。如图4所示：

  • 令牌吞吐量（左图）：随着LP规模从1增加，吞吐量因多GPU利用率提高而上升。但当LP规模过大（如5）时，推测结果过于模糊，导致吞-吐量明显下降。最佳LP规模受树宽影响：更宽的树增加了存在正确令牌的概率，因此允许更模糊、更快的推测（即更大的LP规模）。即使没有树状注意力（宽度=1），EasySpec也能通过层并行实现加速。

  • 推测准确率（右图）：奖励校准显著提高了所有树宽和LP规模组合下的接受率。同时，它也减缓了随着LP规模增加而导致的接受率下降趋势。这些结果证明，在KV值的精度和额外延迟之间进行权衡是合理且有效的。

图4：在MMLU上，Llama-3-70B(8B)-Instruct在不同EasySpec配置下的令牌吞吐量和接受率。左图：令牌吞吐量。右图：接受率。
图4：在MMLU上，Llama-3-70B(8B)-Instruct在不同EasySpec配置下的令牌吞吐量和接受率。左图：令牌吞吐量。右图：接受率。

A7 补充细节

5 相关工作

关注草稿生成阶段的方法大致可分为训练草稿模型、自推测或算法优化三类。Specinfer【索引16，Specinfer: Accelerating large language model serving with tree-based speculative inference and verification，2024，ASPLOS】使用经过增强训练的草稿模型和树状结构注意力进行高效推测。Medusa【索引7，Medusa: Simple llm inference acceleration framework with multiple decoding heads，2024，arXiv】利用原始LLM的特征训练一组额外的MLP头来预测未来令牌。Self-Speculative Decoding【索引27，Draft & verify: Lossless large language model acceleration via self-speculative decoding，2023，arXiv】和LayerSkip【索引28，Layerskip: Enabling early exit inference and self-speculative decoding，2024，ACL】通过跳过目标LLM的某些层来进行自推测。Lookahead【索引26，Break the sequential dependency of llm inference using lookahead decoding，2024，ICML】使用n-gram Jacobi解码来提高接受率。REST【索引29，REST: Retrieval-based speculative decoding，2023，arXiv】利用数据存储进行检索式推测。近期工作还探索了优化的树状结构注意力机制。SpecExec【索引30，Specexec: Massively parallel speculative decoding for interactive llm inference on consumer devices，2024，NeurIPS】从草稿模型中提取最可能的续写来构建缓存树，提高了在资源受限GPU上的推理效率。OPT-Tree【索引31，Opt-tree: Speculative decoding with adaptive draft tree structure，2025，TACL】提出了一种自适应且可扩展的草稿树构建算法，以最大化预期接受长度。

除了草稿模型本身，其他研究还探讨了阶段级并行和推测解码的理论解释。例如，PEARL【索引32，PEARL: Parallel speculative decoding with adaptive draft length，2025，ICLR】通过重叠草稿和验证阶段来实现阶段级流水线，而SpecTr【索引33，Spectr: fast speculative decoding via optimal transport，2023，NeurIPS】则采用最优传输理论为推测解码提供了理论解释和优化。

关于草稿模型和基础模型最优分布规模之间的差异，最早由【索引12，Accelerating large language model decoding with speculative sampling，2023，arXiv】的作者发现，他们建议训练一个更宽但更浅的草稿模型以在张量并行下获得更好性能。目前关注此问题的工作很少，但该问题重要且值得解决。据我们所知，我们是首个有效优化现有可用草稿模型在多GPU上利用率的工作。

A5 结论

本文提出了EasySpec，一种用于在推测解码的草稿生成阶段高效利用多GPU的层并行推测策略。EasySpec对原始推测过程引入了两项修改：模糊推测和奖励校准。模糊推测打破了顺序的层执行顺序，实现了多层并行化；而奖励校准则将令牌并行应用于草稿生成阶段，以消除长期的误差累积。EasySpec无需训练，并且在不同模型和数据集上均能稳定地实现加速。

A6 附录

A 推测解码

推测解码由草稿生成和验证的迭代组成。单次迭代的详细算法如算法3所示。

算法3：一次推测解码的迭代
算法3：一次推测解码的迭代

对于任何分布p'，最终的输出令牌序列等同于直接从p中采样。其证明已在许多现有工作（如【索引11，Fast inference from transformers via speculative decoding，2023，ICML】）中详细说明。

B MLP层的策略

在EasySpec中，MLP层是顺序执行的，而不是像注意力层那样并行化。这一设计选择源于经验结果，即并行化所有层（包括注意力和MLP）会引入大量的近似误差，导致整体性能下降。如表5所示，即使只并行化N=2个完整层，其令牌吞吐量和推测准确率也低于最优的仅注意力并行策略，并且随着N的增加，性能会进一步恶化。

表5：不同并行策略下模型在MMLU上的令牌吞吐量和推测准确率。
表5：不同并行策略下模型在MMLU上的令牌吞吐量和推测准确率。

C 树状注意力

在推测解码的背景下，可以利用树状注意力在一次前向传播中对多个可能的令牌序列进行推测和验证。通过将一个二维注意力掩码应用于扁平化的令牌树，具有相同前缀的令牌序列可以共享KV项，从而增加令牌级别的并行度。如图5所示，通过两次前向传播，生成了6个草稿令牌序列。

图5：树状注意力图示。在两次前向传播中生成了6个草稿令牌序列，这增加了并行性并可能提高接受率。整个树可以在基础模型的一次前向传播中进行验证。
图5：树状注意力图示。在两次前向传播中生成了6个草稿令牌序列，这增加了并行性并可能提高接受率。整个树可以在基础模型的一次前向传播中进行验证。

D 额外结果

D.1 不同TP规模下的结果

表6展示了一些典型测试模型在不同TP规模下的令牌吞吐量。基准测试包含MMLU中的40个数据项。虽然大型模型受益于多GPU分布，但当TP规模大于1时，草稿模型会因方程(3)而出现减速，这证实了多GPU资源利用率不足的问题无法通过传统TP解决。我们还在另一个广泛使用的推理引擎vLLM【索引34，Efficient memory management for large language model serving with pagedattention，2023，SOSP】中测试了吞吐量，结果与上述结论一致。这种一致性支持了加速和减速确实归因于多GPU配置，而非其他可能因素。基于这些发现，我们在测试时对所有小型和大型模型分别使用1和8的分布规模，而不影响泛化性。

需要注意的是，Qwen-2-7B-Instruct不支持TP=8，因为其内部维度大小不能被8整除。这也是并行性不足的一个方面（除了推理速度减慢之外），而EasySpec也可以解决这个问题。

表6：不同TP规模下模型在MMLU上的令牌吞吐量。'*'表示结果来自vLLM。
表6：不同TP规模下模型在MMLU上的令牌吞吐量。'*'表示结果来自vLLM。

D.2 任务特定模型的结果

表7展示了两个任务特定模型在相应数据集MATH和HumanEval上的结果。在这些数据集上，草稿模型的原始推测结果已经非常准确（接受率大于97%），表明推测本身相对简单。因此，模糊推测几乎没有导致准确率下降，实现了理想的草稿阶段加速。

表7：两个任务特定模型在不同数据集上的每100个token的运行时间和接受率。符号与表2一致。
表7：两个任务特定模型在不同数据集上的每100个token的运行时间和接受率。符号与表2一致。

D.3 其他草稿和基础模型组合的结果

我们进一步在其他草稿和基础模型组合上测试了EasySpec。表8中的结果表明，EasySpec在所有模型组合和数据集上都显示出加速效果，突显了我们方法的广泛适应性。这里的“不使用EasySpec”基线与表2中的“+tree”基线相同。“L3”和“Q2”分别代表Llama-3和Qwen2。为获得最佳整体性能，Llama-3.2-1B-Instruct的层并行大小为2，其他模型的大小为4。

如4.1.2节所述，使用更小的草稿模型不一定能带来最佳的整体性能。此外，需要注意的是，Qwen-2-1.5B和Qwen-2-0.5B仅支持TP规模≤4和2，因为其维度大小不能被更大的数整除（类似于Qwen 7B模型），而EasySpec通过引入额外的层并行性增加了并行规模。

表8：使用和不使用EasySpec的其他模型组合的令牌吞吐量。温度为0。左侧和右侧的数字分别代表不使用和使用EasySpec的吞吐量。“不使用EasySpec”与表2中的“+tree”基线相同。
表8：使用和不使用EasySpec的其他模型组合的令牌吞吐量。温度为0。左侧和右侧的数字分别代表不使用和使用EasySpec的吞吐量。“不使用EasySpec”与表2中的“+tree”基线相同。

D.4 结果的详细数据

表9显示了EasySpec及各基线的详细结果数据。与表2相比，这里列出了每100个token的总运行时间，以及原始解码的运行时间。

表9：在温度=0和0.8下，两个主流模型在不同数据集上的草稿和验证阶段每100个token的运行时间、总运行时间和加速比以及接受率的详细结果。符号与表2基本一致。
表9：在温度=0和0.8下，两个主流模型在不同数据集上的草稿和验证阶段每100个token的运行时间、总运行时间和加速比以及接受率的详细结果。符号与表2基本一致。

E 理论支持

E.1 余弦相似度与近似精度之间的联系

当前的Transformer模型通常在注意力块之前应用层归一化（LayerNorm）操作。给定两个输入隐藏状态$h_1$和$h_2$（其中$h_2$是$h_1$的近似），LayerNorm算子将它们转换为归一化表示$h'_1$和$h'_2$，然后用于后续的注意力计算。重要的是，LayerNorm操作保持余弦相似度，即$cos(h_1, h_2) = cos(h'_1, h'_2)$。当这个余弦相似度接近1时，对应的向量在欧几里得距离上也很接近($|h_1 - h_2| \approx 0$)，确保了下游计算保持高数值精度。对于没有LayerNorm操作的模型，图3中的演示也意味着近似的隐藏状态可以（原则上）保持精确。

E.2 推测延迟与准确率之间的平衡

推测解码（SD）的策略是在草稿模型的运行时间$T_{draft}$和推测准确率$\alpha$之间找到一个最佳平衡。记基础模型的运行时间为$T_{base}$，没有推测解码时解码N步的总运行时间为$N * T_{base}$。使用SD的总运行时间可以看作是$T_{draft}$和$\alpha$的函数：$T_{all}(T_{draft}, \alpha) = N * T_{draft} + (N/(n * \alpha)) * T_{base}$（n是推测长度），因为基础模型只需要运行$N/(n * \alpha)$次。我们方法的关键在于大幅降低$T_{draft}$同时保持高$\alpha$，这将导致总时间的减少。

记模糊推测的运行时间为$T_{fuzz}$，原始顺序运行时间为$T_{ori}$。如果每n个令牌都只运行模糊推测，时间将是$T'_{draft} = n * T_{fuzz}$。奖励校准在每次推测开始时运行。由于应用了令牌并行（就像在正常SD中应用于验证一样），它的运行时间基本与$T_{ori}$相同。完整的推测时间则变为$T_{draft} = (n - 1) * T_{fuzz} + T_{ori} > n * T_{fuzz} = T'_{draft}$。如果仅是这样，$T_{all}$会增加。然而，与此同时，$\alpha$通过校准得到了大幅提升（见图4），因为缓存中模糊的KV项被精确的项所取代。因此，奖励校准有助于在$T_{draft}$和$\alpha$之间找到一个更好的平衡点。

上述理论分析补充了经验证据，进一步支持了我们的算法能有效加速推理过程。

F 局限性

本实验是在特定的硬件配置（8×A100 GPU）和特定的推理系统上进行的，加速率可能会因不同平台的计算速度和张量并行效率的差异而有所不同。然而，随着硬件的不断改进，计算资源可能会变得越来越冗余，从而进一步降低张量并行的好处。由于这个问题会随着硬件和算子的加速而变得更加严重，我们相信层并行通过优化当前和未来的多GPU资源利用，提供了一个更有效的解决方案。

G 更广泛的影响

我们的工作通过一种层并行的方法加速了推测解码，从而实现了更快的推理。我们相信它对LLM推理加速领域具有积极的广泛影响。据我们所知，没有需要讨论的特定负面影响。