2.2 · 细粒度索引: 用轻量 selector 召回关键 token

粗粒度路由按块计算, 容易形成规则 kernel, 也会把关键 token周围的低价值位置一并读入. 细粒度索引把最终候选压到 token级: 先用比主 attention便宜的代理分数扫描或检索历史, 再对少量候选执行完整 QK与 Value聚合. 节省来自昂贵表示少读, 风险来自代理排序漏掉主 attention需要的位置.

一条完整路径包含四个张量. 主 query与主 KV负责精确 attention; index query与 index key负责候选打分. selector输出整数位置 $S\in\mathbb{N}^{B\times n\times k}$ , Sparse Attention据此 gather主 KV. indexer分数通常不进入主 softmax, 所以「召回」与「精排」是两个算子. 这一区分决定训练目标、量化误差和 kernel验证方法.

细粒度索引的目标是在固定读取预算内覆盖主 attention的高价值 token, 同时让索引成本显著低于被省掉的主计算. 一个 selector即使 recall很高, 若先做完整高维 QK就没有性能意义; 一个 kernel即使很快, 若漏掉唯一证据也没有质量意义.

1. 从 DSA 建立共同坐标

1.1. 低维全扫与 token top-k

DSA 与索引器给出本路线的基准. Lightning Indexer为每个 query生成多个低维 index heads, 历史 token只保存共享 index key. 各 head经过 ReLU匹配和 query相关权重后聚合成一个标量, selector沿历史维取 top-k. DeepSeek-V3.2 的稀疏训练使用 $k=2048$ , 最终集合供全部主 attention heads共享.

共享候选让 MLA latent KV能跨 heads复用, 也把不同 heads的偏好压进一个集合. indexer内部需要多头表达多种相关性, 这又成为后续 MISA压缩的对象. DSA主 attention降为 $O(nk)$ , indexer仍做低维 $O(n^2)$ 配对; 判断端到端收益时必须把两项分开.

训练不是给 top-k 补一个 loss。 DeepSeek-V3.2 先保持稠密 attention, 冻结主模型, 用跨主 heads聚合的教师分布和 KL训练 indexer. 随后启用 top-k, 让全部主模型参数在稀疏图上继续训练, indexer仍由独立 KL优化且输入 detach. 前一阶段解决随机 selector, 后一阶段解决主模型不适应缺边.

这套过程说明 selector质量与模型适应是两个变量. 只训练 indexer可能获得较高 recall, 模型仍会因 softmax定义域改变而退化; 只让主模型接受随机稀疏图, 又会破坏远程信息. 迁移到其他架构时, 训练 token、上下文长度和教师构造都属于方法的一部分.

索引器的三种降本方向。 HISA 缩小 token搜索域。 HISA 先对块代表打分, 剪掉无关块, 再只在候选块内运行原 DSA token indexer. 最终仍输出相同接口的 token top-k, 可免训练替换 DSA indexer. 块大小控制粗筛数量, 候选块数控制召回余量; 两者共同决定成本 $O(n/b+rb)$ .

粗筛错过的块无法在精排恢复, 因而 HISA 主要用与原 DSA候选的 IoU检验保真. 官方论文报告在 DeepSeek-V3.2和 GLM-5 上直接替换且无需微调, 并取得很高集合重合. 物理收益还取决于块代表读取、块 top-k和候选 token布局.

MISA 缩小 indexer head集合。 MISA 把 DSA indexer heads视为专家池. 轻量 router读取块池化 index keys, 为每个 query选少量 active heads, 只有这些 heads扫描全部候选 token. 最终 token数与 Sparse MLA接口不变, 省的是式中跨 head的计算.

层次变体先用 active heads产生放大候选, 再让原 DSA完整 heads只在候选内精排. 它把召回与精度分开: 第一遍便宜但宽, 第二遍昂贵但窄. 官方论文报告 8 active heads、超过92%的逐层候选恢复以及 H200上的 kernel加速, 数字均应保留模型与硬件口径.

LISA 增加线性全局通道。 LISA 将线性 attention和 indexer引导的 sparse self-attention并联, 由 gate融合. 线性状态以固定规模聚合全历史, sparse分支保留少量 token的精确 softmax交互. selector漏掉分散信息时, 线性分支可能提供补充; 精确复制和唯一证据仍主要依赖 sparse候选.

LISA 先训练线性分支与滑动窗口逼近冻结教师, 再引入 indexer替换静态窗口, 使用 per-head KL对齐. 它改变的信息通道多于 DSA/HISA/MISA, 质量与速度不能只归因于 selector. cache核算也必须包含线性状态、sparse KV和 indexer K.

读这些方法时先看什么。质量轴。统一指标包括教师 top- $m$ 的 recall@k、教师概率质量覆盖、与基线 selector的集合 IoU、top-k cutoff margin和任务质量. recall与 IoU回答不同问题: 前者对照主 attention, 后者对照另一 indexer. 平均值之外还要看唯一证据、多证据和长距离位置的尾部.

量化、分块粗筛、稀疏 heads和跨层复用都会改变候选. 多次近似应逐级测重合, 否则最终退化无法归因. kernel对固定候选的误差又是独立一层, 应先与朴素 masked attention对照.

性能轴。总时间至少拆成 index投影/cache读取、代理打分、top-k、索引整理、主 KV gather和 Sparse Attention. selector不物化 $n^2$ 分数, 通常要融合分块打分与 top-k. element-wise候选地址离散, 主 kernel可能受 HBM事务而非 FLOPS限制.

Prefill与 Decode分别报告. Prefill有大矩阵并行, Decode每步扫描增长中的 index K cache. KV量化降低扫描字节, cutoff附近的排序更敏感. 多卡还要说明候选集合在哪里合并、远程 KV如何搬运.

状态与共享轴。共享可能发生在 indexer heads、相邻 layers、index K或整数 top-k. 每一种省掉的工作不同. MISA是同层按 query少激活 heads; 跨层复用 top-k则直接省掉某些层selector, 但承担层间偏好变化. 共享 index K只省投影或缓存, 仍要每层打分.

本路线聚焦 token级 selector本身. 先沿 DSA 与索引器掌握完整数据路径，再进入 MISA、LISA 与 ResMiSA，比较 HISA、MISA 与 LISA 分别剪掉哪一段。CSA、HCA 以及进一步的跨层索引复用放在下一条路线，避免把不同共享对象压成一个模糊概念.

