Accelerating Large-Scale Reasoning Model Inference: Self-Speculative Decoding with Sparse Attention
Accelerating Large-Scale Reasoning Model Inference: Self-Speculative Decoding with Sparse Attention
文章标题:加速大规模推理模型推断:基于稀疏注意力的自推测解码
作者/机构:Yilong Zhao 1 *, Jiaming Tang 2 *, Kan Zhu 3, Zihao Ye 4, Chi-Chih Chang 5, Chaofan Lin 6, Jongseok Park 1, Guangxuan Xiao 2, Mohamed S. Abdelfattah 5, Mingyu Gao 6, Baris Kasikci 3, Song Han 2 4, Ion Stoica 1
A1 主要贡献
核心问题: 推理语言模型(Reasoning Language Models, RLMs)通过生成详尽的思维链(Chain-of-Thought, CoT)解决方案,在复杂任务上展现出卓越的能力。然而,这种长文本生成方式将推理瓶颈从计算密集型(compute-bound)转变为内存密集型(memory-bound)。为了生成每个词元(token),模型需要对所有先前生成的词元进行全量注意力计算,这需要访问一个不断增大的键值缓存(KV-Cache)。因此,生成序列越长,每一步的内存访问量就越大,给内存带宽带来巨大压力。例如,在使用H100 GPU、批处理大小为128、输出长度为8192的情况下,为Qwen3-8B模型加载KV-Cache平均每步耗时21毫秒,占端到端延迟的70%以上。
研究目标: 本文旨在解决推理语言模型(RLMs)长文本生成中的内存带宽瓶颈问题,提出一个无需训练、无损的加速框架,以提升推理吞吐量。
创新点:
为了应对上述挑战,本文提出了SparseSpec,一个专为RLMs推理设计的无损、免训练的加速框架。其核心思想是利用同一模型作为草稿模型和目标模型(即自推测),并引入新颖的动态稀疏注意力机制和系统级协同设计来释放自推测解码的潜力。
-
PillarAttn动态稀疏注意力机制: 作为SparseSpec的核心,PillarAttn是一种专为推测解码设计的动态稀疏注意力机制。它通过在验证阶段优雅地复用全量注意力计算得到的确切注意力分数,来精确选择并仅加载和计算那些最关键的词元。这些被识别出的关键词元随后被用于后续的草稿生成步骤,从而在适应不同上下文动态性的同时,实现了高推测准确率。与现有的动态稀疏方法相比,PillarAttn识别关键词元的过程实现了零内存开销。
-
系统级协同设计与优化: SparseSpec将PillarAttn与三项系统级优化相结合,以解决现有方法在应用于RLMs时面临的系统性挑战:
- 统一批处理调度器 (Unified Scheduler): 将稀疏的草稿阶段和密集的验证阶段的请求合并到同一个批次中进行处理。这不仅可以分摊模型权重的加载成本,还能在草稿和验证任务间均匀分配请求,使得每次迭代的资源使用保持均衡,从而缓解工作负载波动问题,最大化并行性。
- 延迟验证 (Delayed Verification): 将验证过程推迟一个迭代周期,从而将CPU的验证操作移出关键路径。这使得上一个迭代周期的CPU验证工作可以与当前迭代周期的GPU操作重叠执行,避免了因显式同步造成的性能瓶颈。
- 动态KV-Cache管理器 (Dynamic KV-Cache Management): 通过异步、分块的方式将KV-Cache在GPU和主机内存之间进行换入/换出。这种方式能够最大化GPU内存的利用率,避免因RLMs输出长度不可预测而导致的内存浪费或过度重计算。
主要贡献总结:
- 分析了批量化RLMs推理的性能瓶颈,指出了内存带宽是关键,并从理论上阐述了稀疏自推测解码的潜在收益。
- 提出了一个名为SparseSpec的无损、免训练的加速框架,其核心包括新颖的PillarAttn稀疏注意力机制,以及统一批处理调度器、延迟验证和动态KV-Cache管理器三项系统优化。
- 在真实世界的工作负载上对该框架进行了全面评估,结果表明与最先进的推理框架相比,SparseSpec实现了高达2.13倍的加速。
表1. 三个推理评估数据集上的平均(均值±标准差)词元长度。两种模型的输入长度相同;输出长度包括其标准差。
A3 背景知识与关键洞察
2.1 基于Transformer的大型语言模型
MLP与Attention的性能差异。基于Transformer的LLM主要由多层感知机(MLP)和自注意力模块构成。在批量推理中,MLP和注意力模块表现出不同的性能特征:MLP通常是计算密集型的,而注意力模块是内存密集型的【Tang et al., 2024, Quest: Query-aware sparsity for efficient long-context llm inference, arXiv】。对于MLP中的通用矩阵乘法(GEMM)操作,批处理中的请求共享相同的模型权重,这可以分摊内存加载成本,从而实现高算术强度。因此,延迟函数$T_{GEMM}(B)$(将批大小B映射到延迟)在达到某个阈值$B̂$之前几乎保持不变,在GPU完全饱和后随B线性增加【Zhao et al., 2024a, Atom: Low-bit quantization for efficient and accurate llm serving, arXiv】。相比之下,不同请求具有独立的上下文(即KV-Cache),无法在批处理中分摊。因此,内存流量随序列长度和批大小线性扩展,导致算术强度持续较低,延迟函数$T_{Attn}(M)$呈线性,其中M是KV-Cache的总内存大小。本文重点优化内存密集型的注意力模块,这是长输出RLMs推理的主要瓶颈。
2.2 RLM不断演变的工作负载
RLM通过强化学习生成更长的思维链。推理语言模型(RLMs)是一类旨在通过增强的推理能力解决复杂任务的LLM,如OpenAI-o1【OpenAI, 2024, Openai o1 system card, arXiv】和DeepSeek-R1【DeepSeek-AI, 2025a, Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning, arXiv】所示。与传统LLM不同,它们通过强化学习(RL)被明确激励遵循审慎的思维链(CoT)【Kimi, 2025, Kimi k1.5: Scaling reinforcement learning with llms, arXiv】。因此,RLMs在推理过程中通常会生成比非推理模型多得多的词元。为了证明这一点,我们收集了几个推理任务的平均输入和输出长度。如表1所示,RLM Qwen3-14B在AIME数据集上平均生成13542个词元,比非推理模型Qwen-2.5-32B的2593个词元多出近7倍。
长输出对训练和部署均构成挑战。如此长的输出对RLMs的训练和部署阶段都构成了重大挑战。在RL后训练阶段,RLMs被部署以生成多个轨迹(rollouts),然后由奖励模型评分以演进模型。这种轨迹生成是面向吞吐量的离线推理,可能占整个RL训练时间的90%以上【Luo et al., 2025, Deepscaler: Surpassing o1-preview with a 1.5b model by scaling rl, Notion Blog】。此外,包括聊天机器人和代理工作流在内的RLMs部署是面向延迟的在线推理,同样受到内存密集型注意力的瓶颈制约。本文专注于可应用于训练和部署阶段的无损加速。
2.3 推测解码
推测解码通过计算换内存。为了缓解内存效率低下的问题,推测解码被提出作为一种通过计算换取内存的无损加速技术【Chen et al., 2023, Accelerating large language model decoding with speculative sampling, arXiv】、【Fu et al., 2024, Break the sequential dependency of llm inference using lookahead decoding, arXiv】。其核心思想是利用一个轻量级的草稿模型顺序提出k个候选词元,然后由原始的目标模型并行验证这些词元。与目标模型一致的候选词元被接受,其余的则被丢弃并通过相同的工作流程重新生成,从而确保原始模型的生成质量无损。
推测解码的加速原理。假设草稿模型和目标模型的延迟分别为$T_{draft}$和$T_{target}$,加速来源于$T_{draft} \ll T_{target}$。因此,尽管只有$\alpha \cdot k + 1$个候选词元被接受(其中$\alpha$是接受率,并有一个额外的奖励词元),启用推测解码时的端到端延迟$(k \cdot T_{draft} + T_{verify})$小于不使用推测解码时的延迟$(\alpha k + 1) \cdot T_{target}$。
3.1 批量RLMs推理是注意力密集型的
长文本生成导致注意力成为瓶颈。RLMs的长文本生成在批量推理中引入了显著的内存瓶颈。与短文本生成任务不同,每个推理请求在完成前都会累积大量的KV-Cache,这会迅速耗尽GPU内存并严重限制并发请求的数量(即批大小)。例如,在H100上服务Llama-3-8B,对于8K的生成长度,只能容纳64个请求。在这种小批大小下进行推理,MLP和注意力模块都受到内存限制,而由于KV-Cache通常比模型权重更大,注意力成为主要的运行时组件。
性能剖析证实了内存瓶颈。为了证明这一点,我们对在vLLM上运行的RLMs的端到-端执行过程中的带宽和计算利用率进行了剖析。我们在图2中可视化了单次迭代内的利用率。即使在执行MLP时,计算利用率也始终低于50%;相比之下,内存带宽在整个迭代过程中被大量使用,表明这是一个内存密集型的工作负载。尽管MLP和注意力都是内存密集型的,但注意力占主导地位,占用了超过77%的端到端执行时间。
3.2 稀疏自推测解码的帮助
稀疏自推测解码的概念。如§3.1所述,由于批量场景下KV-Cache的过度内存访问,RLMs的推理是注意力密集型的。幸运的是,先前关于稀疏注意力的研究表明,KV-Cache中一小部分词元(例如5%)主导了注意力的输出,可以提供近乎无损的加速【Tang et al., 2024, Quest: Query-aware sparsity for efficient long-context llm inference, arXiv】、【Lin et al., 2025, Twilight: Adaptive attention sparsity with hierarchical top-p pruning, arXiv】。因此,通过选择性地加载和计算最关键的词元,可以大幅减少内存访问,以准确性换取效率。为确保无损加速,这种方法可以作为推测解码中的草稿模型,与全量目标模型协同工作(即稀疏自推测解码【Sun et al., 2024, Triforce: Lossless acceleration of long sequence generation with hierarchical speculative decoding, arXiv】)。
加速的理论分析。我们现在建立稀疏自推测解码的理论加速比公式,以估算潜在的提升。我们将加速比η定义为使用和不使用推测解码时平均每个词元的生成延迟之比(即$T_{base}/T_{spec}$)。具体来说,设M为KV-Cache的总内存大小,B为并发请求数。使用§2.1中的相同符号,忽略其他小操作,则$T_{base} = T_{GEMM}(B) + T_{Attn}(M)$。
推导过程。假设$k$为草稿长度,$\alpha$为接受率,$s$为草稿注意力的稀疏度,我们通过考虑一个包含$k$个草稿阶段和1个验证阶段的完整推测轮次来计算$T_{spec}$。因此,我们有$(k\alpha + 1) \cdot T_{spec} = k \cdot T_{draft} + T_{verify}$,其中$T_{draft} = T_{GEMM}(B) + T_{Attn}(M \cdot s)$且$T_{verify} = T_{GEMM}((k + 1) \cdot B) + T_{Attn}(M)$。由于B个并发请求随机分布在k个草稿和1个验证阶段,GEMM的平均输入大小为$\frac{2k+1}{k+1}B$,而注意力的平均输入大小为$\frac{ks+1}{k+1}M$。因此,我们可以将先前推导的$(k\alpha + 1) \cdot T_{spec}$方程简化为:
延迟近似。如§2.1所述,当GPU未完全饱和时,$T_{GEMM}$是非线性函数,而$T_{Attn}$近似线性。因此,$T_{spec}$可以近似为:
加速比公式。给定每个接受词元的延迟$T_{spec}$和普通生成中每个词元的延迟$T_{base}$,加速比$\eta$可导出为$\eta = T_{base}/T_{spec}$。
实际性能影响。由于GEMM和注意力的特性截然不同,我们分别分析$\eta$对GEMM和注意力的影响。对于注意力,由于$T_{Attn}(M)$在RLMs推理中占主导地位(如§3.1所述),KV-Cache访问的减少比例为$(ks + 1)/(k\alpha + 1)$,这对加速贡献最大。由于实践中$\alpha \gg s$,注意力延迟可以大幅降低。例如,当$k = 16$, $\alpha = 0.75$, $s = 0.05$时,注意力延迟减少了6.78倍。然而,这种内存减少是以额外的GEMM计算为代价的,即用计算换取内存。在实践中,为避免不可容忍的延迟增加,批大小通常应限制在某个阈值$B̂$以下,以使$T_{GEMM}$不会过饱和。例如,在Hopper GPU上,$B̂ = 256$仅导致最小的延迟增加。此外,$\alpha$对于限制延迟增加也至关重要,因为加速比$\eta$与$\alpha$成反比。
3.3 现有方法在RLMs上的不足
然而,我们发现现有方法在RLMs推理场景下未能提供§3.2模型所预期的加速效果。
现有方法的性能差距。如图3所示,当在H100上使用AIME数据集服务Qwen3-8B时,MagicDec未能接近预言(oracle)top-k注意力【Lin et al., 2025, Twilight: Adaptive attention sparsity with hierarchical top-p pruning, arXiv】的接受率,因此远未达到理论上的最优加速比。这是因为现有方法中的草稿模型不能适应RLMs中上下文的动态性,导致生成的草稿词元不准确。此外,即使接受率相同,由于一些系统性挑战,现有方法与预期加速比之间仍然存在不可忽视的差距。
上下文动态性。对于RLMs的长输出,稀疏模式会根据生成过程中的语义上下文发生显著变化。如图4所示,KV-Cache中的关键词元随时间显著不同。因此,需要一种能够适应上下文动态性的动态稀疏注意力机制,而不是使用静态的稀疏模式(例如,滑动窗口注意力),才能在草稿阶段实现高接受率$\alpha$。此外,这种稀疏注意力的设计应引入最小的计算或存储开销,以避免增加系统复杂性。
工作负载波动。由于草稿和验证阶段的资源使用情况不同,将这两个阶段分开进行朴素调度会导致迭代间的工作负载波动,从而造成严重的硬件资源未充分利用。这种异构性主要来自GEMM操作。例如,给定批大小B和推测步数k,顺序执行模式(所有草稿阶段后跟一个验证阶段)需要:(1) k次输入大小为B的GEMM操作,和(2) 一次输入大小为(k+1)B的验证GEMM。这导致草稿阶段利用率不足,而验证阶段过饱和,总时间为$k T_{GEMM}(B) + T_{GEMM}((k+1)B)$。相反,均匀调度总耗时为$(k+1)T_{GEMM}(\frac{2k+1}{k+1}B)$。考虑到$T_{GEMM}$在饱和前的非线性特性,当B远低于饱和点时,$T_{GEMM}(\frac{2k+1}{k+1}B) < T_{GEMM}(2B) \approx T_{GEMM}(B)$。因此,朴素调度会带来高达2倍的执行时间。
显式同步。推测解码的范式,即先生成候选词元再进行验证,引入了CPU和GPU之间的显式同步,这阻碍了昂贵的CPU操作与GPU操作的重叠。具体来说,GPU必须等待CPU完成验证过程,该过程会从请求中清除被拒绝词元的元数据,并为下一步生成重置草稿状态。因此,最先进的推理框架【Zheng et al., 2024, Sglang: Efficient execution of structured language model programs, arXiv】、【Kwon et al., 2023, Efficient memory management for large language model serving with pagedattention, arXiv】在启用推测解码时必须在关键路径上运行CPU操作,从而产生显著开销。在稀疏自推测解码中,由于统一调度下每次迭代都有验证请求,这种低效率问题被进一步加剧。
KV-Cache利用率不足。RLMs输出长度的巨大差异(§2.2)使得完美管理KV-Cache变得困难,因为输出长度在生成前是未知的。我们从一次真实世界试验中收集了内存利用率和重计算率(即重计算词元与生成的唯一词元之比)。如图5所示,现有方法要么未能充分利用KV-Cache容量,要么因预测失误导致过度重计算。这种KV-Cache利用率不足严重限制了自推测的加速效果,因为效率指标$\eta$会随着KV-Cache尺寸的减小而降低(§3.2)。
3.4 SparseSpec:通过高效的算法-系统协同设计释放稀疏自推测的潜力
机遇。在本文中,我们的目标是利用稀疏自推测解码来缓解批量RLMs推理中长文本生成所带来的内存带宽瓶颈,提供无损且免训练的加速。
挑战。然而,要实现自推测的理论加速比,需要通过算法-系统协同设计来解决以下几个挑战:(1) 如何设计一种能以最小开销适应上下文动态性的稀疏注意力机制;(22) 如何调度草稿和验证阶段以避免工作负载波动;(3) 如何处理显式同步以实现CPU/GPU重叠;以及(4) 如何在不发生重计算的情况下充分利用KV-Cache容量。
A2 方法细节
本节我们描述SparseSpec各组件的设计,包括一个动态稀疏注意力算法(§4.1)、一个统一的批处理调度器(§4.2)、延迟验证(§4.3)和一个动态KV-Cache管理器(§4.4)。我们在图6中可视化了SparseSpec的设计概览。
4.1 PillarAttn: 为自推测解码量身定制的动态稀疏注意力
性能关键在于高稀疏度和高接受率。如§3.2所述,高稀疏度s和高接受率α对稀疏自推测的性能至关重要。因此,一种能够准确高效地识别稀疏模式的稀疏注意力机制是必要的。为此,SparseSpec提出了PillarAttn,一种专为自推测解码量身定制的稀疏注意力机制。
动态更新稀疏模式以适应上下文变化。为了适应§3.3中讨论的上下文动态性,PillarAttn并非使用固定的稀疏模式(如滑动窗口或来自提示的静态模式),而是在草稿阶段周期性地重新识别和更新所使用的稀疏模式。基于上下文语义具有空间局部性的假设,PillarAttn以一个小的步长更新稀疏模式,在该步长内稀疏模式是固定的。因此,识别的开销可以被分摊到一个步长的迭代中。
与自推测协同设计实现无开销识别。这个步长与自推测解码的范式协同设计,其值与推测步数k相同。在每k个使用稀疏注意力的草稿步骤之后,会执行一个使用全量注意力的验证步骤,该步骤会计算KV-Cache中所有词元的注意力分数作为中间结果。为了避免额外的计算和存储开销,PillarAttn通过定制的注意力核在验证过程中动态地转储注意力分数。因此,通过复用并对转储的注意力分数应用Top-K操作,可以直接识别出稀疏模式。如果应用了分组查询注意力,这些分数会先在k个草稿词元和查询头(在同一组内)上进行平均。具体来说,在验证期间会缓存注意力logit和指数的对数和,这些值用于重新物化注意力分数以进行识别。我们在图7中展示了PillarAttn的工作流程,其中稀疏模式在每3个草稿阶段后的验证阶段进行更新。
4.2 统一批处理调度器
统一抽象与感知负载的调度。如§3.3所述,草稿和验证阶段具有异构的资源使用情况,这导致硬件利用率低下,从而性能不佳。为了解决这个问题,SparseSpec引入了一个统一的批处理调度器,为这两个阶段提供统一的抽象和基于工作负载感知的调度。此外,为了提高硬件效率,SparseSpec引入了一个融合的注意力核,可以高效地将稀疏和全量注意力批处理到单个核中。
为草稿和验证阶段提供统一抽象。由于在自推测中草稿模型和目标模型共享相同的模型权重,除了稀疏注意力和全量注意力之外,这两个阶段在GPU上经历完全相同的数据和控制流。幸运的是,现代推理框架已经集成了PagedAttention【Kwon et al., 2023, Efficient memory management for large language model serving with pagedattention, arXiv】,它本质上是一种页面粒度的稀疏注意力实现。因此,通过使用大小为1的页面,稀疏和全量注意力都可以用相同的流水线进行统一,这极大地简化了系统设计,并释放了在这两个阶段进行调度的灵活性。
实现工作负载感知的调度。如§3.3所示,顺序执行所有草稿阶段后再进行一次验证阶段会导致硬件利用率低下,因为工作负载极度不平衡。因此,鉴于统一抽象所带来的调度灵活性,SparseSpec在每个批次(在每个生成步骤)中均匀地混合来自这两个阶段的请求。为了实现完美平衡,SparseSpec引入了一种贪婪的箱装策略,以工作负载感知的方式将新进入的请求分配到不同的草稿阶段。具体来说,SparseSpec维护k个桶来跟踪每个草稿阶段的请求数量,并将每个进入的请求分配给负载最轻的桶。这种调度策略如图8所示。
融合稀疏和全量注意力。即使在理想的调度下,由于草稿和验证阶段的异构性,注意力操作仍然存在硬件利用率低的问题。这是因为草稿和验证中的注意力由于输入词元数量不同,具有不同的算术强度,导致最佳的核实现方式也不同(例如,tile大小和MMA指令)。为了证明这一点,我们对最先进的FlashInfer【Ye et al., 2025, Flashinfer: Efficient and customizable attention engine for llm inference serving, arXiv】在稀疏和全量注意力上进行了性能剖析。结果显示,为验证(全量)注意力优化的核实现了接近85%的带宽,但对草稿(稀疏)注意力的带宽不到50%;而为稀疏注意力优化的核提供了超过80%的带宽,但在全量注意力上降至50%。为了解决这个问题,SparseSpec引入了一种持久化核风格的融合注意力核,它在芯片上将验证和草稿注意力分派到它们各自的最佳模板,而不是启动两个核,从而实现了两全其美。我们在§5.5中评估了融合核。
4.3 延迟验证
将验证移出关键路径。如§3.3所述,验证阶段在CPU和GPU之间引入了显式同步,这阻碍了昂贵的CPU操作与GPU操作的重叠。例如,第i次迭代的执行依赖于第(i-1)次迭代的验证结果,特别是关于被拒绝的词元和更新后的稀疏模式。这种顺序执行可能占到端到端延迟的20%以上【Srivatsa et al., 2024, Can scheduling overhead dominate llm inference performance? a study of cpu scheduling overhead on two popular llm inference systems】。
利用验证请求占比小的特点。我们的关键观察是,这种同步仅适用于处于验证阶段的请求。在§4.2的平衡调度下,这类请求仅占批处理B的一小部分($\frac{1}{k+1}$)。因此,SparseSpec不是让整个批次停顿,而是允许CPU上非验证请求的元数据准备直接进行,而无需等待GPU的验证结果。例如,只有来自第(i-1)次迭代的验证请求被暂停并从第i次迭代中移除。在从GPU获得验证结果后,这些请求在第(i+1)次迭代时被提交到GPU执行。我们用图9说明了延迟验证。
4.4 动态KV-Cache管理
积极的CPU换出策略。鉴于RLMs中输出长度的高方差(表1),在不进行重计算的情况下实现高KV缓存利用率是具有挑战性的(图5)。因此,SparseSpec不优化输出长度预测,而是倾向于积极增加请求并发度以充分利用KV-Cache,并在接近内存不足时将KV-Cache换出到主机内存,以避免重计算。请注意,换出和加载都遵循FIFO顺序,以确保公平性和避免饥饿。
开销分析。首先,SparseSpec通过将大型KV-Cache分解成块,并以异步、逐块的方式进行换出,实现了开销可忽略的换出操作。例如,当在单张H100 GPU上以128的批大小运行Qwen3-8B时,每个解码步骤仅生成128个新词元,只需要18MB的KV-Cache内存。由于每次迭代的GPU延迟在10毫秒量级,只需18 GB/s的带宽即可将换出操作与GPU计算重叠,这远低于PCIe带宽限制。此外,只要GPU有可用内存,SparseSpec就会优先调度已换出的请求。该策略将最坏情况下的CPU使用量限制在GPU容量内,例如,对于一台8xH100服务器为640GB,这在数据中心典型的CPU DRAM限制内是完全可行的。
A4 实验
5.1 实验设置
模型与硬件。我们使用三款最先进的开源RLM来评估SparseSpec:(1) Qwen3-1.7B,(2) Qwen3-8B,和(3) Qwen3-14B,涵盖了广泛的模型尺寸和架构。我们使用张量并行(TP)在多个GPU上划分模型,选择能实现最高单GPU吞吐量的并行化方案,分别为1.7B/8B/14B模型配置TP1/2/4。我们使用NVIDIA DGX-H100-SXM5 GPU作为我们的测试平台。
数据集。我们使用以下数据集作为目标推理工作负载,涵盖数学、通用知识和编程领域:(1) AIME【Veeraboina, 2023, Aime problem set 1983-2024, Kaggle】:来自美国数学邀请赛的数学问题;(2) OlympiadBench【He et al., 2024, Olympiadbench: A challenging benchmark for promoting agi with olympiad-level bilingual multimodal scientific problems, arXiv】:奥林匹克水平的双语科学问题,横跨STEM领域,需要高级多步推理,我们使用其文本问答子集进行评估;以及(3) LiveCodeBench【Jain et al., 2024a, Livecodebench: Holistic and contamination free evaluation of large language models for code, arXiv】:从LeetCode、AtCoder和Codeforces收集的编程问题。对于每个工作负载,我们随机抽样2048个请求以使流水线饱和,用于后续评估。我们按照惯例将温度设置为0.65。
基线系统。我们将SparseSpec与以下服务框架结合不同的推测解码算法进行比较:(1) vLLM:我们默认启用vLLM-V1以获得最佳性能;(2) vLLM-NGram【Leviathan et al., 2023, Fast inference from transformers via speculative decoding, arXiv】:一个集成到vLLM中的免训练推测解码方法;(3) MagicDec【Chen et al., 2024, Magicdec: Breaking the latencythroughput tradeoff for long context generation with speculative decoding, arXiv】:由于原始开源版本硬编码为短输出长度且不易适配,我们在我们的框架内严格按照其论文复现了MagicDec;(4) TriForce【Sun et al., 2024, Triforce: Lossless acceleration of long sequence generation with hierarchical speculative decoding, arXiv】:与MagicDec类似,我们基于vLLM复现了TriForce;以及(5) vLLM-EAGLE3【Li et al., 2025a, Eagle-3: Scaling up inference acceleration of large language models via training-time test, arXiv】:一个由vLLM集成的最先进的基于草稿模型的推测解码算法。我们采用了腾讯开源的所有EAGLE3草稿模型【Contributors, 2025, AngelSlim, GitHub】。对于所有基线,我们都启用了分块预填充、CUDA图和连续批处理。根据我们的测试,我们为NGram(k=4)和EAGLE3(k=3)设置了最佳配置,并遵循【Li et al., 2025a】的建议,在批量推理中禁用了树状推测。我们将所有框架的最大批大小设置为256,这足以在我们的设置中饱和GPU内存。
5.2 端到端性能
与免训练推测解码方法的比较。SparseSpec在设计上是一个免训练的推测解码系统。因此,我们将其与包括vLLM-NGram、MagicDec和TriForce在内的免训练推测解码方法进行比较。具体来说,作为一个分层框架的TriForce包含三个推测层,其中一层依赖于一个独立的草稿模型;由于近期RLMs没有可用的小型草稿模型,我们用免训练的NGram替换了该层。对于中间层,我们使用与MagicDec相同的滑动窗口注意力作为草稿模型。我们为第一层和中间层设置k分别为1和6,这在我们的设置中表现最佳。如图10所示,SparseSpec在所有基线上持续表现优异,与最先进的框架vLLM相比,实现了高达2.13倍的加速。与现有的推测解码框架相比,SparseSpec实现了高达1.36倍的吞吐量提升。有趣的是,TriForce的性能一直落后于MagicDec,主要是因为额外的NGram层接受率较低(图12左),导致了过多的计算。此外,所有推测解码方法在模型更大、TP度更高时,加速比较低。这是因为虽然可用的GPU内存增加了,但每个词元的内存使用量增加得少得多。例如,Qwen3-8B和-14B共享相同的头维度和键/值头数量。因此,请求的最大批大小B增加了很多,接近饱和点$B̂$,在推测中产生了更多的计算开销。
与基于草稿模型的推测解码方法的比较。我们还将SparseSpec与最先进的基于草稿模型的推测解码方法EAGLE3进行了比较,该方法需要额外的训练。在此设置下,尽管无需额外训练,SparseSpec在所有数据集和模型上仍然取得了更好的性能。SparseSpec在完全避免了与草稿模型微调和部署相关的额外成本和工程努力的同时,提供了相似或更高的吞吐量。
执行时间分解。为了分解性能提升的来源,我们剖析了SparseSpec在Qwen3-8B模型和AIME数据集上的CPU操作、注意力和GEMM的执行延迟,并与vLLM在表2中进行了比较。结果显示,SparseSpec成功将注意力的执行时间减少了3.29倍,而GEMM的时间仅略有增加(1.7ms),这与我们在§3.2中的估算一致。此外,SparseSpec保持了极低的CPU开销(< 1ms),从而实现了高GPU利用率。
表2. Qwen3-8B执行时间分解 (ms)
5.3 推测接受率
我们通过测量在不同模型和数据集上草稿生成8个词元(即k=8)时的平均接受长度来评估PillarAttn的接受率。所有方法我们都不计算奖励词元。如图12(左)所示,PillarAttn在8个词元中平均接受了6.16个,超过了所有其他草稿方法。相比之下,NGram和EAGLE3都只能生成少于2个被接受的词元。我们推测这是因为这些推理任务超出了EAGLE3训练数据集的分布范围,表明其泛化能力较差。
5.4 敏感性测试
我们对三个数据集进行了敏感性测试,以量化推测步数k和稀疏度s对PillarAttn性能的影响,如图12(右)所示。尽管在我们的评估中固定了s=0.05和k=8,但SparseSpec在运行时对这些超参数没有严格的限制,实际上允许任意组合和异构的请求配置。
推测步数k。固定稀疏度s=0.05,我们将草稿长度从4变化到20。我们设置草稿长度k=8以平衡接受率和验证开销。
稀疏度s。我们改变PillarAttn中的稀疏度,并测量生成8个词元时的接受率。我们将稀疏度设置为0.05,因为随着所选词元数量的进一步增加,性能趋于饱和。
5.5 消融研究
为了独立评估每个设计组件的贡献,我们从一个稀疏自推测解码的朴素实现开始,然后逐步启用统一批处理调度器、动态KV-Cache管理器和延迟验证。在AIME数据集上使用Qwen3-8B的吞吐量如图13所示。我们的实验表明,这三个设计分别将性能提升了1.23倍、1.61倍和1.12倍,最终实现了2.22倍的总吞吐量增益。我们还为每个组件提供了详细的性能剖析:
批大小和计算利用率。图14展示了有无统一批处理的系统在GEMM输入批大小和硬件利用率(以TFlops计)方面的比较分析。启用统一批处理后,输入批大小在各次迭代中保持稳定,而传统调度则导致显著的方差,并在草稿阶段造成硬件利用率不足。
融合注意力。我们评估了融合的稀疏和全量注意力核的性能,并与两个基线进行了比较:顺序启动两个核(Sequential)和一个朴素地将两种注意力类型共同计算的核(Naive Batch)。结果如图15所示。我们的融合核相比顺序运行实现了1.3倍的加速,相比朴素批处理实现了1.8倍的加速,这得益于以下两个因素:(1) 为草稿和验证注意力分派了最佳的核配置;(2) 在单个核内拥有更多的事务字节,从而重叠了流水线延迟,实现了更高的带宽利用率。
内存容量利用率。我们将SparseSpec的内存利用率和重计算情况与§4.4中描述的基线(包括oracle和抢占)进行了比较。如图5所示,SparseSpec几乎利用了所有可用的GPU内存,且没有产生重计算。我们进一步通过比较启用换出操作的执行时间与将这些操作替换为空核的基线来量化换出的开销。结果表明,换出操作平均仅使周期时间延长了0.5%,这在实践中可以忽略不计。
A7 补充细节
6. 讨论
SparseSpec在混合专家(MoE)模型上的应用。SparseSpec可以无缝应用于MoE模型,因为它只涉及注意力模块,而没有修改FFN。此外,由于在推理过程中只激活一部分专家,每个专家的输入词元大小显著减少,这增加了稀疏化FFN计算的饱和点$B̂$。因此,由于计算开销的减少,自推测具有更高的潜力。
SparseSpec与多词元预测(MTP)的结合。SparseSpec可以与其他轻量级草稿方法(包括EAGLE3和MTP)结合,形成一种如TriForce【Sun et al., 2024, Triforce: Lossless acceleration of long sequence generation with hierarchical speculative decoding, arXiv】中提出的分层推测方法。例如,使用MTP一次性生成k1个词元,这些词元由PillarAttn验证后存入缓冲区;当累积了k2个被接受的词元后,再使用底层的全量注意力进行验证。这种分层方法除了减少KV-Cache加载外,还可以减少FFN的计算量,为进一步加速带来了巨大机遇。
局限性。所提出的方法专注于提高长文本生成工作负载的内存效率。对于短上下文任务,最大并发批处理大小足以饱和GPU计算,使得整体工作负载变为计算密集型【Zhu et al., 2024, Nanoflow: Towards optimal large language model serving throughput, arXiv】。
7. 相关工作
长上下文稀疏注意力。先前的工作广泛探索了通过利用注意力的内在稀疏性来加速注意力计算。静态方法【Xiao et al., 2024b, Efficient streaming language models with attention sinks, arXiv】、【Zhang et al., 2023, H2o: Heavy-hitter oracle for efficient generative inference of large language models, arXiv】、【Li et al., 2024a, Snapkv: Llm knows what you are looking for before generation, arXiv】采用固定的KV-cache剪枝策略,因此无法适应不断变化的推理上下文。而查询感知方法【Tang et al., 2024, Quest: Query-aware sparsity for efficient long-context llm inference, arXiv】、【Zhu et al., 2025, Tactic: Adaptive sparse attention with clustering and distribution fitting for long-context llms, arXiv】、【Xiao et al., 2024a, Infllm: Training-free long-context extrapolation for llms with an efficient context memory, arXiv】、【Wu et al., 2025, Tokenselect: Efficient long-context inference and length extrapolation for llms via dynamic token-level kv cache selection, arXiv】、【Lin et al., 2025, Twilight: Adaptive attention sparsity with hierarchical top-p pruning, arXiv】在生成过程中动态识别重要词元,但它们在估计词元重要性时会产生不可忽略的开销。相比之下,PillarAttn自然地复用推测解码验证阶段的注意力分数,为下一轮草稿阶段动态选择关键词元,开销极小。
推测解码。推测解码为长输出任务提供了无损加速。基于训练的方法——包括EAGLE系列【Li et al., 2025b, Eagle: Speculative sampling requires rethinking feature uncertainty, arXiv】、【Li et al., 2024b, Eagle-2: Faster inference of language models with dynamic draft trees, arXiv】、【Li et al., 2025a, Eagle-3: Scaling up inference acceleration of large language models via training-time test, arXiv】、Hydra【Ankner et al., 2024, Hydra: Sequentiallydependent draft heads for medusa decoding, arXiv】、多词元预测【DeepSeek-AI, 2025b, Deepseek-v3 technical report, arXiv】和EESD【Liu et al., 2024a, Speculative decoding via early-exiting for faster llm inference with thompson sampling control mechanism, arXiv】——通过不同的草稿模型设计或学习方法有效提高了接受率,但代价是部署复杂性更高。免训练方法如N-gram【Leviathan et al., 2023, Fast inference from transformers via speculative decoding, arXiv】、Lookahead Decoding【Fu et al., 2024, Break the sequential dependency of llm inference using lookahead decoding, arXiv】和SAM Decoding【Hu et al., 2024, Sam decoding: Speculative decoding via suffix automaton, arXiv】从当前上下文中预测未来词元,但在推理任务上准确性下降。MagicDec【Chen et al., 2024, Magicdec: Breaking the latencythroughput tradeoff for long context generation with speculative decoding, arXiv】和TriForce【Sun et al., 2024, Triforce: Lossless acceleration of long sequence generation with hierarchical speculative decoding, arXiv】采用静态稀疏注意力作为长输入场景的草稿,但在处理长而动态的推理输出时遇到困难。RhymeRL【He et al., 2025, History rhymes: Accelerating llm reinforcement learning with rhymerl, arXiv】和SuffixDecoding【Oliaro et al., 2024, Suffixdecoding: A model-free approach to speeding up large language model inference, arXiv】中的后缀树方法在RL轨迹生成中有效,但不适应于每个文档只出现一次的服务场景。
相比之下,SparseSpec将注意力识别为推理过程中的瓶颈,并无缝地将原始模型与一个准确的动态稀疏注意力模块PillarAttn集成为草稿模型,实现了显著的端到端加速,且无需额外的训练或存储。
A5 结论
由于输出序列长,推理模型的推断过程严重受到内存限制。我们提出了SparseSpec,一个针对推理模型的无损且免训练的服务框架,它采用稀疏自推测解码。SparseSpec在验证阶段的全量注意力计算过程中识别关键词元,并利用它们来指导草稿阶段的稀疏注意力计算。通过系统级优化——包括一个统一且资源感知的批处理调度器、延迟验证和动态KV-Cache管理器——SparseSpec相比现有的服务框架和推测解码基线,实现了1.36倍的吞吐量提升。
💬 评论讨论
欢迎在这里分享您的想法和见解!