DoSA 完全拆解:用可证明的提前拒绝加速精确稀疏注意力索引

9110 字
46 分钟
DoSA 完全拆解:用可证明的提前拒绝加速精确稀疏注意力索引

稀疏注意力省下了 KV 访问,却留下了索引账单

长上下文稀疏注意力通常被概括为“先选少量 token,再只对这些 token 做注意力”。这个说法容易掩盖一个越来越重要的事实:选择本身也可能是一遍稠密计算。如果每个 query 都要用若干索引头扫描全部历史 key,计算完整分数后再取 Top-K,那么真正的注意力虽然只访问 2048 个 token,索引器仍然访问了十万到百万个 token。

小红书 Dots Studio 的 Fanxu Meng 与 Taosong Fang 于 2026 年 9 月发布 DoSA: Dots Sparse Attention。DoSA 不改变模型学到的索引分数,也不把精确 Top-K 换成近似 Top-K;它改变的是“什么时候必须把一个候选的所有索引头都算完”。其关键观察可以压缩成一句话:淘汰一个候选只需要证明它不可能越过当前门槛,不需要知道它的精确分数。

长上下文下 Indexer 的时间与 FLOPs 占比
长上下文下 Indexer 的时间与 FLOPs 占比

图 1 来自 DoSA 论文。左图统计 query/key 准备、打分与 Top-K 在内的完整 Indexer 时间占比,右图统计理论矩阵 FLOPs 占比。在约 100 万 token 上下文中,Indexer 占 GLM-5.2 预填充 GPU 时间的 73.0%,在 DeepSeek-V3.2 上达到 90.8%;对应矩阵 FLOPs 占比分别为 47.4% 和 83.2%。这说明瓶颈不只来自乘加,还来自 score 写回、候选整理和 Top-K 选择。

DoSA 的答案不是“少算几个头然后祈祷结果相同”,而是构造一个对每个候选都成立的分数上界 UiU_i。只要 UiU_i 已低于一个可靠的获胜阈值下界 τ\tau,就可以严格断定该候选不在 Top-K 中。只有上界仍有机会越过门槛的候选,才执行全部索引头并计算原始分数。

这个思路把稀疏注意力的索引阶段再稀疏化了一次:第一层稀疏决定注意力读取哪些 KV,第二层条件执行决定哪些候选值得算完整索引分数。后者就是 DoSA 的“Dots”——完整计算集中在少量仍有竞争力的 query-key 点上。

从 DSA 索引器到 DoSA 的问题定义

DoSA 针对的是 DeepSeek Sparse Attention 一类带多头索引器的细粒度稀疏注意力。关于 DSA 的结构、选择分支与 Kernel 背景,可参考 NSA 原生稀疏注意力(二);这里只保留理解 DoSA 必需的打分定义。

对固定 query,索引器有 HH 个头。第 hh 个 query 向量记为 qh∈RD\mathbf q_h\in\mathbb R^D,共享候选 key 记为 ki∈RD\mathbf k_i\in\mathbb R^D,索引头的有符号权重为 whw_h。候选 ii 的完整索引分数是

si=∑h=1Hwh[qh⊤ki]+,[x]+=max⁡(x,0).s_i=\sum_{h=1}^{H}w_h\left[\mathbf q_h^\top\mathbf k_i\right]_+, \qquad [x]_+=\max(x,0).

其中 [⋅]+[\cdot]_+ 是 ReLU。一个头先做点积,再截去负值,最后乘上可能为正也可能为负的 whw_h。可见这里不是普通的单向量最大内积搜索:正权重头会把分数向上推,负权重头会把分数向下拉,而且 ReLU 使“先把向量相加再点积”与“分别点积再相加”通常不相等。

设因果遮罩后可见 key 的集合为 Ω\Omega,目标是返回原始分数最高的 KK 个 token。GLM-5.2 与 DeepSeek-V3.2 的论文实验均取 K=2048K=2048,但前者有 32 个索引头,后者有 64 个索引头。直接实现大致包含三步:

  1. 对每个 i∈Ωi\in\Omega 计算所有 HH 个点积和完整分数;
  2. 把分数写入显存或直接流入融合选择器;
  3. 在全部候选中执行精确 Top-K。

SGLang 基线会计算完整多头分数再做 Top-K。LiteTopK 已经把打分、候选过滤和选择融合起来,显著减少中间分数的全量写回;但它仍然为每个 key 计算所有索引头。DoSA 要削掉的正是这部分仍然稠密的多头乘加。

这里还要区分三种常被混用的“无损”:

  • 模型层面的无损:不改参数,不重新训练,不替换索引函数;
  • 选择层面的无损:返回与全域原始打分相同的 Top-K,或在并列分数下返回一个合法 Top-K;
  • 数值实现层面的无损:低精度与舍入误差不能让理论上应保留的候选被错误剪掉。

DoSA 的理论覆盖前两项,并在附录中给出实现第三项所需的保守数值条件。它不是把误差“做得很小”,而是要求误差余量也进入上界,使剪枝方向始终安全。

核心思想:先证明候选不可能入选,再决定是否算完

DoSA 首先选出一个至少包含 KK 个可见 key 的种子集合 S⊆Ω\mathcal S\subseteq\Omega,对种子使用原始 HH 头打分。令

τ=τK(S)\tau=\tau_K(\mathcal S)

表示种子分数中的第 KK 大值。因为 S\mathcal S 只是全体候选的子集,所以全域第 KK 大分数 τK(Ω)\tau_K(\Omega) 不可能比它更低:

τK(S)≤τK(Ω).\tau_K(\mathcal S)\leq\tau_K(\Omega).

这就是“获胜阈值的下界”。它未必等于真实门槛,但绝不会高估门槛。接着,若能为候选 ii 构造满足 si≤Uis_i\leq U_i 的上界,就有

Ui<τ⟹si<τ≤τK(Ω).U_i<\tau \quad\Longrightarrow\quad s_i<\tau\leq\tau_K(\Omega).

因此 ii 不可能进入 Top-K。注意剪枝条件必须是严格小于:当 Ui=τU_i=\tau 时,候选可能恰好落在并列边界,仍须保留。

DoSA 的四阶段条件执行流程
DoSA 的四阶段条件执行流程

图 2 来自 DoSA 论文,按编号可分为四个阶段:

  1. 筛选点积:只计算一组保留头与少量辅助方向;
  2. 上界判定:由已算点积、key 范数余量与负头聚合项构造 UiU_i,剪掉 Ui<τU_i<\tau 的候选并压紧 survivor ID;
  3. 完整细化:只对 survivor 计算全部 HH 个原始头;
  4. 稀疏回写:恢复原始 ReLU 与有符号加权,仅回写 si≥τs_i\geq\tau 的精确分数,再与种子一起执行最终 Top-K。

灰色块表示被 certificate 省掉的完整打分,紫色块表示被精确阈值过滤省掉的回写。Tensor Core 负责规则的矩阵乘阶段,SIMT 线程从 Tensor Memory 中取累加器,完成上界、压紧、ReLU 和归约。

这套流程有两个彼此独立的筛子。第一道筛子用便宜上界减少完整算分;第二道筛子用已经得到的精确分数减少写回和 Top-K 输入。只做第一道而仍然密集写分数,会被 HBM 流量拖住;只做第二道则不会减少多头点积。论文的消融实验正好验证了二者互补。

一个最小数值例子

假设要选 Top-2,种子集合中两个精确分数为 8.1 和 5.7,于是 τ=5.7\tau=5.7。扫描到三个新候选时,便宜 certificate 分别给出上界 4.9、6.2 和 9.0。

  • 第一个候选即使取到上界也只有 4.9,可以立刻拒绝,无需知道它究竟是 4.8 还是负数;
  • 第二、第三个候选的上界超过门槛,只表示“尚不能证明落选”,并不表示它们一定获胜;
  • 完整细化后,若二者精确分数为 5.1 和 7.0,5.1 还会被第二道阈值过滤挡住,最终只需把 7.0 写给选择器;
  • 最后在种子 8.1、5.7 与新候选 7.0 中取 Top-2,得到 8.1 和 7.0,与三名新候选全部完整打分的结果一致。

这个例子也说明了 UiU_i 的含义:它是淘汰证书,不是近似 score。系统不会用 6.2 或 9.0 参与最终排序;凡是可能入选的候选,最终都必须恢复原始精确分数。

Certificate 的数学构造

真正困难的部分,是用少于 HH 个点积得到一个足够紧、又永远不低估 sis_i 的上界。DoSA 把索引头按权重符号分成

H+={h:wh>0},H−={h:wh<0}.\mathcal H_+=\{h:w_h>0\}, \qquad \mathcal H_-=\{h:w_h<0\}.

再选一个预算为 rr 的保留头集合 RR,直接计算

SR(ki)=∑j∈Rwj[qj⊤ki]+.S_R(\mathbf k_i)=\sum_{j\in R}w_j \left[\mathbf q_j^\top\mathbf k_i\right]_+.

未保留的正头记作 P=H+∖RP=\mathcal H_+\setminus R,未保留的负头记作 N=H−∖RN=\mathcal H_-\setminus R。正头不能直接丢弃,因为它们可能抬高真实分数;负头若直接丢弃,上界仍安全,却会因为少减了一大截而过松。DoSA 分别处理这两个问题。

正头:投影到已计算的锚点

对每个未保留正头 h∈Ph\in P,从 RR 中挑一个方向最相近的锚点 jh⋆j_h^\star,并只允许非负投影:

jh⋆=argmax⁡j∈R[qh⊤qj]+∥qj∥2,j_h^\star=arg\max_{j\in R} \frac{[\mathbf q_h^\top\mathbf q_j]_+}{\|\mathbf q_j\|_2},βh=[qh⊤qjh⋆]+∥qjh⋆∥22,dh=qh−βhqjh⋆.\beta_h= \frac{[\mathbf q_h^\top\mathbf q_{j_h^\star}]_+} {\|\mathbf q_{j_h^\star}\|_2^2}, \qquad \mathbf d_h=\mathbf q_h-\beta_h\mathbf q_{j_h^\star}.

βh≥0\beta_h\geq0 很关键,因为 ReLU 只对非负缩放满足正齐次性。将 qh\mathbf q_h 写成“锚点方向 + 残差”后,利用 ReLU 的次可加性和 Cauchy–Schwarz 不等式可得

[qh⊤ki]+≤βh[qjh⋆⊤ki]++∥dh∥2∥ki∥2.[\mathbf q_h^\top\mathbf k_i]_+ \leq \beta_h[\mathbf q_{j_h^\star}^\top\mathbf k_i]_+ +\|\mathbf d_h\|_2\|\mathbf k_i\|_2.

第一项复用已计算的锚点点积;第二项只依赖 query 侧残差范数和每个 key 的范数摘要。该 βh\beta_h 还是所有单锚点、非负投影形式中使统一残差系数最小的选择:几何上,它把 qh\mathbf q_h 投影到由锚点生成的非负射线,∥dh∥2\|\mathbf d_h\|_2 就是到这条射线的最短距离。