1.2. 为什么 indexer 逐渐成为主角

早期稀疏注意力常把 mask 当成既定结构: 局部窗口、固定全局 token 或规则步长先写好, attention 只负责执行. 动态稀疏把问题向前推了一步: 每个 query 需要的历史位置不同, mask 本身必须由模型计算. 当主 attention 被压到很低的候选预算后, 生成 mask 的 indexer 不再是附属小模块, 它开始决定质量上限、端到端复杂度与 KV 的物理访问.

这个变化可以从时间分解看得很清楚. 令 $T_{main}^{dense}$ 为原主 attention 时间, 稀疏后变为 $T_{main}^{sparse}$ , 新增 indexer 和选择成本 $T_{idx}+T_{topk}$ . 整层能得到的节省为:

\Delta T=T_{main}^{dense}- (T_{idx}+T_{topk}+T_{main}^{sparse}). \tag{1}

最初 $T_{main}^{dense}$ 很大, 一个不够精致的索引器也可能划算. kernel 与主候选继续压缩后, $T_{main}^{sparse}$ 下降, $T_{idx}$ 的占比会越来越高. HISA 从 token 搜索域入手, MISA 从 indexer head 数入手, 跨层共享从调用次数入手, 都是在处理式 (1) 中同一项. 稀疏率继续提高时, 系统的中心问题会从「怎样少算主 attention」转成「怎样用更低成本找对位置」.

Indexer 还定义了稀疏 attention 的错误形态. 主 kernel 对给定候选计算得再准确, 也无法恢复候选外 token. 因而它既像检索器, 又不同于普通语义检索. 普通检索只需返回语义相关片段; attention selector 要逼近当前层、当前 query、当前模型的内部计算偏好. 同一文本位置在不同层承担的作用不同, 表面语义相似也不保证主 QK 分数高.

2. 一个统一的 selector 数学接口

2.1. 代理空间、候选集合与主空间

设主 attention 的第 $l$ 层 query、key 和 value 为 $q^{(l)}_t,k^{(l)}_s,v^{(l)}_s$ . Indexer 从 hidden state 或主投影的某个中间量生成低成本表示:

\tilde q^{(l)}_t=f_q^{(l)}(h_t),\qquad \tilde k^{(l)}_s=f_k^{(l)}(h_s). \tag{2}

代理分数为:

r^{(l)}_{t,s}=g^{(l)}(\tilde q^{(l)}_t,\tilde k^{(l)}_s), \tag{3}

选择器输出:

\mathcal S^{(l)}_t=\operatorname{Select} (r^{(l)}_{t,0:t},\mathcal B^{(l)}_t), \tag{4}

$\mathcal B^{(l)}_t$ 表示预算或阈值规则. 主 attention 再在 $s\in\mathcal S_t^{(l)}$ 上计算原分数. 式 (2) 到式 (4) 把几条路线放到同一接口里: DSA 的 $g$ 是多 index-head 的 ReLU 匹配加权和; HISA 在式 (3) 前增加块级过滤; MISA 对参与 $g$ 的 index heads 做路由; 近似最近邻方法用离散桶或向量检索实现 Select; 跨层共享则让某些层复用别层的 $\mathcal S$ 或 $\tilde k$ .

这个接口里有三种容量. 表示容量由 $\tilde q,\tilde k$ 的维度和 head 数决定; 搜索容量由参与打分的历史范围决定; 输出容量由 $|\mathcal S_t|$ 决定. 降低其中任意一项都能省成本, 误差却不同. 低维代理会改变排序, 缩小搜索域会造成不可恢复的粗召回错误, 减少输出数则在排序正确时仍会截断长尾.

共享候选与逐 head 候选。设主 attention 有 $H_q$ 个 query heads. 最自由的选择是每个 head 各自输出 $\mathcal S_{t,h}$ , 候选总数上限为 $H_qk$ . 共享选择只输出一个 $\mathcal S_t$ , 所有 heads 从同一批 KV 中精排. 共享集合有利于 MLA、GQA 和 KV gather, 也需要覆盖不同 head 的偏好并集.

可以按组折中. 把 $H_q$ 个 heads 分成 $G$ 组, 每组一个 selector:

\mathcal S_{t,g}=\operatorname{TopK} \left(\operatorname{Agg}_{h\in g}r_{t,h,:},k_g\right). \tag{5}

$G=1$ 是全共享, $G=H_q$ 是逐 head. 增大 $G$ 提高表达容量, 也增加索引输出、KV 地址集合和 kernel 分支. 合理组数取决于 heads 偏好的重合度, 不能只由主 head 数决定. 聚类 head 的教师分布、测量候选并集扩张, 比随意等分更有依据.

MISA 稀疏激活的是 indexer 内部 heads, 与式 (5) 的主候选分组不是一回事. 它仍可以为主 attention 输出共享 token 集合. 同一个「head 稀疏」说法可能指少算代理 head, 也可能指为不同主 head 生成不同候选; 两者影响的张量和 kernel 完全不同.

Select 不等于一次普通 top-k。直接 materialize 全分数再调用 top-k, 数学最直观, 内存写入是平方级. 长序列实现会把历史切成片, 每片计算局部候选, 再不断合并当前 top-k. 对分片 $c$ 的候选 $(v_c,i_c)$ 与已有候选 $(v,i)$ , 合并只需从 $2k$ 项中保留 $k$ 项. 最终索引与全局 top-k 一致, 中间不保存完整分数矩阵.

因果 Prefill 还要求每个 query 的有效历史长度不同. 一个二维 tile 内, 后半 query 能看到更多 key. kernel 可以给未来位置打负无穷, 但局部 top-k 的有效项数会随行变化. 序列开头 $t<k$ 时, 只能返回 $t+1$ 个真实位置, 其余槽位必须显式标记无效, 不能用位置 0 重复填充后让主 softmax算多次.

分数并列要规定稳定性. 两个位置代理分数完全相同, 按较小位置、较大位置或实现线程顺序打破平局, 都可能得到不同集合. 浮点误差在 cutoff 附近也会产生类似替换. 训练与推理只要求集合质量时, 平局顺序未必重要; 多卡各 rank 必须得到同一集合时, 不确定 tie-break 会直接造成结果分叉.

