DOSA: Dots Sparse Attention

发表时间: 2026-10 · arXiv:2610.13990 (Xiaohongshu)

原文: https://arxiv.org/abs/2610.13990

Fanxu Meng, Taosong Fang (Dots Studio, Xiaohongshu Inc.)

速读

一句话结论 提出了一种名为 DOSA 的无损稀疏注意力机制，通过构建严格的得分上界来跳过大量不必要的全头注意力计算，在保持精确 Top-$K$ 选择的前提下，将百万级上下文大模型的首字生成延迟降低了近一半。

要解决什么问题 在长上下文模型中，稀疏注意力机制（如 DeepSeek Sparse Attention）虽然最终只访问少数键，但为了找出这些键，现有的路由机制必须用所有的索引头对每一个键进行密集的打分。这个打分过程包含点积、ReLU 激活以及带符号的加权求和。随着上下文长度的增加，这种全量扫描的计算量急剧膨胀。例如在处理约 1M 词元的上下文时，索引器占据了 DeepSeek-V3.2 预填充阶段 90.8% 的 GPU 时间。现有的优化路线（如 LiteTopK）主要通过算子融合来减少显存读写流量和排序开销，但依然无法避免对每一个键执行全头计算的算力瓶颈；而其他采用部分头进行粗筛的方法（如 MISA）则会带来精度损失，无法保证选出的候选者与原模型完全一致。因此，如何在不丢失任何原本应该入选 Top-$K$ 的键的前提下，安全地跳过绝大多数键的全头打分计算，是当前长文本稀疏注意力面临的核心卡点。

怎么做的 DOSA 的核心思路是条件计算：如果能证明一个候选键的得分上限，低于当前已知的 Top-$K$ 门槛下限，就可以直接将其淘汰，而无需计算其完整的全头得分。为了以极低的计算成本构造出这个紧致且安全的得分上界，DOSA 设计了三个关键部件。首先，利用一组随机采样的种子键计算出真实得分，将其中的第 $K$ 大得分作为全局 Top-$K$ 阈值的下界 $\tau$。其次，对于候选键，DOSA 只保留少数几个难以估计的头（保留头集合 $R$）进行真实点积计算。对于被省略的正权重头，DOSA 利用非负投影和联合残差边界，将其贡献转化为对保留头点积的复用以及一个基于键向量 $L_2$ 范数的统一余量。对于被省略的负权重头（它们会降低总分），如果简单地将其全部聚合，未激活的头会抵消掉激活头的负向贡献从而导致上界变松；因此 DOSA 利用种子键的激活频率，选出一个多数激活的子集 $A$ 进行聚合，以此安全地估算负向扣分下界。最终，候选键 $i$ 的得分上界定义为：

$$ U_i = \sum_{j \in R} \widetilde{w}_j [\mathbf{q}_j^\top \mathbf{k}_i]_+ + E_P \|\mathbf{k}_i\|_2 - [\mathbf{a}_A^\top \mathbf{k}_i]_+ $$

其中 $\widetilde{w}_j$ 是折叠后的权重，$E_P$ 是正向残差余量，$\mathbf{a}_A$ 是负头聚合向量。在执行时，GPU 先用极少的算力计算 $U_i$，只有当 $U_i \geq \tau$ 时，该候选键才会被保留并送入下一阶段计算完整的原模型得分。这种机制将原本密集的矩阵运算转化为了低成本粗筛加稀疏全量计算的两段式流水线，并配合稀疏写回机制，大幅削减了计算量和显存流量。

效果如何 实验在 NVIDIA GPU 上搭建，评测了 GLM-5.2（32 个索引头）和 DeepSeek-V3.2（64 个索引头）两款模型，测试数据涵盖 128K 到约 1M 词元的英文、中文和代码长文本。对比基线有两个：一是 SGLang，代表计算全头得分后进行常规 Top-$K$ 排序的传统路线；二是 LiteTopK，代表将打分与过滤融合以消除中间显存读写的当前最优工程优化路线。在微基准测试中，DOSA 在所有 60 种配置下均超越基线，在 1M 上下文时，完整索引器速度最高达到 SGLang 的 2.74 倍、LiteTopK 的 1.69 倍。在端到端生成首字时间测试中，面对 1M 上下文，DOSA 在 GLM-5.2 和 DeepSeek-V3.2 上分别比 SGLang 提速 1.86 倍和 1.94 倍，比 LiteTopK 提速 1.26 倍和 1.50 倍。作者也指出了该方法的局限性：DOSA 依然需要扫描每一个键的向量，如果上下文较短，或者得分上界过于松散导致大量键存活，节省下来的全头计算量可能不足以抵消粗筛阶段引入的额外开销。

主要贡献

在大型语言模型中，条件计算（Conditional computation）能够在每次处理输入时仅激活模型的一部分，从而有效扩展模型容量。稀疏注意力（Sparse attention）机制虽然只需要访问少量的键（keys），但在定位这些键时，通常仍需要使用每个索引头（indexer head）对每个键进行密集的评分计算。例如，在DeepSeek Sparse Attention (DSA) 中，通过多头索引器计算分数后进行Top-$K$选择，这一过程包含了点积、ReLU激活以及带符号的加权归约，导致稀疏注意力内部存在密集的评分计算瓶颈。

