ASPIRE：面向长上下文大语言模型推理的异步批处理自推测解码技术

# ASPIRE：用于长上下文大语言模型推理的异步批处理自推测解码
来源：https://arxiv.org/html/2609.17943
ASPIRE：用于长上下文大语言模型推理的异步批处理自推测解码

###### 摘要

长上下文大语言模型的推理受限于注意力机制，其重复的KV缓存读取使得解码过程成为内存瓶颈。自推测解码通过使用稀疏注意力草拟令牌，再用全注意力进行验证，缓解了这一问题。然而，现有的批处理方法仍然是*同步的*：即使每个请求的最优草拟长度差异巨大且在每个请求过程中动态变化，一个批次中的所有请求仍共享单一的草拟-验证调度。我们提出了ASPIRE，一个基于三个组件构建的非同步批处理自推测解码框架。首先，*统一混合前向传播*允许草拟和验证请求在同一次批处理前向传播中共存，消除了对全局草拟-验证阶段的需求。其次，轻量级的*在线推测调度器*利用每个请求的接受率估计和一个具有批次感知能力的成本模型，使每个请求能独立决定何时进行验证。第三，*草拟内刷新层*在草拟过程中在单个指定层执行全注意力，在每个草拟步骤更新稀疏上下文，以减少草拟过程中的过时问题。在三个模型和五个推理及长上下文基准测试中，ASPIRE实现了相对于自回归基线1.70-4.58倍的解码吞吐量加速，并将平均加速比提高了约27%，优于之前最强大的自推测基线。

## 1 引言

