论文大模型重点稀疏注意力KV cache推理加速长上下文

OpenAttention: One Percent Is Enough with Tail Proxies

AI 解读已生成 · AI 生成，数字以原文为准

AI 全文总结值得深入基于论文全文 3.8 万字 · 完整阅读

OpenAttention 用尾部代理在 1% read 下接近密集精度，RULER-32K HARD 得 86.30。

长上下文解码受KV cache内存流量限制，稀疏注意力在极低读取时丢弃尾部会偏置、采样尾部又易漏掉重击者；作者发现单一计数感知均值代理比采样权重误差更低。据此提出 OpenAttention：QGC按查询二阶矩主特征向量投影后聚类，缓存均值/方差token；ErrorGate按估计权重误差分配top-k预算，剩余成员折叠为均值代理（logit加log n）；QGS低比特编码召回高分token，由OpenAttention.flash做一次在线softmax。在Llama-3.1-8B的RULER-32K HARD上，1%读取平均精度86.30，接近稠密87.04，作者报告保留 99% 稠密精度，优于HashAttention、vAttention等基线；跨其他模型和更长上下文仍接近稠密，解码注意力比稠密快 10×。局限是ErrorGate以对角方差近似簇内方差，相关性高但不能恢复精确误差。

1尾部改用 count-aware 均值代理，替代采样，避免低密度错过 heavy hitters。
2QGC 在 16 维查询特征空间聚类，QGS 单层选择 0.22ms。
33% read 下三模型距 dense 不足 0.7 点，96K 仍稳定。

为什么重要　给做长上下文推理优化的读者：一条 1% read 下接近密集精度的完整管线，含开源 kernel 与多模型验证。适合　做长上下文 LLM 推理、稀疏注意力、KV cache 优化的工程师与研究者。

来源: arxiv.org
发布: 2026/10/23

中文（AI 翻译）

Sparse attention 只读取 key-value cache 的一小部分，如何处理未读取的尾部仍是一个关键挑战。丢弃尾部会使输出产生偏差，而以往方法的采样在 5% 及以下读取率时会失败，因为尾部可能包含选择器遗漏的 heavy hitters。

我们的关键观察是，最简单的代理——一个将整个尾部折叠为其 count-aware mean 的 token——已经能在实用选择器背后的每个读取率下产生比采样更低的输出误差。我们进一步以闭式形式表明，在低读取率下，采样几乎使丢弃的权重误差翻倍，而单个代理仍低于该误差。

基于这些发现，我们探索用少量针对尾部的代理来近似稠密注意力分布。我们提出 OpenAttention，一个由聚类、ErrorGate 和 top-k 选择组成、随后接 count-aware mean proxy 的流水线。

为了高效且有效地实现聚类和选择，我们设计了无需训练的 Query-Geometry Clustering（QGC）和 Query-Geometry Selector（QGS），它们在压缩 keys 的同时尊重 query 方向。此外，与以往仅按 centroid 排序 cluster 的方法不同，ErrorGate 根据每个 cluster 的估计权重误差分配精确读取量。

在 Llama-3.1-8B 上的 RULER-32K HARD 中，OpenAttention 在 1% 读取率下保留了稠密精度的 99%。它还在多个模型、上下文长度和数据集上取得了最先进精度。

我们定制的注意力 kernel OpenAttention.flash 在 1% 读取率下，使 decode attention 阶段比稠密注意力快 10×。

原文

Sparse attention reads only a small subset of the key-value cache, and how to treat the unread tail remains a key challenge.

Discarding the tail biases the output, and sampling it as previous approaches do fails at reads of 5% and below, where the tail likely holds heavy hitters the selector missed.

Our key observation is that the simplest possible proxy, one token that folds the entire tail into its count-aware mean, already yields lower output error than a sample at every read behind a practical selector.

We further show in closed form that at low reads sampling nearly doubles the weight error of discarding while a single proxy stays below it.

Building on these findings, we explore approximating the dense attention distribution with a few proxies that target the tail.

We propose OpenAttention, a pipeline of clustering, ErrorGate and top-k selection followed by a count-aware mean proxy.

To implement clustering and selection both efficiently and effectively, we design training-free Query-Geometry Clustering (QGC) and a Query-Geometry Selector (QGS), which compress the keys while respecting the query directions.

Additionally, unlike prior methods that rank clusters by centroid alone, ErrorGate allocates the exact read by the estimated weight error of each cluster.

On RULER-32K HARD with Llama-3.1-8B, OpenAttention retains 99% of dense accuracy at a 1% read. It also achieves state-of-the-art accuracy across models, context lengths and datasets.

Our customized attention kernel, OpenAttention.flash, runs the decode attention stage 10 $\times$ faster than dense attention at a 1% read.

86.30

OpenAttention

RULER-32K 均值

87.04

Dense 注意力

同基准

48.2

最佳基线 1%

采样/丢弃方法

10×

解码加速

batch 32, 32K

注释 · 基准，10× 加速来自定制 Triton kernel。出自 §摘要, §4.1

无组件
67.35
仅 proxy
77.91
proxy+
85.90
proxy+
86.73
全部三件套
86.82

注释 · 无组件配置等同 vAttention on HashAttention selector。出自表 4

图 1 展示 Llama-3.1-8B 在 RULER-32K HARD 的每任务精度、总读取量与解码加速；1% read 下 OpenAttention 接近密集精度并 10× 加速。

图 5 展示 OpenAttention 概览：QGC 按查询几何聚类，ErrorGate 用 QGS 读高误差簇 top tokens，其余成 count-aware 均值 proxy。

图 7 在 137,216 个 decode 行上对比对角与全方差 logit 方差估计，Spearman 0.92，支持对角近似。

模型: Llama-3.1-8B-Instruct bf16
QGC rank: $r = 16$
QGS rank: $r = 32$
聚类数: $C = 64$
量化: int8 每 head/坐标

深度总结

1 / 6

OpenAttention 四阶段方法· 流程

四组件顺序：聚类 → 预算 → 选 token → 余量成 count-aware 均值代理。

01
聚类缓存
将 keys 投影到 top-16 查询特征向量做 k-means，缓存每簇数量、均值 token 与对角方差 token。
02
分配预算
按估计权重误差 $\hat{{} ϵ}_{b}$ 将 token 预算 E 按比例分给各簇，最小化残余误差。
03
读取 top tokens
用 int8 编码 keys 在 $r = 32$ 子空间评分，每簇读取得分最高的 $q_{b}$ 个 token。
04
剩余成员成代理
每簇未读 token 合并为一个虚拟 token，logit 为 $\overset{ˉ}{{} z}_{b} + log$ $n_{b}$ ，value 为均值。
05
OpenAttention.flash 融合
Triton 版 FlashAttention-2 decode kernel 在单次 online softmax 中处理 exact tokens 与 proxies。

注释 · 聚类每请求一次，选择每解码步复用；proxy 等价于一个带 log n 的虚拟 token。出自 §3.1–3.4

OpenAttention 用尾部代理在 1% read 下接近密集精度，RULER-32K HARD 得 86.30。

1尾部改用 count-aware 均值代理，替代采样，避免低密度错过 heavy hitters。

2QGC 在 16 维查询特征空间聚类，QGS 单层选择 0.22ms。

33% read 下三模型距 dense 不足 0.7 点，96K 仍稳定。

来源

arxiv.org

发布

2026/10/23