如图1所示，当上下文长度达到约100万个token时，索引器（indexer）在GLM-5.2和DeepSeek-V3.2的预填充（prefill）GPU时间中分别占据了$73.0\%$和$90.8\%$，而其在逻辑矩阵浮点运算（FLOPs）中的比例分别为$47.4\%$和$83.2\%$。这表明，随着上下文的增长，决定“关注哪里”的计算成本已经主导了稀疏注意力本身的成本。

为了解决这一问题，本文提出了DOSA (Dots Sparse Attention)，这是一种无损（lossless）的优化方法，旨在使Token的选择计算条件化，同时严格保留原始的精确Token选择规则。本文的核心创新点与贡献如下：

1. 基于充分证据的Top-$K$拒绝机制：观察到Top-$K$选择只需要足够的证据来拒绝一个候选者，而不必计算其完整的分数。通过一个种子集合（seed set）提供获胜阈值的下界，如果一个候选者的得分上界低于该下界，则可以直接拒绝。
2. 针对带符号DSA分数的定制化证书（Certificate）：构建了Token级别的得分上界。通过保留部分非负的注意力头投影（Nonnegative head projections）和联合残差界限（joint residual bounds）来限制省略的正向头的贡献上限；同时，利用基于激活信息的聚合（activation-informed aggregation）来限制省略的负向头所减去的幅度下限。
3. 减少计算与内存流量的GPU实现：该证书机制引导系统仅对符合条件的候选者进行完整的评分计算，精确的分数过滤进一步减少了数据写回（writeback）和Top-$K$选择的工作量。
4. 显著的端到端性能提升：在英语、中文和代码工作负载的60个微基准测试配置中，DOSA的GPU实现使完整索引器的速度比SGLang最高提升了$2.74\times$，比LiteTopK最高提升了$1.69\times$。在GLM-5.2和DeepSeek-V3.2的端到端评估中，首字生成时间（TTFT）分别获得了最高$1.94\times$和$1.50\times$的加速，几乎将端到端延迟减半。

方法细节

问题形式化

索引器分数的数学定义：对于具有$H$个索引头的固定查询（query），设$\mathbf{q}_h \in \mathbb{R}^D$为第$h$个查询向量，$w_h \in \mathbb{R}$为其权重，$\mathbf{k}_i \in \mathbb{R}^D$为token $i$的键（key）。记$[x]_+ = \max(x, 0)$。索引器的分数计算公式为：

$$ s_i = \sum_{h=1}^H w_h \left[ \mathbf{q}_h^\top \mathbf{k}_i \right]_+, \qquad i \in \Omega, $$

其中$\Omega$仅包含可见的键。

精确选择的目标：对于$|\Omega| \geq K$的情况，我们的目标是恢复出得分最高的$K$个可见键，同时避免对那些不可能进入该集合的键进行全头（full-head）评分计算。

种子阈值与保留的注意力头

种子阈值的设定与拒绝条件：首先选择不同的可见键子集$\mathcal{S} \subseteq \Omega$，要求$|\mathcal{S}| \ge K$，并将其第$K$大的完整分数设为阈值$\tau = \tau_K(\mathcal{S})$。该子集提供了$K$个见证（witnesses）：

$$ \tau_K(\mathcal{S}) \leq \tau_K(\Omega), \qquad s_i \leq U_i < \tau_K(\mathcal{S}) \implies i \notin \mathrm{TopK}(\Omega). $$

因此，$\tau_K(\mathcal{S})$是真实Top-$K$阈值的下界，而$U_i$是token得分的上界：它们之间的严格分离构成了拒绝该候选者的证书（certifies rejection）。一个较差的种子会削弱过滤效果，而增加额外的精确分数只会提高阈值的下界。

保留选定头的必要性：如果对每一个头都进行边界估计，可能会导致所有候选者的上界都高于阈值。因此，我们选择直接评估那些其贡献难以被紧密界定的选定头。

保留头的选择策略：设$R$为保留的$1 \leq r < H$个头的集合。定义正权重头集合$\mathcal{H}_+ = \{h : w_h > 0\}$和负权重头集合$\mathcal{H}_- = \{h : w_h < 0\}$。如果$|\mathcal{H}_+| \leq r$，则保留所有正向头，并用负向头填充剩余位置；负向头的优先级根据$|w_h| \|\mathbf{q}_h\|_2 \min(p_h, 1 - p_h)$确定，其中$p_h$是它们在采样的种子键上的激活频率。这种方式倾向于选择那些激活状态难以预测的强权重头。否则，根据$w_h \|\mathbf{q}_h\|_2$保留最强的$r$个正向头。权重为零的头不需要贡献。记被省略的正向头为$P = \mathcal{H}_+ \setminus R$，被省略的负向头为$N = \mathcal{H}_- \setminus R$。保留部分的贡献将被直接评估：

$$ S_R(\mathbf{k}_i) = \sum_{j \in R} w_j \left[ \mathbf{q}_j^\top \mathbf{k}_i \right]_+ $$