Indexer 怎样学习「主 attention 需要什么」。排序监督、分布监督与输出监督。最直接的教师是主 attention logits 或 softmax. 分布监督用 KL 让代理分布逼近聚合后的主 attention; 排序监督只要求正样本排在负样本之前; 集合监督把教师 top- $m$ 当作多标签目标. 三者强调的误差不同. KL 对高概率质量敏感, 低概率但任务关键的位置可能权重很小; 排序损失可针对 cutoff, 需要设计负样本; 集合监督容易解释, 教师第 1 名和第 $m$ 名却常被同等处理.

还可以直接蒸馏主 attention 输出. 给定候选后计算稀疏输出 $o_t^{sp}$ 与稠密输出 $o_t^{dense}$ , 最小化二者距离. 这会把 Value 方向考虑进来: 漏掉两个权重较高但 Value 互相抵消的位置, 输出误差可能仍小. 其训练成本更高, 也难定位 selector 究竟学会了排序还是利用了抵消. 实际训练常把代理分布损失与主任务或输出对齐结合.

监督聚合方式必须对应候选共享方式. 若所有主 heads 共享 $\mathcal S_t$ , 教师需要形成一个能覆盖多 head 的目标. 对 softmax 权重直接求和偏向总质量大的位置; 对每个 head 先取 top- $m$ 再求并集, 能照顾专用 head, 教师集合可能迅速膨胀. 用每个位置在 heads 中的最大权重可保留单 head 尖峰, 也会让噪声 head 主导. 这不是数据预处理小事, 它决定共享候选想保存哪类依赖.

离散选择为何需要分阶段。随机初始化的 indexer 在启用 top-k 后会向主模型提供近乎随机历史. 主模型收到的语言模型梯度也无法穿过离散集合, selector 很难靠最终 loss 自己纠正. 稠密教师预热先让式 (3)具有合理排序, 再启用稀疏图让主模型适应, 是 DSA 两阶段路径背后的优化原因.

可微 top-k、Gumbel 噪声或直通估计能让主任务梯度近似影响 selector, 但梯度偏差、温度和训练稳定性会成为新变量. 另一条路线保留稠密或局部分支作为信息保险, 让 selector逐渐承担更多预算. MoBA 从 full 到 sparse 的过渡、NSA 的多分支、LISA 的线性通道都能从这个角度理解: 离散路由切换时, 模型仍有一条连续的信息路径.

训练数据决定 selector 见过什么依赖. 如果长样本主要是简单拼接, 教师会大量强调局部与文档边界, 不会自动学会跨文档多跳. 代码、表格、长对话和多语言的 QK 分布不同. 训练 token 总量重要, 长度和依赖结构同样重要. 只报告 128K 样本数量, 仍不足以说明样本中是否存在需要远程精确召回的信号.

正负样本与 cutoff 附近的难例。对 top-k selector, 真正决定集合的是第 $k$ 名附近. 教师第一名与随机负样本差距很大, 很容易学; 第 $k$ 与第 $k+1$ 的细小差别才影响候选替换. 排序训练可以对这一区域做 hard-negative mining. 设教师正样本 $s^+$ 、难负样本 $s^-$ , margin 损失为:

\mathcal L_{rank}=\max(0,\gamma-r_{t,s^+}+r_{t,s^-}). \tag{6}

若 $s^+$ 只是教师尾部小权重项, 过大的 $\gamma$ 会浪费表示容量; 若唯一证据在教师聚合后恰好位于尾部, 忽略它又可能伤害任务. 因此 hard negative 最好与概率质量、head 专用性和任务标注联合选择. 评测时保存 cutoff margin, 可以解释量化或少 head近似为什么在某些层更容易替换候选.

计算量减少不等于访存减少。 Index K 全扫的字节下界。 Decode 第 $t$ 步扫描全部 index keys. 假设 $\tilde k_s\in\mathbb{R}^{d_i}$ 以 $b_i$ 字节存储, 仅 key 主体读取量约为 $td_ib_i$ . 当 batch 为 $B$ 时, 不同请求的 key cache 通常不能共享, 读取近似乘以 $B$ . 低维和 FP8 很关键, 因为它们同时降低计算和 HBM 字节.

若 $d_i=32$ , FP8 每元素 1 字节, 128K 历史的 index K 主体约 4 MiB. 单个 query 扫一遍可能落在 L2/HBM 的边界, 多层、多 batch 会迅速放大. 主 latent KV 即使更宽, 稀疏后只读 2048 项; index K 反而成为连续全扫. 这解释了为什么 HISA、量化与跨层 index K共享会在 DSA 之后受到关注.

Prefill 的 index Q 和 index K能形成矩阵乘, 算术强度高于 Decode 扫描. 同一 indexer 在两个阶段的最佳 kernel 可能完全不同: Prefill偏向 tiled GEMM 加流式 top-k, Decode偏向连续 key 读取、少量 query和融合 reduce. 只优化其中一条路径, 不能推断另一条也快.

主 KV gather 的物理代价。 Token top-k 在逻辑上只读 $k$ 个位置. 如果这些位置分散在 $k$ 个 cache pages 或 cache lines, 每次事务会带入未使用字节. 将索引按物理地址排序能合并部分读取, 但 Paged Attention 的逻辑位置要先经过 block table. 选择器输出的是序列位置, attention kernel消费的是物理地址, 中间映射也有成本.

同一组 token 供多个主 heads共享时, K/V可以加载一次后被多个 heads使用. 逐 head selector虽然提高召回自由度, 也会扩大候选并集. 若每个 head各取 256, 32 个 heads的并集可能远大于共享的 2048; 最终读字节未必更少. 因此候选共享是算法表达力与硬件复用的直接交换.

候选位置排序不会改变 softmax, 前提是索引和值保持对应. kernel 可以先按 page 分组, 页内按 offset排序, 再逐段执行在线 softmax. 若同一位置因多路召回出现两次, 必须去重. 重复项会把概率质量算两遍, 不是单纯的性能浪费.

top-k 自身也有带宽与同步。流式 top-k需要维护每个 query 的候选缓冲. $k=2048$ 时, 候选值与索引远超单个 warp 的寄存器容量, 实现会分层选择: 每个线程块保留局部较小候选, 写入中间缓冲, 后续归并. 候选越大, 中间写入和归并轮数越多. 所以把主候选从 1024 增到 2048, 影响的不只 Sparse Attention 边数, indexer 选择阶段也会变重.

