超越 Flash：利用注意力稀疏性实现高效长上下文解码

# 快于Flash：利用注意力稀疏性实现高效长上下文解码
来源：https://arxiv.org/html/2609.00097
1\]复旦大学，上海，中国 2\]上海创新研究院，上海，中国 3\]哈尔滨工业大学，哈尔滨，中国\\contribution\[†\]通讯作者\\correspondence253108120105@sii\.edu\.cn,zhangmin2021@hit\.edu\.cn,ziwei\.he@sii\.edu\.cn,xpqiu@fudan\.edu\.cn

宁志远 李瑞晓 刘小然 宋月蓉 张敏 何子威 丘锡鹏
单位：\单位：\单位：\[

###### 摘要

长上下文大语言模型（LLMs）的发展受限于解码过程中注意力机制的内存带宽瓶颈与二次复杂度。为克服基于元数据指标的内存开销与自适应选择策略计算效率低下之间的固有权衡，我们提出**更快Flash解码（FFD）**，这是一种新颖的硬件-算法协同设计框架，旨在突破长上下文解码的内存墙。FFD将选择器与计算器整合到一个完全融合的内核中，通过低比特量化实现内容感知扫描，取代了外部元数据索引。此外，我们引入了**top-δ策略**，该策略可动态过滤块以实现分布自适应的稀疏性，且无需全局同步。FFD提供了一种无需训练、即插即用的解决方案，并允许扫描结果复用于计算，最终实现了高达11.6倍的内核级加速，并可扩展至256K上下文长度，端到端吞吐量提升达2.37倍。在RULER和LongBench上的实证验证表明，FFD在保持模型精度的同时实现了高比率稀疏性，代码已开源：[https://github\.com/qluoluo/faster\-flash\-decoding (https://github.com/qluoluo/faster-flash-decoding)]。

## 1 引言

请参见标题图1：选择策略比较。(a) Top-k（固定预算）快速但无法适应分布形状。(b) Top-p提供动态基数，但由于全局同步和排序带来显著开销，破坏了流式管道。(c) Top-δ（我们提出）使用从局部最大值获得的相对阈值δ，实现了top-p的适应性与top-k的硬件效率。
处理长上下文的能力已成为现代大语言模型（LLMs）的一个决定性特征\[1 (https://arxiv.org/html/2609.00097#bib.bib1),2 (https://arxiv.org/html/2609.00097#bib.bib2),3 (https://arxiv.org/html/2609.00097#bib.bib3),4 (https://arxiv.org/html/2609.00097#bib.bib4)\]。然而，长上下文能力在解码阶段带来了高昂的代价。随着序列长度增加，标准注意力机制表现出二次计算复杂度和键值（KV）缓存的线性内存增长\[5 (https://arxiv.org/html/2609.00097#bib.bib5)\]，在解码阶段形成了严重的内存墙瓶颈\[6 (https://arxiv.org/html/2609.00097#bib.bib6),7 (https://arxiv.org/html/2609.00097#bib.bib7)\]。与预填充阶段（计算主导且查询被并行处理）不同，自回归解码阶段是内存带宽受限的：每个token生成都需要从高带宽内存（HBM）重新加载整个KV缓存，使得IO成为主要瓶颈，而非浮点运算。因此，稀疏注意力机制，特别是无需训练、即插即用的方法，已成为一个关键研究方向，旨在减少计算开销而无需重新训练模型\[8 (https://arxiv.org/html/2609.00097#bib.bib8),9 (https://arxiv.org/html/2609.00097#bib.bib9),10 (https://arxiv.org/html/2609.00097#bib.bib10),11 (https://arxiv.org/html/2609.00097#bib.bib11),12 (https://arxiv.org/html/2609.00097#bib.bib12)\]。

尽管稀疏注意力技术多种多样，大多数方法遵循**解耦的选择器-计算器**范式。在这种标准工作流中，选择器首先基于指标过滤索引，然后计算器检索数据进行计算。这种分离阻止了计算器复用选择器的中间计算，导致效率瓶颈。此外，该范式内的两个核心组件——重要性指标和选择策略——各自面临困境。一方面，为确定保留哪些token，现有方法通常依赖额外元数据，如部分维度、聚类质心或几何边界\[9 (https://arxiv.org/html/2609.00097#bib.bib9),13 (https://arxiv.org/html/2609.00097#bib.bib13),8 (https://arxiv.org/html/2609.00097#bib.bib8)\]。这些方法导致了**指标困境**：它们要么为存储这些元数据结构产生额外的内存开销，要么遭受信息失真。另一方面，选择最佳token子集仍然具有挑战性。如图1 (https://arxiv.org/html/2609.00097#S1.F1)所示，top-k选择计算效率高，但无法适应注意力分布的不同熵\[14 (https://arxiv.org/html/2609.00097#bib.bib14),15 (https://arxiv.org/html/2609.00097#bib.bib15)\]。相比之下，top-p选择在理论上具有优势，因为它能适应分布，但需要全局softmax操作来计算累积概率，从而限制了其效率\[10 (https://arxiv.org/html/2609.00097#bib.bib10)\]。

本文提出**FFD（更快Flash解码）**，这是一种新颖的稀疏注意力框架，通过硬件-算法协同设计应对这些挑战。通过整合选择器与计算器，FFD使计算阶段能够直接复用选择阶段的扫描结果。我们从基于元数据的索引转向内容感知扫描。我们不是存储额外元数据，而是将K缓存拆分为低比特量化部分和半精度残差部分，以前者作为KV缓存的压缩表示。接着，我们引入了基于数值精度分析的**top-δ策略**。top-δ不是强制固定预算（k）或计算全局概率（p），而是根据块对注意力分数和的估计贡献进行动态过滤，并且得益于LLMs中注意力的特性\[16 (https://arxiv.org/html/2609.00097#bib.bib16)\]，可以实现近似并行化。我们将FFD实现为一个完全融合的内核，消除了规约气泡，实现了11.6倍的内核级加速。我们的贡献总结如下：

- •**通过低比特量化实现内容感知扫描**。我们用内容感知扫描取代基于元数据的索引，解决了“指标困境”。通过将K缓存划分为低比特量化表示和8位残差，我们在保持高信息保真度的同时消除了额外元数据的内存开销。
- •**Top-δ选择策略**。我们引入了top-δ，这是一种新颖的选择机制，不同于top-k（固定预算）或top-p（高开销），top-δ根据块对注意力总和的贡献动态过滤注意力块，提供了分布自适应的稀疏性和并行执行效率。
- •**更快Flash解码（FFD）**。我们提出FFD，这是一种为带宽关键的解码阶段量身定制的硬件-算法协同设计框架。我们实现了一个融合的选择器-计算器内核，其中计算阶段复用选择器的量化扫描，优化了内存访问模式。我们的实现实现了高达11.6倍的内核级加速，并可扩展至256K上下文，端到端吞吐量最高提升2.37倍，同时保持即插即用、无需训练的特性。
- •**广泛的实证验证**。我们证明FFD在RULER\[17 (https://arxiv.org/html/2609.00097#bib.bib17)\]和LongBench\[18 (https://arxiv.org/html/2609.00097#bib.bib18)\]等各种短上下文和长上下文基准测试中保持了LLM的性能，证明了在不重新训练或牺牲精度的情况下可以实现高比率稀疏性。

请参见标题图2：我们的FFD实现概述。输入包括16位查询、值缓存，以及拆分为2位量化缩略图和8位残差的键缓存。在K量化、伪最大估计、top-δ选择、基于top-δ策略的过滤和细粒度优化之后，获得注意力输出。
#### 利益冲突声明

我们声明与本工作没有财务利益冲突。

## 2 相关工作

#### 硬件感知精确注意力。

标准注意力机制在预填充期间带来二次成本，在解码期间每一步带来线性成本，为长上下文应用造成了显著瓶颈。FlashAttention\[6 (https://arxiv.org/html/2609.00097#bib.bib6),19 (https://arxiv.org/html/2609.00097#bib.bib19)\]通过引入IO感知优化，采用分块和内核融合来最小化高带宽内存（HBM）访问，彻底改变了这一领域。FlashDecoding和FlashDecoding++\[20 (https://arxiv.org/html/2609.00097#bib.bib20)\]进一步将这些原则适配到生成阶段，通过在序列维度上并行化注意力计算。虽然FFD继承了这些工作中内核融合和硬件-算法协同设计的精神，但现有的基于Flash的方法本质上仍然是精确注意力机制。随着上下文长度增加，它们仍然必须加载和扫描整个KV缓存，导致延迟不可避免地线性增长。FFD通过将稀疏性直接集成到融合内核中解决了这个问题，打破了线性扫描屏障，同时保持了硬件效率。

#### 动态稀疏注意力

为规避精确注意力的成本，已提出各种稀疏注意力机制。这些方法在重要性估计（指标困境）和子集过滤（选择困境）方面普遍面临特定挑战。关于指标困境，诸如H2O\[14 (https://arxiv.org/html/2609.00097#bib.bib14)\]、Scissorhands\[21 (https://arxiv.org/html/2609.00097#bib.bib21)\]和SnapKV\[22 (https://arxiv.org/html/2609.00097#bib.bib22)\]等方法依赖历史积累来识别token驱逐的“重要条目”。然而，纯粹基于历史的指标有丢弃那些仅在未来上下文中才变得相关的token的风险。为解决此问题，诸如SparQ\[9 (https://arxiv.org/html/2609.00097#bib.bib9)\]和Quest\[8 (https://arxiv.org/html/2609.00097#bib.bib8)\]等动态检索方法使用辅助元数据（例如，均值向量或最小-最大边界）来估计块的重要性。同样，InfLLM\[13 (https://arxiv.org/html/2609.00097#bib.bib13)\]利用代表性块进行邻域搜索。虽然这些基于元数据的方法降低了扫描成本，但它们遭受信息失真并需要额外的内存管理。关于选择困境，大多数方法（例如，Quest, SparQ）强制使用静态top-k预算，这在变化的注意力熵分布中缺乏灵活性。Twilight\[10 (https://arxiv.org/html/2609.00097#bib.bib10)\]试图适应各种注意力分布，但引入了全局softmax操作，从而限制了其效率。与我们工作同时期，SALE\[23 (https://arxiv.org/html/2609.00097#bib.bib23)\]独立探索了基于锚点的低比特估计选择用于稀疏注意力。然而，SALE针对预填充阶段，并采用4位量化，与我们针对带宽关键解码设计的2位设计相比，扫描带宽翻倍。相比之下，FFD通过融合的选择器-计算器内核执行内容感知扫描，绕过了对辅助结构的需求，并实现了动态的top-δ选择策略，可实时适应熵的变化。

#### KV缓存量化

KV缓存量化主要用于缓解内存容量约束。诸如KIVI\[24 (https://arxiv.org/html/2609.00097#bib.bib24)\]、Atom\[25 (https://arxiv.org/html/2609.00097#bib.bib25)\]和KVQuant\[26 (https://arxiv.org/html/2609.00097#bib.bib26)\]等方法将键和值压缩为低精度格式，以将更长的上下文放入GPU内存。通常，这些框架在执行注意力计算之前将数据反量化回更高精度（FP16/BF16）以保持精度。FFD从根本上重新利用了量化。我们不是将低比特表示仅视为存储压缩技术，而是利用量化的INT2/4数据作为重要性指标本身的高速代理。这使得我们的选择器能够以高保真度和最小延迟扫描完整上下文，避免了标准量化框架中典型的解压缩开销，同时与基于元数据的启发式方法相比保持了更高的精度。

## 3 方法

本节介绍**FFD（更快Flash解码）**，这是一种硬件感知的稀疏注意力框架，旨在突破长上下文解码的内存墙。FFD建立在硬件-算法协同设计之上，通过两项核心创新系统地解决了稀疏注意力中的权衡问题：通过低比特量化实现内容感知扫描和top-δ选择策略。FFD利用K缓存的低比特量化缩略图进行高保真度重要性估计，而无需额外的内存开销。为适应不同的注意力分布，引入了top-δ准则，为僵化的top-k或计算昂贵的top-p选择提供了一种理论依据充分、可并行的替代方案。我们为选择器和计算器实现了一个融合算子以优化执行；具体来说，计算器复用选择器的估计结果继续计算，从而在高效支持自适应选择的同时加速过程。整体流程总结如图2 (https://arxiv.org/html/2609.00097#S1.F2)和算法1 (https://arxiv.org/html/2609.00097#alg1)所示。

本节其余部分组织如下。首先，我们详细阐述高效的基于量化的扫描机制，该机制实现了高保真度的重要性估计。接下来，我们介绍top-δ选择的公式化与并行实现。最后，我们描述支撑FFD高速执行的系统级优化和融合内核设计。

### 3.1 通过低比特量化实现内容感知扫描

受KV缓存量化进展的启发\[24 (https://arxiv.org/html/2609.00097#bib.bib24),26 (https://arxiv.org/html/2609.00097#bib.bib26),27 (https://arxiv.org/html/2609.00097#bib.bib27)\]，我们将检索保真度与存储开销解耦，将K缓存分解为k→\{kquant,kres\}。其中，kquant表示用于高吞吐量相似性扫描的2位量化缩略图，而kres表示8位残差。

为在有限的位宽内最大化信息熵，我们采用对称零中值（mid-rise）量化方案，确保即使对于小幅度特征也能保留方向信息（详见附录A (https://arxiv.org/html/2609.00097#A1)）。

这种架构确保扫描阶段以最低的算术强度运行，同时识别出的候选被重建至接近FP16精度进行精确计算。扫描误差的形式化集中界在附录B (https://arxiv.org/html/2609.00097#A2)中提供。

算法1：使用混合量化的FFD解码
1：输入：查询q<sub>i</sub>\