在筛选期间，集合$P$和$N$中被省略的头不会被单独评估；它们的贡献将在接下来的两个小节中被界定。该证明对任何集合$R$都成立；其选择仅影响效率和边界的紧密性。

确定省略的正向头的上界

正向头的非负投影与残差分解：对于每个$h \in P$，选择一个保留的锚点（anchor）和一个非负投影：

$$ j_h^\star = \arg\max_{j \in R} \frac{\left[ \mathbf{q}_h^\top \mathbf{q}_j \right]_+}{\left\| \mathbf{q}_j \right\|_2}, \quad \beta_h = \frac{\left[ \mathbf{q}_h^\top \mathbf{q}_{j_h^\star} \right]_+}{\left\| \mathbf{q}_{j_h^\star} \right\|_2^2}, \quad \mathbf{d}_h = \mathbf{q}_h - \beta_h \mathbf{q}_{j_h^\star}. $$

范数为零的锚点具有匹配的零分数和零系数。

基于子可加性的正向界限：根据ReLU的子可加性、正齐次性以及柯西-施瓦茨不等式（Cauchy–Schwarz），可以得到：

$$ \left[ \mathbf{q}_h^\top \mathbf{k}_i \right]_+ \leq \beta_h \left[ \mathbf{q}_{j_h^\star}^\top \mathbf{k}_i \right]_+ + \left\| \mathbf{d}_h \right\|_2 \left\| \mathbf{k}_i \right\|_2. $$

投影部分复用了已经计算过的点积。非负投影最小化了残差范数，从而在这种形式的单锚点边界中产生了最小的一致范数容限（详见附录B.1）。

折叠权重与正向容限的定义：定义折叠权重和正向残差容限如下：

$$ \widetilde{w}_j = w_j + \sum_{\stackrel{h \in P}{j_h^\star = j}} w_h \beta_h, \qquad E_P = \sum_{h \in P} w_h \left\| \mathbf{d}_h \right\|_2. $$

保留与省略正向项的综合上界：保留的和被省略的正向项加在一起，其上界被限制为$\sum_{j \in R} \widetilde{w}_j [\mathbf{q}_j^\top \mathbf{k}_i]_+ + E_P \|\mathbf{k}_i\|_2$。折叠权重仅属于筛选计划；它们不会修改模型参数。如果$P$为空，则$E_P = 0$且$\widetilde{w}_j = w_j$。

联合残差容限的优化：被省略的残差也共享相同的键。令$\mathbf{u}_h = w_h \mathbf{d}_h$对于$h \in P$。一个联合容限（joint allowance）为：

$$ E_P^{\mathrm{joint}} = \frac{1}{2} \left( \left\| \sum_{h \in P} \mathbf{u}_h \right\|_2 + \sqrt{\sum_{h, \ell \in P} |\mathbf{u}_h^\top \mathbf{u}_\ell|} \right) \leq E_P. $$

该联合容限将$\sum_{h \in P} [\mathbf{u}_h^\top \mathbf{k}_i]_+$限制在$E_P^{\mathrm{joint}} \|\mathbf{k}_i\|_2$以内，因此可以用来替代$E_P$。这种改进利用了残差的几何性质，且不需要额外的键读取操作。附录A.1证明了该陈述及其分组形式。

确定省略的负向贡献的下界

负向头的聚合下界：负向头会从总分中减去数值。为了构建一个上界，我们聚合一个子集$A \subseteq N$并忽略剩余的头；省略负向贡献只会使上界变高。定义：

$$ \mathbf{a}_A = \sum_{h \in A} |w_h| \mathbf{q}_h, \qquad L_N(\mathbf{k}_i) = \left[ \mathbf{a}_A^\top \mathbf{k}_i \right]_+ \leq \sum_{h \in N} |w_h| \left[ \mathbf{q}_h^\top \mathbf{k}_i \right]_+. $$

避免全头聚合的原因：ReLU的子可加性确保了$L_N$不会超过被省略负向头真实减去的数量。为什么不简单地将每个头都包含在$A$中呢？一个未激活的头在自身的ReLU之后贡献为零，但其负的点积可能会在聚合ReLU之前抵消掉已激活头的正点积。因此，包含这样的头可能会削弱负向贡献的估计，从而导致得分上界变得松弛。

基于激活频率的多数支持选择：我们使用种子的激活模式来减少这种抵消。对于采样的种子键，记录$M_i = \{h \in N : \mathbf{q}_h^\top \mathbf{k}_i > 0\}$，并令$p_h$为头$h$的频率。我们选择逐分量的多数支持（componentwise majority support），即$A = \{h \in N : p_h \geq 1/2\}$。这最小化了归一化负向松弛的一个可分离上界；联合来看，前文所述的保留负向优先策略消除了该替代目标中剩余的最大项（详见附录B.2）。如果一个新的键其激活集合恰好是$A$，则聚合是精确的。其他模式仍然是安全的，只是可能会给出一个较松的界限。这些频率用于指导效率，而不是正确性。

分数边界与候选者细化