如果 indexer 分数、top-k 和主 attention 分成多个 kernel, 中间存在启动与全局同步. 融合能减少写回, 又增加寄存器压力和代码特化. HISA 两级选择、MISA 路由加精排还会增加阶段. 判断改进时要看完整链条, 不能把每个单核的最好数字相加.

2.2. 层次索引不只是「先粗后细」

两级召回的条件概率。设教师重要 token 集合为 $G$ , 粗筛保留区域 $C$ , 细排输出 $S\subseteq C$ . 最终召回可拆成:

P(s\in S\mid s\in G)= P(s\in C\mid s\in G) P(s\in S\mid s\in C,s\in G). \tag{7}

粗筛召回为 0.99, 细排条件召回也为 0.99, 两级相乘约为 0.9801. 再叠加量化和跨层复用, 每一级看起来都很高, 最终尾部仍会明显下降. 因而层次系统要报告条件召回, 不能只比较最终集合与基线的平均 IoU.

HISA 的块代表承担第一项. 块越大, 代表更难保留单 token尖峰, 候选块数较少时风险更高; 块越小, 块级打分数量增加. MISA 的 active heads也可以视为另一种粗筛: 只有被路由到的代理子空间参与广域扫描. 层次变体随后让全 heads在放大的 token候选上精排, 对应式 (7) 的第二项.

粗筛单位与物理页可以协同。如果 HISA 块与 KV cache页对齐, 选中一个粗块恰好对应连续页, 精排读取容易合并. 若算法块跨物理页, 一个候选块会拆成多个事务; 若多个算法块落在同页, 第一阶段可以共享读取. 对齐能提高工程效率, 但页大小由内存管理决定, 算法块大小由召回决定. 最优值不一致时, 可让一个算法块由若干完整小页组成, 避免跨边界.

粗筛摘要也可随 cache页写入. 均值、最大值或学习型代表的更新规则不同. 学习型代表如果依赖整个块, 当前未满块需要单独处理. Decode 对当前块通常直接保留局部注意力, 历史完整块才进入层次检索, 这样可避免不断重算未满摘要.

失败样本怎样定位到具体层级。保存教师 token后, 先检查它所在粗块是否被选. 若没有, 问题在块摘要或粗预算; 若块被选但 token未进入最终集合, 问题在细排; 若 token已进入集合而输出仍错, 再检查主 logits、Value和 kernel. 这条诊断链能避免一看到任务下降就全局增加 $k$ .

对 MISA, 先检查负责教师 token的 index heads是否被激活, 再检查活跃 heads生成的放大候选, 再看全 head精排. 某些 query需要罕见 head, router平均准确仍可能稳定漏掉. 统计每个 index head被激活频率与它独占召回的教师 token数, 可以区分冗余 head和低频关键 head.

跨层共享到底共享什么。共享 index key。相邻层可复用同一 $\tilde k_s$ , 省掉 key投影和缓存容量, 每层仍生成自己的 index query并扫描. 这种共享保留层级不同的 query偏好, 代理 key表达却固定. 若深层需要的语义特征在共享 key里没有编码, 再强的 query也无法找回. 可以用分组层共享: 每若干层重新生成一次 key, 在容量和表达之间取点.

共享候选索引。直接复用 $\mathcal S_t$ 能省掉完整 selector, 收益更大, 假设也更强. 相邻层 attention偏好常有相关性, 不代表集合完全相同. HCA、CSA一类路线关注层间可复用性, 核心问题是由哪一层产生候选、复用多少层、何时刷新, 以及共享的是完整集合还是候选的一部分.

一种稳妥组合是共享基础候选 $S^{base}$ , 每层再追加较小的残差候选 $S_l^{res}$ :

\mathcal S_t^{(l)}=S_t^{base}\cup S_{t,l}^{res}. \tag{8}

基础集合承担稳定远程依赖, 残差集合修正层特有偏好. 这与共享专家加稀疏残差的思路相似, 实际价值取决于 $S^{res}$ 多小以及两集合如何去重. 若残差仍需全历史扫描, 省下的只是一部分代理维度; 若残差在基础块附近搜索, 成本低但纠错范围受限.

共享 main KV读取与共享 selector 是两回事。多层 attention通常各有自己的 K/V, 即使 token位置相同也不能直接复用 Value. 共享索引只省选择, 不自动省每层主 KV读取. CSA、CSA2、HCA、HySparse2等名称中经常同时出现跨层、共享、残差或缓存概念, 阅读时要逐项问清: 共享的是 index key、整数位置、主 K/V的线性项、还是 attention输出. 同样写「跨层共享」, 张量不同, 对质量和带宽的影响就不同.

跨层共享按这些对象继续展开, CSA 与 CSA2和 HySparse负责具体方法. 细粒度索引的共同接口只要求接收候选并返回 token位置; 是否复用由上层状态协议决定.

它与向量数据库检索有多像。相似之处是「便宜召回, 昂贵精排」。两者都把候选生成和高成本判断分开. 向量检索先用压缩向量、倒排表或图索引召回文档, 再用交叉编码器重排; token selector先用低维 index Q/K召回位置, 再用主 Q/K与 Value计算. coarse-to-fine、召回率、cutoff margin和分层索引等概念可以共用.

向量数据库的 IVF 思路也能帮助理解 HISA. 历史 token先被分到若干区域, query只探测少数区域, 然后在区域内精排. 若探测区域数为 $n_{probe}$ , 增大它会提高召回并增加扫描. 图索引则从少量入口沿邻接关系搜索, 适合静态或缓慢更新的库. KV Cache 每步追加、每层表示不同, 建图成本和更新一致性会削弱这类结构的吸引力.

区别在于索引更新频率。文档库可以离线构建索引, 多次 query摊销. 自回归 KV 每生成一个 token就增长, index key也是模型逐层产生的内部状态. Decode 第 $t$ 步刚写入的位置, 第 $t+1$ 步可能立即需要. 索引若有延迟合并区, 新 token必须通过局部窗口或增量缓冲保持可见.

Prefill 又会一次写入成千上万 token. 逐 token更新树或图可能比直接矩阵扫描更慢. 分层块摘要适合批量生成, 结构简单; 复杂 ANN只有在超长 cache、许多后续 Decode query复用同一历史时才可能摊销. 是否值得建立, 取决于请求生成长度, 不能只看 prompt长度.