大语言模型正越来越多地应用于需要长上下文处理的任务，包括代理工作流、长期推理和检索增强生成（OpenAI, 2025 (https://arxiv.org/html/2609.17943#bib.bib1); Google DeepMind, 2025 (https://arxiv.org/html/2609.17943#bib.bib2); Yang et al., 2025a (https://arxiv.org/html/2609.17943#bib.bib3)）。随着上下文窗口的扩展，在大语言模型服务中同时实现高吞吐量和低延迟变得越来越困难（Liu et al., 2025 (https://arxiv.org/html/2609.17943#bib.bib4)）。这一挑战源于Transformer解码两个主要组成部分的不同性能特性：多层感知机和自注意力。MLP主要受模型权重访问限制，通常是计算密集型的，而自注意力取决于上下文长度，通常是内存密集型的（Tang et al., 2024 (https://arxiv.org/html/2609.17943#bib.bib9)）。

批处理主要通过分摊批次内请求的模型权重加载来提高吞吐量，从而增加MLP的算术强度，并将其执行从类似矩阵-向量乘法转变为类似矩阵-矩阵乘法（Kwon et al., 2023 (https://arxiv.org/html/2609.17943#bib.bib5); Agrawal et al., 2024b (https://arxiv.org/html/2609.17943#bib.bib6); Zheng et al., 2024 (https://arxiv.org/html/2609.17943#bib.bib7)）。然而，注意力机制并不能以同样方式受益：每个请求都有一个独立的键值缓存，因此其上下文无法在批次间分摊。对于每个新生成的令牌，模型必须读取该请求完整先前上下文的键和值。因此，随着序列长度和批处理大小的增加，注意力机制的内存访问量增长，而其算术强度保持较低（Agrawal et al., 2024a (https://arxiv.org/html/2609.17943#bib.bib8)）。

为了缓解这种内存低效性，自推测解码通过将高效的草拟与精确的验证配对，用计算换取更低的内存访问量。模型使用稀疏注意力自回归地草拟多个令牌，然后通过一次全注意力的前向传播并行验证它们。由于对完整KV缓存的昂贵访问仅在验证期间发生一次，该成本可以被多个被接受的令牌分摊。因此，自推测解码在长上下文场景中特别具有吸引力，因为重复的KV缓存读取主导了运行时间。最近的系统如MagicDec（Sadhukhan et al., 2025 (https://arxiv.org/html/2609.17943#bib.bib10)）和Vegas（Yue et al., 2026 (https://arxiv.org/html/2609.17943#bib.bib11)）展示了这一方向的有效性。

MagicDec（Sadhukhan et al., 2025 (https://arxiv.org/html/2609.17943#bib.bib10)）是一种使用窗口注意力进行草拟的批处理自推测解码方案。Vegas（Yue et al., 2026 (https://arxiv.org/html/2609.17943#bib.bib11)）通过根据验证令牌的注意力分数选择高分上下文令牌来提高草拟准确性。然而，Vegas在草拟过程中保持部分KV上下文固定，仅在验证后才刷新它。结果，随着生成的进行，草拟上下文变得越来越过时，限制了草拟长度和质量，最终限制了加速比。为了解决这个问题，我们引入了轻量级的草拟内上下文刷新。通过利用同一令牌在不同层之间以及连续令牌之间注意力模式的相似性（Zarch et al., 2026 (https://arxiv.org/html/2609.17943#bib.bib12); Yang et al., 2025b (https://arxiv.org/html/2609.17943#bib.bib13)），我们使用一个刷新层在草拟过程中执行全注意力，并以最小的开销更新下一个令牌的部分KV。

尽管有这些进展，现有的批处理推测解码系统在很大程度上仍然以同步方式运行：同一批次中的请求通常在共享的推测调度下一起经历草拟和验证。虽然这种设计简化了执行，但它与实际的批处理工作负载匹配不佳。不同的请求通常具有不同的最优草拟长度，即使在单个请求内，随着生成的进行，最优草拟长度也可能发生显著变化。因此，同步批处理迫使异构请求同步推进，导致一些请求草拟不足，而另一些请求草拟过度。因此，批处理推测解码的收益不仅受到次优推测长度的限制，还受到同步本身的限制。

为了解决这些限制，我们提出了ASPIRE，一个用于大语言模型推理的动态、非同步批处理推测解码框架。ASPIRE围绕一个统一的混合批次构建，该批次允许同一批次中的请求在单次模型前向传播中保持不同的推测状态：一些请求使用部分KV上下文草拟，而另一些则使用完整KV进行验证。这消除了对同步的批处理范围草拟和验证阶段的需求，同时保留了共享批处理执行的效率优势。在这个执行模型之上，ASPIRE采用了一个细粒度的在线调度器，让每个请求遵循自己的推测轨迹。使用每个请求的接受率估计和具有批次感知能力的草拟/验证成本比，调度器决定每个请求在每个模型批次前向传播中是草拟还是验证。最后，ASPIRE集成了一个轻量级的全注意力刷新层，在草拟过程中以最小的开销更新部分KV上下文，提高了草拟质量并支持更长的草拟长度。因此，ASPIRE提高了推测利用率，并提升了批处理长上下文大语言模型推理的吞吐量。

总之，我们的贡献如下：
- • 我们对批处理推理中的草拟长度行为进行了实证研究，表明最优草拟长度在请求间高度异构，且在每个请求的生成过程中高度动态。
- • 我们提出了ASPIRE，一个动态非同步批处理推测解码框架，它在单个请求的粒度上调度草拟和验证，允许同一批次内的请求遵循各自的推测轨迹。
- • 我们引入了一种轻量级的上下文刷新机制，在草拟过程中利用跨层和连续令牌之间的注意力相关性，以最小的开销更新部分KV集。
- • 我们在三个模型和五种工作负载上进行了广泛的实验，表明ASPIRE显著提高了相对于自回归解码的吞吐量，最高可达4.58倍加速，并且与之前的推测解码基线相比，始终实现了最佳的整体性能。

## 2 背景

基于Transformer的大语言模型推理。仅解码器的大语言模型由堆叠的Transformer块组成，每个块主要包含一个自注意力模块和一个MLP。大语言模型推理包括一个*预填充*阶段，后跟一个*解码*阶段。在预填充阶段，提示被并行处理，生成的每一层的键和值存储在KV缓存中。在解码阶段，模型一次生成一个令牌；对于每个新令牌，注意力必须读取该请求完整先前上下文的缓存KV。在长上下文推理中，这种重复的KV缓存访问主导了注意力成本。这在批处理服务中造成了一个重要的不对称：MLP层主要受共享模型权重访问的限制，通常从批处理中受益，而注意力仍然与每个请求的缓存上下文相关联，因此更受内存限制。相应地，解码可以使用*全上下文注意力*（关注整个缓存上下文）或*稀疏上下文注意力*（仅关注该上下文的一个选定的稀疏子集）。对于请求\(i\)，我们将这个选定的稀疏上下文记为\(\mathcal{S}_i\)。稀疏上下文注意力降低了注意力成本，但当\(\mathcal{S}_i\)遗漏了对当前令牌重要的上下文位置时，可能会降低准确性。

稀疏上下文自推测解码。自推测解码在*草拟*和*验证*之间交替进行。在草拟期间，请求\(i\)使用对\(\mathcal{S}_i\)的稀疏上下文注意力生成推测令牌。在验证期间，相同的目标模型对推测前缀运行一次全上下文的前向传播，并将草拟的令牌与相应的全上下文输出进行比较。令\(d_i\)表示当前的*草拟长度*，即尚未验证的缓冲草拟令牌数。如果请求\(i\)在草拟长度\(d_i\)时进行验证，并且接受了\(a_i \in \{0, \dots, d_i\}\)个草拟令牌，则被接受的草拟前缀以及验证的下一个令牌会被附加到请求的上下文后。然后请求重置其草拟长度，并开始累积新的推测前缀。现有的批处理推测解码系统通常使用同步执行：批次中的请求在共同的推测调度下共享全局草拟和验证阶段。所有请求一起草拟，直到批次级条件触发验证，然后所有请求一起验证。这种设计简化了执行，但它迫使异构请求即使在最优验证深度不同时也同步推进。

## 3 动机观察

#### 最优草拟长度取决于请求。

我们首先探究单一的共享草拟长度是否能很好地服务于所有请求。为了回答这个问题，我们在256个前缀长度在\([16\mathrm{k},20\mathrm{k}]\)范围内的LongBench样本上进行了一项研究。对于每个请求和每轮验证，我们记录了在草拟20个令牌的情况下被接受的草拟令牌数，最多生成1k个令牌。图1(a)显示了每个请求的平均接受草拟长度。分布很广，范围从2.56到20.00，平均值为11.51，表明存在显著的异构性，说明不同的请求偏好不同的草拟长度，因此单一的共享草拟长度本质上是次优的。

#### 最优草拟长度在单个请求内也会随时间变化。

图1(b)显示了每个请求在每轮验证中被接受的草拟令牌数。行按每个请求的平均接受草拟长度降序排列，以突出跨请求的异构性。明显的递减趋势是这种排序的结果。平均接受草拟长度较高的请求出现在顶部，每轮验证提交更多令牌，因此在更少的轮次内达到最终答案或生成长度上限。然而，在单个行内，接受的草拟长度在不同验证轮次间波动很大。因此，单一的共享固定草拟长度或固定的请求特定草拟长度都不足够。结合图1(a)中的跨请求异构性，这为每个请求、步骤级的调度而非同步的批处理范围草拟和验证提供了动力。

参考图1：(a) 平均接受草拟长度在请求间差异很大。(b) 每个请求内接受的长度在不同验证轮次间变化。行按其均值降序排列，产生了整体的递减趋势；白色单元格表示完成后的轮次。(c) 用于时间、跨层和单刷新层源的稀疏上下文预测器召回率。固定源使用前一个令牌的倒数第二层。

#### 跨层信号比时间重用更稳定。

图1(c)评估了不同的注意力图源在多大程度上能够预测目标令牌在目标层所需的稀疏上下文。对于每个候选源，我们使用其注意力图选择前100个KV页面，并衡量目标令牌/层的总注意力质量有多少落在这些页面上，以整体质量召回率@100报告。预言机使用目标令牌/层自身的注意力图选择前100个页面，因此代表了在相同页面预算下可达到的最佳召回率。

第一组，时间\(\Delta t\)，使用同一层的早期令牌的注意力图来预测当前令牌。召回率随着令牌滞后\(\Delta t\)的增长而迅速下降，表明来自早期草拟令牌的注意力图在草拟过程中变得过时。第二组，层\(\Delta \ell\)，使用同一令牌的前\(\Delta \ell\)层的注意力图。例如，当\(\Delta \ell=1\)时，每个目标层使用其紧前一层的注意力图。随着\(\Delta \ell\)增加，召回率的下降速度比时间重用慢得多，表明附近的跨层信号是目标稀疏上下文的更稳定预测器。然而，这种比较主要是诊断性的，因为为每个目标层使用固定的层偏移将需要网络中多个层特定的刷新信号。

#### 固定的深层源提供了最佳的质量-开销权衡。

第三组，刷新\(i\)，评估了一个实用的设计，即单个相同的令牌刷新点。所有直到层\(i\)的层都执行全注意力，之后层\(i\)的注意力图用于为剩余层选择稀疏KV页面。将刷新点移得更深可以提供更强的预测器，但会直接增加每个草拟步骤中执行的全注意力层数。在图1(c)中，从刷新层0到2召回率有所提高，然后变化很小，表明增加全注意力层的好处有限。

这些观察结果促使ASPIRE采用来自前一个令牌的固定深层源。该源只有一个令牌的年龄，且来自较深层，结合了\(\Delta t=1\)的低时间过时性和强大的跨层信号。与相同令牌的刷新\(i\)不同，