最终的分数证书公式：结合公式(5)–(8)，我们得到了分数证书（score certificate）：

$$ \boxed{U_i = \sum_{j \in R} \widetilde{w}_j \left[ \mathbf{q}_j^\top \mathbf{k}_i \right]_+ + E_P \left\| \mathbf{k}_i \right\|_2 - \left[ \mathbf{a}_A^\top \mathbf{k}_i \right]_+, \qquad s_i \le U_i.} $$

一个聚合方向取代了任意数量的被省略的负向头。其效用取决于激活的一致性，而正向容限则取决于方向的相似性。

证书执行的四个阶段：图2将证书映射为四个阶段。

• 阶段1（$\bullet$）：在键块（key tiles）之间复用筛选行$\widetilde{Q}$（保留的和辅助的）；外层和内层循环分别遍历键和查询行。
• 阶段2（$\otimes$）：形成$U = p + R$：其中$p$结合了公式(9)中的保留项和聚合项，标量$R = E_P \|\mathbf{k}_i\|_2$是残差容限（与头集合$R$区分开来）。FP8精度下会添加数值保护（详见附录C）。
• 阶段3（$\otimes$）：丢弃$U_i < \tau$的键，并压缩幸存者的ID。收集这些幸存者的键，并评估所有$H$个原始头。
• 阶段4（$\bullet$）：应用公式(1)中的原始ReLU和权重，仅保留$s_i \geq \tau$的分数，以便与种子记录一起进行最终选择。
  在阶段$\otimes$和$\bullet$的数据块中，灰色标记表示避免的计算，紫色标记表示避免的写回。相等的分数能在两次过滤中存活；边界永远不会替代最终分数。

算法的精确性保证：在拥有有效的种子阈值且不丢失存活候选者的前提下，该过程保留了精确的Top-$K$选择；附录A给出了完整的算法、形式化保证（定理1）及其证明。

GPU内核设计

执行架构的设计选择：经过认证的修剪减少了算术运算，但不规则的幸存者集合可能会限制GPU利用率并引入分散的键访问。在预填充期间，我们将DOSA组织为一个规则的筛选遍（screening pass），随后是稀疏的全头细化（sparse full-head refinement）（图3）。三个设计选择将证书转化为高效的执行：对查询进行分组以进行Tensor Core计算、在细化期间共享键加载、以及减少传递给选择阶段的记录。

查询分组与流水线筛选：我们将具有相似保留头预算的查询分组为固定宽度的块（tiles），使得一小部分筛选内核能够使用密集的Tensor Core MMA。保留的查询行、折叠权重和两个FP8辅助行只需准备一次，即可在键块之间复用。可见种子键的精确分数初始化了每个查询的阈值，该阈值在扫描期间保持固定。专用的warp发出异步键加载和MMA操作，而消费者warp则从TMEM累加器评估边界。缓冲的键块将加载与计算重叠。仅当$U_i^{\mathrm{num}} < \tau$时，键才会被拒绝，并且warp ballots将幸存者决策打包成紧凑的掩码（masks）以供细化使用。

细化阶段的共享键访问：独立的每查询收集（per-query gathers）会重复加载被多个查询选中的键。因此，我们将幸存者掩码压缩为共享的（shared）和每查询的（per-query）细化流。在四个查询的组内，如果一个键在至少两个边界中存活，并且位于该组的公共可见范围内，它就会进入共享流；其他候选者使用每查询流。共享键加载服务于该组的全头评分，随后是原始的ReLU、权重和归约。这种组织方式用一些额外的查询-键评估换取了更少的键传输。每个分数都会根据其查询的阈值进行过滤，从而保留精确的选择。

紧凑写回与精确选择：细化阶段将全分计算与阈值过滤融合在一起：只有$s_i^{\mathrm{ref}} \geq \tau$的分数才会产生输出记录。这些记录存储了带有精确键ID的紧凑单调分数代码，减少了高带宽内存（HBM）流量和选择输入。在细化之后，Top-$K$选择保留严格高于边界代码的记录，并使用其原始分数解析边界上的候选者。综合来看，筛选减少了全头算术运算，共享细化复用了键数据，而紧凑记录降低了写入和选择分数的成本。

解码阶段的优化：在解码时，由于每步的查询很少，查询组之间键复用的机会有限。我们保留了完整的FP8评分，并利用时间相干性（temporal coherence）来减少选择工作量。在最近步骤中选出的不同键提供了见证，它们在当前查询中的分数产生了一个Top-$K$阈值的保守下界。达到或超过该下界的分数被压缩以进行精确选择，该选择过程同时映射选定的索引并更新见证。

参考文献引用汇总（方法细节部分）

• 关于精确检索原则的背景，引用了【18，Query evaluation: Strategies and optimizations+1995+Information Processing & Management】和【3，Efficient query evaluation using a two-level retrieval process+2003+CIKM】。原文描述为："A subset of at least $K$ visible keys supplies the latter through its $K$th largest score, following established exact-retrieval principles (Turtle & Flood, 1995; Broder et al., 2003)."

背景知识/关键Observation/设计原则