区别还在分数函数与层级。通用 ANN常假设欧氏距离、余弦或内积. Lightning Indexer 的多 head ReLU加权分数随 query改变各子空间权重, 未必能直接放入标准最大内积索引. 即使每个 head分别建索引, 合并后的全局 top-k也需要考虑 query权重, MISA选择 active heads正是沿这个维度降本.

每一 Transformer层还有自己的表示. 为 60 层分别构建复杂 ANN, 状态和更新成本很大; 共享一个索引又要求层间偏好足够一致. CSA/HCA等跨层路线之所以重要, 正因为它们尝试用模型结构上的相关性减少重复索引, 而非假定一个外部语义向量库能服务所有层.

此外, attention候选受因果、padding、文档边界和位置编码约束. 检索到未来 token必须屏蔽, packed batch不能跨样本, 相同内容在不同位置的 RoPE主分数也会变化. Indexer若不含位置信息, 召回后仍由主 attention精排, 但候选预算可能被位置上不合适的相似 token占用.

3. 数值精度怎样改变一个离散集合

3.1. cutoff margin 比平均分数误差更有用

令精确代理分数排序后为 $r_{(1)}\ge\cdots\ge r_{(n)}$ , 预算为 $k$ . 定义边界间隔:

\delta_k=r_{(k)}-r_{(k+1)}. \tag{9}

当 $k$ 小于有效候选数、 $\delta_k>0$ ，且每个分数的误差绝对值都小于 $\delta_k/2$ 时，整个 top-k 集合保持不变。任一集合内分数与集合外分数的原始间隔至少为 $\delta_k$ ，两侧误差之和不足以跨过这个间隔。集合内部的排序仍可能变化。 $\delta_k$ 很小时，FP8 舍入、不同 reduce 顺序甚至并列规则都会改变集合。平均分数均方误差可能很小，top-k IoU 仍会下降。

这也是扩大粗候选再精排有效的原因. 第一阶段取 $k'>k$ , 只要教师重要 token进入较宽集合, 第二阶段用更高精度或完整 heads重排. 代价从扫描全历史转为扫描 $k'$ 项. $k'$ 应通过 recall—时延曲线选择, 不能默认取 $2k$ 就足够.

ReLU、scale 与异常值。 ReLU打分在零点存在离散变化. 一个略小于零的点积量化后变成正数, 会新增贡献; 略大于零变负则整项消失. 若大多数 head点积位于零附近, indexer分数易受量化影响. 监控 ReLU零比例、每 head输出范围和权重分布, 能发现某些 heads已经坍缩或 scale失配.

分块 scale要与 cache分页一起设计. 每 token一组 scale精度高、元数据多; 每页一组 scale遇到异常 key时会压低其他 token的有效分辨率. 对 selector而言, 异常值不仅增加数值误差, 还可能长期占据 top-k. 可以对 index Q/K做归一化或裁剪, 前提是训练与推理使用相同规则.

NaN 必须在 top-k前处理. 不同排序库对 NaN 的位置约定可能不同, 某些实现会把它当成极大值. 一个 NaN若进入候选, 主 attention读取的位置合法, 错误不会表现为越界. kernel入口应检查 index表示和 scale为有限值, 训练日志则追踪首次出现异常的层与 token.

多卡一致性。复制 indexer时, 各 rank理论上产生相同分数. 混合精度归约次序、不同 kernel版本和并列 tie-break仍可能让 cutoff附近集合分叉. 如果每个 rank只消费自己主 heads, 分叉会使同一 query的 heads读取不同 token, 破坏训练时的共享候选假设.

确定性方案可以在比较键中加入位置, 以 (score,index) 做词典序 top-k; 对量化相等值按位置稳定选择. 这不会恢复分数精度, 但让多个设备结果一致. 另一种方案由单 rank选择后广播索引, 增加通信. 哪一种更快, 取决于索引大小、并行规模与现有同步点.

建立一条可复现实验曲线。固定候选预算扫描质量。先在同一模型、同一上下文集上扫描 $k\in{256,512,1024,2048,4096}$ . 每个 $k$ 记录教师 top- $m$ recall、概率质量覆盖、hidden state误差和任务结果. $m$ 不必等于 $k$ : 例如固定 $m=256$ 观察候选扩张能否覆盖最重要位置, 再用 $m=k$ 比较同预算排序.

每条曲线按长度和任务分组. 32K上的 $k=2048$ 占历史比例与 256K差八倍; 如果质量随长度下降, 要判断是固定容量不足、代理排序退化, 还是模型本身没有在该长度训练. 按相对位置分桶还能看出 selector偏向近邻还是远程.

固定质量扫描系统成本。从质量曲线选择满足阈值的最小预算, 再测 index投影、扫描、top-k、索引映射、主 gather和 attention. HISA、MISA与平坦 DSA应在相近质量点比较, 而非各自选最有利的不同 $k$ . 同时报告读取字节、唯一 KV页数和候选连续段数, 解释延迟差来自少算还是更规则的访存.

Prefill使用多个长度与 batch, Decode使用多个 cache长度与并发. 首 token延迟和单 token延迟分开. PD分离系统还需在 Prefill节点和 Decode节点分别部署相应 kernel; 只在一个阶段有收益时, 不必让另一阶段承担同样复杂度.

消融近似叠加。从 FP16平坦 indexer开始, 依次加入 FP8、块粗筛、少 heads、跨层复用. 每加一项都保存上一阶段和当前阶段的候选, 计算条件召回与新增时间节省. 若最终质量下降, 可以沿这条链找到主要损失来源.

消融顺序也应反向跑一次. 例如 HISA 与 MISA同时启用时, 先块后 head和先 head后块可能得到不同候选与成本. 前者对较少 token运行少 heads, 后者先用少 heads筛全历史. 两条路径的理论乘法量相近, cache访问和误差传播不同. 实测以后再决定组合, 不把两个独立论文倍数直接相乘.

从研究结果进入工程决策。先确定瓶颈在哪一段。如果 profiler显示主 Sparse Attention仍占大头, 继续压缩 indexer不会带来明显端到端收益, 应优先改善 KV gather、候选排序与 attention kernel. 如果 indexer随上下文增长成为主耗时, HISA的候选块、MISA的 active heads或跨层复用才是直接方案. 如果两者都不高而通信占主导, 单卡 kernel倍数也无法改变集群延迟.