把所有未保留正头的投影项折叠进锚点权重,定义

w~j=wj+∑h∈Pjh⋆=jwhβh,EP=∑h∈Pwh∥dh∥2.\widetilde w_j=w_j+ \sum_{\substack{h\in P\\j_h^\star=j}}w_h\beta_h, \qquad E_P=\sum_{h\in P}w_h\|\mathbf d_h\|_2.

则保留头与省略正头的总正向贡献至多为

∑j∈Rw~j[qj⊤ki]++EP∥ki∥2.\sum_{j\in R}\widetilde w_j [\mathbf q_j^\top\mathbf k_i]_+ +E_P\|\mathbf k_i\|_2.

w~j\widetilde w_j 只是本次筛选计划的折叠系数,不会修改模型权重。若所有正头均已保留,则 PP 为空,正头残差余量也退化为零。

联合残差:不要把相关方向当成各自最坏

逐头使用三角不等式会假设所有残差都能同时取到各自最坏方向,通常偏保守。令 uh=whdh\mathbf u_h=w_h\mathbf d_h,DoSA 给出联合余量

EPjoint=12(∥∑h∈Puh∥2+∑h,ℓ∈P∣uh⊤uℓ∣).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).

它满足

∑h∈P[uh⊤ki]+≤EPjoint∥ki∥2,EPjoint≤EP.\sum_{h\in P}[\mathbf u_h^\top\mathbf k_i]_+ \leq E_P^{\mathrm{joint}}\|\mathbf k_i\|_2, \qquad E_P^{\mathrm{joint}}\leq E_P.

直觉是:所有残差共享同一个 ki\mathbf k_i,不同残差方向之间会相互制约,不可能都把同一个 key 当成自己的最优方向。联合界用残差 Gram 矩阵中的方向关系收紧余量,却不增加 key 读取。代价发生在 query 侧规划阶段,需要处理省略正头之间的二次配对。

精确 Top-K 为什么仍然无损

在解释负权重头之前,先把 DoSA 的正确性闭环写清楚。算法可以抽象成七步:

输入:可见 key 集合 Ω、Top-K 大小 K、原始 H 头索引器
1. 选择至少 K 个不同的可见种子 S
2. 用原始索引器计算 S 的精确分数,取其中第 K 大值 τ
3. 建立保留头、锚点、负头聚合支持集和安全上界 U
4. V = {i ∈ Ω \ S : U_i ≥ τ}
5. 对 V 中每个候选计算原始 H 头精确分数 s_i
6. C = {i ∈ V : s_i ≥ τ}
7. 在 S ∪ C 的原始分数上执行 Top-K

对任何未进入 S∪C\mathcal S\cup C 的候选只有两种情况:要么它在筛选阶段满足 si≤Ui<τs_i\leq U_i<\tau;要么它经过完整细化,但原始分数 si<τs_i<\tau。两者都得到

si<τK(S)≤τK(Ω).s_i<\tau_K(\mathcal S)\leq\tau_K(\Omega).

因此,被丢弃的候选不可能严格高于全域第 KK 大分数。最终集合至少包含 KK 个种子,又包含所有可能获胜者,在同一 tie-break 全序下重新选择就会得到与全量打分完全相同的 token ID。若系统不要求并列项的 ID 也完全相同,则结果至少是原始分数下的一个合法 Top-K 集合。

这里有四个容易破坏“无损”的工程细节:

  • 种子必须是至少 KK 个不同且可见的 key,不能让重复 ID 虚增见证数量;
  • Ui=τU_i=\tau 的候选必须保留,只有严格小于才可剪枝;
  • survivor 或精确分数缓冲区溢出时不能直接丢记录,必须扩容、分块做精确 Top-K 合并,或回退到全域原始打分;
  • 一旦缩放、RMS 摘要或数值上界失效,回退必须覆盖所有可见 key,不能只在尚未丢失的局部缓冲区重选。

尤其是缓冲区容量。最终只返回 KK 个 token,不代表中间候选最多也只有 KK 个。极端情况下所有 score 都等于 τ\tau,每个候选都必须存活。论文附录报告,为修复约 1M-key 容量下的候选存储,额外预留了每个 query 1024 个 candidate-ID page 与 114688 条输出记录,相对原分配增加 9.94 GiB workspace。这是 exactness 的真实系统成本,而不是可以忽略的实现脚注。

种子门槛也体现了正确性与效率的分离。只要 S\mathcal S 合法,再差的 seed 也不会导致错误;它只会让 τ\tau 太低、更多候选进入细化。增加高质量见证或更新门槛可以提升效率,但也会带来额外原始打分与同步。DoSA 选择先从可见种子建立固定下界,再让大规模筛选保持规则。换句话说,certificate 的安全性不依赖 seed“代表整体分布”,性能才依赖。

从实数证明走到 FP8

理论式 si≤Uis_i\leq U_i 发生在实数域,而 GPU 上的原始索引器本身可能使用 FP8、缩放因子与特定舍入顺序。实现真正需要保证的是

siref≤Uinum,s_i^{\mathrm{ref}}\leq U_i^{\mathrm{num}},

其中 sirefs_i^{\mathrm{ref}} 是原始 Kernel 实际算出的参考分数,UinumU_i^{\mathrm{num}} 是包含量化、重构和舍入余量的数值上界。论文把 FP8 key 写成 ki=αikˉi\mathbf k_i=\alpha_i\bar{\mathbf k}_i,并保存满足