注：本文在Introduction章节阐述了背景与观察，此处根据要求进行缩写。

稀疏注意力中的密集评分瓶颈：条件计算（Conditional computation）可以通过仅激活部分模型来扩展模型容量，例如混合专家模型（MoE）在FFN中的应用【8，Switch transformers+2022+JMLR】。随着上下文的增长，注意力机制也提供了类似的机会。稀疏注意力虽然只访问少数键，但识别这些键仍然需要密集的扫描。例如，DeepSeek Sparse Attention (DSA)【5，DeepSeek-V3.2+2025+arXiv】通过多头索引器选择token，然后进行Top-$K$选择；每个分数的计算都需要点积、ReLU激活和带符号的加权归约。这在稀疏注意力内部留下了密集评分的瓶颈，限制了MoE风格的容量与计算分离。

现有加速方法的局限性：最近的系统试图加速选择过程：Guess-Verify-Refine (GVR)【4，Guess-Verify-Refine+2026+arXiv】利用解码中的时间相关性；DeepSelect【15，DeepSelect+2026+GitHub】通过有界缓冲区流式处理物化分数；LiteTopK【21，LiteTopK+2026+arXiv】将评分与候选者过滤融合以避免密集的得分写回。然而，这些路径仍然要求对每个键进行全头评估。MISA【27，MISA+2026+arXiv】则使用较少的头进行筛选并使用完整索引器重新排序候选者，但重新排序无法恢复被排除的获胜者。

核心观察与设计原则：本文提出，一个token如果其得分上界低于真实Top-$K$阈值的下界，就不可能进入Top-$K$。通过利用已知可见键子集的第$K$大得分，可以获得该下界。面临的挑战是如何以远低于计算完整分数的成本构建紧密的上界。对于带有ReLU的带符号索引器，直接丢弃正权重头是不安全的；丢弃负权重头是安全的但会导致界限松弛。因此，一个有用的证书必须利用头的几何形状和激活结构。

实验环境

• 数据集：使用来自InfiniteBench【25，∞Bench+2024+arXiv】的英语、中文和代码（Code）输入，并辅以Project Gutenberg（英文书籍）和SQLite源代码以构建长上下文。提取了128K到约1M token的连续上下文。
• 模型架构关键参数：评估了GLM-5.2（32个索引头）和DeepSeek-V3.2（64个索引头）。Top-$K$均设置为2048。
• 硬件配置：NVIDIA GPU。微基准测试在单张GPU上运行GLM-5.2；端到端推理使用8路张量并行（TP8），GLM-5.2的预填充块大小为32K token，DeepSeek-V3.2为16K。
• 软件配置：基线方法包括SGLang v0.5.18【26，SGLang+2023+arXiv】和LiteTopK【21，LiteTopK+2026+arXiv】。
• 评测指标：完整索引器延迟（CUDA events测量），首字生成时间（TTFT）。

实验结果

1. 预填充性能（微基准测试）

• 实验内容：在128K、256K、512K、768K和1M token上下文长度下，以及4K、8K、16K和32K预填充块大小下测试GLM-5.2。重放所有21个生成索引器的层的最终查询块。
• 实验结果：在所有60个配置中，DOSA都是最快的方法（图4）。相比SGLang，加速比在$1.12\times$到$2.74\times$之间（在512K及以上均超过$2\times$）。相比LiteTopK，加速比在$1.17\times$到$1.69\times$之间。在1M token时，相较于两个基线的加速比分别达到$2.47\mathrm{-}2.74\times$和$1.47\mathrm{-}1.69\times$。
• 分析结论：上下文长度对增益的驱动作用强于查询块大小。DOSA从减少长键维度上的工作中获益最多。与LiteTopK的对比表明，仅融合评分和选择而不减少计算是不够的，DOSA通过认证拒绝移除了剩余的计算工作。

2. 解码性能

• 实验内容：构建长约1M token的KV cache，使用HELMET【20，HELMET+2024+arXiv】、InfiniteBench、LongBench v2【1，LongBench v2+2024+arXiv】和SQLite代码进行解码测试。
• 实验结果：在256K、512K和约1M请求长度下，DOSA分别将解码索引器加速了$1.06\times$、$1.11\times$和$1.26\times$（图5）。
• 分析结论：增益来自于减少选择工作：最近选择的键为紧凑选择提供了阈值，使解码受益。

3. 各层的修剪率（Pruning across layers）

• 实验内容：定义修剪率$p = 1 - N_{\mathrm{keep}} / N_{\mathrm{screen}}$，统计筛选前后的查询-键对数量。
• 实验结果：在三个领域的21层中，有20层的修剪率超过$95\%$，领域平均值为$97.2\% - 97.5\%$（图6）。所有领域在第1层的修剪率最低（$89\% - 93\%$）。
• 分析结论：这种变化跟随模型深度，表明层依赖的分数几何形状比输入领域的影响更大。大多数键远低于Top-$K$边界，证书能够成功拒绝它们。

4. 端到端推理性能