长度门槛必须实测. 短上下文上稠密 kernel规则且无需 top-k, token稀疏常常得不偿失; 超长上下文上主 attention节省变大, indexer平方项又会暴露. 一套部署可以按长度选择稠密、平坦 DSA或层次 indexer, 但不同路径要用同一模型质量协议验收.

再决定接受哪种近似。 HISA的近似发生在 token搜索域, MISA发生在 indexer head集合, LISA还改变信息通道. 若要求最大限度复现原 DSA候选, 两阶段扩大召回后精排更合适; 若硬件瓶颈来自 head维度, 稀疏激活 heads更直接; 若任务需要分布式长程聚合, 线性并联提供另一条路径.

近似可以串联, 误差也会串联. 块粗筛后再选少量 heads, 然后量化打分, 每一步单独的高 recall不保证最终乘积仍高. 应保存逐阶段候选并计算条件召回: 在上一步仍包含教师 token的前提下, 下一步保留了多少. 这样能定位预算应该加在哪一级.

状态与生命周期。训练保存 indexer中间量和重计算元数据, Prefill批量构建 index K, Decode逐 token追加, 服务重排或分页 KV. 四个阶段对索引表示的要求不同. 一个论文级 PyTorch mask能验证数学, 不会自然变成不物化分数的融合 top-k kernel.

量化 scale、position、请求边界和 cache页表必须与 index K同步. 主 KV发生驱逐或迁移时, selector返回的位置仍要能映射到有效物理地址. 跨层共享则要注明索引何时刷新, 模型切换、prefix复用和 speculative分支如何失效. 这些生命周期规则决定动态路由能否安全进入长期运行的推理服务.

落到具体实现时, 一次 code review可以沿六个对象核对. index_q 的 shape决定 query侧并行, index_k_cache 决定每 token额外状态, score 的生成方式决定是否物化平方矩阵, selected_idx 的 dtype与排序决定地址开销, block_table 完成逻辑位置到物理页映射, 主 attention的在线 softmax统计保证多个离散片段共享分母. 任何一个对象没有明确 shape和生命周期, 复杂度分析都可能漏掉真实工作.

线上异常也可沿同一顺序定位. Recall下降而 index分数范围正常, 先检查训练域与共享层跨度; 多卡结果分叉, 检查 tie-break、scale和全局位置; selector时间上升, 检查 cache长度、量化字节与 top-k归并; 主 attention时间上升, 检查候选页分散度和重复索引; 只有特定 batch出错, 检查 sequence length、packed边界与 block table. 这种定位方式比统一增加候选预算更快找到真正瓶颈.

读完这条路线后, 可以把任何 token selector还原成一条可测的数据路径: 从 hidden state生成代理 Q/K, 产生候选, 用何种目标对齐教师, 将整数位置送入 Sparse Attention, 并在 cache和设备拓扑上找到对应 KV. 这条路径完整以后, recall、复杂度与真实时延才处于同一口径.

3.2. Token Indexer 的概率对象

主注意力在第 $\ell$ 层、query 位置 $t$ 对历史 token $i$ 的 logit 为

z_{t,i}^{(\ell,h)}= \frac{\langle q_t^{(\ell,h)},k_i^{(\ell,h)}\rangle}{\sqrt{d_h}}.

细粒度 indexer 用更便宜的表示 $\bar q_t,\bar k_i$ 计算代理分数 $r_{t,i}$ ，只把候选集合 $S_t$ 交给主 attention。它不需要精确重建所有 $z_{t,i}$ ，只需保证会影响输出的 token 进入 $S_t$ 。这把回归问题改成了极端尾部排序问题：绝大多数低分 token 排序错一点没有关系，cutoff 附近和少数高质量 token 的错误才危险。

若教师分布为 $p_i=\operatorname{softmax}(z)_i$ ，候选覆盖质量

C(S_t)=\sum_{i\in S_t}p_i

比普通 top-k overlap 更贴近输出。教师前 100 个 token 中，第一名可能占 40% 概率，其余各自很小；漏第一名和漏第 100 名不应同罚。另一方面，主模型在稀疏候选内会重新 softmax，概率质量高仍不保证 Value 方向正确，最终还需比较 attention 输出和任务。

Indexer 的目标也可以写成学习一个 inclusion probability $P(i\in S_t\mid h_t,h_i)$ 。实际 top-k 施加固定容量，token 之间相互竞争，单点概率并不独立。一个 token 分数没有下降，只因加入更多高分干扰也会掉出集合。长度增加会改变候选基数和极值分布，训练长度上的分类阈值不能直接外推。

低维代理保留什么。令投影矩阵 $W_q,W_k\in\mathbb R^{d_r\times d}$ ，最简单代理为

r_{t,i}=\langle W_qh_t,W_kh_i\rangle,\qquad d_r\ll d.

若主 attention 有多个 head，单一代理相当于要求它们的相关性能够压缩进同一个低维空间。Lightning Indexer 一类多头代理允许

r_{t,i}=\sum_{g=1}^{G}w_{t,g} \operatorname{ReLU}(\langle \bar q_{t,g},\bar k_{i,g}\rangle),

不同 indexer head 捕获不同关系，query 权重 $w_{t,g}$ 决定当前使用哪些子空间。表达能力提高，扫描成本从 $d_r$ 变成多个 head 的总维度；MISA 进一步选择少量活跃 indexer head，正是在这个成本项上做稀疏。

ReLU 截掉每个 head 的负点积，也带来零点敏感。若聚合权重非负，多个弱正相关 head 可以累积成高分，负点积不会拉低总分；若聚合权重允许为负，正的 ReLU 响应仍可能产生负贡献。截断是否损失任务所需的有符号关系，要结合权重、分数分布与消融一起看。

低维投影的可行性依赖相关结构的有效秩。若所有主 head 的重要方向彼此独立，需要很大 $d_r$ 才能同时保留；若关系共享少数主题、实体和位置方向，小 indexer 足够。可以对教师 QK 或重要 token 对做谱分析，也可直接扫 $d_r$ —recall 曲线。维度翻倍后 recall 几乎不变，瓶颈可能在训练目标或 top-k；仍持续上升则说明代理容量不足。

排序损失与 cutoff 学习。对每个 query 构造教师正集合 $P_t$ 与负集合 $N_t$ ，pairwise 排序损失可写为