ρiD≥∥kˉi∥2\rho_i\sqrt D\geq\|\bar{\mathbf k}_i\|_2

的保守 RMS 摘要。最终上界形式是

Uinum=Eval⁡up[αi(U^i+ρiC+εraw)+εout].U_i^{\mathrm{num}} =\operatorname{Eval}_{\mathrm{up}} \left[ \alpha_i\left(\widehat U_i+\rho_i C+\varepsilon_{\mathrm{raw}}\right) +\varepsilon_{\mathrm{out}} \right].

U^i\widehat U_i 是筛选点积得到的主体,CC 汇总正残差、辅助向量重构和权重误差,εraw\varepsilon_{\mathrm{raw}} 与 εout\varepsilon_{\mathrm{out}} 分别覆盖缩放前后的绝对误差,Eval⁡up\operatorname{Eval}_{\mathrm{up}} 表示用定向上舍入或显式误差界向上包住结果。普通“就近舍入”不具备这个保证;RMS 摘要若只是最近舍入,也可能低估真实范数。

当所有正权重头都已保留时,还可以利用 [x]+=(x+∣x∣)/2[x]_+=(x+|x|)/2 改写上界。令

V=aA−12∑j∈Rwjqj,\mathbf V=\mathbf a_A- \frac{1}{2}\sum_{j\in R}w_j\mathbf q_j,

则可使用

Uabs(k)=12∑j∈Rwj∣qj⊤k∣−V⊤k.U_{\mathrm{abs}}(\mathbf k) =\frac{1}{2}\sum_{j\in R}w_j |\mathbf q_j^\top\mathbf k| -\mathbf V^\top\mathbf k.

它相对基础式会多出 [−aA⊤k]+[-\mathbf a_A^\top\mathbf k]_+ 的 slack,但便于把聚合方向编码成两个缩放 FP8 辅助行。论文的 32 头实现对 r=14,16r=14,16 使用这一绝对值路径,r=20r=20 则保留带正头投影的基础路径。两种形式都是 r+2r+2 个逻辑筛选点积,体现了数学紧度与 Kernel 形状之间的权衡。

所以 DoSA 的无损不是“实测选出来都一样”,而是由集合阈值、逐候选 certificate、保留边界等号、上舍入数值余量与失败回退共同组成的契约。

Activation-informed aggregation 如何约束负向贡献

负权重头是 DoSA certificate 最容易被低估、也最关键的一部分。真实分数中,未保留负头贡献为

−∑h∈N∣wh∣[qh⊤ki]+.-\sum_{h\in N}|w_h| [\mathbf q_h^\top\mathbf k_i]_+.

若把它们全部忽略,相当于用零代替一个非正量,得到的当然仍是上界。但这个上界会高得几乎剪不掉任何候选。DoSA 从 NN 中挑一个聚合支持集 AA,构造聚合向量

aA=∑h∈A∣wh∣qh.\mathbf a_A=\sum_{h\in A}|w_h|\mathbf q_h.

ReLU 次可加性给出

[aA⊤ki]+≤∑h∈N∣wh∣[qh⊤ki]+.[\mathbf a_A^\top\mathbf k_i]_+ \leq \sum_{h\in N}|w_h| [\mathbf q_h^\top\mathbf k_i]_+.

左侧是“至少会被负头减掉多少”的下界。因此,从正向上界中减去它仍然安全,而且任意多的未保留负头只需增加一个聚合方向的点积。把正头与负头合并,得到 DoSA 的核心 certificate:

Ui=∑j∈Rw~j[qj⊤ki]++EP∥ki∥2−[aA⊤ki]+,si≤Ui.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]_+, \qquad s_i\leq U_i.

为什么不是把所有负头都聚合起来

看似 A=NA=N 应该减得最多,事实恰好不是。假设两个负头分别产生点积 55 与 −4-4。逐头 ReLU 后,真实扣减包含 [5]++[−4]+=5[5]_++[-4]_+=5;若先聚合再 ReLU,只有 [5−4]+=1[5-4]_+=1。未激活头的负点积会在聚合 ReLU 之前抵消激活头的正点积,让下界变小、最终上界变松。

DoSA 用种子 key 上的激活模式估计哪些负头经常同时激活。对每个负头记录激活频率 php_h,选择逐头多数支持集

A={h∈N:ph≥1/2}.A=\{h\in N:p_h\geq 1/2\}.

这个选择不是经验拍脑袋。令 ah=∣wh∣∥qh∥2a_h=|w_h|\|\mathbf q_h\|_2,可把负侧归一化 slack 的可分上界写成

J(A)=∑h∈Aah(1−ph)+∑h∈N∖Aahph.J(A)= \sum_{h\in A}a_h(1-p_h) +\sum_{h\in N\setminus A}a_hp_h.

每个头若加入 AA,付出的是“它不激活”的概率代价 ah(1−ph)a_h(1-p_h);若不加入,付出的是“漏掉它激活”的概率代价 ahpha_hp_h。逐头比较两者,阈值自然就是 ph=1/2p_h=1/2。

若正头数量不超过预算 rr,DoSA 先保留全部正头,再用剩余名额保留那些最难由多数模式预测的负头。优先级为

∣wh∣∥qh∥2min⁡(ph,1−ph).|w_h|\|\mathbf q_h\|_2\min(p_h,1-p_h).

权重大、query 范数大且激活概率接近一半的负头最危险:它既影响大,又没有稳定的开关模式,适合直接计算。若正头数量超过预算,则先按 wh∥qh∥2w_h\|\mathbf q_h\|_2 保留最强正头,因为遗漏正头必须用范数残差兜底。

