SpecLA：面向线性注意力模型的高效推测解码

# SpecLA: 面向线性注意力模型的高效推测解码
来源: https://arxiv.org/html/2607.16673 ,Xuying Han 南京大学计算机软件新技术国家重点实验室 南京 中国, Zhaohua Yang 南京大学计算机软件新技术国家重点实验室 南京 中国, Fuliang Liu 南京大学计算机软件新技术国家重点实验室 南京 中国, Xue Li 阿里巴巴集团 杭州 中国, Rong Gu 南京大学计算机软件新技术国家重点实验室 南京 中国, Sheng Zhong 南京大学计算机软件新技术国家重点实验室 南京 中国, Chen Tian 南京大学计算机软件新技术国家重点实验室 南京 中国

###### 摘要。线性注意力模型用循环状态取代了不断增长的 KV 缓存，但自回归解码仍然每次处理一个 token，读取、更新和写入这些状态。推测解码通过在单次目标推理中验证多个草稿 token，可以降低这一成本。然而，现有的推测系统是为 Transformer KV 缓存设计的。对于有状态的线性注意力目标而言，验证必须遵循跨链和分支的循环依赖关系，接受过程必须仅更新已接受的状态轨迹，并且草稿模型必须避免提交会浪费有状态验证工作的候选 token。本文提出了 SpecLA，一个面向有状态线性注意力模型的推测解码运行时系统。SpecLA 使用拓扑感知的内核来验证链和树，存储验证过程中产生的紧凑因子以恢复已接受状态，并采用置信度剪枝加目标对齐的 EAGLE 风格草稿器，为验证器提供有用的候选 token。在配备公共 GDN-1.3B 目标的 NVIDIA H100 上，SpecLA 相比自回归解码实现了高达 1.70× 的端到端加速。

## 1.引言