• 实验内容：在TP8下测量GLM-5.2和DeepSeek-V3.2在128K到1M token书籍和代码工作负载上的首字生成时间（TTFT）。
• 实验结果：在1M token时，DOSA几乎将两个模型的TTFT减半（图8）。在GLM-5.2和DeepSeek-V3.2上，相较于SGLang分别获得$1.86\times$和$1.94\times$的加速；相较于LiteTopK分别获得$1.26\times$和$1.50\times$的加速。
• 分析结论：随着上下文增长，索引器所占的时间比例增加（根据阿姆达尔定律），因此加速索引器能够移除更大比例的请求延迟，证明了针对长上下文服务优化索引器的价值。

5. 消融实验

• 实验内容：在LOFT【11，Can Long-Context Language Models Subsume...+2024+arXiv】的SciFact检索示例上，独立启用筛选（Dense）和稀疏写回（Full），并与Base（均不启用）和DOSA（均启用）对比。
• 实验结果：筛选单独降低$37.4\%$延迟，稀疏写回单独降低$23.4\%$。两者结合（DOSA）降低$47.3\%$，带来$1.90\times$加速（图7）。
• 分析结论：Full方法避免了密集写回但仍为每个头评估付费；Dense方法减少了评分评估但密集的中间结果消耗了HBM带宽。这表明修剪需要匹配的执行策略才能实现其收益。

(注：论文中Figure 5, 6, 7共用排版，此处引用原始图片路径)

结论

DOSA证明了精确的Token选择所需的工作量可以少于计算每个Token的完整分数。通过带符号头的证书（Signed-head certificates）、原始分数补全和稀疏写回，DOSA将这一观察转化为注意力索引器内部的条件计算。查询的几何形状决定了有多少工作可以被证明是不必要的；头预算分析使得获取该证书的成本变得明确。TTFT结果显示了在长上下文下不断增加的优势，而工作负载分析解释了筛选成本与避免的全头评分之间的权衡。未来的工作可能包括：在训练时对证书松弛度进行正则化，以鼓励学习出更容易安全跳过计算的索引器；以及将每键证书扩展为使用紧凑摘要的键块（key-block）边界，从而在上下文增长时同时降低评分和键访问成本。

附录细节

A 证书有效性与精确选择

本节在实数算术下建立分数边界，并证明固定查询在可见键上的精确选择。

A.1 分数证书的有效性

• 命题1（正向头投影界限）：对于任意$\beta_h \geq 0$和分解$\mathbf{q}_h = \beta_h \mathbf{q}_j + \mathbf{d}_h$，有
  

  $$ [\mathbf{q}_h^\top \mathbf{k}]_+ \leq \beta_h [\mathbf{q}_j^\top \mathbf{k}]_+ + \left\| \mathbf{d}_h \right\|_2 \left\| \mathbf{k} \right\|_2. $$

  
  证明：对于实数$a, b$和$c \geq 0$，有$[a + b]_+ \leq [a]_+ + [b]_+$以及$[ca]_+ = c[a]_+$。因此$[\beta_h \mathbf{q}_j^\top \mathbf{k} + \mathbf{q}_h^\top \mathbf{k}]_+ \le \beta_h [\mathbf{q}_j^\top \mathbf{k}]_+ + [\mathbf{q}_h^\top \mathbf{k}]_+ \le \beta_h [\mathbf{q}_j^\top \mathbf{k}]_+ + \|\mathbf{d}_h\|_2 \|\mathbf{k}\|_2$。乘以$w_h > 0$并在$P$上求和即得到公式(9)的正向部分。任何有限的非负系数只要其残差被保守地界定，就是有效的。

• 命题2（负向头支持界限）：对于任意$A \subseteq N$，$[\mathbf{a}_A^\top \mathbf{k}]_+ \leq \sum_{h \in N} |w_h| [\mathbf{q}_h^\top \mathbf{k}]_+$。
  证明：对于$c_h \geq 0$，ReLU子可加性给出$\big[ \sum_{h \in A} c_h x_h \big]_+ \leq \sum_{h \in A} c_h \big[ x_h \big]_+ \leq \sum_{h \in N} c_h \big[ x_h \big]_+$。代入$c_h = |w_h|$和$x_h = \mathbf{q}_h^\top \mathbf{k}$。该结果对任何数据依赖的支持（包括$A = \emptyset$）均逐点成立。

• 分组正向残差：将坐标划分为$G$个不相交的组。在实数算术中，令$E_{P,g} = \sum_{h \in P} w_h \left\| \mathbf{d}_h^{(g)} \right\|_2$。然后
  

  $$ \sum_{h \in P} w_h [\mathbf{d}_h^\top \mathbf{k}_i]_+ \leq \sum_{g=1}^G E_{P,g} \left\| \mathbf{k}_i^{(g)} \right\|_2 \leq \left( \sum_g E_{P,g}^2 \right)^{1/2} \left\| \mathbf{k}_i \right\|_2. $$

  
  闵可夫斯基不等式（Minkowski’s inequality）给出$\left( \sum_g E_{P,g}^2 \right)^{1/2} \leq E_P$。这通过每个键一个RMS值紧缩了残差系数。