三种 slack 决定剪枝能力

对一个 key 记 xh=qh⊤kx_h=\mathbf q_h^\top\mathbf k、ch=∣wh∣c_h=|w_h|。论文把 U−sU-s 精确分成三个非负项:

U(k)−s(k)=δP(k)+∑h∈N∖Ach[xh]++min⁡{ZA+(k),ZA−(k)}.U(\mathbf k)-s(\mathbf k) =\delta_P(\mathbf k) +\sum_{h\in N\setminus A}c_h[x_h]_+ +\min\{Z_A^+(\mathbf k),Z_A^-(\mathbf k)\}.

第一项 δP\delta_P 是正头投影与范数余量造成的松弛;第二项是没有进入聚合集合、却在当前 key 上激活的负头;第三项是已聚合负头之间正负点积抵消造成的损失。这里

ZA+=∑h∈Ach[xh]+,ZA−=∑h∈Ach[−xh]+.Z_A^+=\sum_{h\in A}c_h[x_h]_+, \qquad Z_A^-=\sum_{h\in A}c_h[-x_h]_+.

这个分解解释了激活感知聚合为何不可替代:它直接控制后两种 slack。论文消融中,移除负头聚合后,候选存活率从 2.65% 暴涨到 95.06%,Indexer 时间增至 DoSA 的 9.51 倍;把所有负头无差别聚合,存活率也会上升到 4.06%,延迟增加 12.0%。相比之下,正头投影只在评测 query-layer 行中的 0.42% 被需要,平均影响较小,但它仍是预算不足以保留全部正头时保证正确性的必要部件。

Certificate 各组成部分对存活率与延迟的影响
Certificate 各组成部分对存活率与延迟的影响

图 3 来自论文附录的 certificate 消融。它传达的重点不是“负头越多越好”,而是负头聚合必须匹配激活支持。聚合向量的价值取决于多个头在同一 key 上是否同号激活;错误地加入经常不激活的头,会在 ReLU 之前制造抵消。

从公式到 GPU:条件执行的索引 Kernel

有了数学上很紧的上界,不代表 GPU 上一定更快。若 survivor 分布不规则,第二阶段会退化成大量离散 gather;若每个 query 独立读取 key,同一个 key 可能被相邻 query 反复搬运;若最终仍写回密集 score,省下的乘加会被 HBM 流量吃掉。DoSA 的 Kernel 因此围绕三个目标设计:保持筛选阶段规则、在细化阶段共享 key、在输出阶段只保留必要记录。

DoSA 筛选、细化与稀疏回写的数据流
DoSA 筛选、细化与稀疏回写的数据流

图 4 来自 DoSA 论文。蓝色和橙色表示相邻 tile,G2S 是 global-to-shared 传输,T2R 是 Tensor Memory 到寄存器传输。加载下一块 key 的同时计算当前块,实现流水重叠;灰色 MMA 区域是 certificate 判死后无需执行的完整多头矩阵乘。

规则的筛选通道

DoSA 按相近的保留头预算对 query 分组,装进固定宽度 tile。保留 query 行、折叠权重与两个 FP8 辅助行只准备一次,然后跨 key tile 复用。专用 warp 异步加载 key 并发起 Tensor Core MMA,consumer warp 从 TMEM 读取累加结果,计算 UinumU_i^{\mathrm{num}} 并与固定阈值 τ\tau 比较。

这里“固定阈值”是重要的工程取舍。理论上,扫描到更高精确分数时不断抬高 τ\tau 可以增强后续剪枝;但动态门槛会引入跨 tile 的同步与状态依赖。固定种子阈值较弱,却让大量 tile 独立流水化。只要 seed 足够好,较小的阈值缺口仍能换来高剪枝率。

warp ballot 把每个候选的生存布尔值压成位掩码,再将 survivor ID 压紧。筛选阶段仍扫描每个 key,但每个 key 只支付 r+dar+d_a 个逻辑点积,其中 dad_a 是辅助行数,论文 FP8 路径通常为 2;只有 survivor 才再支付完整的 HH 头。

四 query 共享细化 key

多个相邻 query 可能让同一个 key 通过 certificate。如果各自 gather,一份 key 会从全局显存重复加载。DoSA 在四个 query 的小组内把 survivor 分成两条流:若 key 至少通过两个 query 的上界,并位于共同可见范围,就进入 shared stream;其余进入 per-query stream。

shared stream 让一次 key 加载服务四个 query 的完整打分。代价是某些 query 可能额外计算一个自己并未单独选择的 query-key 对,但这笔冗余计算可能比重复搬运 key 更便宜。这个设计体现了 GPU 稀疏计算的常见规律:算术最少不等于时间最短,适度扩大规则计算常能换取更少的数据传输与更高的 Tensor Core 利用率。

精确分数过滤与压缩编码

细化阶段恢复原始 HH 头、原始 ReLU 和权重归约。得到 sis_i 后立即与各 query 的 τ\tau 比较,只有 si≥τs_i\geq\tau 才生成输出记录。记录包含单调的紧凑 score code 与精确 key ID;最终选择先保留严格高于边界 code 的候选,再对边界候选用原始分数解开歧义。

压缩编码本身不能改变次序,它的角色只是减少 HBM 写流量和 Top-K 的输入规模。真正的 exactness 仍由原始分数与统一 tie-break 保证。

Decode 为什么采用另一条路

