Distributed Speculative Inference (DSI): Speculation Parallelism for Provably Faster Lossless Language Model Inference
Distributed Speculative Inference (DSI): Speculation Parallelism for Provably Faster Lossless Language Model Inference
作者/机构: Nadav Timor∗ (Weizmann Institute of Science), Jonathan Mamou (Intel Labs), Daniel Korat (Intel Labs), Moshe Berchansky (Intel Labs), Oren Pereg (Intel Labs), Moshe Wasserblat (Intel Labs), Tomer Galanti (Weizmann Institute of Science), Michal Gordon (Texas A&M University), David Harel (Weizmann Institute of Science)
A1 主要贡献
本文介绍了一种名为分布式推测式推理(Distributed Speculative Inference, DSI)的新型推理算法。该算法通过并行化推测式推理(Speculative Inference, SI),克服了现有SI方法的根本局限性。传统的SI方法依赖于一个顺序的“起草-验证”过程,其中每次验证必须完成后才能起草新的令牌。这导致当起草模型(drafter)速度过慢或准确率不足时,SI的性能甚至可能劣于标准的自回归推理。DSI通过引入一种名为“推测并行(Speculation Parallelism, SP)”的新型任务并行机制,将SI转变为一个非阻塞算法,通过重叠验证和起草过程来有效隐藏验证延迟。
核心问题与研究目标:
核心问题在于现有SI算法的顺序性限制,即验证过程会阻塞后续的起草过程,导致在起草模型不够理想的情况下性能下降。研究目标是设计一种新的推理算法,能够克服这一瓶颈,无论使用何种起草模型,都能保证比标准自回归推理和传统SI更快,从而拓宽SI技术的适用范围并充分利用可用硬件资源。
创新点与主要贡献:
* 引入推测并行(SP): 提出了一种新颖的任务并行类型,通过使用目标模型和起草模型的多个实例,实现验证和起草过程的并发执行,从而消除了SI的阻塞特性。
* 可证明的加速效果: 从理论上证明了DSI的速度总是至少与非SI相当,并且在期望上严格快于SI和非SI。这解决了SI可能比非SI更慢的问题,保证了在任何起草模型下的性能增益。
* 更广泛的适用性: DSI能够加速那些对于传统SI方法无效的起草模型,从而使其能够应用于更广泛的语言模型。
* 对可用硬件的可扩展性: DSI可以通过调整其lookahead超参数来协调任意数量(≥ 2)的GPU,为利用更多计算资源换取更低延迟提供了新的权衡基础。
* 实证验证: 在模拟的单节点、多GPU真实设置中,实验表明DSI在多种现有语言模型和任务上比SI快1.29至1.92倍。
A3 背景知识与关键观察
推测式推理(Speculative Inference, SI)
SI基本原理。推测式推理(SI)是一种加速目标语言模型 $f_m$(例如GPT系列成员)推理的方法。这类方法使用比目标模型 $f_m$ 更快的语言模型 $f_1, \dots, f_{m-1}$ 来近似 $f_m$,以减少总推理时间。例如,【28,Fast inference from transformers via speculative decoding,2023,ICML】和【11,Accelerating large language model decoding with speculative sampling,2023】等工作可以通过批处理(batching)来减少从目标模型 $f_2$ 生成 $N>1$ 个词元所需的时间。推理过程首先使用一个更快的起草模型 $f_1$ 起草k个词元 $x'i := f_1(x' \oplus x'}) := f_1(x_{\le 01 \oplus \dots \oplus x')$,其中 $i \in [k]$ 且 $1 \le k < N$。然后,算法将提示序列 ${x'{\le i}}^k$ 作为一个输入批次一同发送给目标模型 $f_2$。其核心思想是利用现代GPU提供的数据并行能力,并行计算与这些提示 ${x'{\le i}}^k$ 对应的logits,这比顺序计算这 $k+1$ 个单独的logits要快。得到logits后,算法可以在没有额外前向传递的情况下生成 [1, k+1] 个词元。通过重复此过程,算法可以生成超过 $k+1$ 个词元。
SI的无损特性。简单的推测式推理算法通常是期望无损的(lossless in expectation),即它们生成的词元与目标模型在没有推测的情况下生成的词元遵循相同的分布。一些朴素的推测算法能保证返回与目标模型完全相同的词元(【20,Assisted generation: a new direction toward low-latency text generation,2023】;【53,Accelerating LLM inference with staged speculative decoding,2023】;【60,Accelerating llm inference with lossless speculative decoding algorithms for heterogeneous vocabularies,2025】)。而更复杂的推测算法可能会生成不同的词元,但这些生成的词元遵循目标模型的分布(【28,Fast inference from transformers via speculative decoding,2023,ICML】;【11,Accelerating large language model decoding with speculative sampling,2023】;【38,Specinfer: Accelerating large language model serving with tree-based speculative inference and verification,2024】;【58,Block verification accelerates speculative decoding,2025】;【60,Accelerating llm inference with lossless speculative decoding algorithms for heterogeneous vocabularies,2025】)。
分布式算法的实现与符号。为了实现分布式推测式推理算法,我们使用多个处理器或服务器,这些是能够执行线程的硬件组件。处理器可以计算前向传播并从输出概率向量中采样词元,我们假设线程可以并行运行。在使用DSI时,我们会运行一系列起草模型 $f_{j_1}, f_{j_2}, \dots, f_{j_k}$,其中第一个模型接收 $x_{\le 0}$ 并返回某个词元 $x_{j_1}^1$,第二个模型接收 $x_{\le 0} \oplus x_{j_1}^1$ 作为提示并返回 $x_{j_1, j_2}^2$,依此类推。因此,为了表示一个给定线程正在计算序列 $x_{j_1,\dots,j_{k-1}}^{\le k-1} := x_{\le 0} \oplus x_{j_1}^1 \oplus \dots \oplus x_{j_1,\dots,k-1}^{j_{k-1}}$ 上 $f_{j_k}$ 的输出,我们用 $C_J$ 表示,其中 $J = (j_1, \dots, j_k)$。当线程 $C_J$ 计算一个LM时,我们用 $C_J[\text{prob}]$ 表示输出的概率向量。如果 $C_J$ 从 $C_J[\text{prob}]$ 中采样一个新词元,我们用 $C_J[\text{new}]$ 表示这个词元。例如,实现公式3的线程 $C_J$ 将有 $C_J[\text{prompt}] := x_{\le i}$,$C_J[\text{prob}] := f(C_J[\text{prompt}])$ 以及 $C_J[\text{new}] \sim C_J[\text{prob}]$。
线程输出的表示。一旦线程 $C_J$ 完成新词元的采样,该线程会输出 $C_J[\text{prompt}]$ 和 $C_J[\text{new}]$ 的拼接结果。沿用公式3的例子,我们有:
$$C_{J}[\text{return}] := C_{J}[\text{prompt}] \oplus (C_{J}[\text{new}]) := (x_{\leq 0}, x_{1}, \dots, x_{i+1}).$$由 $C_J$ 启动的新线程表示为 $C_{J\oplus(j)}$,其中 $J \oplus (j)$ 是 $J$ 和 $(j)$ 的拼接。所有源自 $C_J$ 的线程集合为 ${C_{J\oplus J'} : J' \text{是一个非空元组}}$。我们假设终止一个并发线程会终止所有源自它的线程。
时间度量。本文中的时间指的是挂钟时间(wall time)。我们测量从任务启动到其终止所经过的时间。任务是一个非空的线程集合,表示为 ${C_J : J \in \mathcal{J}}$。其时间为 $T_{\text{wall}}[{C_J : J \in \mathcal{J}}] := \max_{J \in \mathcal{J}} T_{\text{wall}}[C_J]$。
单线程时间。当任务仅包含单个线程时,我们省略花括号,即:
$$T_{\text {wall }}\left[C_J\right]:=T_{\text {wall }}\left[\left\{C_J\right\}\right] \text { where }\left|\left\{C_J\right\}\right|=1 .$$需要注意的是,两个线程 $C_J$ 和 $C_{J'}$ 可能并发运行并在时间上重叠。因此,可能出现 $\max {T_{\text{wall}}[C_J], T_{\text{wall}}[C_{J'}]} \le T_{\text{wall}}[{C_J, C_{J'}}] < T_{\text{wall}}[C_J] + T_{\text{wall}}[C_{J'}]$ 的情况。但是,如果 $C_J$ 和 $C_{J'}$ 在时间上没有重叠,则 $T_{\text{wall}}[{C_J, C_{J'}}] \ge T_{\text{wall}}[C_J] + T_{\text{wall}}[C_{J'}]$。
现有推测式推理方法的局限性
DSI框架概述。本节提出了一个理论上完善的编排框架,用于并行化SI(【28,Fast inference from transformers via speculative decoding,2023,ICML】;【11,Accelerating large language model decoding with speculative sampling,2023】;【38,Specinfer: Accelerating large language model serving with tree-based speculative inference and verification,2024】;【60,Accelerating llm inference with lossless speculative decoding algorithms for heterogeneous vocabularies,2025】),该框架在本质上与底层的正向传播计算解耦。如后文所示,我们的方法适用于任意固定数量的处理器(≥ 2)。首先,我们介绍一个该方法的简化版本,它假设有足够多的处理器,以确保线程永远不需要等待。
SI的顺序性瓶颈。在介绍我们的算法之前,我们首先讨论现有SI方法的局限性。每当一个草稿词元被接受时,SI就会减少一次目标模型的正向传播。使用一个准确且快速的起草模型,SI可以显著减少目标模型正向传播的次数,从而可能比非SI方法更快地完成推理。例如,如果平均每次迭代接受一个草稿词元,那么目标模型正向传播的次数将减少一半,因为平均一次目标模型正向传播会生成两个词元。只要总的起草延迟小于因减少目标模型正向传播而节省的延迟,SI就能提供比非SI更快的速度。现有SI方法的主要局限性在于其顺序性。每个SI迭代都需要一次目标模型正向传播,而下一次迭代只有在当前迭代完成后才能开始。因此,如果起草模型速度过慢或不准确,SI会比非SI更慢,即使减少了目标模型正向传播的次数。然而,我们观察到每个SI迭代的验证过程本质上并非顺序的,可以被并行化。
DSI的理论加速潜力。以往关于SI的研究通过使用运行时间仅为目标模型1-5%的起草模型,以及2-5个草稿词元的迭代(即lookahead ∈ [2, 5])来展示加速效果。我们可以根据阿姆达尔定律(Amdahl’s law)计算我们提出的方法相对于SI的最大加速比。假设起草模型是完美准确的,在这种情况下,我们的方法隐藏了所有的验证延迟,使得整体端到端延迟仅剩下起草延迟。对于延迟为1-5%且lookahead ∈ [2, 5]的起草模型,我们提出的并行化方法理论上可以比SI快4倍到50倍。
DSI、SI与非SI的时间线对比。图1和表1说明了我们提出的方法相对于SI和非SI的潜在加速效果,假设起草模型延迟为14%且lookahead=1。在实践中常用的更大lookahead值会产生更大的理论加速。
A2 方法细节
3.1 方法概述
任务设定。考虑这样一个任务:给定一个提示 $x_{\le 0}$,需要从一个目标模型 $f_m$ 自回归地计算出 $N$ 个输出词元。我们拥有一组比 $f_m$ 更快的起草模型 $f_1, \dots, f_{m-1}$(如假设2中定义)。我们的目标是计算所有的 $x_i = f_m(x_{\le i-1})$,其中 $i \in [N]$。附录C提供了算法1的详细分步解释。
接受率(Acceptance rate)。算法1的第8行和第10行会终止任何返回与当前验证器返回的词元不匹配的线程(及其后代线程)。我们称这个草稿词元被“拒绝”了。给定一个目标模型、一个起草模型和一个输入提示,我们将接受率定义为接受草稿词元的概率。为了提高接受率,我们可以用更宽松的方法来替代严格的精确匹配(第8、10行)来拒绝草稿。例如,应用【28,Fast inference from transformers via speculative decoding,2023,ICML】、【11,Accelerating large language model decoding with speculative sampling,2023】、【38,Specinfer: Accelerating large language model serving with tree-based speculative inference and verification,2024】和【60,Accelerating llm inference with lossless speculative decoding algorithms for heterogeneous vocabularies,2025】中提出的程序,可以在保持目标模型输出分布(即期望上无损)的同时提高接受率。
推测并行(Speculation parallelism, SP)。DSI的核心是提供了一种我们称之为推测并行(SP)的新型任务并行机制,它能够协调目标模型和起草模型的实例在时间上重叠。推测并行度(SP degree)是指在同一时间使用的目标服务器(即专用于计算目标模型的服务器)的最大数量。DSI并行化了SI中所有非顺序的部分。与SI中当前迭代的草稿验证过程是顺序的,并且会阻塞算法进入下一迭代的起草阶段不同,DSI在额外的线程上运行验证,以隐藏其延迟。在DSI中,只有当验证拒绝一个词元时,才会对整体延迟产生影响。在DSI中拒绝一个词元会触发一个同步机制,终止那些接收了被拒绝词元的线程(第8行),以确保输出的词元都是被接受的。随着生成的词元数量趋于无穷,DSI所并行化的推理部分比例趋向于期望接受率。例如,在一个简单的案例中,我们只有一个接受率为80%的起草模型,DSI有效地并行化了80%的工作,根据阿姆达尔定律,在生成 $N \to \infty$ 个词元时,期望加速比为5倍。
算法1: 生成N个词元的分布式推测式推理(DSI)
--------------------------------------------------------------------------------------------------------------------
要求: 一个提示 x≤0,以及 m 个自回归模型 f1, f2, . . . , fm。
1: v = 1
2: 启动 m 个线程 C(1), . . . , C(m),使得 C(j1) 并发地为所有 j1 ∈ [m] 生成词元 xj1 ∼ fj1(x≤0)。
3: 将线程 C(m) 标记为当前验证器。
4: 一旦任何线程 CJ⊕(j) 完成生成一个词元,即采样了 CJ⊕(j)[new] ∼ fj(CJ⊕(j)[prompt]):
5: 如果 |J| + 1 < N 那么
6: 启动 m 个线程 CJ⊕(j,1), CJ⊕(j,2), . . . , CJ⊕(j,m),分别从 f1, f2, . . . , fm 并发地生成一个词元。
7: 如果 CJ⊕(j) 是当前验证器线程 那么
8: 终止所有采样了与 CJ⊕(j) 不同词元的所有线程 CJ⊕(j') (及其后代线程)。
9: 令 j* = arg min {j' | CJ⊕(j')[new] = CJ⊕(j)[new]}。
j'∈[m]
10: 终止所有 j' > j* 的线程 CJ⊕(j') (及其后代线程)。
11: 将 CJ⊕(j*,m) 标记为当前验证器。
12: 更新 v = v + 1。
13: 如果 CJ⊕(j*,m) 已经完成 那么
14: 返回第7步,令 J = J ⊕ (j*, m)。
15: 结束如果
16: 结束如果
17: 否则如果 J ⊕ (j) 的最后一个条目等于 m (即 j = m) 那么
18: 返回 CJ⊕(j)[return]。
19: 结束如果
20: 结束一旦
前瞻(Lookahead)。虽然算法1中描述的DSI抽象版本利用了足够多的服务器,但在实践中,我们通常只有固定数量的服务器。通过选择一个足够大的lookahead超参数,我们可以在任意数量的服务器(≥ 2)上部署DSI,具体细节在附录D中阐述。lookahead被定义为发送给目标服务器的每个验证任务中所包含的草稿词元数量。为了简化,算法1中的lookahead设置为1,但它可以是任意大的值。例如,在实践中,SI方法通常设置为每验证五个草稿词元(lookahead=5)而不是一个(lookahead=1)。在DSI中,更大的lookahead值会降低向目标服务器发送验证任务的频率,从而降低所需的SP度。给定一个SP度,lookahead必须足够大以满足以下不等式,确保验证任务永远不会等待目标服务器处理。
较小的SP度或较快的起草模型需要选择较大的lookahead值。例如,给定一个延迟为5%的单个起草模型和SP=4,设置lookahead=5就足够了。假设单个起草模型在单个处理单元上运行,所需的最大处理单元数量为 $1 + 5 \cdot 0.05 = 5$。如果有超过五个处理单元可用,我们可以选择一个更小的lookahead值,从而更频繁地产生验证任务。通常,选择满足公式1的最小lookahead值是最佳选择,这使得DSI能够更早地检测到拒绝(第8行)。使用SP度使得 SP = $\frac{\text{目标延迟}}{\text{起草延迟}}$ 可以达到最大的期望加速,任何更大的SP度都无法再加速推理,因为目标服务器的数量将超过可以并行处理的验证任务数量。
资源竞争。在实践中,当多个线程竞争相同的硬件资源,如内存带宽、数据传输通道或用于编排的CPU核心时,可能会发生资源竞争。通过选择满足可用资源支持的SP度(即满足公式1)的最小lookahead,我们可以保证算法在实践中高效运行,没有显著的资源竞争,因为验证请求在不同时间发送到目标服务器,并且响应预计在交错的时间点返回。
模型并行(Model parallelism, MP)。DSI引入了一种新的并行化方式,不同于张量并行(TP)和流水线并行(PP)。DSI可以采用TP、PP或两者的组合。为简化起见,我们称任何此类并行组合为MP,包括TP、PP或TP+PP。MP加速了前向计算,而SI减少了目标前向传播的次数。它们的组合(SI+MP)既减少了目标前向传播的次数,又加速了每一次目标前向传播。DSI进一步减少了对延迟有贡献的目标前向传播次数,因为它通过并行计算来隐藏验证任务的延迟。下面我们比较SP和MP,并提供一个简单的例子证明在相同计算预算下SP优于MP。
SP与MP的对比。在DSI中,目标模型的前向传播只有在拒绝一个草稿时才会产生延迟。下面是一个比较MP和SP的简单例子。给定一个延迟为10%的起草模型,我们可以设置lookahead=2,这样DSI就能在一个只有6个GPU的单节点上运行(5个用于目标模型,1个用于起草模型)。设a为起草模型的接受率。DSI通过隐藏操作所消除的目标模型前向传播的期望数量是 $a^2$。例如,对于一个接受率为80%的起草模型,只有36%的目标模型前向传播会产生延迟(通常是 $1-a^{\text{lookahead}}$)。在相同的计算预算下(MP=5),MP必须将目标模型前向传播加速2.78倍或更多才能比DSI快。然而,对于某些硬件设置、模型架构和大小,MP是无效的,而DSI仍然有效。
DSI与MP的结合。DSI可以很自然地与MP结合以加速底层的正向传播,这不需要对算法进行任何修改,因为DSI提供了一个与底层正向传播计算无关的编排算法。由于DSI可以协调多个节点,它解锁了具有足够大SP度的设置,从而在理论上DSI和MP之间没有权衡。将DSI与MP方法结合可能在单节点和多节点设置中进一步加速推理。本文涵盖了实现这种DSI的所有基础概念。
KV缓存。我们可以将DSI看作一个动态构建和验证词元树的编排算法。无论在理论上还是实践上,DSI都与底层的正向传播计算(包括KV缓存管理)解耦。每个服务器维护自己的KV缓存。这些服务器协同处理一个具有共享前缀的词元树。每次草稿被拒绝时都会发生同步。
高效的KV缓存管理。SpecInfer中已经开发了高效的词元树KV缓存管理技术,其中树的路径可以共享公共前缀(【38,Specinfer: Accelerating large language model serving with tree-based speculative inference and verification,2024】)。实践者可以直接应用SpecInfer的KV缓存管理来实现本文报告的预期加速。虽然实现SpecInfer的KV缓存管理可能需要一些工程努力,但这已是一个已解决的研究问题,并且已被证明增加的延迟可以忽略不计。
算法的分步解释 (附录C)
任务设定。考虑这样一个任务:给定一个提示 $x_{\le 0}$,需要从一个目标模型 $f_m$ 自回归地计算出 $N$ 个输出词元。我们拥有一组比 $f_m$ 更快的起草模型 $f_1, \dots, f_{m-1}$(如假设2中定义)。我们的目标是计算所有的 $x_i = f_m(x_{\le i-1})$,其中 $i \in [N]$。
初始化与并行起草。为了实现这一目标,我们启动 m 个线程,$C_{(1)}, \dots, C_{(m)}$(算法1的第2行)。每个表示为 $(j_1)$ 的线程负责计算 $x^1_{j_1} = f_{j_1}(x_{\le 0})$。
递归地生成与验证。一旦一个线程 $C_{(j_1)}$ 完成计算,我们实例化 m 个新线程 $C_{(j_1, j_2)}$,为所有 $j_2 \in [m]$ 计算 $x^2_{j_1, j_2} = f_{j_2}(x_{\le 0} \oplus x^1_{j_1})$。通常,一旦我们计算出 $x^{r-1}{j_1, \dots, j}}$,我们就会启动 m 个新线程 $C_{(j_1, \dots, j_{r-1}, 1)}, \dots, C_{(j_1, \dots, j_{r-1}, m)}$,为所有 $j_r \in [m]$ 计算 $x^r_{j_1, \dots, j_r} = f_{j_r}(x_{\le 0} \oplus x^1_{j_1} \oplus \dots \oplus x^{r-1{j_1, \dots, j)$。这体现在第4行和第6行。}
验证与剪枝。一旦 $C_{(m)}$ 完成计算并提供了第一个输出词元的正确值 $x_1^m = x_1$,我们就可以验证其他哪些线程 $C_{(j_1)}$ 准确地计算了 $x_1$。任何 $x^1_{j_1} \ne x_1$ 的线程 $C_{(j_1)}$ 及其后代进程都会被立即终止(第8行)。
选择最优路径并更新验证器。对于每个正确计算出 $x^1_{j_1} = x_1$ 的 $j_1 \in [m]$,我们继续为所有 $j_2 \in [m]$ 计算 $x_{j_1, j_2} = f_{j_2}(x_{\le 0} \oplus x^1_{j_1})$。然而,由于所有线程都在计算同一组词元,我们只保留满足 $x^1_{j_1} = x_1$ 的最小 $j_1$ 对应的线程,并终止其他线程(第10行)。本质上,$C_{(m)}$ 充当验证器,识别出错误计算自回归初始部分的起草者。一旦我们保留了一个有效的 $j_1$,我们就将 $C_{(j_1, m)}$ 重新标记为新的验证器线程。我们知道,由于 $C_{(j_1)}$ 返回了正确的词元 $x^1_{j_1} = x_1$ 并且 $x_2 = f_m(x_{\le 1})$,所以 $C_{(j_1, m)}$ 的输出必然是正确的。当该线程完成时,在剩下的线程 $C_{(j_1, j_2)}$ 中,我们终止那些错误计算了 $x_2 = x^2_{j_1, m}$ 的线程(第8行),并只保留其索引 $j_2$ 最小且满足 $x^2_{j_1, j_2} = x^2_{j_1, m} = x_2$ 的线程。重新标记验证器线程和终止无关线程的过程在第8、10和11行中进行了概述。
处理已完成的新验证器。第13行考虑了新标记的线程可能已经完成的情况。如果是这样,在第14行,我们带着新的验证器线程返回到第7行。
Lookahead (附录D)
适配固定数量的服务器。通过选择一个足够大的lookahead超参数,我们可以在任意数量的服务器上部署DSI。算法1的第6行在生成任何词元(除了位置为N的词元)后会立即调用一个新的进程来计算目标模型 $f_m$。特别是,在从一个起草模型 $f_j$ (其中 $j < m$) 生成一个词元后。这样的“草稿”词元可能会在第8行被拒绝,否则被接受。我们可以将此过程视为生成一个草稿词元并将其发送进行“验证”。
Lookahead=1的特例。发送单个草稿词元的验证任务只是DSI的一个特例,即lookahead=1。对于一个足够大的SP度,可以并行运行的此类验证任务数量是无限的。然而,在实践中,可用的处理器数量是给定的,因此SP度必须是固定的。例如,给定一个单个起草模型(即m=2)和一个拥有8个GPU的单节点,DSI必须确保SP ≤ 7,假设1个GPU足以计算起草模型。
通过Lookahead控制SP度。为了控制SP度,我们引入一个新的超参数lookahead,即发送给目标服务器的每个验证任务中的草稿词元数量。当lookahead > 1时,第2行和第6行会相应调整:启动 m-1 个线程 $C_{J\oplus(j,1)}, C_{J\oplus(j,2)}, \dots, C_{J\oplus(j,m-1)}$,分别从 $f_1, f_2, \dots, f_{m-1}$ 并发地生成lookahead个词元,并启动1个线程 $C_{J\oplus(j,m)}$ 从 $f_m$ 生成1个词元。这里为了简化,我们重载了符号,允许 $J\oplus(j)$ 为空,使得 $J\oplus(j, j') = (j')$。这一改变减少了所需的调用次数。
Lookahead的选择。对于任何给定的模型 $f_1, f_2, \dots, f_m$ 和一个SP度,存在一个足够大的lookahead值,使得在任何验证任务被发送时,至少有一个可用的目标服务器,从而验证任务永远不会等待被目标服务器处理。因此,lookahead超参数允许调整DSI以使用任意最大数量的可用处理单元。
理论上的可扩展性。理论上,如果不按比例 lookahead $\propto \frac{\text{目标延迟}}{\text{起草延迟}}$ 来扩展lookahead,SP度可能会增长到无穷大。例如,如果计算起草模型一次前向传播的时间趋于0,或者计算目标模型一次前向传播的时间趋于无穷大。
3.2 理论结果
接下来,我们正式声明DSI(算法1)是严格无损的,总是返回正确的词元序列 $x_1, \dots, x_N$,其运行速度至少与非SI相当,并且在期望上比非SI和SI都快。在介绍我们的主要理论结果之前,我们概述了分析中使用的假设。证明在附录E中提供。
假设1: 我们假设存在一个(通用)常数 $c>0$,使得对于任何输入提示 $x_{\le 0}$ 和模型索引 $j \in [m]$,我们有:
$T_{\text{wall}}[\text{计算 } f_j(x_{\le 0})] \in (0, c)$ 并且 $T_{\text{wall}}[\text{采样 } x \sim f_j(x_{\le 0})] = 0$。
假设2: 我们假设对于所有的 $j \in [m-1]$,$f_j$ 都比 $f_m$ 快(表示为 $f_j \preceq f_m$),具体意义如下:
$\max_{x_{\le 0}} T_{\text{wall}}[\text{计算 } f_j(x_{\le 0})] \le \min_{x_{\le 0}} T_{\text{wall}}[\text{计算 } f_m(x_{\le 0})]$。
假设3: 我们假设 $T_{\text{wall}}[{C_{(j_1, \dots, j_i)}}{i=1}^k] = \sum]$。}^k T_{\text{wall}}[C_{(j_1, \dots, j_i)
第一个假设断言,计算任何模型的输出都需要一个非零且有界的时间,而从输出概率中采样一个词元的时间可以忽略不计。第二个假设断言,对于任何给定的输入提示,每个起草模型都比目标模型运行得快。第三个假设断言,计算 $x^k_{j_1, \dots, j_k}$ 的时间是先计算 $x^1_{j_1}$,然后是 $x^2_{j_1, j_2}$,依此类推,直到 $x^k_{j_1, \dots, j_k}$,中间没有延迟。
以下定理表明,我们的方法返回的词元与目标模型在没有推测的情况下生成的词元具有相同的分布,并且其运行速度至少与迭代应用目标模型本身一样快。
定理1: 在假设1、2和3下,算法1返回与不使用推测式推理(SI)的情况下运行目标模型本身相同的输出,并且运行速度至少一样快。
定理2: 在假设1、2和3下,算法1在期望上运行速度至少与SI一样快。
算法1的优势在于其并发性。下面的命题展示了DSI如何利用一个比目标模型快且能以高概率返回正确输出的起草模型来加速给定目标模型的推理。
命题1: 假设我们有一个起草模型 $f_1$,一个目标模型 $f_2$ 和一个提示 $x_{\le 0}$。假设 $f_1$ 计算每个输出需要 $t_1$ 个时间单位,$f_2$ 需要 $t_2$ 个时间单位,其中 $t_2 > t_1$。假设给定提示 $x_{\le i} = x_{\le 0} \oplus x_1 \oplus \dots \oplus x_i$, $f_1$ 返回(正确的)词元 $x_{i+1}$ 的概率是 $p$。那么,算法1计算正确输出的期望时间至多为 $t_1 p(N-1) + t_2((1-p)(N-1)+1)$ 个时间单位,而如果我们不使用推测式推理来计算 $f_2$,则需要 $t_2 N$ 个时间单位。
A4 实验环境
-
模型架构:
- 目标/起草模型对: 使用了多个模型家族进行配对,包括:
- StarCoder-15B (目标) / StarCoder-168M (起草)
- Phi3-14B (目标) / Phi3-4B (起草)
- Vicuna-13B (目标) / Vicuna-68M (起草)
- Vicuna-7B (目标) / Vicuna-68M (起草)
- 目标/起草模型对: 使用了多个模型家族进行配对,包括:
-
数据集:
- CNN Daily Mail: 用于文本摘要任务。
- Alpaca: 用于指令遵循任务。
- MBPP 和 HumanEval: 用于代码生成任务。
-
硬件配置:
- 由于预算限制,实验并未在物理的多GPU节点上运行。
- 实验在一个单块NVIDIA A100 80GB GPU上进行。
- 为了评估DSI在多GPU环境下的性能,实验模拟了一个拥有最多8个GPU的单节点设置。DSI被配置为协调这8个模拟的GPU,其中SP degree(推测并行度)最高为7(即最多7个GPU用于目标模型),1个GPU用于起草模型。
-
软件配置:
- DSI 使用 Python 实现。
- 利用Python内置的线程池(thread pool)来管理操作系统(OS)级别的线程,以模拟真实的并发执行环境并计入线程管理开销(如上下文切换、创建和调度延迟)。
- 实验代码与DSI实现均已开源。
-
实验模拟方法:
- 在模拟环境中,对模型前向传播的调用被替换为
wait命令。 wait的等待时间基于在真实A100 GPU上独立测量的首次令牌生成时间(TTFT)和后续每令牌输出时间(TPOT)。- 模型对的接受率也是通过独立实验预先测量的,以确保模拟的真实性。
- 在模拟环境中,对模型前向传播的调用被替换为
A4 实验结果
实验结果证实了理论分析,表明DSI在各种实际场景中均优于传统的推测式推理(SI)。
主要实验结果
该实验通过一个多线程系统在线模拟DSI,并测量了包含真实世界多线程延迟的端到端时间。
实验内容:
- 比较DSI与SI在多种模型(StarCoder, Phi3, Vicuna)和任务(HumanEval, MBPP, Alpaca, CNN-DM)上的性能。
- 实验生成50个词元,
lookahead在{1, 5, 10}中选择,并确保配置能够在模拟的8-GPU单节点(SP degree ≤ 7)上运行。 - 性能指标为DSI相对于SI的端到端延迟加速比,该延迟包括预填充和解码时间。
实验结果 (表2):
- DSI在所有测试的模型和任务组合中始终优于SI。
- 加速比范围为1.29倍到1.92倍。
- 例如,在HumanEval任务上,使用Starcoder-15B/168M对,DSI比SI快1.92倍;在使用Vicuna-13B/68M对和CNN-DM数据集时,DSI比SI快1.47倍。
- 结果表明,即使起草模型的延迟相对较高(例如Phi3-4B相对于Phi3-14B的延迟为65%-67%),DSI依然能取得显著的加速效果。
消融研究:离线模拟
该实验通过离线模拟进行,忽略了多线程的实际开销,仅累加前向传播的延迟,旨在从理论层面验证算法的性能,并探索更广泛的配置空间。
实验内容:
- 在一个大的配置空间(起草模型延迟和接受率的笛卡尔积)上模拟非SI、SI和DSI的性能。
- 比较三者之间的成对加速比(或减速比)。
实验结果与结论 (图2):
- (a) SI vs. 非SI: 证实了SI的一个关键弱点:当起草模型速度慢或准确率低时(图中的粉色区域),SI的性能会劣于非SI。
- (b) DSI vs. SI 和 (c) DSI vs. 非SI: 结果显示,在所有具有非零接受率的配置中,DSI始终比SI和非SI更快。DSI在任何配置下都不会比SI或非SI慢。
- (d) DSI vs. 基线: 将DSI与每个配置下SI和非SI中较快者的性能进行比较。结果显示,DSI比最优基线算法快高达1.6倍。这突显了DSI的稳健性,它不仅避免了SI的性能陷阱,还在SI有效的区域提供了进一步的加速。
结论: 离线模拟强有力地支持了理论分析,即DSI通过并行化克服了SI的顺序瓶颈,无论起草模型的性能如何,都能提供稳定且显著的加速。
A7 补充细节
5 相关工作
对现有推测式推理(SI)工作的扩展。对SI的研究已经扩展了该框架,包括动态控制每次迭代的草稿词元数量——这项技术在实践中被广泛采用(【37,Dynamic speculation lookahead accelerates speculative decoding of large language models,2024】;【34,Optimizing speculative decoding for serving large language models using goodput,2024a】;【20,Assisted generation: a new direction toward low-latency text generation,2023】),以及探索其他各种方向(【31,EAGLE: Speculative sampling requires rethinking feature uncertainty,2024】;【10,Medusa: Simple LLM inference acceleration framework with multiple decoding heads,2024】;【57,Spectr: Fast speculative decoding via optimal transport,2024b】;【35,Online speculative decoding,2024b】;【71,Distillspec: Improving speculative decoding via knowledge distillation,2024】;【69,Fastdraft: How to train your draft,2024】;【40,Faster cascades via speculative decoding,2025】;【4,Judge decoding: Faster speculative sampling requires going beyond model alignment,2025】)。然而,在先前的工作中,计算目标模型的前向传播仍然是一个阻塞操作,限制了算法处理后续位置的词元,使得SI作为一个顺序算法的根本局限性没有得到解决。
与PEARL的比较。PEARL是最近对SI的一项扩展,它展示了在验证的同时进行起草可以将非SI和原生SI的速度分别提升高达4.43倍和1.5倍(【33,PEARL: Parallel speculative decoding with adaptive draft length,2025】)。他们的实证结果凸显了DSI打破SI“起草-然后-验证”顺序性的优势。然而,PEARL存在一个根本性限制:它仍然是一个顺序算法,因为它只能处理下一个SI迭代的词元,而DSI可以处理任何未来迭代的词元。PEARL的算法采用了一种启发式方法(控制是否验证每次迭代的第一个草稿词元),并且没有理论保证能够加速SI或非SI。实际上,如果起草模型速度过慢或不准确,PEARL会比非SI更慢,这与DSI不同。由于PEARL无法协调超过一个目标模型实例和一个起草模型实例,它对硬件设置的可扩展性有限,而DSI可以协调任意数量的GPU(≥ 2)。因此,除非PEARL运行时的lookahead = $\frac{\text{目标延迟}}{\text{起草延迟}} > 1$,否则它会未充分利用可用硬件,因此在期望上,它比使用更小lookahead的DSI要慢。
A 其他相关工作
通过增加计算能力加速推理。除了SI,近期减少LM推理延迟的其他工作可以根据其硬件利用方式分为两大类。第一类专注于通过使用更多的计算能力来加速推理。这包括数据并行(DP)和各种类型的模型并行(MP),如流水线并行(PP)、张量并行(TP)(【41,Efficient large-scale language model training on gpu clusters using megatron-lm,2021】)、上下文并行(【30,Sequence parallelism: Long sequence training from system perspective,2023b】;【67,Context parallelism for scalable million-token inference,2024】)和专家并行(【44,DeepSpeed-MoE: Advancing mixture-of-experts inference and training to power next-generation AI scale,2022】)。这种在多个处理器上的分区可以加速受内存限制的推理设置,例如通过避免内存卸载。它们通常在提高推理吞吐量和支持更长上下文长度方面是有效的。然而,典型的LM架构需要自回归推理。这种推理本质上是顺序的,只有有限(如果有的话)的MP机会,具体取决于LM架构。由于在典型LM中可以并行的推理部分很小(如果有的话),根据阿姆达尔定律(【1,Validity of the single processor approach to achieving large scale computing capabilities,1967】;【45,Improvements in multiprocessor system design,1985】),MP方法的潜在加速是有限的。因此,虽然DP方法可以提高推理吞吐量,但它们本质上仍然是顺序的。此外,随着上述并行性在更多处理器上分片,通信开销增加,可能导致收益递减,限制了它们可以有效利用的处理器数量。
通过减少或优化计算资源加速推理。第二类专注于让LM使用更少的计算资源或更好地利用相同的资源。这包括通过剪枝进行的训练后压缩(例如,(【17,Sparsegpt: Massive language models can be accurately pruned in one-shot,2023】;【56,A simple and effective pruning approach for large language models,2024a】;【36,LLM-pruner: On the structural pruning of large language models,2023】))、知识蒸馏(例如,【24,Distilling the knowledge in a neural network,2015】;【22,MiniLLM: Knowledge distillation of large language models,2024】)、量化(例如,(【26,Quantized neural networks: Training neural networks with low precision weights and activations,2018】;【18,OPTQ: Accurate quantization for generative pre-trained transformers,2023】;【32,Awq: Activation-aware weight quantization for llm compression and acceleration,2024】;【68,Exploring post-training quantization in llms from comprehensive study to low rank compensation,2024】;【15,Gpt3. int8 (): 8-bit matrix multiplication for transformers at scale,2022】))、低秩分解(例如,(【25,Language model compression with weighted low-rank factorization,2022】;【66,Tensorgpt: Efficient compression of the embedding layer in llms based on the tensor-train decomposition,2023】))、提前退出(例如,(【50,Confident adaptive language modeling,2022】;【27,Speculative decoding with big little decoder,2023】;【16,Depth-adaptive transformer,2020】;【5,Controlling computation versus quality for neural sequence models,2020】;【49,Consistent accelerated inference via confident adaptive transformers,2021】))以及替代架构(例如,(【10,Medusa: Simple LLM inference acceleration framework with multiple decoding heads,2024】;【31,EAGLE: Speculative sampling requires rethinking feature uncertainty,2024】;【70,The hedgehog & the porcupine: Expressive linear attentions with softmax mimicry,2024b;a】;【65,Efficient streaming language models with attention sinks,2024】;【21,Mamba: Linear-time sequence modeling with selective state spaces,2024】))。虽然这些解决方案在实践中很有用,但它们通常需要修改模型架构、改变训练过程和重新训练模型,而不能保证输出相同。尽管减少了推理时间,这些方法通常有一个显著的缺点,即降低输出质量。也有一些解决方案可以保持输出质量,例如内核优化(【14,Flashattention: Fast and memory- efficient exact attention with io-awareness,2022】;【13,Flashattention-2: Faster attention with better parallelism and work partitioning,2024】),但它们高度依赖硬件,因此并不总是可用甚至可行。
A5 结论
本文提出了一种通过利用任意数量的额外多个处理单元(例如GPU)来减少推测式推理算法运行时间的方法。我们已经表明,尽管SI算法被广泛采用,但在各种实际应用场景中,当起草模型不够准确或速度不够快时,它们最终可能会减慢语言模型的推理速度。我们证明了,通过利用至少一个额外的GPU,我们可以设计一个可证明能减少非SI和SI算法推理时间的推测式推理算法。我们的模拟证实了我们的理论,表明在给定一个拥有最多八个GPU的单节点的所有可能配置下,相对于SI和非SI都有显著的加速。本质上,这项工作为通过推测并行(SP)来同时协调多个处理单元的其他SI算法铺平了道路。
我们引入了分布式推测式推理(DSI),并通过理论分析和实验证明,它在所有可能的配置下都比SI和非SI更快。虽然理论保证对单节点和多节点设置都成立,但我们的实验侧重于具有最多八个GPU且SP度≤7的单节点场景。由于预算限制,我们调整了DSI的实现,以模拟对这样一个节点的访问,而不是在物理的八GPU节点上运行。尽管如此,实证结果是现实的,因为DSI是作为一个多线程系统实现的,会产生这类系统的所有真实世界延迟(例如,上下文切换),并且所有与GPU相关的真实世界延迟都基于访问GPU的独立实验。
未来工作。DSI的关键优势在于它利用推测并行来通过并行化验证减少延迟,以及它能够在不修改模型架构的情况下实现无损加速。在此基础上,未来的工作应侧重于评估DSI在需要多个GPU以避免内存卸载(即MP度 > 1)的LM上的表现。虽然大多数最先进的模型可以在单个GPU上运行(通过压缩、使用较低精度等),但我们确实看到了更大LM的趋势。由于更大的LM通常更慢,DSI的并行性可能特别有效。尽管多节点推理还不是一个常见的设置,但在现实的多节点环境中测试DSI,尽管存在通信延迟,也可能释放其潜力。
A6 附录
B 扩展预备知识
自回归语言模型 (LMs) 是确定性的、实值多变量函数。LM的输入是一个维度为 $n_{\text{vocab}}$ 的向量序列。我们称这些向量为词元(token),序列为提示(prompt)。LMs输出一个维度为 $n_{\text{vocab}}$ 的实值向量,也称为logits。由于提示的长度可能不同,我们将前向传播的符号简化如下:$f : \mathbb{R}^{* \times n_{\text{vocab}}} \to \mathbb{R}^{n_{\text{vocab}}}$。
自注意力 (Self-Attention) LMs 是具有预定义上下文长度 $n_{\text{ctx}}$ 的LMs (【61,Attention is all you need,2017】)。因此,我们将这类LMs的前向传播表示为:$f : \mathbb{R}^{n_{\text{ctx}} \times n_{\text{vocab}}} \to \mathbb{R}^{n_{\text{vocab}}}$。例如,GPT-2和GPT-3是Transformer模型,其 $n_{\text{vocab}} = 50257$,上下文长度分别为 $n_{\text{ctx}} = 1024$ 和 $n_{\text{ctx}} = 2048$ (【43,Language models are unsupervised multitask learners,2019】; 【7,Language models are few-shot learners,2020】)。在本文中,所有的LMs都是具有预训练(冻结)参数的自注意力模型。
提示表示法。我们扩展提示的表示法,使其长度 $l \le n_{\text{ctx}}$。自注意力LMs通过在输入序列的开头加上一个长度为 $n_{\text{ctx}} - l$ 的前缀,然后跟上给定的 $l$ 个词元,来处理长度为 $l < n_{\text{ctx}}$ 的提示。LMs会忽略这个前缀,要么通过将与前缀对应的注意力部分置零(masking),要么通过使用专门的词元进行左填充(left-padding)。在本文中,长度为 $l < n_{\text{ctx}}$ 的提示是输入序列中长度为 $n_{\text{ctx}}$ 的非掩码、非填充的后缀。
生成下一个词元。生成下一个词元是自回归LMs的主要应用。这个过程包括两个步骤:计算LM的前向传播,然后根据输出选择下一个词元。选择可以是确定性的或非确定性的。非确定性的选择过程在LM的前向传播后应用softmax函数,并从得到的概率向量中采样:
softmax: $\mathbb{R}^{n_{\text{vocab}}} \to [0, 1]^{n_{\text{vocab}}}$,使得所有条目之和为1。
输出概率向量。为了方便,我们用 $f(x_{\le i})$ 表示输出的概率向量:
$$x_{i+1} \sim f\left(x_{\leq i}\right):=\operatorname{softmax}\left(f\left(x_{\leq i}\right)\right):=\operatorname{softmax}\left(f\left(x_{\leq 0} \oplus x_{1} \oplus \cdots \oplus x_{i}\right)\right),$$其中 $a \oplus b = (a, b)$ 是向量 $a$ 和 $b$ 的拼接,$x_{\le i} := x_{\le 0} \oplus x_1 \oplus \dots \oplus x_i$。对于确定性选择过程,组合像softmax这样的单调函数通常是不必要的。例如,最可能的下一个词元是logits和softmax输出的 arg max。尽管如此,为了方便,我们假设LMs总是输出概率向量。公式3中的采样过程要么是确定性的(即,$x_{i+1}$ 是具有最大概率的词元),要么是随机的(通过从分布 softmax$(f(x_{\le i}))$ 中随机选择 $x_{i+1}$ 来实现)。
E 证明
定理1: 在假设1、2和3下,算法1返回与不使用推测式推理(SI)的情况下运行目标模型本身相同的输出,并且运行速度至少一样快。
证明: 我们首先证明算法的无损性。我们希望证明当 $v=k$ 时,存在一个线程 $C_{J_k}$,它是唯一被标记为验证器的线程,并且它正确地计算了下一个词元,且 $J_k = J' \oplus (m)$,其中 $J' = (j_1, \dots, j_{k-1})$ 是一个长度为 $k-1$ 的序列,且对于所有 $i \in [k-1]$,都有 $x_{j_1, \dots, j_i}^i = x_i$。我们将通过对 $v$ 的值进行归纳来证明这一点。此外,我们注意到如果算法遵循这个模式,那么它显然是一个无损算法。
基础情况 (v=1): 最初,当 $v=1$ 时,只有一个验证器 $C_{(m)}$,它运行目标模型 $f_m$。因此,当它完成时,它将返回正确的词元 $x_1$。由于验证器只有在 $v$ 的值改变时才会被重新标记(见第11-12行),所以只要 $v=1$,唯一被标记为验证器的线程就是 $C_{(m)}$。
归纳假设: 假设只要 $v=k$,就只有一个线程 $C_J$ 被标记为验证器,它返回正确的词元 $x_k$,并且 $J_k = J' \oplus (m)$,其中 $J' = (j_1, \dots, j_{k-1})$ 是一个长度为 $k-1$ 的序列,且对于所有 $i \in [k-1]$,都有 $x_{j_1, \dots, j_i}^i = x_i$。
归纳步骤: 当 $v$ 从 $k$ 更新到 $k+1$ 时,这种变化只在满足第7行的条件时发生。这个条件表明,长度为 $|J_k| = k$ 的单个验证器线程 $C_J$ 已经完成了其输出词元的计算。根据归纳假设,这个线程返回 $x_k$ 作为其输出。由于 $f_m$ 比所有起草模型 $f_1, \dots, f_{m-1}$ 都慢,所有线程 $C_{J' \oplus (i)}$ 都已经完成了它们的输出计算。因此,当执行第8、10和11行时,唯一保持活跃的线程是 $C_{J' \oplus (j^)}$ 的后代,唯一充当验证器的线程是 $C_{J' \oplus (j^, m)}$。由于对于所有 $i \le k-1$,都有 $x_{j_1, \dots, j_i}^i = x_i$,并且 $x_{j_1, \dots, j_{k-1}, j_k^}^k = x_k$,那么 $C_{J' \oplus (j^, m)}$ 只是在正确的序列 $x_{\le 0} \oplus x_1 \oplus \dots \oplus x_k$ 上计算目标模型 $f_m$ 的输出。因此,它正确地返回第 $(k+1)$ 个词元 $x_{k+1}$,符合要求。
时间: 我们注意到,算法一旦计算出 $C_{J_N}$ 的输出就会终止。根据假设3,我们有 $T_{\text{wall}}[\text{算法1}] = \sum_{i=1}^N T_{\text{wall}}[\text{计算 } f_{j_i}(x_{\le i})]$,并且根据假设2,我们有 $T[\text{计算 } f_{j_i}(x_{\le i})] \le T[\text{计算 } f_m(x_{\le i})]$。综合起来,我们得到 $T_{\text{wall}}[\text{算法1}] \le \sum_{i=1}^N T_{\text{wall}}[\text{计算 } f_m(x_{\le i})]$,这正是不使用推测式推理计算输出词元所需的时间。□
命题1: 假设我们有一个起草模型 $f_1$,一个目标模型 $f_2$ 和一个提示 $x_{\le 0}$。假设 $f_1$ 计算每个输出需要 $t_1$ 个时间单位,$f_2$ 需要 $t_2$ 个时间单位,其中 $t_2 > t_1$。假设给定提示 $x_{\le i} = x_{\le 0} \oplus x_1 \oplus \dots \oplus x_i$, $f_1$ 返回(正确的)词元 $x_{i+1}$ 的概率是 $p$。那么,算法1计算正确输出的期望时间至多为 $t_1 p(N-1) + t_2((1-p)(N-1)+1)$ 个时间单位,而如果我们不使用推测式推理来计算 $f_2$,则需要 $t_2 N$ 个时间单位。
证明: 为了理解其工作原理,令 $j_1 \in {1, 2}$ 为使得 $x_{j_1}^1 = x_1$ 的最小索引,并且对于所有 $i \in [N-1]$,我们递归地定义 $j_i \in {1, 2}$ 为使得 $x_{j_1, \dots, j_i} = x_i$ 的最小索引。我们还固定 $j_N=2$。此外,令 $i_0 = 0$ 且 $i_r$ 为 $[N]$ 中第 $r$ 个使得 $j_i = 2$ 的索引。我们注意到计算 $x_{j_1, \dots, j_{i_1}}^{i_1}$ 的值需要 $t_1(i_1-1)+t_2$ 时间单位。这是因为我们首先计算 $x_1^1$,然后是 $x_{1,1}^1$,一直到 $x_{1, \dots, 1}^{i_1-1}$,最后是 $x_{1, \dots, 1, 2}^{i_1}$。前 $(i_1-1)$ 个词元每个需要 $t_1$ 时间单位,而最后一个词元需要 $t_2$ 时间单位。在 $t_1(i_1-1)+t_2$ 时间单位后,我们将计算出 $x_2^1, x_{1,2}^2, x_{1,1,2}^3$ 等等,直到 $x_{1, \dots, 1, 2}^{i_1}$。由于 $f_1$ 一直准确生成到索引 $i_1-1$ 的词元,一旦我们观察到 $x_2^1$ 与 $x_1^1$ 匹配,我们就知道 $x_{1,2}^2=x_2$,然后可以验证 $x_{1,1}^2=x_2$,$x_{1,1,1}^2$ 等等。我们注意到计算所有这些词元直到计算出 $x_{1, \dots, 1, 2}^{i_1}$ 最多需要 $t_1(i_1-1)+t_2$ 时间单位。因此,我们可以验证 $x_{1, \dots, 1, 2}^{i_1}=x_{i_1}$,计算 $x_{j_1, \dots, j_N}^N$ 的值(并验证其正确性)最多需要 $\sum_{r=1}^k (t_1((i_r-i_{r-1})-1)+t_2)$ 时间单位。我们注意到 $Q = \sum_{i=1}^{N-1} \mathbb{I}(j_i=1)$ 是 $[N-1]$ 中使得 $j_i=1$ 的索引数量。由于 $E[Q]=p(N-1)$,我们有 $E[\sum_{r=1}^k (t_1((i_r-i_{r-1})-1)+t_2)] = t_1 p(N-1) + t_2((1-p)(N-1)+1)$。□
定理2: 在假设1、2和3下,算法1在期望上运行速度至少与SI一样快。
证明: 假设我们有一个起草模型 $f_1$,一个目标模型 $f_2$ 和一个提示 $x_{\le 0}$。假设 $f_1$ 计算每个输出需要 $t_1$ 时间单位,$f_2$ 需要 $t_2$ 时间单位,其中 $t_1 < t_2$。假设给定提示 $x_{\le i} = x_{\le 0} \oplus x_1 \oplus \dots \oplus x_i$,$f_1$ 返回(正确的)词元 $x_{i+1}$ 的概率是 $p$。考虑使用lookahead=k的SI(或DSI)算法从 $f_2$ 生成 $N > k+1$ 个词元。在时间=0时,根据SI的定义,SI开始生成草稿词元。在时间=$kt_1$时,SI完成了前k个草稿词元 $x_1^1, x_{1,1}^2, \dots, x_{1, \dots, 1}^k$ 的生成。在时间=$kt_1+t_2$时,SI完成了对前k个词元 $x_1^1, x_{1,1}^2, \dots, x_{1, \dots, 1}^k$ 的验证。令 $A_1, A_2, \dots, A_{k+1}$ 为如下采样的指示变量。对于所有 $i \in [k+1]$,$A_i=1$ 的概率为p,否则 $A_i=0$。令 $n := \min{i | A_i=0}-1$。注意n服从SI(或DSI)前k个草稿中被接受草稿数量的分布。根据SI的定义,对于任何 $n \in {0, 1, \dots, k}$,SI在时间=$kt_1+t_2$时完成前n+1个词元的生成。根据SI的定义,SI的第一次迭代无法输出位置 > k+1 的词元。SI能完成生成 $x_{k+2}$ 的最早时间是在其第二次迭代结束时。因此,SI在时间 $\ge 2(kt_1+t_2)$ 时完成生成 $x_{k+2}$。考虑在至少 $\frac{t_2}{kt_1}$ 个服务器上使用相同 $f_1, f_2$ 和lookahead的DSI。我们证明DSI可以在时间 $\le 2(kt_1+t_2)$ 内完成生成 $x_{k+2}$,并且在期望上,时间 < $2(kt_1+t_2)$。根据DSI的定义,DSI从不抢占当前的验证器。在时间=$kt_1$时,DSI并发地调用 (i) 验证包含前k个尚未验证词元 $x_1^1, x_{1,1}^2, \dots, x_{1, \dots, 1}^k$ 的批次,以及 (ii) 起草 $x_{1, \dots, 1}^{k+1}, x_{1, \dots, 1}^{k+2}, \dots, x_{1, \dots, 1}^{2k}$。根据DSI的定义,我们使用... 两种SI和DSI都在时间=$kt_1+t_2$时完成第(k+1)个首个词元 $x_{k+1}$ 的生成。此时,DSI要么调用一个新的当前验证器线程,要么将一个现有线程标记为当前验证器(取决于 $\frac{t_2}{t_1}$ 和lookahead)。因此,DSI在时间 $\le kt_1+2t_2$ 时完成生成 $x_{k+2}$,这恰好是当前验证器线程完成其验证的时间。对于所有 $t_1, t_2, k$,DSI都比SI快,因为 $kt_1+2t_2 < 2(kt_1+t_2)$。否则,两种算法都在时间 $\le kt_1+t_2$ 内接受前n+1个词元。此时,证明对 $N-(n+1)$ 重复。□
F 实验细节
F.1 首次令牌生成时间(TTFT)和每输出令牌时间(TPOT)
本节详细说明了我们用于估计预期TTFT和TPOT的独立实验。
为了确保我们模拟中的等待时间是现实的,我们区分了首次令牌生成时间(TTFT)和每输出令牌时间(TPOT)。我们在单个NVIDIA A100 80GB GPU上独立地为每个模型和数据集的组合估计了TTFT和TPOT延迟。
对于数据集d和相应目标或起草模型f的每个组合⟨d, f⟩,我们按以下方式估计f的平均延迟。首先,我们从d中均匀随机选择50个提示,并为每个提示使用f生成20个词元,测量每个词元的延迟(以毫秒为单位)。遵循先前的工作,我们区分了首次令牌生成(TTFT)和所有后续19个词元的每输出令牌生成时间(TPOT)。最后,我们计算每个模型-数据集对在所有提示上的平均TTFT和TPOT,以估计单次前向传播的预期延迟。在我们所有测量前向计算延迟(即TTFT和TPOT)的实验中,模型以MP度为一且无内存卸载的方式运行。我们使用单个NVIDIA A100 80GB GPU,被测模型完全加载到GPU内存中。
我们的主要实验(表2)和消融实验(§4.1)都按如下方式使用估计的TTFT和TPOT:生成第一个词元会增加TTFT的等待时间,而生成每个后续词元会增加TPOT的等待时间。由于TTFT通常比TPOT(主导整个序列生成时间)长得多,为简洁起见,表2中的所有延迟数据均指TPOT。目标模型和起草模型的估计TPOT延迟分别显示在“Target Latency (ms)”和“Drafter Latency (ms)”中。我们还报告了目标和起草模型延迟之间的比率,并以百分比形式呈现在“Drafter Latency (%)”中。每个模型和数据集对的有效预填充-解码延迟比率在表3中提供。
F.2 接受率
本节详细说明了我们用于估计预期接受率的独立实验。
为确保我们的评估是现实的,我们使用了真实世界的接受率,计算方法如下。对于任何给定的⟨目标模型, 起草模型, 数据集⟩组合,我们独立地估计它们的接受率。对于数据集中的每个输入提示,我们同时从起草模型和目标模型生成词元。然后我们考虑目标模型和起草模型之间最长精确词元匹配序列的长度。下面是一个简化的例子,其中词元被计为英文单词。如果目标模型生成“We can only see a short distance ahead, but we can see plenty there that needs to be done. [...]”而起草模型生成“We can only see a short distance ahead, we done. [...]”,那么最长精确匹配序列的长度是8个词元。接受的草稿的期望数量是 $\bar{n} := \frac{1}{N} \sum_{i=1}^N n_i$,其中 $n_i$ 是第i个提示接受的草稿词元数量。然后从几何分布计算接受率:$(\text{接受率}) := 1 - \frac{1}{1+\bar{n}}$。总的来说,当 $N \to \infty$ 时,报告的接受率保证收敛于期望接受率(附录F.2.1)。在本实验中,报告的接受率是基于为每个提示生成256个词元计算的。
这个实验表明,像StarCoder(【29a,Starcoder: may the source be with you!,2023a】)或Vicuna(【70,Judging LLM-as-a-judge with MT-bench and chatbot arena,2023】)这样的现成模型“家族”可以组成很好的目标和起草模型对,因为它们的接受率相对较高,如表2所示。这些家族由不同大小的模型组成,它们以相似的方式在相似的数据集上训练。我们观察到,即使是相对较小的起草模型也与同一家族的较大模型表现出良好的一致性。例如,Starcoder-168M(起草模型)和Starcoder-15B(目标模型)产生的接受率为93%。
我们的主实验(表2)和消融实验(§4.1)都测量了Non-SI、SI和DSI的总延迟。总延迟主要由目标和起草模型的前向传播决定。所需的目标和起草模型前向传播次数取决于接受率。为了计算前向传播的次数,我们可以通过从模型中采样来模拟算法,执行验证过程(对于SI和DSI),并计算所需的前向传播次数。由于从模型中采样本质上是随机的,总延迟T是一个随机变量,我们关注其期望值E[T]。
存在一个接受率p,使得期望延迟满足E[T] = $\hat{T}$。这个接受率可以通过两种方式估计。第一种方法直接计算一个词元被验证器接受的似然。或者,p可以通过将几何分布拟合到接受过程来估计。具体来说,这涉及到计算每次迭代接受的平均词元数,并提取p作为几何分布的参数。
两种方法是一致的。随着迭代次数趋于无穷大,估计的接受率收敛于真实的经验接受率p。因此,基于p估计的总延迟收敛于实际的总延迟。这保证了报告的延迟和计算的延迟在期望上是相等的。
该领域的先前工作已经基于词元接受是独立同分布(i.i.d.)的简化假设发展了理论和方法;例如,参见【28,Fast inference from transformers via speculative decoding,2023,ICML】。在本文中,从拟合的几何分布估计接受率是基于相同的假设。【37,Dynamic speculation lookahead accelerates speculative decoding of large language models,2024】中的结果表明,这种假设实际上是现实的。【37,Dynamic speculation lookahead accelerates speculative decoding of large language models,2024】中的图4显示,每次迭代接受的词元数遵循几何分布。
F.3 消融研究:离线模拟
本节提供了关于消融分析的更多实现细节,该分析涉及一个补充性的“离线”实验(§4.1)。
我们对所有可能的⟨起草模型延迟, 接受率⟩配置进行了非SI的模拟,这些配置是 {0.01, 0.02, ..., 1} × {0, 0.01, 0.02, ..., 1} 的笛卡尔积。我们对所有可能的⟨起草模型延迟, 接受率, lookahead⟩配置进行了SI的模拟,其中 lookahead ∈ {1, 2, ..., 200}。然而,给定任意数量的可用服务器,可以调整lookahead超参数,使得DSI仅协调可用的服务器。为确保DSI可以在一个最多有八个GPU的单节点上部署,我们对所有满足SP=7的公式1的SI配置进行了DSI模拟,假设起草模型在单个GPU上运行。对于每个⟨起草模型延迟, 接受率, lookahead⟩的组合,我们将算法(SI或DSI)运行五次并对结果取平均值。由于lookahead是一个可调参数,我们的实验假设用户会对其进行优化,以便为每个配置优化SI。为了让SI选择其最佳的lookahead,每个⟨起草模型延迟, 接受率⟩配置的期望延迟是所有lookahead值中的最小平均延迟。
为了计算算法X相对于算法Y在每个⟨起草模型延迟, 接受率⟩下的加速比,我们将Y的延迟除以X的延迟。由于lookahead超参数的离散化,当起草模型延迟 < 20% 时,加速比并不平滑。例如,对于lookahead=5,SI和DSI的加速比都是平滑的,如附录F.7所示。
与在线实验一样,离线实验中的每次前向传播都被替换为将现实的等待时间加到总延迟中。因此,离线模拟中非SI的延迟就是目标模型前向延迟(即,计算目标模型单次前向传播的平均时间,这是在上一节描述的独立实验中测量的)乘以要生成的词元数。例如,如果目标模型前向延迟是30ms,要生成的词元数是100,那么非SI的延迟是3000ms。为了估计SI的期望延迟,我们必须引入lookahead超参数和起草模型的接受率。在接受率的随机性下,我们计算目标模型和起草模型前向传播的次数。每个被接受的草稿词元相当于生成了另一个输出词元,因此将剩余的目标模型前向传播次数减一。例如,如果接受率对应于每次迭代接受1.5个草稿,那么目标模型前向传播的期望次数是 $\frac{100}{2.5} = 40$ 次,因为每次目标模型前向传播预计会产生 1.5 + 1 = 2.5 个输出词元(即,接受的草稿加上一个额外的词元)。起草模型前向传播的次数取决于lookahead超参数。例如,如果lookahead=5,那么起草模型前向传播的次数是 $40 \cdot 5 = 200$ 次,因为每次目标模型前向传播都有5次起草模型前向传播。SI的期望延迟就是目标模型和起草模型前向传播的期望延迟之和。例如,如果起草模型前向延迟是6ms,那么SI的期望延迟是 $200 \cdot 6 + 40 \cdot 30 = 2400$ms,比非SI快1.25倍。DSI的模拟与SI的模拟类似。在这两种模拟中,唯一的随机性是起草模型的接受率。然而,由于DSI可以成功地隐藏几乎所有目标模型前向传播的延迟,简单地累加前向传播延迟无法准确估计其总延迟。相反,DSI的模拟基于一种分析,该分析仅在目标模型前向延迟未被推测并行性隐藏时才对其进行累加。
F.4 SI模拟
我们开源了DSI在单节点设置下的实现以及本文所有实验的代码。以下是估算SI算法端到端延迟下限的代码。端到端延迟是总的挂钟时间,包括预填充和解码延迟,但不包括任何分词延迟。这是一个下限,因为我们忽略了除前向传播延迟之外的任何真实世界延迟。例如,我们忽略了所有与模型切换、CPU与GPU之间通信等相关的延迟。
def si(target_latency: float, drafter_latency: float, lookahead: int, N: int) -> float:
total_cost: float = 0
total_toks: int = 0
while total_toks < N:
num_accepted: int = get_num_accepted()
total_toks += num_accepted + 1
total_cost += lookahead * drafter_latency + target_latency
return total_cost
F.5 模型
对于所有模型,我们从Hugging Face获取模型权重。为清晰和可复现性,我们提供了每个使用模型的URL:
* Vicuna-13B: https://huggingface.co/lmsys/vicuna-13b-v1.3,根据非商业许可分发 。
* Vicuna-7B: https://huggingface.co/lmsys/vicuna-7b-v1.3,根据非商业许可分发 。
* Vicuna-68M: https://huggingface.co/double7/vicuna-68m,根 据Apache License 2.0分发。
* Starcoder-15B: https://huggingface.co/bigcode/starcoder,根 据Responsible AI License分发。
* Starcoder-168M: https://huggingface.co/bigcode/tiny_starcoder_py,也根 据Responsible AI License分发。
* Phi3-14B: https://huggingface.co/microsoft/Phi-3-medium-128k-instruct,根 据MIT license分发。
* Phi3-4B: https://huggingface.co/microsoft/Phi-3-mini-128k-instruct,根 据MIT license分发。
F.6 数据集与提示
我们使用来自Hugging Face的标准数据集和来自最先进技术的标准提示。
F.6.1 MBPP
MBPP数据集包含众包的Python编程问题,并根据cc-by-4.0许可证分发。
关于提示,我们遵循(【6,A framework for the evaluation of code generation models,2022】; 【19,Incoder: A generative model for code infilling and synthesis,2023】)的做法,包含了编程任务的描述和用于验证解决方案的单个测试,以帮助模型捕捉函数的签名(见图3)。
F.6.2 HUMANEVAL
HumanEval数据集包含编程问题,并根据MIT许可证分发。
提示仅包含数据集中的prompt字段。
F.6.3 CNN-DM
CNN-DM包含新闻文章,并根据Apache License 2.0分发。
我们在提示中包含了文章字段,如图4所示。
F.6.4 ALPACA
Alpaca数据集包含指令和演示。它根据cc-by-nc-4.0许可证分发。
我们遵循【59,Alpaca: A strong, replicable instruction-following model,2023】来定义提示。对于输入字段不为空的样本,我们使用如图5所示的提示;而对于输入字段为空的样本,我们使用如图6所示的提示。
F.7 LOOKAHEAD=5时的加速比
图7展示了在lookahead=5的静态设置下,推测式推理(SI)、非推测式推理(non-SI)和动态推测式推理(DSI)的性能比较。该图包含三个热力图,每个图都代表了这些算法在不同起草模型速度和准确率下的成对比较。
在图7(a)中,我们比较了SI和non-SI。粉色区域表示SI比non-SI慢的场景,这发生在起草模型速度过慢或不够准确时。这突显了SI在特定条件下的局限性。
图7(b)和(c)表明,DSI在所有起草模型配置下都持续优于SI和non-SI。这一经验证据支持了我们的理论发现,即DSI在期望上比SI和non-SI都快,无论起草模型的速度或准确率如何。
这些结果强调了DSI相比现有推理方法的鲁棒性和效率,特别是在SI可能因起草模型性能不佳而失败的场景中。
💬 评论讨论
欢迎在这里分享您的想法和见解!