L_{rank}=\sum_{i\in P_t}\sum_{j\in N_t} \log(1+e^{-(r_{t,i}-r_{t,j})}).

全量正负配对成本高，容易负样本失衡。随机负例大多非常容易，梯度很快趋近零；最有价值的是教师 cutoff 附近、语义相似但不重要、位置很近却无关的 hard negatives。难例过度集中又可能把模型训练成只区分微小教师噪声，应混合随机、近邻与边界负例。

分布蒸馏用代理 softmax $\hat p$ 拟合教师 $p$ ：

L_{KL}=\sum_i p_i\log\frac{p_i}{\hat p_i}.

它保留概率质量，却要求全历史归一，训练成本接近平坦扫描。教师分布极尖时，梯度集中在少数 token；分散时需要代理重建大量尾部，而部署只保留 top-k。可以截断教师质量到较宽候选 $k'>k$ ，在集合内蒸馏并增加集合外 margin。

集合监督直接优化 top-k overlap 的连续近似，目标更贴近部署，但可微排序开销与数值稳定性更复杂。实践中常先用稠密或较宽监督预训练 indexer，再用硬选择进入语言模型训练。第一阶段学会基本召回，第二阶段让主模型适应剩余近似。若一开始就硬 top-k，未选 token 没有梯度，早期随机错误容易固化。

Cutoff margin 是部署稳定性的核心。设第 $k$ 、 $k+1$ 名教师重要性很接近，它们互换可能无害；若第 $k+1$ 实际承载关键证据，固定预算本身已不够。训练既要提高代理对教师排序的 margin，也要判断教师边界是否有真实质量断层。报告代理 margin、教师概率差和替换后的输出变化，能区分「集合不稳定但等价」与「关键 token 反复被挤出」。

固定 $k$ 的 loss 容易忽略 query 难度。简单 query 只需几十个 token，复杂多跳可能需要数千。可以预测每个 query 的预算或按累积置信度停止，训练时再惩罚候选成本：

L=L_{task}+\lambda\,\mathbb E[|S_t|].

$\lambda$ 控制质量—成本工作点。预算预测若没有校准，困难 query 可能反而过度自信地选择更少；需用 oracle 曲线监督不确定性，并设置最低局部与全局保障。

主任务梯度的归因边界。若 indexer 只把硬 top-k 的整数索引交给主 attention，语言模型 loss 不会沿这些索引传回评分参数，选中 token 也不会自动建立这条梯度路径。训练 indexer 需要额外的评分监督，或让软门控、连续松弛、自定义反向等机制真正进入输出计算。Straight-through estimator 拼接硬选择前向与近似反向，梯度有偏；Gumbel 与软松弛可以增加探索，但是否节省训练计算取决于实现。判断端到端训练，要沿输出回到评分参数检查计算图。

主模型还能适应一个有偏 indexer：把重要信息复制到更容易被召回的位置，或在局部 token 中形成摘要。最终任务变好，不代表 selector 更接近稠密 attention。对原生稀疏模型，稠密教师 recall 只是诊断，不能当唯一目标；还应做候选干预、信息搬运分析和任务证据覆盖。

反过来，强制 indexer 完全复现稠密 QK 也可能限制新模型。稠密 attention 中大量低权重连接没有必要保留，某些高权重模式只是训练历史的副产物。目标应是以较小候选保持或提高任务能力，而非把每层热图复制成缩小版。Oracle 需要同时包含「模仿稠密」与「对任务 loss 的反事实贡献」两种口径。

HISA：先缩小搜索域。平坦 token indexer 对每个 query 扫描全部历史 index K，主 attention 已从 $N$ 降到 $k$ ，selector 仍保留 $O(Nd_r)$ 的逐 token读取。HISA 用层次结构先选择较粗区域，再在区域内做细粒度排序。设历史分成 $B=N/b$ 块，先选 $k_b$ 块，再扫描其中 $k_bb$ 个 token：

C_{HISA}=O(Bd_c+k_bbd_r),

其中 $d_c$ 是粗索引维度。相对平坦 $O(Nd_r)$ ，收益取决于 $k_bb\ll N$ 和粗层足够便宜。

两层 recall 满足条件分解。目标 token 被最终选中的概率为

R_{total}=P(\text{粗层命中}) P(\text{细层命中}\mid\text{粗层命中}).

若两层各 95%，总 recall 约 90.25%。因此粗层通常要取较宽候选，让第二层负责精排。只报告细层在已命中块内的 99% recall，会隐藏粗筛损失。

粗单元可以与 KV 页对齐，选中块后读取连续 index K 与主 K/V，减少随机访问。算法块太大，块内扫描浪费；太小，粗索引和 top-k 变长。最佳粒度同时由语义聚集、分页大小和 GPU tile 决定。把论文块数照搬到另一种 KV page 或硬件上，成本曲线会变化。

层次索引也有更新优势。Decode 每形成一个 token，只追加细 index K；块完成后才生成一个粗摘要。未满块由局部或单独 open-block 路径处理。Prefix cache 可以同时复用粗、细索引，前提是路由器版本和块边界一致。

HISA 的失败定位很清楚：目标不在粗候选，问题属于区域摘要或粗预算；粗候选包含目标、细层没选中，问题属于 token indexer；最终 token 已选中、主 attention 仍错，问题在主模型或 Value 竞争。逐层保存候选能让质量下降有可操作解释。

MISA：先缩小评分子空间。多头 indexer 的成本随 $G$ 增长，但不同 query 未必需要全部 heads。MISA 为每个 query 选择活跃 head 集合 $A_t$ ，代理分数变为

r_{t,i}=\sum_{g\in A_t}w_{t,g} \operatorname{ReLU}(\bar q_{t,g}^T\bar k_i).

这里的 index key $\bar k_i$ 在 heads 间共享。若 $|A_t|=g\ll G$ ，逐 token 点积从 $G$ 个 head 减到 $g$ 个，但仍要扫描历史共享 key，常驻 index K 缓存不会随 head 数按 $G/g$ 缩小。路由器还需用全部 heads 读取块摘要。算术减少能换来多少显存读取和延迟收益，取决于 tile 内复用与实际 kernel。