Prefill 一次有成千上万个 query,适合分组、共享 key 和 Tensor Core tile;decode 每步 query 很少,难以摊薄规划与稀疏 gather 开销。因此 DoSA 在 decode 保留完整 FP8 多头打分,主要利用时间局部性减少选择工作:最近若干步选中的不同 key 作为见证,对当前 query 重新打分后给出保守门槛,只把不低于门槛的 score 压紧送入精确选择。

这也解释了性能差异。论文中 decode Indexer 在 256K、512K 和约 1M 请求长度下分别加速 1.06、1.11 和 1.26 倍,远小于 prefill 的最高 2.74 倍。DoSA 的主体收益来自有足够 query 并行度时跳过完整多头打分,而不是一个对所有阶段都同样有效的通用 Top-K 技巧。

实验结果该怎样解读

DoSA 的实验对象是 GLM-5.2 与 DeepSeek-V3.2,Top-K 均为 2048;前者有 32 个索引头,后者有 64 个。微基准在单卡上运行 GLM-5.2,端到端推理使用 8 路张量并行,GLM-5.2 的 prefill chunk 为 32K,DeepSeek-V3.2 为 16K。输入覆盖 InfiniteBench 的英文、中文与代码,以及 Project Gutenberg 长文本和 SQLite 源码,长度从 128K 到约 1M token。

基线有两条:SGLang v0.5.18 代表完整打分再 Top-K,LiteTopK 代表融合打分、候选过滤与选择,但仍计算全部索引头。微基准统计的是包括 query/key 准备在内的 complete Indexer 延迟;端到端统计清空 prefix cache 后单请求的 TTFT。二者不能混为同一指标。

不同上下文、输入域与 chunk 下的完整 Indexer 加速
不同上下文、输入域与 chunk 下的完整 Indexer 加速

图 5 来自 DoSA 论文。在上下文长度 128K、256K、512K、768K、1M,prefill chunk 4K、8K、16K、32K,以及三类输入域构成的 60 个组合中,DoSA 全部快于两条基线。相对 SGLang 的加速为 1.12–2.74 倍,相对 LiteTopK 为 1.17–1.69 倍;到 1M token 时,这两个区间分别收敛到 2.47–2.74 倍和 1.47–1.69 倍。

上下文长度比 chunk 大小更能拉开差距。以英文、16K query 为例,从 128K 扩到 1M,SGLang 对比加速由 1.19 倍升到 2.58 倍;固定 1M 上下文,把 query chunk 从 4K 放大八倍到 32K,加速只从 2.47 倍变为 2.69 倍。原因是 certificate 的 query 侧准备可以沿很长的 key 维度复用;key 越多,被避免的完整打分越多。

97% 剪枝率不是一个孤立数字

论文定义 certificate 剪枝率

p=1−NkeepNscreen.p=1-\frac{N_{\mathrm{keep}}}{N_{\mathrm{screen}}}.

在 GLM-5.2 的 21 个产生索引的层中,三种输入域的平均剪枝率为 97.2%–97.5%;每种输入中都有 20 层超过 95%。第一个索引层最低,但仍拒绝 89%–93% 的 query-key 对。这种跨英文、中文和代码都相似、而随层深变化明显的模式,说明 head 几何与层结构比文本域更决定 certificate 是否紧。

为什么上界即使不等于真实分数,也能淘汰绝大多数 key?令全域第 KK 大真实分数为 τ∗\tau^*,定义候选距边界的 margin、种子门槛缺口和上界 slack:

Δi=τ∗−si,b=τ∗−τ,δi=Ui−si.\Delta_i=\tau^*-s_i, \qquad b=\tau^*-\tau, \qquad \delta_i=U_i-s_i.

剪枝条件可以等价改写为

Ui<τ⟺Δi>δi+b.U_i<\tau \quad\Longleftrightarrow\quad \Delta_i>\delta_i+b.

这条式子比“上界越紧越好”更具体。一个 key 能否被剪掉,由它离真实边界有多远,与 certificate 松多少、种子门槛低估多少共同决定。百万候选里,大多数 key 的分数离第 2048 名很远,即便上界有一定余量,仍足以证明它们无望入选。反过来,若大量分数挤在 Top-K 边界附近,或者 seed 很差,survivor 就会变多。

保留头不是越多越好

设 N0N_0 个 seed key 完整打分,NsN_s 个 key 进入筛选,MM 个候选存活,辅助行数为 dad_a。基线与 DoSA 的逻辑点积数近似为

Dbase=H(N0+Ns),D_{\mathrm{base}}=H(N_0+N_s),DDoSA=HN0+(r+da)Ns+HM.D_{\mathrm{DoSA}}=HN_0+(r+d_a)N_s+HM.

令 f=N0/(N0+Ns)f=N_0/(N_0+N_s)、m=M/Nsm=M/N_s,归一化工作量是

DDoSADbase≈f+(1−f)(r+daH+m).\frac{D_{\mathrm{DoSA}}}{D_{\mathrm{base}}} \approx f+(1-f)\left(\frac{r+d_a}{H}+m\right).

筛选只有在 m<1−(r+da)/Hm<1-(r+d_a)/H 时才减少逻辑点积。多保留一个头会收紧上界、降低 mm,却也让每个 key 的筛选更贵。若预算从 rr 增到 r+Δrr+\Delta r,新增头值得的条件是额外拒绝率超过 Δr/H\Delta r/H。因此 rr 是一个需要按头几何和 Kernel tile 共同选择的系统参数,不是越大越保险。