• 联合正向残差：对于一个坐标组，写$\mathbf{u}_h = w_h \mathbf{d}_h^{(g)}$并定义$B_g = \frac{1}{2} \left( \left\| \sum_{h \in P} \mathbf{u}_h \right\|_2 + \sqrt{\sum_{h, \ell \in P} |\mathbf{u}_h^\top \mathbf{u}_\ell|} \right)$。

• 命题3（联合残差容限）：对于每个$\mathbf{k}^{(g)}$，$\sum_{h \in P} [\mathbf{u}_h^\top \mathbf{k}^{(g)}]_+ \leq B_g \left\| \mathbf{k}^{(g)} \right\|_2$，且$B_g \leq \sum_{h \in P} \left\| \mathbf{u}_h \right\|_2 = E_{P,g}$。
  证明：使用$[x]_+ = (x + |x|) / 2$展开，绝对值和等于$\max_{\sigma_h \in \{-1, 1\}} (\sum_h \sigma_h \mathbf{u}_h)^\top \mathbf{k}^{(g)}$。对于每个符号向量，$\left\| \sum_h \sigma_h \mathbf{u}_h \right\|_2^2 \leq \sum_{h, \ell} |\mathbf{u}_h^\top \mathbf{u}_\ell|$。柯西-施瓦茨不等式证明了第一个主张。三角不等式证明了第二个。

A.2 固定阈值算法与精确Top-$K$选择

• 命题4（子集阈值）：对于任意$\mathcal{S} \subseteq \Omega$且$|\mathcal{S}| \ge K$，$\tau_K(\mathcal{S}) \leq \tau_K(\Omega)$。

• 算法1步骤：

  1. 选择不同的可见种子键$\mathcal{S} \subseteq \Omega$，$|\mathcal{S}| \ge K$。
  2. 使用原始评分器计算分数$s$；设$\tau = \tau_K(\mathcal{S})$。
  3. 构建$R$、锚点、$A$和公式(9)中的分数边界。
  4. 筛选：$V \gets \{i \in \Omega \setminus \mathcal{S} : U_i \geq \tau\}$。
  5. 使用原始$H$头评分器对每个$i \in V$进行评分。
  6. 细化：$C \gets \{i \in V : s_i \geq \tau\}$。
  7. 返回$\mathrm{TopK}(\mathcal{S} \cup C)$。

• 定理1（精确Top-$K$选择）：在算法1中，假设在筛选期间被拒绝的每个键满足$s_i \leq U_i < \tau$，且没有丢失存活的候选者。在有共同的总平局顺序下，该算法返回与全域评分相同的Top-$K$。
  证明：考虑任何键$i \notin \mathcal{S} \cup C$。如果筛选拒绝了它，则$s_i \leq U_i < \tau$。否则它属于$V \setminus C$且原始分数满足$s_i < \tau$。在两种情况下，$s_i < \tau_K(\mathcal{S}) \leq \tau_K(\Omega)$。因此被丢弃的键不可能属于全域Top-$K$。

B 证书紧密性与计算成本

B.1 命题5（最优单锚点容限）：

• 对于向量$\mathbf{q}, \mathbf{p} \in \mathbb{R}^D$，考虑所有满足$[\mathbf{q}^\top \mathbf{k}]_+ \leq \beta [\mathbf{p}^\top \mathbf{k}]_+ + c \|\mathbf{k}\|_2$的对$\beta, c \geq 0$。可实现的最小系数是$c^* = \min_{\beta \geq 0} \|\mathbf{q} - \beta \mathbf{p}\|_2$。对于$\mathbf{p} \neq 0$，它由$\beta^* = \frac{[\mathbf{q}^\top \mathbf{p}]_+}{\|\mathbf{p}\|_2^2}$和$c^* = \|\mathbf{q} - \beta^* \mathbf{p}\|_2$达到。
  证明：通过反证法，代入特定方向的$\mathbf{k}$，证明不存在更小的$c$。这优化了一致范数容限。

B.2 证书松弛与感知激活的头选择

• 命题6（精确证书松弛）：公式(9)中的界限包含三个非负松弛分量：正向投影松弛$\delta_P(\mathbf{k})$，错过的活跃负向头松弛$\sum_{h \in N \setminus A} c_h [x_h]_+$，以及聚合抵消松弛$\min\{Z_A^+(\mathbf{k}), Z_A^-(\mathbf{k})\}$。
  证明：保留的贡献在$U - s$中抵消。正向剩余为$\delta_P \geq 0$。负向剩余为$\delta_N(\mathbf{k}) = \sum_{h \in N} c_h [x_h]_+ - [Z_A^+ - Z_A^-]_+$，利用$z - [z - y]_+ = \min(z, y)$展开即得。

• 支持与保留头选择：对于任何支持，负向松弛服从$0 \leq \delta_N(\mathbf{k}) \leq \|\mathbf{k}\|_2 \sum_{h \in A \triangle M(\mathbf{k})} a_h$，其中$a_h = c_h \|\mathbf{q}_h\|_2$。对其求平均得到可分离的替代目标$J(A) = \sum_{h \in A} a_h (1 - p_h) + \sum_{h \in N \setminus A} a_h p_h$。逐分量多数支持$A^* = \{h : p_h \geq 1/2\}$最小化了$J$。