MISA 用块平均 index key 估计各 head 对当前前缀的贡献：将加权 ReLU 响应的绝对值在块间取平均，再选择分数最高的 heads，无需另训一个 head selector。独立衡量各 head 仍可能低估互补性：某 head 单独无法召回目标，却在与另一 head 聚合后改变排序。成对消融可以检查这种交互，代价是更多评测计算。

ResMiSA 则先利用 $\operatorname{ReLU}(z)=z+\operatorname{ReLU}(-z)$ ，把全部 heads 的线性项汇成一次共享点积，再只计算部分 heads 的非线性残差。Indexer LISA 进一步只保留共享线性项，并用候选精排或跨层复用改善效率。这些操作发生在索引评分里，与下面并联线性 attention 的同名 LISA 是两项不同工作。具体分解与计算流见共享线性项与残差索引。MISA 的共享 key、块摘要路由和活跃 head 评分对应原论文 §3–4。

Head 坍缩有两种表现。一种是永远选同几个 head，MISA 退化成静态裁剪；如果质量不降，这可能说明其余 head 冗余。另一种是 query 权重很分散却硬选很少，重要子空间随样本丢失。应报告 head 使用频率、条件熵、不同领域的活跃组合与 oracle head 上限，而非把均衡使用当作唯一健康标准。

MISA 与 HISA 可以组合：先选区域再用少 heads扫 token，或先用廉价公共分数全扫再对少数区域启用多头精排。两条顺序的乘法量看似相似，内存路径不同。前者避免读取无关区域的所有 head，后者保留全局公共召回。组合实验要分别记录条件 recall，不能把两篇论文各自的速度倍数直接相乘。

并联 LISA：线性状态与稀疏分支。纯 top-k 将未选 token 排除在当前主 attention 之外，容易损失广泛但单点权重不高的信息；这些 token 的缓存可以仍然保留。以非负特征映射 $\phi$ 、非零归一化分母为例，归一化线性注意力通过前缀统计聚合历史：

o_t^{lin}=\frac{\phi(q_t)^T\sum_{i\le t}\phi(k_i)v_i^T} {\phi(q_t)^T\sum_{i\le t}\phi(k_i)}.

前缀统计可增量更新，Decode 每步不必扫描全部 token。线性通道提供低分辨率全局背景，稀疏 selector 再读取少量高精度原始 KV，结构上类似「全局基底加稀疏残差」。

线性聚合的误差和 top-k 不同。它不会因候选外 token完全消失，却受特征映射近似和有符号 softmax困难影响；精确实体、极值和少量关键 token 可能被平均。稀疏残差应优先召回线性通道难以表达的部分，而非机械模仿完整 attention top-k。训练中若两路同时存在，selector 学到的目标会随线性分支能力变化。

可把理想输出分解为

o_t=o_t^{lin}+\Delta o_t.

这里将 $\Delta o_t$ 定义为目标输出与线性输出之差，用来分析两路应如何分工；并联 LISA 实际采用门控融合，未必显式计算这个差值。用有限原始 token 补足线性通道难以表达的信息，是一种可检验的训练思路。若仍只模仿稠密 attention 的高权重位置，selector 可能重复选择线性通道已经覆盖的背景；可以通过加入某 token 后的误差变化和分支干预检查实际分工。

线性状态也进入缓存生命周期。每层保存前缀统计矩阵与归一化向量，Prefix cache 命中时一起复用；分支生成从同一前缀复制状态；滑窗若要删除旧 token，普通累加状态还需支持减法或分段统计。它省去扫描，不等于没有状态成本。

LISA 类混合路线对长程多证据可能更稳：大量背景经线性通道聚合，top-k 留给尖锐依赖。若任务需要精确恢复多个独立值，线性混合会发生信息碰撞，最终仍依赖稀疏候选全部命中。评测要区分主题聚合、精确检索和多跳组合，单一 needle 不能验证两路分工。

4. 跨层共享为什么成为第二条主线

每层独立构建 index Q/K、扫描历史和执行 top-k，selector 成本会乘层数。相邻层表示具有相关性，CSA、HCA 等路线尝试复用某种索引状态。共享对象不同，误差来源也不同：共享 index K 保留各层 query，但历史 token 的代理表示较旧；共享候选连 query差异也一起忽略；共享主 KV 则属于注意力表示压缩，不能与 selector 复用混为一谈。

设第 $\ell$ 层独立候选为 $S_t^{(\ell)}$ ，复用上一刷新层 $r$ 的候选。可用 Jaccard 和教师质量衡量：

J_{\ell,r}=\frac{|S_t^{(\ell)}\cap S_t^{(r)}|} {|S_t^{(\ell)}\cup S_t^{(r)}|},\qquad C_{\ell\leftarrow r}=\sum_{i\in S_t^{(r)}}p_{t,i}^{(\ell)}.

高 Jaccard 说明集合相似，高 $C$ 才说明复用集合覆盖当前层重要质量。两个集合可能只差少数 token，恰好那几个在当前层概率很大；只看 overlap 会低估风险。

固定每 $m$ 层刷新一次最简单。浅层、中层、深层的语义变化速度不同，统一 $m$ 未必合理。可以根据 hidden state变化、候选 margin 或预测的不确定性触发刷新。触发器也有成本，若它需要计算当前层完整 indexer，便失去节省；应使用便宜统计或定期探测。

共享 index K 比共享候选更柔和。历史 token 的低维 key 在刷新层生成，后续层用自己的 query 扫描这些 key，仍能随当前 query 改变排序。假设是历史表示在若干层内变化不大。可测同一 token 的 index K 余弦和用旧 K 得到的 teacher recall；若深层语义旋转明显，刷新跨度要缩短。

共享候选省去扫描和 top-k，收益更大。它假设相邻层既关注相似 token，也允许主 QK 在同一集合内重新分配权重。为了防止关键新 token无法进入，可在复用集合外加入小额探索候选，或让部分 head刷新。组合候选

S_t^{(\ell)}=S_t^{(r)}\cup E_t^{(\ell)}

用较小 $E$ 修补层间变化。探索如何生成决定成本，可从块粗筛、局部异常或少 heads打分获得。

跨层共享还影响梯度。多个层依赖同一 index K 时，更新信号汇到刷新层或共享投影；它可能学到跨层折中表示，也可能被深层大梯度主导。若使用 stop-gradient，后续层只能适应固定索引；若梯度全部回传，需要观察尺度和冲突。论文应说明共享张量的计算图，而不只画推理缓存箭头。