这个代价模型也给出了可操作的调优顺序。先在真实激活上测每层 m(r)m(r),排除 certificate 几乎不剪枝的层;再把候选的 rr 对齐硬件支持的固定 tile,而不是逐个整数搜索;最后才比较端到端 TTFT,因为 query/key 准备、候选压紧和 Top-K 不在逻辑点积模型中。若某个更大预算只让拒绝率小幅提高,却引入新的 Tensor Core padding 或辅助行,它即使在公式中接近盈亏平衡,真实 Kernel 也可能更慢。

还可以把门槛质量与上界质量分别诊断。记录 b=τ∗−τb=\tau^*-\tau 能看出 seed 是否拖累筛选,记录 δi=Ui−si\delta_i=U_i-s_i 的分布能看出 head 投影与负头聚合是否过松。如果 bb 大而 δi\delta_i 小,应改进见证集合;如果 bb 小而 δi\delta_i 大,应调整保留头、锚点或激活支持;如果二者都小但 mm 仍大,问题就在真实分数集中于 Top-K 边界,继续收紧 certificate 的回报有限。这种拆分比只看最终剪枝率更容易定位层间差异。

从 Indexer 到 TTFT

GLM-5.2 上 DoSA 的端到端 TTFT
GLM-5.2 上 DoSA 的端到端 TTFT

DeepSeek-V3.2 上 DoSA 的端到端 TTFT
DeepSeek-V3.2 上 DoSA 的端到端 TTFT

图 6 与图 7 来自论文端到端实验。在约 1M token 上,DoSA 相对 SGLang 将 GLM-5.2 与 DeepSeek-V3.2 的 TTFT 分别加速 1.86 倍和 1.94 倍;相对 LiteTopK 分别为 1.26 倍和 1.50 倍。一次性反馈中所概括的“相比分别加速 1.94、1.50 倍”,对应的是 DeepSeek-V3.2 上相对两条基线的峰值;不能把它误读为两个模型各自只有一个统一加速值。

Indexer 加速高于端到端加速符合 Amdahl 定律。若原始 prefill 中 Indexer 时间占比为 ff,Indexer 自身加速为 ss,理想端到端上限为

S=11−f+f/s.S=\frac{1}{1-f+f/s}.

GLM-5.2 的 ff 从 128K 时的 23.6% 增到 1M 时的 73.0%,DeepSeek-V3.2 在 1M 时达到 90.8%,所以长上下文越长,DoSA 越接近决定总延迟的关键路径;注意力主体、MoE 与其他层仍然存在,TTFT 不会等于 Indexer 的 2.74 倍峰值。

筛选与稀疏回写各自贡献多少

论文在 SciFact 检索输入上使用 16K query、896K key,对 21 个索引层做消融。只启用 certificate 筛选,延迟下降 37.4%;只启用稀疏回写,下降 23.4%;二者合用下降 47.3%,即 1.90 倍加速。给只有稀疏回写的方案再加筛选,继续节省 31.1%;给只有筛选的方案再加稀疏回写,继续节省 15.7%。

这组数据说明 DoSA 不是单纯的数学剪枝,也不是另一个融合 Top-K Kernel。上界减少计算,精确过滤减少写回和选择,两者必须在同一数据流里配合,才能把“少算了”转化成真实延迟下降。

论文附录还给出更细的逐层复测:在 21 层、12 组输入组成的 252 个 cell 中,DoSA 相对 SGLang 有 251 个更快,相对 LiteTopK 有 250 个更快;逐层 retrieval 峰值分别达到 3.74 倍和 2.08 倍。三个变慢 cell 的回退幅度为 0.7%–9.0%。这比只报告峰值更有用:它说明收益广泛存在,但并非每一层、每一输入都严格占优,部署时仍应保留按长度、层或预算切换路径的余地。

与 LiteTopK、Quest 和近似索引的边界

DoSA、LiteTopK 与其他稀疏注意力方案解决的问题处在不同层级,不能只看最终加速数字横向排序。

方法主要减少什么是否扫描全部 key是否计算全部索引头Top-K 语义
SGLang 基线无,完整打分与选择是是精确
LiteTopKscore 写回、同步与选择流量是是精确
DoSA完整多头打分、score 写回与选择输入是否,只对 survivor 算全头精确
MISA 类少头筛选索引头计算是只对候选重排可能漏掉初筛淘汰的赢家
Quest 页级筛选KV 页读取与后续注意力计算读取页级摘要,跳过部分 KV 页不适用对注意力贡献做近似筛选

LiteTopK 用采样估计分数范围、分桶并融合 candidate filtering,解决“所有分数都算了以后,怎样少写、少同步、少选择”。DoSA 在这个基础上向前推进一步:能否连完整分数也不算。两者并非完全互斥,DoSA 的细化与最终精确选择同样需要紧凑写回技术;论文以 LiteTopK 为基线,是为了分离“融合收益”与“certificate 条件计算收益”。

Quest 用页级 min/max 元数据估计一个 query 对 KV 页可能产生的注意力贡献,再只读取有希望的页。它关注的是注意力本体的 KV 流量,接受页级筛选所对应的模型近似;DoSA 针对已经存在的 DSA 索引器,要求恢复其原始精确 Top-K。二者共同使用“上界低于门槛即可跳过”的思想,但被界定的分数、摘要粒度与正确性目标不同。

MISA 一类方法用少数索引头先选候选,再对候选做全头 rerank。若真正赢家在第一阶段已被排除,rerank 无法把它找回来。DoSA 的关键差别不是“也有第二阶段”,而是第一阶段排除每个 token 时都有 si≤Ui<τs_i\leq U_i<\tau 的证据。