B.3 选择边界与筛选成本

• 拒绝条件：设$\tau^*$为全域第$K$大分数，$\Delta_i = \tau^* - s_i$为选择裕度（margin），$b = \tau^* - \tau \geq 0$为阈值赤字，$\delta_i = U_i - s_i \ge 0$为边界松弛。则$U_i < \tau \Longleftrightarrow \Delta_i > \delta_i + b$。
• 算术成本与头预算：设$N_0$个键在初始化时被完全评分，$N_s$个被筛选，$M$个存活。归一化工作量为$D_{\mathrm{DOSA}} / D_{\mathrm{base}} \approx f + (1 - f) \left( \frac{r + d_a}{H} + m \right)$。只有当存活率$m < 1 - (r + d_a) / H$时，筛选才能节省逻辑点积。

C 数值条件与恢复

C.1 参考分数与FP8证书

• 必须独立建立数值支配关系$s_i^{\mathrm{ref}} \leq U_i^{\mathrm{num}}$。
• 在基本形式中，两个辅助行近似负向聚合：$\mathbf{a}_A = \mathbf{v}_0 + \mathbf{v}_1 + \mathbf{e}_A$。筛选表达式为$\widehat{U}_i = \mathrm{fl}_{\mathrm{screen}} \left( \sum_{j \in R} \omega_j [\mathbf{q}_j^\top \bar{\mathbf{k}}_i]_+ - [(\mathbf{v}_0 + \mathbf{v}_1)^\top \bar{\mathbf{k}}_i]_+ \right)$。
• 绝对值形式：当保留所有正权重头时，设$\mathbf{V} = \mathbf{a}_A - \frac{1}{2} \sum_{j \in R} w_j \mathbf{q}_j$。使用绝对值形式，额外的松弛为$[-\mathbf{a}_A^\top \mathbf{k}]_+$。32头实现使用此形式（$r=14,16$）。
• 数值证书形式为：$U_i^{\mathrm{num}} = \mathrm{Eval}_{\mathrm{up}} \Big[ \alpha_i \big( \widehat{U}_i + \rho_i C + \varepsilon_{\mathrm{raw}} \big) + \varepsilon_{\mathrm{out}} \Big]$。

C.2 误差容限与保守键摘要

• 令$E_g^+$为有效正向残差容限的上限包络。假设存在非负坐标包络$\mathbf{J}$和绝对容限，满足$s_i^{\mathrm{ref}} \leq \alpha_i \Biggl( \widehat{U}_i + \sum_g E_g^+ \left\| \bar{\mathbf{k}}_i^{(g)} \right\|_2 + \sum_d J_d |\bar{k}_{i,d}| + \varepsilon_{\mathrm{raw}} \Biggr) + \varepsilon_{\mathrm{out}}$。
• 充分的RMS系数为$C \geq \sqrt{D} \left( \sum_g \big[ E_g^+ + \Big\| \mathbf{J}^{(g)} \Big\|_2 \big]^2 \right)^{1/2}$。
• 保守键RMS：存储的摘要必须满足$\rho_i \sqrt{D} \geq \|\bar{\mathbf{k}}_i\|_2$，需要向上覆盖。

C.3 候选者存储与全域恢复

• 缓冲区容量不能限制满足$U_i \geq \tau$的键的数量。如果缓冲区满，必须通过精确的Top-$K$合并进行处理，丢弃溢出的候选者会导致丢失ID。无效的缩放、RMS摘要或非有限边界会阻止修剪，查询级错误会触发全域恢复（重新评估所有可见键）。

D 额外实验分析

D.1 索引器分析方法

• 图1使用真实模型权重和文本，禁用了DOSA和LiteTopK。分析排除了CPU分词和请求传输。FLOPs计算因果逻辑矩阵工作。

D.2 证书组件消融实验

• 实验内容：在GLM-5.2的21层中重放最后4K个查询。对比聚合所有省略的负向头、移除负向聚合、仅按强度选择保留的负向头、移除正向头投影。
• 实验结果：移除负向聚合会使候选者存活率从$2.65\%$飙升至$95.06\%$，索引器时间增加$9.51\times$（图9）。聚合所有负向头也会削弱过滤（存活率$4.06\%$，延迟增加$12.0\%$）。
• 结论：负向聚合对于获得带符号索引器分数的紧密证书是必不可少的。

D.3 逐层证书行为

• 实验结果：移除负向聚合在大多数层中几乎留下所有候选者（图10）。聚合所有负向头的成本在早期层最为明显：在第2层，其相对索引器延迟达到$3.41\times$。

D.4 逐层检索加速

• 实验内容：重放最后32K查询跨越所有21个索引器层（252个单元格）。
• 实验结果：在252次比较中，DOSA在251次（$99.6\%$）中快于SGLang，在250次（$99.2\%$）中快于LiteTopK。所有21层在至少一个输入上相较于SGLang实现了至少$2\times$的加速，最高达$3.74\times$；相较于LiteTopK的加速最高达$2.08\times$（图11和图12）。