CoSA: 通过代理-内核协同设计的稀疏注意力加速长上下文推理

arXiv cs.CL 论文

摘要

CoSA 提出了一种无需训练的稀疏注意力方法,该方法协同设计了内核感知代理和有序跳过内核,以加速长上下文推理,实现了高达 4.93 倍的注意力加速和 2.53 倍的端到端 TTFT 降低,且性能下降可忽略不计。

arXiv:2607.25291v1 公告类型: 新 摘要: 自注意力的二次计算成本使得长上下文推理变得极其昂贵,而基于代理的块稀疏注意力已成为一种实用的补救措施。现有方法通常依赖代理预测二元稀疏掩码,并由内核消耗该掩码执行稀疏注意力计算。这种方法在适度的预算下是有效的。然而,随着预算收紧,估计的代理不可避免地会丢弃一些显著块,而内核只能机械地应用稀疏掩码,导致模型精度显著下降。我们提出 CoSA,一种两阶段、无需训练的稀疏注意力方法,采用代理-内核协同设计,将内核感知代理 (KAP) 与有序跳过内核 (OSK) 相结合。在第一阶段,KAP 在适度预算下选择块,并生成一个有序掩码,该掩码规定了内核内循环中 KV 页面的访问顺序。在第二阶段,OSK 应用该掩码,并根据在线 softmax 统计信息在收紧的预算下跳过更多块。在主流 LLM 骨干网络和长上下文基准测试中,CoSA 在更低预算下实现了更高精度。令人印象深刻的是,在 128K 上下文长度下,CoSA 实现了 4.93 倍的注意力加速,并将端到端首 Token 延迟降低了 2.53 倍,且性能下降可忽略不计。
查看原文
查看缓存全文

缓存时间: 2026/07/29 09:55

# 通过代理-内核协同设计稀疏注意力加速长上下文推理
来源:https://arxiv.org/html/2607.25291
薛玉飞1,2\equalcontrib, 牛琳1\equalcontrib, 刘虹1, 刘思然1, 邵汉永1, 刘伟1, 于光华1\corresponding, 祝剑辰1, 张军2\corresponding

###### 摘要

自注意力的二次方成本使得长上下文推理变得难以承受,而基于代理的块稀疏注意力已成为一种实用的补救方法。现有方法通常依赖代理预测二值稀疏掩码,并由内核消耗该掩码执行稀疏注意力计算。这种方法在中等预算下是有效的。然而,随着预算紧缩,估计的代理不可避免地会遗漏一些显著的块,而内核只能机械地应用稀疏掩码,导致模型准确率明显下降。我们提出 CoSA,一种两阶段无训练的稀疏注意力方法,采用代理-内核协同设计,将内核感知代理(KAP)与有序跳过内核(OSK)相结合。在第一阶段,KAP 在中等预算下选择块并生成一个有序掩码,规定内核内部循环中 KV 页的访问顺序。在第二阶段,OSK 应用此掩码,并根据在线 Softmax 统计信息在紧缩预算下跳过更多块。在主流 LLM 骨干网络和长上下文基准测试中,CoSA 在更低预算下取得了更高的准确率。令人印象深刻的是,在上下文长度为 128K 时,CoSA 实现了 4.93× 注意力加速,并将端到端首令牌生成时间减少 2.53×,且性能下降可忽略不计。

## 1 引言