自回归解码仍然是大型语言模型服务中的核心瓶颈。每生成一个 token 都需要调用一次目标模型，而对于受内存限制的解码路径，延迟不仅由算术运算主导，还受到缓存访问、状态移动、同步和内核启动开销的影响 (Kwon et al., 2023 (https://arxiv.org/html/2607.16673#bib.bib6); Qwen Team, 2025a (https://arxiv.org/html/2607.16673#bib.bib23))。推测解码通过将生成过程分为快速的草稿阶段和目标验证阶段来解决这个瓶颈：一个草稿模型提出几个未来的 token，然后目标模型一起验证它们，同时通过后验接受规则保持目标分布 (Leviathan et al., 2023 (https://arxiv.org/html/2607.16673#bib.bib1); Chen et al., 2023 (https://arxiv.org/html/2607.16673#bib.bib2))。这已成为 Transformer 服务中一种实用的加速策略，后续系统使用了特征级草稿、辅助草稿头和树结构验证 (Li et al., 2024b (https://arxiv.org/html/2607.16673#bib.bib7), a (https://arxiv.org/html/2607.16673#bib.bib8), 2025 (https://arxiv.org/html/2607.16673#bib.bib9); Cai et al., 2024 (https://arxiv.org/html/2607.16673#bib.bib10); Ankner et al., 2024 (https://arxiv.org/html/2607.16673#bib.bib11); Miao et al., 2024 (https://arxiv.org/html/2607.16673#bib.bib15); Wu et al., 2025 (https://arxiv.org/html/2607.16673#bib.bib16))，因为提交的草稿可以被表示为对现有 KV 缓存历史的 token 索引扩展 (Kwon et al., 2023 (https://arxiv.org/html/2607.16673#bib.bib6))。

与此同时，线性注意力和有状态序列模型正在改变服务的基础设施。这些模型不是存储不断增长的 KV 缓存，而是将前缀压缩为循环状态，并用每个新 token 更新该状态 (Gu and Dao, 2024 (https://arxiv.org/html/2607.16673#bib.bib12); Yang et al., 2024a (https://arxiv.org/html/2607.16673#bib.bib18); Dao and Gu, 2024 (https://arxiv.org/html/2607.16673#bib.bib22); Yang et al., 2024b (https://arxiv.org/html/2607.16673#bib.bib20))。Gated DeltaNet (GDN) 进一步将门控与 Delta 规则更新相结合，并出现在最近的混合和开源模型栈中 (Yang et al., 2025 (https://arxiv.org/html/2607.16673#bib.bib25); Qwen Team, 2025b (https://arxiv.org/html/2607.16673#bib.bib26); NVIDIA Research, 2026 (https://arxiv.org/html/2607.16673#bib.bib27))。这些层的高效内核已经需要专门的预填充和解码实现 (Yang and Zhang, 2024 (https://arxiv.org/html/2607.16673#bib.bib21); Qin et al., 2024 (https://arxiv.org/html/2607.16673#bib.bib19); Beck et al., 2025 (https://arxiv.org/html/2607.16673#bib.bib24); Yang et al., 2024b (https://arxiv.org/html/2607.16673#bib.bib20))，而设计的有状态特性并没有让解码变得免费：每个 token 仍然需要读取、更新和写入一个密集的循环状态 (Gupta et al., 2026 (https://arxiv.org/html/2607.16673#bib.bib29))。因此，线性注意力解码仍然是推测的一个自然目标，因为一次性验证多个草稿 token 可以将循环状态移动和固定的启动开销分摊到多个接受的 token 上。

两种直接的适配方法无法将推测解码迁移到有状态的线性注意力目标上。
*1) 解码内核重放* 保留了循环性，但也保留了推测旨在分摊的 per-token 循环状态往返开销；对于分支提议，不同分支还需要不同的中间循环状态 (Yang and Zhang, 2024 (https://arxiv.org/html/2607.16673#bib.bib21); Gupta et al., 2026 (https://arxiv.org/html/2607.16673#bib.bib29))。
*2) 预填充内核复用* 暴露了 token 并行性，但推测窗口太短，无法分摊长预填充的设置成本，并且分支提议不遵循预填充内核假定的规则因果依赖关系 (Yang and Zhang, 2024 (https://arxiv.org/html/2607.16673#bib.bib21); Beck et al., 2025 (https://arxiv.org/html/2607.16673#bib.bib24); Yang et al., 2024b (https://arxiv.org/html/2607.16673#bib.bib20), 2025 (https://arxiv.org/html/2607.16673#bib.bib25))。

在本文中，我们旨在使推测解码对有状态的线性注意力目标有效，这需要解决三个耦合的挑战。
*1) SRAM 驻留验证：* 目标内核必须为链状候选 token 保持循环状态分片驻留，为分支候选 token 保留仅祖先依赖关系，并避免在短推测窗口上使用预填充风格设置成本。
*2) 密集状态接受：* 后验选择必须仅推进已接受的循环状态更新，而无需完整状态快照、token 重放或每次迭代边界的额外状态访问步骤。
*3) 目标对齐的草稿：* 提议路径必须通过剪枝低置信度路径来提交有用的候选工作，将草稿特征与循环目标动态对齐，并将保留的候选 token 打包到验证器使用的拓扑表示中。

我们提出 SpecLA，一个推测解码运行时系统，它保留了外部的草稿-验证-接受循环，但用三种目标感知机制取代了 KV 缓存后缀操作：拓扑感知验证，使用提交的拓扑作为目标内核调度；接受因子状态管理，记录紧凑的验证因子而不是密集的状态端点；以及目标对齐的草稿集成，将提议生成视为验证器的成本感知调度。

**拓扑感知验证（第 4 节 (https://arxiv.org/html/2607.16673#S4)）。** SpecLA 将提交的拓扑视为目标内核的调度信号，而不是强制每个提议通过解码重放或预填充复用。由于在目标执行之前已知草稿 token，链状候选 token 可以逐层验证，同时将 V 分块的循环状态切片保留在片上。相反，分支提议需要仅祖先状态流：树掩码分解的 GDN 路径给出了完全并行的端点，而链分解的混合路径在尊重依赖关系的链内保持低开销的串行内核，并跨就绪链并行化。

**接受因子状态管理（第 5 节 (https://arxiv.org/html/2607.16673#S5)）。** 由于无法通过截断循环状态来移除被拒绝的候选 token，运行时将推测进度与持久状态分离。关键观察是，验证已经产生了重建已接受状态更新所需的紧凑层特定因子，因此 SpecLA 缓存因子而不是完整状态或仅 token 日志。在后验选择之后，它仅收集已接受的路径并延迟状态更新，使得待定因子缓冲区在下一个验证内核内部应用，而不是通过独立的状态更新步骤。

**目标对齐的草稿集成（第 6 节 (https://arxiv.org/html/2607.16673#S6)）。** 草稿端被调整为适应循环目标，而不是被视为 Transformer 插件。由于目标端成本取决于提交的拓扑，提议生成也充当调度问题：置信度引导的剪枝在验证前移除低概率路径，而一个在循环目标特征上训练的 EAGLE 风格草稿器使得这些剪枝分数反映目标动态 (Li et al., 2024b (https://arxiv.org/html/2607.16673#bib.bib7), a (https://arxiv.org/html/2607.16673#bib.bib8))。保留的候选 token 随后作为父指针拓扑打包，供目标内核使用。

**评估（第 8 节 (https://arxiv.org/html/2607.16673#S8)）。** 我们在 PyTorch (Paszke et al., 2019 (https://arxiv.org/html/2607.16673#bib.bib4)) 和 Triton (Tillet et al., 2019 (https://arxiv.org/html/2607.16673#bib.bib5)) 中实现了 SpecLA，并在配备公共 GDN-1.3B 目标 (m-a-p, 2026 (https://arxiv.org/html/2607.16673#bib.bib28)) 的 NVIDIA H100 GPU 上进行了评估。端到端地，SpecLA 在混合套件、GSM8K 和 HumanEval 上相比自回归解码分别实现了 1.42×、1.70× 和 1.06× 的加速比。微基准测试表明，链分解混合验证相比根到叶重放提升了 1.80–7.11×，因子缓冲区复用将已接受状态恢复延迟降低了 2.74–4.28×，延迟状态更新将提交边界延迟降低了 1.15–1.44×。

## 2.背景与动机

在本节中，我们首先回顾两个技术领域：推测解码和线性注意力。然后，我们通过一个简单的适配，说明为什么推测解码不能直接迁移到有状态的线性注意力模型。最后，我们确定了从内核设计到状态管理再到运行时栈集成，在将推测解码有效地与线性注意力结合使用时面临的三个挑战。

### 2.1.推测解码

推测解码是一种用于自回归生成的推理时加速技术：一个快速的近似模型草拟几个未来 token，然后目标模型一起验证它们，使得一次昂贵的目标调用可以产生多个已接受的 token，同时通过类似拒绝采样的接受规则保持目标模型的输出分布 (Leviathan et al., 2023 (https://arxiv.org/html/2607.16673#bib.bib1); Chen et al., 2023 (https://arxiv.org/html/2607.16673#bib.bib2))。这对于受内存限制的解码尤其有用，因为目标端的缓存访问、状态移动、同步和启动开销可以为一组候选 token 支付一次，而不是为每个生成的 token 支付一次。图 1 (https://arxiv.org/html/2607.16673#S2.F1) 展示了工作流程，我们在本文中将其分为三个阶段。我们以 Transformer 架构下的推测解码作为运行示例。

**草稿构建。** 运行时从一个已接受的前缀开始，并基于该前缀构建推测工作。在最初的公式中，已接受的历史由 token 索引的 KV 缓存条目表示，一个更快但能力较弱的草稿模型提出一个短 token 延续，近似目标模型将生成的内容 (Leviathan et al., 2023 (https://arxiv.org/html/2607.16673#bib.bib1); Chen et al., 2023 (https://arxiv.org/html/2607.16673#bib.bib2))。后来的系统使用特征级草稿 (Li et al., 2024b (https://arxiv.org/html/2607.16673#bib.bib7), a (https://arxiv.org/html/2607.16673#bib.bib8), 2025 (https://arxiv.org/html/2607.16673#bib.bib9))、辅助草稿头 (Cai et al., 2024 (https://arxiv.org/html/2607.16673#bib.bib10); Ankner et al., 2024 (https://arxiv.org/html/2607.16673#bib.bib11)) 或树结构提议 (Miao et al., 2024 (https://arxiv.org/html/2607.16673#bib.bib15); Wu et al., 2025 (https://arxiv.org/html/2607.16673#bib.bib16)) 来实例化相同的步骤。这个阶段适合 Transformer 运行时，因为草稿可以表示为额外的 token 位置，或 token 位置的路径，扩展了一个已经由 token 索引的历史。

**目标验证。** 目标模型在一次批处理前向传递中对提交的候选 token 进行评分，而不是逐个解码，产生接受规则使用的目标端概率。这将会调用一次昂贵的目标模型的开销分摊到多个可能的输出 token 上：对短延续的并行评分可能与单次目标解码调用相当，同时可能确认多个未来 token (Chen et al., 2023 (https://arxiv.org/html/2607.16673#bib.bib2))。基于树的验证通过在同一目标 pass 中检查多个候选延续来扩展相同的想法 (Miao et al., 2024 (https://arxiv.org/html/2607.16673#bib.bib15); Wu et al., 2025 (https://arxiv.org/html/2607.16673#bib.bib16))。这个阶段很适合 Transformer，因为已提交的 KV 缓存可以作为共享前缀重用，而草稿 token 的 KV 条目形成一个紧凑的推测后缀，在一次批处理验证 pass 中生成和访问。对于受内存限制的 Transformer 解码，验证将重复的 per-token KV 缓存访问转变为规则的批处理访问模式，使得缓存流量更容易分摊到接受的 token 上。

**接受与缓存提交。** 后验选择应用修改后的拒绝采样规则，以保留提议的已接受前缀，同时保持目标模型的输出分布 (Leviathan et al., 2023 (https://arxiv.org/html/2607.16673#bib.bib1); Chen et al., 2023 (https://arxiv.org/html/2607.16673#bib.bib2))。当一个候选 token 被拒绝时，生成回退到修正后的目标端分布用于下一个输出；运行时提交已接受的输出 token 并从 KV 缓存中移除被拒绝的后缀。这很廉价，因为推测窗口存储为已提交前缀之后的 token 位置：拒绝是一个后缀截断操作，而不是模型状态的重建 (Kwon et al., 2023 (https://arxiv.org/html/2607.16673#bib.bib6))。因此，这个缓存提交步骤对于 Transformer 尤其方便，因为其生成进度已经组织为 token 索引的缓存历史。

参见图注
图 1. Transformer 导向的推测解码中的后缀截断。接受的草稿 token 保留在 KV 缓存中，而被拒绝的 token 作为缓存后缀被丢弃。一个存储在 KV 缓存中的前缀被扩展了四个提出的 token。目标验证接受前两个提出的 token，拒绝剩余的两个。已接受的前缀被保留，被拒绝的后缀从 KV 缓存中被截断。

### 2.2.线性注意力与有状态序列模型

长上下文服务场景，包括长文档理解、流式或多轮生成，以及具有大前缀的检索或代码密集型工作负载，给全注意力模型带来了持续压力 (Kwon et al., 2023 (https://arxiv.org/html/2607.16673#bib.bib6); Qwen Team, 2025a (https://arxiv.org/html/2607.16673#bib.bib23))。全注意力形成了逐对 token 交互