从信息检索视角看,DoSA 与 MaxScore、WAND、Block-Max WAND 一脉相承:维护 Top-K 阈值,用贡献上界跳过不可能获胜的文档。新难点来自 DSA 的有符号多头 ReLU 分数。文档检索中的项贡献上界不能直接处理负权重头的扣减与跨头激活抵消,DoSA 的正头投影、联合残差和激活感知负头聚合正是为这种 scorer 定制的。

局限与工程使用条件

DoSA 很强,但它没有把索引从 O(∣Ω∣)O(|\Omega|) 变成次线性复杂度。筛选仍须扫描每个可见 key,读取 key payload 或摘要并计算 r+dar+d_a 个点积;它省掉的是大多数候选的剩余 HH 头计算、密集 score 写回和后续选择工作。因此它最适合 key 维度很长、原始头数较多、Top-K 边界之外存在大批低分候选的场景。

具体限制包括:

  1. 短上下文可能不划算。query 侧规划、seed 完整打分、筛选 Kernel 和 survivor 压紧都有固定开销。论文在 128K 的最低加速只有 1.12 倍,继续缩短上下文可能触发 crossover。
  2. certificate 紧度依赖模型几何。未保留正头若与锚点方向差异大,EPE_P 会变大;负头激活若频繁变动,聚合会产生更多漏项与抵消。DoSA 不是对任意 scorer 都保证高剪枝率,只保证不错误剪枝。
  3. 边界分数拥挤会留下大量 survivor。当许多 key 的 Δi\Delta_i 很小,哪怕 δi\delta_i 与 bb 都不大,也无法证明它们落选。最坏情况下所有 key 都存活,筛选变成额外工作。
  4. 仍然读取所有 key。未来若构造 key-block 的保守摘要,可先用块级上界跳过整块,从而连逐 key 读取都避免;但一个块上界必须覆盖块内最坏 key,往往比逐 key certificate 更松。
  5. 数值与恢复路径复杂。FP8 辅助向量、key scale、保守 RMS、上舍入误差、候选溢出和 nonfinite 值都要有清晰契约。任一环节错误都可能把“高概率一致”伪装成“精确一致”。
  6. 论文验证的是计算收益。实验没有证明任务准确率超出模型已声明的上下文范围;DeepSeek-V3.2 超出声明长度的测量只能用于观察计算缩放,不能外推模型能力。
  7. 显存换时间不可忽略。为保证候选不丢失,极端 survivor 数量需要充分 workspace 或可靠的分块合并方案。部署时应把这一容量纳入并发数与 KV Cache 预算,而不是只看 Kernel 峰值。

工程接入前至少应检查四件事:原始 scorer 的精确算术定义是否稳定;每个模型层的正负头比例与激活模式是否让界足够紧;候选缓冲溢出是否有全域恢复;端到端收益能否覆盖额外 workspace 对并发的影响。只有这四项同时成立,微基准中的索引器优势才会变成服务级收益。

小结

DoSA 的价值不在于又设计了一种稀疏注意力模式,而在于重新界定了“精确选择必须付出的计算”。完整 Top-K 只要求每个赢家被保留,并不要求每个输家都得到精确分数。

它的技术闭环由五部分组成:种子集合给出真实 Top-K 门槛的下界;保留头直接提供部分精确贡献;非负投影与联合残差安全覆盖未算正头;激活感知聚合给出负头扣减量的下界;GPU Kernel 把 survivor 压紧、共享 key,并只回写仍有资格的精确分数。最终 Top-K 始终在原始分数上完成,certificate 只负责证明谁可以提前离场。

更普遍的启发是:在推理系统中,“近似计算更少”与“用界证明无需计算”是两条不同路线。前者用误差换速度,后者用额外元数据和控制逻辑换来同一结果。DoSA 证明,当候选离选择边界足够远、query 侧几何能够构造紧界时,第二条路线也能在 GPU 上获得可观收益。

它也把后续优化方向划得很清楚:更好的见证降低阈值缺口,更紧的几何界降低 certificate slack,块级摘要减少 key 扫描,而更稳健的稀疏流水线负责把理论节省兑现为端到端延迟。四者分别作用于不同成本,不能用单一剪枝率替代完整评估。

参考资料

  1. DoSA: Dots Sparse Attention
  2. DoSA 官方项目与论文仓库
  3. DoSA SGLang 实现仓库
  4. LiteTopK: Exploiting the Curse of Dimensionality for a Fused Indexer-TopK Kernel in Long-Context Sparse Attention
  5. DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models
  6. Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
  7. Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference
  8. Efficient Query Evaluation Using a Two-Level Retrieval Process
  9. Faster Top-k Document Retrieval Using Block-Max Indexes

文章分享

如果这篇文章对你有帮助,欢迎分享给更多人!

DoSA 完全拆解:用可证明的提前拒绝加速精确稀疏注意力索引
https://pinghaoyang.com.cn/articles/insights/posts/dosa-sparse-attention/
作者
平昊阳
发布于
2026-10-16
许可协议
CC BY-NC-SA 4.0

评论区

访客信息
--°天气加载中

正在获取所在地区…

💨 风力--
💧 湿度--

正在读取天气

Profile Image of the Author
平昊阳
乘长风,破巨浪, 展鸿图于未央!
公告
欢迎来到我的个人博客!欢迎关注交流吖!
更多相关公告,见
「社交-留言」。
音乐
封面

音乐

暂未播放

0:000:00
暂无歌词
站点统计
文章
448
分类
29
标签
451
总字数
6,178,652
运行时长
0 天
最后活动
0 天前

文章目录