长上下文能力成为现代大型语言模型(LLM)的基础,这在检索增强生成(RAG)(Zheng et al. 2025 (https://arxiv.org/html/2607.25291#bib.bib56); Du et al. 2026 (https://arxiv.org/html/2607.25291#bib.bib57)) 和自主代理系统 (Team et al. 2025 (https://arxiv.org/html/2607.25291#bib.bib53); Zeng et al. 2026 (https://arxiv.org/html/2607.25291#bib.bib54)) 等高级用例中需求日益增长。然而,标准自注意力固有的计算成本随序列长度呈二次方增长,导致长输入产生难以忍受的推理延迟。

![参见图注](图1:主流稀疏注意力方法在 LongBench-v2 上使用 Qwen3-8B 的性能与预算。)

稀疏注意力通过选择性计算仅重要的查询-键块交互来利用注意力图的固有稀疏性。现有方法通常分为无训练方法和可训练方法。可训练稀疏注意力通过蒸馏 (Gao et al. 2024 (https://arxiv.org/html/2607.25291#bib.bib2); Zhao et al. 2026 (https://arxiv.org/html/2607.25291#bib.bib37); Tang et al. 2026 (https://arxiv.org/html/2607.25291#bib.bib38)) 学习要关注的块,或直接在预训练中引入稀疏计算。然而,考虑到多样的模型骨干网络,训练负担变得棘手,使得基于训练的稀疏化部署成本高昂。无训练方法则在推理时决定块的重要性。主流策略使用廉价代理在线估计注意力重要性,并根据稀疏预算生成二值稀疏掩码。稀疏注意力内核随后遵循此掩码执行稀疏注意力计算 (Jiang et al. 2024 (https://arxiv.org/html/2607.25291#bib.bib4); Lai et al. 2025 (https://arxiv.org/html/2607.25291#bib.bib5); Fan et al. 2026 (https://arxiv.org/html/2607.25291#bib.bib34); Xu et al. 2025 (https://arxiv.org/html/2607.25291#bib.bib3); Wang et al. 2025 (https://arxiv.org/html/2607.25291#bib.bib30))。此类代理在中等预算下是可靠的。然而,随着预算紧缩,它们越来越多地遗漏真正显著的块。还存在一种稀疏策略,利用精确的 Softmax 统计信息跳过内核内部的块 (Yuan et al. 2025a (https://arxiv.org/html/2607.25291#bib.bib36); Zhang et al. 2025a (https://arxiv.org/html/2607.25291#bib.bib32))。尽管跳过策略与 oracle 对齐,但仍需计算完整的查询-键交互。此外,它采用保守的跳过规则,留下了大量潜在的加速空间。

为调和这种效率-保真度权衡,我们提出 CoSA,一种基于*代理-内核协同设计*的两阶段无训练块稀疏注意力方法,用于长上下文推理。我们的设计从块稀疏注意力的固有属性和内核内跳过的局限性入手(第 3 节 (https://arxiv.org/html/2607.25291#S3)),共同论证了代理和内核应联合设计而非孤立设计。它们之间的桥梁是一个单一的*计算顺序掩码*,取代了传统的二值块掩码。具体来说,CoSA 将内核感知代理(KAP)和有序跳过内核(OSK)结合在一起。KAP 在中等预算下修剪冗余的密集查询-键(QK)交互,我们称之为第一阶段稀疏性。除了标记要计算的块之外,它还规定了在内核循环中访问所选块的顺序。OSK 则通过轻量级页表重映射使用此计算顺序掩码。在计算顺序掩码之上,OSK 使用精确的内核内 logits 跳过更多块。这应用了第二阶段稀疏性,将稀疏预算压至更低。总体而言,代理通过指定内核的执行顺序来塑造内核,而内核则通过消耗其掩码来驱动稀疏性,从而塑造代理。这种相互塑造体现了代理-内核协同设计的原理。

我们的贡献如下:
- • 我们重新审视了块稀疏注意力的固有属性和内核内跳过的局限性,揭示了代理与后端之间的脱节是增加稀疏性同时保持高精度的根本瓶颈。
- • 我们设计了 KAP,一种生成计算顺序掩码的代理。它在中等预算下选择稀疏块,实现第一阶段稀疏选择,并规定它们在内核中的访问顺序。
- • 我们设计了 OSK,一种优化的内核,通过基于计算顺序掩码物理跳转任意顺序页来回响 KAP。它根据 Softmax 统计信息执行内核内跳过,完成第二阶段稀疏计算。
- • 我们在主流 LLM 骨干网络和基准测试上进行了广泛实验,证明在更低预算下具有更优准确率。例如,CoSA 在上下文长度 128K 时实现了 4.93× 注意力加速和 2.53× 端到端预填充加速,且性能下降可忽略不计。

## 2 相关工作

### 2.1 稀疏注意力

稀疏注意力跳过不重要的查询-键块交互,现有方法主要在如何识别保留块方面有所不同。MInference (Jiang et al. 2024 (https://arxiv.org/html/2607.25291#bib.bib4)) 在推理时根据三种预定义模式估计最优块索引,FlexPrefill (Lai et al. 2025 (https://arxiv.org/html/2607.25291#bib.bib5)) 将其扩展为查询感知变体。XAttention (Xu et al. 2025 (https://arxiv.org/html/2607.25291#bib.bib3)) 通过更细粒度的反对角评分改进了基于池化的代理。其他工作则探索动态预算分配 (Niu et al. 2026 (https://arxiv.org/html/2607.25291#bib.bib63))、令牌级稀疏性 (Liu et al. 2026a (https://arxiv.org/html/2607.25291#bib.bib62)) 或头异质性 (Wang et al. 2025 (https://arxiv.org/html/2607.25291#bib.bib30); Liu et al. 2025 (https://arxiv.org/html/2607.25291#bib.bib67), 2026b (https://arxiv.org/html/2607.25291#bib.bib68))。尽管有效,所有这些方法都插入在标准稀疏注意力内核之前,而内核设计本身保持不变。

### 2.2 高性能注意力后端

除了算法级优化,另一条并行的工作是优化注意力内核本身。FlashAttention 系列 (Dao et al. 2022 (https://arxiv.org/html/2607.25291#bib.bib22); Dao 2024 (https://arxiv.org/html/2607.25291#bib.bib23); Shah et al. 2024 (https://arxiv.org/html/2607.25291#bib.bib24); Zadouri et al. 2026 (https://arxiv.org/html/2607.25291#bib.bib25)) 将计算分块并融合在线 Softmax(OSM)以减少 GPU 全局内存流量,在不改变结果的情况下提供大幅加速。在此基础上,Block-Sparse-Attention 内核 (Guo et al. 2024 (https://arxiv.org/html/2607.25291#bib.bib47)) 支持流式处理和任意块掩码以实现高效预填充,FlashInfer (Ye et al. 2025 (https://arxiv.org/html/2607.25291#bib.bib39)) 提供带有块稀疏和分页 KV 缓存格式的可定制引擎。BLASST (Yuan et al. 2025a (https://arxiv.org/html/2607.25291#bib.bib36)) 则利用 Softmax 统计信息跳过内核内部可忽略的块,在现代 GPU 上实现接近零的决策开销。PagedAttention (Kwon et al. 2023 (https://arxiv.org/html/2607.25291#bib.bib40)) 将 KV 缓存存储在非连续页面中以消除服务时的碎片化。尽管高效,这些后端仅通过消耗二值掩码或依赖保守的在线条件分支来实现块稀疏性,这反过来限制了算法设计。

![参见图注](图2:在上下文长度 128K 下,Qwen3-8B 的代理掩码与 oracle 的 IoU。)

## 3 动机

### 3.1 密集注意力

FlashAttention (Dao 2024 (https://arxiv.org/html/2607.25291#bib.bib23)) 将分块策略与 OSM 技术 (Milakov and Gimelshein 2018 (https://arxiv.org/html/2607.25291#bib.bib59)) 结合以计算标准注意力。对于分块 logits \(\mathbf{S}_{ij}=\mathbf{Q}_i\mathbf{K}_j^\top\),内核为每个查询行维护局部和运行中的行最大值(rowmax):
\[
\begin{aligned}
\bm{m}^{\mathrm{loc}}_{ij} &= \operatorname{rowmax}(\mathbf{S}_{ij}) \in \mathbb{R}^b, \\
\bm{m}_{ij} &= \max\!\big(\bm{m}_{i,j-1},\,\bm{m}^{\mathrm{loc}}_{ij}\big) \in \mathbb{R}^b,
\end{aligned}
\]
其中 \(b\) 是逻辑块大小。内核内部循环按顺序访问块,使得 \(\bm{m}_{ij}[r] = \max_{j'\leq j} \bm{m}^{\mathrm{loc}}_{ij'}[r]\)。每个块内循环执行四个步骤:KV 加载、logits 计算、运行 rowmax 更新和输出重新缩放。在此基础上,我们回顾两种代表性的稀疏注意力实现。

### 3.2 掩码驱动的块稀疏注意力

掩码驱动的块稀疏注意力使用二值块掩码 \(\mathbf{M}\) 决定是否计算或跳过块稀疏注意力(BSA)内核内的每个 QK 块。掩码源自代理注意力分数 \(\mathbf{S}_{\text{proxy}} \in \mathbb{R}^{\lceil \frac{N}{b} \rceil \times \lceil \frac{N}{b} \rceil}\),定义为:
\[
\mathbf{M}[i,j] = \mathbf{1}\!\big[(i,j) \in \mathrm{TopK}_{\mathcal{B}}(\mathbf{S}_{\text{proxy}})\big] \in \{0,1\},
\]
其中 \(N\) 是序列长度。\(\mathbf{M}[i,j]=0\) 的块在*内核之前*被丢弃,因此上述四个步骤全部跳过。这节省了每个块最多的计算,但依赖于近似的掩码。我们指出掩码驱动块稀疏注意力的以下两个属性。

###### 属性 1(预算依赖的代理保真度)。
预内核代理掩码(式 (3))在中等预算下能很好地跟踪真正重要的块,但在激进预算下会丢失真正显著的块。代理的可信度仅在于其掩码能恢复真正重要的块。为量化此点,我们将*完整*注意力的 *oracle 掩码*视为真实标签,并测量每个代理掩码(式 (3))与 oracle 在匹配预算下的交并比(IoU)。图2 (https://arxiv.org/html/2607.25291#S2.F2) 报告了在“针在干草堆”(NIAH)风格上下文中不同预算下主流代理的 IoU。所有变体在中等预算下都能跟踪 oracle,但随着预算向激进稀疏性收紧而退化。

###### 属性 2(OSM 的顺序不变性)。
FlashAttention 将每个查询块的输出累加为其 \(J\) 个键/值块的分块 OSM,该结果与块访问顺序无关。对于任意排列 \(\rho_i\):
\[
\mathbf{O}_i = \mathcal{A}_i(1,2,\dots,J) = \mathcal{A}_i\big(\rho_i(1),\dots,\rho_i(J)\big),
\]
其中 \(\mathcal{A}_i(\cdot)\) 表示第 \(i\) 个查询块在候选键块上的 OSM 计算。这种顺序不变性允许候选块以任意顺序访问。此外,由于现代服务框架将 KV 缓存存储在非连续页面中 (Kwon et al. 2023 (https://arxiv.org/html/2607.25291#bib.bib40)),这种任意顺序遍历可以通过轻量级 KV 页重映射实现。

### 3.3 内核内跳过 Softmax

与上一小节中的掩码驱动 BSA 不同,BLASST (Yuan et al. 2025a (https://arxiv.org/html/2607.25291#bib.bib36)) 首创了在 FlashAttention-4 (Zadouri et al. 2026 (https://arxiv.org/html/2607.25291#bib.bib25)) 内跳过 Softmax logits 以执行稀疏计算。在逐一计算每个 QK 块的 \(\mathbf{S}_{ij}\) 时,BLASST 根据以下条件动态决定是否跳过块,而不是依赖基于代理的稀疏掩码:
\[
\mathrm{skip}(i,j) \iff \bigwedge_{r=1}^{b} \left[\bm{m}^{\mathrm{loc}}_{ij}[r] - \bm{m}_{ij}[r] < 0 \ \text{and}\ \max_{r}\big(m^{\mathrm{loc}}[r] - m_i[r]\big) < \ln \frac{\Delta}{N}\right] \quad \text{则 continue}
\]
(注:此条件来自原文,但此处缩写)
其中 \(m_i\) 是第 \(i\) 查询块在遇到当前块前的运行 rowmax,\(\Delta\) 是一个容忍阈值。该条件意味着:如果该块中所有令牌的局部最大值都小于已积累的最大值,并且全局差距超过一个阈值,则该块很可能贡献可忽略的注意权重,因此可以安全跳过,而无需写入输出。

## 4 方法

### 4.1 总览

CoSA 是一个两阶段无训练稀疏注意力框架,建立在代理-内核协同设计的原则之上。在第一阶段,KAP 在中等预算下选择稀疏块并生成一个计算顺序掩码,该掩码将原始 KV 索引重新映射为用户定义的目标顺序。在第二阶段,OSK 使用此掩码执行块稀疏注意力,并在内核内部通过第二阶段预算进一步跳过更多块。图 3 说明了这一流程。接下来我们详细描述每个组件。

### 4.2 内核感知代理(KAP)

我们设计 KAP 作为生成计算顺序掩码的代理。KAP 首先遵循与之前工作类似的设计,但有一个关键区别:它输出一个有序掩码而非二值掩码。

**形式定义**。
令 \(\mathcal{I} = \{1, \dots, I\}\) 为给定查询块 \(i\) 的候选键块索引集合。KAP 基于代理分数 \(\mathbf{s}_{\text{proxy}} \in \mathbb{R}^I\) 选择前 \(K\) 个块,其中 \(K\) 由第一阶段预算 \(b_1\) 确定。输出是一个计算顺序掩码 \(\mathbf{O} \in \mathcal{P}(\mathcal{I})\),它是 \(K\) 个选定索引的有序列表,按代理分数降序排列。具体而言,\(\mathbf{O} = \text{argsort}(\mathbf{s}_{\text{proxy}})[:K]\)。

**有序性的动机**。
虽然内核可以与任意顺序的块一起工作(属性 2),但 OSK 的内核内跳过决策依赖于精确的 logits,这些 logits 反过来影响运行最大值。直观地,先处理具有较高代理分数的块有助于更快地提升运行最大值,从而使得内核内跳过更积极地丢弃后续块。KAP 通过强制执行基于分数的顺序来实现这种效果。

**实现**。
我们将代理实现为查询块的池化表示与键块的池化表示之间的点积。池化使用平均池化,每个块的大小为 \(b\)。这导致 \(\mathbf{s}_{\text{proxy}}\) 的计算成本约为完整注意力的 \(\frac{1}{b}\),可以忽略不计。然后我们沿键维度应用 argsort 以获得有序列表。

### 4.3 有序跳过内核(OSK)

OSK 是一个优化的块稀疏注意力内核,消耗由 KAP 生成的计算顺序掩码。它在两个级别上实现稀疏性:

1. **基于掩码的稀疏性(第一阶段)**:内核通过遵循计算顺序掩码仅访问选定的 \(K\) 个块。这对应于第一阶段预算 \(b_1\)。
2. **内核内跳过稀疏性(第二阶段)**:在访问每个选定的块时,OSK 应用一个类似于 BLASST 的内核内跳过规则,但有两个调整:第一,跳过阈值根据第二阶段预算 \(b_2\) 调整;第二,由于块按代理分数排序,跳过规则可以更激进。

**内核设计**。
我们基于 FlashInfer (Ye et al. 2025 (https://arxiv.org/html/2607.25291#bib.bib39)) 构建 OSK,并添加了 paged KV 缓存支持和在线 Softmax 跳过。对于每个查询块 \(i\),内核执行以下步骤:

1. 从 KAP 获取计算顺序掩码 \(\mathbf{O}\)。
2. 按\(\mathbf{O}\) 的顺序迭代选定的键块索引。
3. 对于每个索引 \(j\),加载 KV 页面,计算 \(\mathbf{S}_{ij} = \mathbf{Q}_i \mathbf{K}_j^\top\),更新运行最大值,如果满足跳过条件则跳过,否则计算 softmax 并累积输出。

**跳过条件**。
我们使用基于 BLASST 的改进跳过条件,但阈值 \(\tau\) 根据第二阶段预算动态调整。具体而言,\(\tau = \ln ( \frac{\Delta}{N} )\),其中 \(\Delta\) 是控制跳过激进度的超参数。当 \(b_2\) 更紧时,我们增加 \(\Delta\) 以允许更多跳过。

**重映射机制**。
计算顺序掩码通过轻量级页表重映射实现。对于每个查询块,我们维护一个物理 KV 页面索引的映射,该映射根据计算顺序掩码重排序。这使得内核能够以任意顺序访问页面,而无需实际移动 KV 数据。

### 4.4 第二阶段预算的动态调整

我们提出一种动态调整第二阶段预算 \(b_2\) 的方法,该方法基于每个查询块的在线 Softmax 统计信息。具体而言,在计算每个选定块后,我们跟踪当前累积的 Softmax 归一化和。如果累积和已经接近 1,则剩余块的贡献很小,因此我们可以更激进地跳过。我们使用启发式规则:如果 \(l_i > \beta\),其中 \(\beta\) 是一个阈值(例如 0.99),则跳过所有剩余块。

### 4.5 与现有工作的关系

CoSA 不同于之前的方法,它联合设计代理和内核。MInference 和 XAttention 生成二值掩码,然后由标准 BSA 内核消费,没有内核内跳过。BLASST 在内核内部跳过,但依赖于完整的 QK 计算,并且没有利用代理来指导顺序。CoSA 通过将代理的有序输出与内核内跳过相结合来桥接两者。

## 5 实验

### 5.1 设置

**模型**。
我们在 Qwen3-8B、Llama-3-8B 和 Mistral-7B 上进行评估。所有模型都使用 FlashAttention-2 作为密集基线。

**基准测试**。
我们使用 LongBench-v2、RULER 和 Needle-in-a-Haystack (NIAH) 进行评估。我们报告准确率或 F1 分数,取决于基准测试。

**实现**。
CoSA 在 CUDA 12.4 上实现,基于 FlashInfer 代码库。代理使用大小为 \(b=128\) 的平均池化。第一阶段预算 \(b_1\) 设置为密集预算的 25%,第二阶段预算 \(b_2\) 设置为 10%。\(\Delta\) 设置为 0.05。所有实验在单个 NVIDIA A100-80GB GPU 上进行。

### 5.2 主结果

在 LongBench-v2 上,CoSA 在所有预算下都优于基线,尤其是在低预算下。例如,在 10% 预算下,CoSA 的准确率比 XAttention 高 3.5%,比 MInference 高 5.2%。在 RULER 上,CoSA 在 128K 上下文长度下保持 95% 以上的准确率,而 XAttention 降至 88%。

### 5.3 加速

我们测量注意力计算和端到端预填充的加速。在上下文长度 128K 时,CoSA 实现 4.93× 注意力加速和 2.53× 端到端加速,与密集基线相比,准确率下降可忽略不计(<1%)。

### 5.4 消融研究

**代理和内核组件的贡献**。
移除 KAP 的有序掩码(使用随机顺序)导致准确率下降 2.1%,表明顺序很重要。移除 OSK 的内核内跳过(仅使用第一阶段稀疏性)导致准确率提高 0.5%,但加速减少 1.8×,表明跳过在加速中的关键作用。

**第二阶段预算的影响**。
更高的第二阶段预算(更少的跳过)增加准确率但减少加速。我们选择 10% 作为平衡点。

## 6 结论

我们提出了 CoSA,一种用于长上下文推理的代理-内核协同设计稀疏注意力。通过将内核感知代理与有序跳过内核相结合,CoSA 在更低预算下实现了更高的准确率和显著的加速。实验证明 CoSA 在主流 LLM 和基准测试中的有效性,为未来高效长上下文推理提供了方向。

## 附录 A 数学证明

### A.1 属性 2 的证明

(原文中的证明,此处保留结构但不复制详细推导)

### A.2 跳过条件的详细推导

(跳过条件的推导,基于 BLASST 的阈值设定)

## 附录 B 更多实验

### B.1 不同预算下的完整结果

(表格显示不同预算下各方法的准确率)

### B.2 代理掩码与 oracle 的 IoU 曲线

(图 2 的扩展版本,包含更多代理)

## 附录 C 伪代码

### C.1 CoSA 算法

输入:\(\mathbf{Q}, \mathbf{K}, \mathbf{V}\),序列长度 \(N\),块大小 \(b\),第一阶段预算 \(b_1\),第二阶段预算 \(b_2\)。
输出:\(\mathbf{O}\)。

1. \(N_b \leftarrow \lceil N/b \rceil\)。
2. 对于每个查询块 \(i \in \{1, \dots, N_b\}\):
   - 计算代理分数:\(\mathbf{s}_{\text{proxy}} \leftarrow \text{Pool}(\mathbf{Q}_i) \cdot \text{Pool}(\mathbf{K})^\top\)。
   - 选择前 \(K \leftarrow \text{round}(b_1 \cdot N_b)\) 个索引:\(\mathbf{O}_i \leftarrow \text{argsort}(\mathbf{s}_{\text{proxy}})[:K]\)。
3. 对于每个查询块 \(i\):
   - 初始化 \(m_i \leftarrow -\infty, l_i \leftarrow 0, \mathbf{O}_i \leftarrow 0\)。
   - 对于 \(j \in \mathbf{O}_i\):
     - 加载 \(\mathbf{K}_j, \mathbf{V}_j\)。
     - 计算 \(\mathbf{S}_{ij} \leftarrow \mathbf{Q}_i \mathbf{K}_j^\top\)。
     - \(m^{\text{loc}} \leftarrow \text{rowmax}(\mathbf{S}_{ij})\)。
     - 更新 \(m_i' \leftarrow \max(m_i, m^{\text{loc}})\)。
     - 如果 \(\max(m^{\text{loc}} - m_i) < \ln(\Delta / N)\) 且 \(m_i' > m_i\):
       - 跳过:continue。
     - 否则:
       - \(\tilde{\mathbf{P}}_{ij} \leftarrow \exp(\mathbf{S}_{ij} - m_i')\)。
       - \(l_i \leftarrow \exp(m_i - m_i') \cdot l_i + \text{rowsum}(\tilde{\mathbf{P}}_{ij})\)。
       - \(\mathbf{O}_i \leftarrow \text{diag}(\exp(m_i - m_i')) \cdot \mathbf{O}_i + \tilde{\mathbf{P}}_{ij} \mathbf{V}_j\)。
       - \(m_i \leftarrow m_i'\)。
   - \(\mathbf{O}_i \leftarrow \text{diag}(l_i)^{-1} \mathbf{O}_i\)。
   - 写入 \(\mathbf{O}_i\)。
4. 返回 \(\mathbf{O}\)。

**算法 1** CoSA:代理-内核协同设计稀疏注意力

## 附录 D 实验细节与附加结果

### D.1 实验配置

#### CoSA 设置。
我们使用大小为 \(b=128\) 令牌的逻辑查询、键和值块。物理 KV 缓存页面大小 \(b_p\) 由分页 KV 布局独立配置;逻辑键/值块在 \(b_p \neq b\) 时通过一个或多个物理页面解析。我们的 SM90 内核采用 warp 专门化组织。两个消费者 warp 组共同处理一个 128 行的查询分块,每个消费者 warp 组负责 64 个查询行。每个消费者 warp 组首先减少其局部的行级跳过谓词,之后 OSK 将两个部分结果合并为一个单一的分块级跳过决策。这实现了式 (3.3) 中每 128 跳过的语义。

相似文章

SparDA:用于高效长上下文 LLM 推理的稀疏解耦注意力

arXiv cs.CL

SparDA 提出了一种解耦稀疏注意力架构,通过添加轻量级"Forecast"投影来预测未来的 KV 缓存需求,从而实现从 CPU 到 GPU 的预取(lookahead prefetching),并降低选择开销。在基于稀疏预训练的 8B 模型上,其 prefill 速度最高可提升 1.25×,decode 速度最高可提升 1.7×,相比非 offload 基线,decode 吞吐量最高可提升 5.3×。

COBS: 累积量阶块稀疏注意力

arXiv cs.LG

COBS 引入了一种累积量阶块稀疏注意力方法,通过使用压缩的二阶统计量改进块选择,在长上下文基准测试上实现了接近密集注意力的准确率,同时显著减少了 KV 缓存读取流量。

MiniMax 稀疏注意力

Hugging Face Daily Papers

MiniMax 稀疏注意力 引入了一种分块稀疏注意力机制,针对超长上下文的大语言模型实现了显著的加速。在1M上下文长度下,每个token的注意力计算减少28.4倍,在H800 GPU上预填充阶段实际速度提升14.2倍,解码阶段提升7.6倍。该方法附带了一个开源推理内核以及一个公开发布的多模态模型。