IndexCache: Accelerating Sparse Attention via Cross-Layer Index Reuse

发表时间: 2026-03 · arXiv:2603.12201 (Tsinghua & Z.ai)

文章标题：IndexCache：通过跨层索引重用加速稀疏注意力
作者/机构：Yushi Bai, Qian Dong, Ting Jiang, Xin Lv, Zhengxiao Du, Aohan Zeng, Jie Tang, Juanzi Li (清华大学, http://Z.ai)

A1 主要贡献

核心问题：自注意力机制的二次方复杂度（$O(L^2)$）是长上下文推理的根本瓶颈。尽管像DeepSeek稀疏注意力（DSA）这样的方法通过为每个查询选择top-k个最相关的token，将核心注意力计算从$O(L^2)$降低到$O(Lk)$，但其用于选择token的索引器（indexer）本身在每一层仍然具有$O(L^2)$的复杂度。随着上下文长度的增加，所有N个层的索引器总成本$O(NL^2)$成为一个不可忽视的延迟来源，尤其是在预填充（prefill）阶段，这限制了长上下文DSA推理的加速效果。

研究目标：本文旨在解决DSA中索引器带来的性能瓶颈。基于一个关键观察——由索引器产生的top-k选择在连续层之间高度相关（即存在跨层冗余），研究目标是能否在不降低模型质量的前提下，移除DSA中的大部分索引器，让多数层重用少数保留索引器层的top-k索引，从而大幅减少索引器计算开销。

创新点 (IndexCache)：本文提出了IndexCache方法，通过跨层索引重用，在DSA中消除了高达75%的索引器计算。其核心思想是将模型层划分为两类：
- F (Full) 层：保留其索引器，独立计算新的top-k索引。
- S (Shared) 层：不运行索引器，直接继承并重用其最近的前一个F层所计算的top-k索引。

为了确定和优化这种F/S层配置，本文提出了两种互补的方法：
1. 训练无关的IndexCache (Training-free IndexCache)：该方法适用于任何现成的DSA模型，无需更新权重。它提出了一种贪心层选择算法，通过在一个小型校准集上直接最小化语言模型损失，来确定保留哪些层的索引器（即最优的F/S模式）。
2. 训练感知的IndexCache (Training-aware IndexCache)：该方法在模型训练阶段优化参数以适应跨层共享。它引入了一种多层蒸馏损失，使得每个保留的索引器能够学习生成对所有其服务的层都有效的注意力分布，从而让简单的均匀交错模式也能达到与完整索引器相当的精度。

实验结果表明，在一个30B的DSA模型上，IndexCache可以在移除75%索引器计算的同时，性能几乎没有下降，与标准DSA相比，在200K上下文长度下实现了高达1.82倍的预填充加速和1.48倍的解码加速。在生产级GLM-5模型上的初步实验也证实了该方法的有效性（图1）。

(b) IndexCache 推理过程

需要: 输入 X, 层 1 ... N, 模式 c
1: for ℓ = 1 to N do
2:   if cℓ = F then
3:     I(ℓ) ← INDEXERℓ(X)
4:     T(ℓ) ← Top-k(I(ℓ))
5:     Tcache ← T(ℓ)
6:   else {cℓ = S}
7:     T(ℓ) ← Tcache ▷ 重用
8:   end if
9:   X ← SPARSEATTN(X, T(ℓ))
10:  X ← FFNℓ(X) ▷ + norm, residual, 等.
11: end for

A3 背景知识

2.1 DeepSeek稀疏注意力 (DSA)

DSA工作原理及瓶颈。DeepSeek稀疏注意力（DSA）【【23】，Deepseek-v3. 2: Pushing the frontier of open large language models，2025】将每个注意力层分解为选择和计算两个阶段。一个轻量级的“闪电索引器”（lightning indexer）首先使用多头ReLU门控点积为当前查询对所有先前的token进行评分，然后选择得分最高的top-k个位置。核心的注意力计算只在这个稀疏子集上进行，从而将每层的核心注意力复杂度从$O(L^2)$降低到$O(Lk)$，其中$k=2048 \ll L$（$L$为序列长度）。索引器为效率而设计：它使用较少的头、低秩投影和FP8算术，使其每个FLOP的成本比主多头潜在注意力（MLA）【【22】，Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model，2024a】便宜一个数量级。

DSA的训练过程。DSA通过两阶段的持续预训练进行实例化。首先，一个简短的密集预热阶段仅通过KL散度蒸馏，针对每一层聚合的完整注意力分布来训练索引器，而所有其他参数都保持冻结。然后，一个更长的稀疏训练阶段激活top-k选择，并联合优化整个模型，索引器在一个分离的计算图上接收蒸馏梯度。

引出核心问题。尽管有这些效率提升，索引器本身仍然以$O(L^2)$的复杂度运行：在每一层，它都必须独立地对所有先前的token进行评分以确定自己的top-k集合。在一个拥有N层的模型中，索引器的总成本为$O(NL^2)$，在长上下文长度下，这成为总注意力预算的一个重要部分。一个自然的问题是，所有N个逐层索引器的计算是否真的都是必需的，以及层间的冗余是否可以被利用。

2.2 Token选择的跨层稳定性

跨层稳定性观察。答案来自一个更广泛的经验性发现：重要token的集合在连续的Transformer层中非常稳定。Deshmukh等人【【7】，Kascade: A practical sparse attention method for long-context llm inference，2025】和Gao等人【【11】，Hysparse: A hybrid sparse attention architecture with oracle token selection and kv cache sharing，2026】都观察到，相邻层共享其绝大部分的top-k注意力权重，并利用这一点，指定少数“锚点层”计算完整注意力，而让中间层重用锚点层的top-k索引。

将稳定性观察应用于DSA的挑战与机遇。关键在于，这两种方法都依赖于完整注意力作为识别重要token的“神谕”（oracle）。而在DSA中，完整注意力被完全消除——由轻量级索引器取而代之。这就提出了一个据我们所知尚未解决的问题：索引器的输出是否也表现出跨层稳定性？如果是，我们就可以应用相同的共享原则来消除冗余的索引器计算，而根本不需要任何完整注意力的神谕。此外，在质量下降之前可实现的最大重用率是多少？我们能否调整模型以弥合由激进的索引重用所带来的性能差距？我们在下一节中提出了两种新颖的互补方法来实现这一点。

A2 方法细节

符号定义。我们用$N$表示Transformer层的数量，用$L$表示序列长度，用$k$表示每个查询选择的token数量。在第$ℓ$层，闪电索引器为查询位置$t$生成一个得分向量$I^{(ℓ)}_t \in R^L$，从中提取top-k索引集$T^{(ℓ)}_t = \text{Top-k}(I^{(ℓ)}_t)$（本文中$k=2048$）。我们用$p^{(ℓ)}_t$表示第$ℓ$层在位置$t$的聚合注意力分布（通过对各头的softmax注意力权重进行平均得到），用$q^{(ℓ)}_t = \text{Softmax}(I^{(ℓ)}_t)$表示索引器的输出分布。

IndexCache核心机制。IndexCache通过将$N$个层划分为两种角色来修改DSA，这由一个二进制模式字符串$c = c_1c_2 \cdots c_N$（其中$c_ℓ \in \{F, S\}$）编码：
- F (Full) 层：该层保留其索引器，对所有先前的token计算新的$T^{(ℓ)}$，并对所选子集执行稀疏核心注意力，与标准DSA相同。
- S (Shared) 层：该层没有索引器。它继承最近的前一个F层的索引集，即$T^{(ℓ)}_t \leftarrow T^{(f(ℓ))}_t$，其中$f(ℓ) = \max\{j < ℓ : c_j = F\}$，并直接使用这些继承的索引应用稀疏核心注意力。

推理过程与设计目标。第一层始终为F层，以提供初始索引。在推理时，S层只需跳过索引器的前向传播，并重用其F层前驱缓存的索引张量。图2展示了这种修改的简单性：对每层推理循环的唯一改变是增加一个条件分支，要么运行索引器，要么复制缓存的索引。关键的设计问题是如何选择模式$c$。如果大多数层可以安全地共享索引，那么总索引器成本$O(NL^2)$的很大一部分可以被消除，而核心注意力成本$O(NLk)$保持不变。我们提出了两种方法：一种是训练无关的方法，通过在已建立的DSA模型上进行贪心搜索来确定$c$（第3.1节）；另一种是训练感知的方法，通过多层蒸馏损失联合优化索引器参数以实现跨层共享（第3.2节）。

3.1 训练无关的IndexCache

训练无关IndexCache的目标。给定一个预训练的DSA模型，我们的目标是找到一个模式$c$，在最小化对模型质量影响的同时，最大化S层的数量。我们首先讨论为什么最直接的方法会失败，然后介绍我们的贪心搜索算法。

3.1.1 为什么均匀交错是次优的

均匀交错策略的局限性。最简单的策略是均匀交错：保留每$r$层的索引器，跳过其余的（例如，对于$r=4$，模式为FSSSFSSS...）。然而，这种策略忽略了索引器的重要性在不同层之间显著变化的事实。我们凭经验观察到，某些层，特别是网络早期和过渡区域的层，对索引器移除的敏感度远高于其他层。均匀交错可能会移除一个关键的索引器，同时保留一个冗余的索引器，导致明显的质量下降（定量比较见第4.3节）。这促使我们采用一种数据驱动的方法：让模型自己告诉我们哪些索引器是可有可无的，从而引出我们的贪心层选择算法。

3.1.2 层选择算法

贪心搜索算法。我们提出一种贪心搜索方法，它逐步将F层转换为S层，使用在小型校准集上的语言建模损失作为下游质量的代理指标。

校准集定义。我们从训练数据中缓存$B$个小批量（mini-batch）。所有候选模式都在完全相同的批次上进行评估，确保损失差异仅反映模式变化的影响，而非数据方差。损失来自于在整个批次$D$上的一次前向传播：$EVAL\_LOSS(M, D, c)$。

搜索过程详解。从全F层的基线（所有$ℓ$的$c_ℓ = F$）开始，算法进行$K$步，其中$K$是目标S层的数量（例如，为保留1/4的索引器，$K = 3N/4$）。在每一步中，我们遍历所有当前的F层（不包括第一层），试探性地将每个F层翻转为S层，评估由此产生的LM损失，并提交导致最低损失的翻转。算法1展示了完整的流程。

算法复杂度与并行优化。从全F层到全S层的完整搜索需要进行$N(N-1)/2$次前向传播。当流水线并行将模型划分为$P$个阶段时，我们通过将层分成$P$个块（每个块的第一层固定为F层）并在每一步内顺序搜索这些块来加速搜索：每个块中的最佳翻转在搜索下一个块之前被提交，因此每步最多放置$P$个层，总前向传播次数大约减少$P$倍。

贪心解的三个良好特性。尽管贪心搜索不保证全局最优，我们一致观察到三个令人满意的特性：
1. 搜索到的模式在相同保留率下优于均匀交错模式（表2）。
2. 如右图所示（标记了1/2、1/4和1/8保留率的步骤），每步的LM验证损失曲线揭示了“简单”层（前20步）和“关键”层（35步之后）之间的明显分离，这表明索引器的重要性存在一个自然排序。
3. 结果在不同的校准集上是稳定的，表明这种重要性排名是模型的内在属性，而不是数据伪影。此外，LM损失可作为下游任务的有效代理，因为较低的LM损失与较好的任务性能正相关。

3.2 带有分层蒸馏的训练感知IndexCache

训练感知IndexCache的动机。训练无关的IndexCache不需要更新权重，但其局限性在于每个索引器最初只为自己的层服务。当从头开始或通过持续预训练来训练DSA模型时，我们可以做得更好：明确地训练每个保留的索引器以同时服务多个层。

多层蒸馏损失函数。在标准的DSA训练中（第2.1节），第$ℓ$层的每个索引器通过KL散度与其自身层的聚合注意力分布$p^{(ℓ)}_t$进行蒸馏：$L_I = \sum_t D_{KL}(p^{(ℓ)}_t || q^{(ℓ)}_t)$。我们将其推广为一个多层目标。假设第$ℓ$层是一个保留的F层，而第$ℓ+1, \ldots, ℓ+m$层是后续将重用其索引集$T^{(ℓ)}_t$的S层。多层蒸馏损失为：

直观上，这鼓励索引器预测一个对其服务的所有层都共同有用的top-k集合，而不是仅仅对第$ℓ$层过拟合。

梯度等价性证明。一个自然的担忧是，优化KL项之和是否会引入意想不到的相互作用。我们证明了多层损失有一个清晰的解释：它产生的梯度与对单一平均目标进行蒸馏完全相同。

定义平均目标$\bar{p}_t = \frac{1}{m+1} \sum_{j=0}^{m} p^{(ℓ+j)}_t$和相应的单目标损失：

命题1及其证明。命题1：$\nabla_{\theta} L_{I_{multi}} = \nabla_{\theta} L_{I_{avg}}$。
证明：由于$q^{(ℓ)}_t$是$D_{KL}(p || q^{(ℓ)}_t)$中唯一依赖参数的项，p的熵在微分下消失：$\nabla_{\theta} D_{KL}(p || q^{(ℓ)}_t) = -\nabla_{\theta} \sum_s p(s) \log q^{(ℓ)}_t(s)$。应用于公式1：

多层蒸馏的物理解释。命题1表明，多层蒸馏不仅仅是一种启发式正则化器——它完全等同于将索引器蒸馏到目标层注意力分布的质心。因此，索引器学会了预测一个能够共同覆盖所有服务层中重要token的共识top-k。

实现上的考量。尽管两种损失公式产生相同的梯度，我们在实践中采用$L_{I_{multi}}$以提高实现效率。当后续层是S层时，它只需要接收当前层预测的$q^{(ℓ)}$。相比之下，使用$L_{I_{avg}}$进行训练需要传递$q^{(ℓ)}$和$p^{(ℓ)}$，这会引入不必要的内存开销和额外的运行时成本。

训练过程。我们遵循标准的DSA两阶段训练程序。在预热阶段，我们使用$L_{I_{multi}}$训练F层中的索引器，同时保持所有其他参数固定。在稀疏训练阶段，我们继续使用$L_{I_{multi}}$（仅在选定的top-k token上计算KL散度）来训练索引器，并额外包含LM损失来训练其余参数。

A4 实验环境

A5 实验结果

4.2 端到端推理加速

实验内容：在30B参数的DSA模型上，比较了原始DSA基线与两种保留率（1/2和1/4）的IndexCache的端到端推理性能。评测了在10K到200K不同上下文长度下的三个指标：预填充延迟、单并发下的每请求解码吞吐量、以及KV缓存完全利用时的总解码吞吐量。

实验结果（见表1和图3）：

分析结论：IndexCache对于长上下文场景尤其有价值，能有效加速预填充和解码阶段。

Table 1: 30B DSA模型在使用IndexCache（两种保留率）时的端到端推理性能。预填充时间：秒（越低越好）。单请求解码：单并发下每秒token数（越高越好）。满载解码：总每秒token数（越高越好）。解码吞吐量按每GPU报告。


图3: 30B模型上，IndexCache相对于DSA基线在三种推理设置下的相对加速比。DSA基线被归一化为100%。

4.3 训练无关的IndexCache结果

实验内容：在30B DSA模型上，比较了训练无关IndexCache在1/2、1/4和1/8三种保留率下的性能。每种保留率都对比了均匀交错模式和贪心搜索模式。

实验结果（见表2）：
- 长上下文任务：激进的均匀交错模式（如1/4和1/8）会导致显著的性能下降。而贪心搜索出的模式能基本弥补这一损失，在1/4保留率下性能与原始DSA相当。这证明了保留哪些层的索引器比保留多少层更重要。但在1/8的极端稀疏度下，即使是搜索模式也无法避免性能的明显下降。
- 长链推理能力：除了1/8均匀交错模式外，所有配置的通用与推理（G&R）平均分都与基线（74.6）持平。值得注意的是，1/4搜索模式在某些任务上甚至略优于DSA，表明移除冗余计算可能起到了轻微的正则化作用。

分析结论：训练无关IndexCache在不牺牲通用推理能力的前提下，通过贪心搜索能有效保持长上下文性能。

Table 2: 训练无关IndexCache在1/2、1/4和1/8索引器保留率下的结果。“Long”和“G&R”分别汇总了基准测试得分。我们比较了均匀交错模式和搜索模式。

4.4 训练感知的IndexCache结果

实验内容：使用多层蒸馏损失进行训练感知的IndexCache实验，保留率为1/2和1/4，均采用均匀交错模式。并进行了两项消融研究：1) 将均匀模式替换为贪心搜索模式；2) 移除跨层蒸馏损失。

实验结果（见表3）：
- 主要结果：训练感知的IndexCache即使在1/4保留率下，使用简单的均匀模式也能达到与DSA基线相当的性能。
- 模式敏感性消失：与训练无关的情况形成鲜明对比，在训练感知设置下，均匀交错模式的表现与贪心搜索模式相当，甚至略优。这表明，当模型经过共享感知的重新训练后，S层学会了适应继承的索引，而F层索引器也学会了生成更具泛化性的选择，从而消除了对特定层的敏感性。
- 跨层蒸馏的益处：移除跨层蒸馏损失会导致长上下文性能下降，证实了多层蒸馏目标的实际效用。它帮助索引器学习一个跨层泛化而非过拟合于单层的共识top-k。

分析结论：训练感知的方法使得IndexCache更加鲁棒，简单的均匀模式即可生效，而多层蒸馏损失是其中的关键。

Table 3: 训练感知IndexCache在1/2和1/4索引器保留率下使用均匀交错模式的结果。w/ searched pattern：使用贪心搜索模式替换均匀交错。w/o cross-layer loss：每个索引器仅针对其自身层进行蒸馏。

4.5 扩展性实验

实验内容：将训练无关的IndexCache应用于默认使用DSA的GLM-5（744B参数）模型。

实验结果（见表4和图1）：
- 趋势与30B模型一致：均匀交错模式在激进保留率下性能下降，而搜索模式能恢复性能。
- 在1/2保留率下，搜索模式的性能略超基线；在1/4保留率下，性能与基线非常接近（仅下降0.4分）。
- 在Artificial Analysis Index的所有测试中，1/2保留率的IndexCache表现与原始GLM-5模型几乎相同。

分析结论：IndexCache方法能有效扩展到生产规模的超大模型。

Table 4: GLM-5 (744B)上使用训练无关IndexCache的初步结果。

A7 补充细节

5.1 高效注意力

高效注意力机制概述。降低自注意力的二次方成本是核心研究方向。训练无关的稀疏方法在推理时通过固定模式、启发式策略或轻量级重要性估计引入稀疏性【【56】，H2o: Heavy-hitter oracle for efficient generative inference of large language models，2023；【49】，Efficient streaming language models with attention sinks，2024；【34】，Sparq attention: Bandwidth-efficient llm inference，2024；【40】，Quest: Query-aware sparsity for efficient long-context llm inference，2024；【18】，Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention，2024；【50】，Duoattention: Efficient long-context llm inference with retrieval and streaming heads，2025；【51】，Xattention: Block sparse attention with antidiagonal scoring，2025；【19】，Flexprefill: A contextaware sparse attention mechanism for efficient long-sequence inference，2025；【8】，Moa: Mixture of sparse attention for automatic large language model compression，2025；【57】，Sampleattention: Near-lossless acceleration of long context llm inference with adaptive structured sparse attention，2025；【52】，Lserve: Efficient long-sequence llm serving with unified sparse attention，2025c；【55】，Spargeattention: Accurate and training-free sparse attention accelerating any model inference，2025】，但可能因训练-推理不匹配导致长上下文误差累积【【16】，Lil: Less is less when applying post-training sparse-attention algorithms in long-decode stage，2026】。相比之下，可训练的稀疏方法在训练阶段就引入稀疏性，例如通过学习门控机制【【9】，Seerattention: Learning intrinsic sparse attention in your llms，2024；【10】，Seerattention-r: Sparse attention adaptation for long reasoning，2025】、端到端稀疏预训练【【54】，Native sparse attention: Hardwarealigned and natively trainable sparse attention，2025】、块级混合路由【【25】，Moba: Mixture of block attention for longcontext llms，2025；【58】，Infllm-v2: Dense-sparse switchable attention for seamless short-to-long adaptation，2025；【43】，Minicpm4: Ultra-efficient llms on end devices，2025b】或全注意力到稀疏注意力的蒸馏【【23】，Deepseek-v3. 2: Pushing the frontier of open large language models，2025；【35】，Ssa: Sparse sparse attention by aligning full and sparse attention outputs in feature space，2025】。本文工作的基础DSA【【23】】通过从全注意力中蒸馏一个轻量级索引器来选择top-k token，将核心注意力复杂度降至$O(Lk)$。除了稀疏性，混合架构通过将二次方复杂度的层与滑动窗口注意力【【1】，gpt-oss-120b & gpt-oss-20b model card，2025；【12】，Gemma 3 technical report，2025；【48】，Mimo-v2-flash technical report，2026；【42】，Longcat-flash-thinking-2601 technical report，2026】、线性注意力【【53】，Gated delta networks: Improving mamba2 with delta rule，2025d；【20】，Minimax-01: Scaling foundation models with lightning attention，2025；【4】，Nemotron-h: A family of accurate and efficient hybrid mamba-transformer models，2025；【41】，Kimi linear: An expressive, efficient attention architecture，2025a】或状态空间层【【13】，Mamba: Linear-time sequence modeling with selective state spaces，2024；【6】，Transformers are ssms: Generalized models and efficient algorithms through structured state space duality，2024；【21】，Jamba: A hybrid transformer-mamba language model，2024】交错来减少昂贵层的数量。

5.2 跨层共享

跨层共享相关工作。最近的研究表明，表征在相邻层之间表现出很强的一致性。这种结构特性常被用来减少计算冗余和加速推理。TidalDecode【【52】，Lserve: Efficient long-sequence llm serving with unified sparse attention，2025a】、LessIsMore【【52】，Lserve: Efficient long-sequence llm serving with unified sparse attention，2025b】、OmniKV【【14】，Omnikv: Dynamic context selection for efficient long-context llms，2025】和DELTA【【55】，Spargeattention: Accurate and training-free sparse attention accelerating any model inference，2025】重用周期性锚点层的top-k索引进行稀疏解码。Kascade【【7】，Kascade: A practical sparse attention method for long-context llm inference，2025】通过在跨层相似度矩阵上进行动态规划来形式化锚点层的选择，并指出头感知的重映射对于保持准确性至关重要。所有这些方法都依赖于锚点层的完整注意力来计算精确的top-k索引。独立地，跨层KV缓存共享通过让多个层重用相同的键值张量来减少内存【【36】，You only cache once: Decoder-decoder architectures for language models，2024；【5】，Reducing transformer key-value cache size with cross-layer attention，2024；【24】，Minicache: Kv cache compression in depth dimension for large language models，2024b；【31】，Swiftkv: Fast prefilloptimized inference with knowledge-preserving model transformation，2025；【59】，Mlkv: Multi-layer key-value heads for memory efficient transformer decoding，2025；【47】，Improving model representation and reducing KV cache via skip connections with first value heads，2025】。HySparse【【11】，Hysparse: A hybrid sparse attention architecture with oracle token selection and kv cache sharing，2026】统一了这两个方向，将全注意力层与继承top-k块索引和KV缓存的稀疏层交错。然而，所有这些方法都需要全注意力层作为神谕，而DSA完全移除了它。

IndexCache与现有工作的区别。IndexCache在两个关键方面有所不同。首先，神谕从根本上更便宜，因为我们共享的是DSA轻量级索引器的输出，而不是完整的$O(L^2)$注意力分数。其次，我们引入了系统性的技术来优化共享配置，包括用于识别最优结构布局的训练无关贪心搜索和用于参数适应的训练感知多层蒸馏损失。尽管我们在DSA上实例化了IndexCache，但其核心原则可以扩展到任何不依赖固定稀疏模式、而是涉及动态token选择步骤的稀疏注意力方法：例如，MoBA【【25】，Moba: Mixture of block attention for longcontext llms，2025】和NSA【【54】，Native sparse attention: Hardwarealigned and natively trainable sparse attention，2025】中的块级选择同样可以从跨层重用中受益。

A6 结论

本文提出了IndexCache，一种通过利用负责token选择的索引器在层间的冗余性来加速稀疏注意力的方法。IndexCache将层划分为少数保留索引器的F层和大多数重用继承top-k索引的S层，仅通过一个条件分支便可消除高达75%的$O(NL^2)$总索引器成本，且几乎没有性能下降。更广泛地说，我们的工作表明，以前仅适用于以全注意力为神谕的跨层共享原则，可以自然地扩展到稀疏注意力模型。随着稀疏注意力成为前沿大语言模型（如DeepSeek-V3.2, GLM-5）的默认配置，我们预计跨层索引重用将成为高效推理流水线的标准组成部分。

A7 附录

A 跨层Top-k索引重叠

跨层Top-k索引重叠率计算方法。为了通过实验验证DSA中top-k索引选择的跨层冗余性，我们计算了每层闪电索引器选择的top-k索引之间的成对重叠率。具体来说，对于每对层$(i, j)$，我们测量$|T^{(i)} \cap T^{(j)}|/k$（其中$k=2048$），并在一个包含768个长度为200K的样本的校准集上取平均。

热力图分析。图4将这个重叠率可视化为47层30B DSA模型的热力图，其中根据贪心搜索的1/4保留率IndexCache模式得到的共享块被标记为红色方框。可以观察到几个明显的模式：
- 对角线附近的高重叠：相邻层的重叠率在0.7-1.0之间，证实了连续层选择的token集合基本相同。
- 块状结构：热力图揭示了具有相互高重叠的明显层簇，例如3-5层、6-8层、17-30层、31-36层等，这表明模型组织成了内部token选择一致的功能块。
- 不均匀衰减：重叠率在块边界处的下降速度比块内部快，表明少数“过渡”层显著地改变了注意力焦点。
- 早期-晚期差异：热力图的左下角和右上角明显较暗（重叠率≤0.4），显示早期和晚期层关注的token子集根本不同。


图4: 30B DSA模型所有层对之间的成对top-k索引重叠率。根据贪心搜索的1/4 IndexCache模式确定的共享块已标记。

贪心搜索块与重叠簇的对比分析。比较红色方框（贪心搜索的共享块）和自然的重叠簇，可以发现一个有信息量的错位。虽然贪心搜索确实在一些视觉上明显的簇边界附近放置了F层，但这两个分区并不完全重合。根本原因在于，重叠是一个聚合指标：它计算了共享了多少token，但没有计算哪些token不同。在训练无关的设置中，权重是固定的，因此即使是一小部分不匹配的关键token也可能扰动一个层的隐藏状态，这种扰动会级联到所有下游层。早期层尤其脆弱，因为它们的扰动传播路径最长。

局部指标的局限性。这一观察也呼应了我们基于相似度的模式搜索的负面结果（附录C）：局部指标——无论是注意力输出的余弦相似度还是top-k索引重叠——都缺乏识别最优共享模式的判别力，因此必须进行端到端评估。

B 搜索到的模式

GLM-4.7-Flash 30B DSA的搜索模式：
- 保留1/2 'F'层: FSFSFSSSSFSFFFFSFFSSFFSFFFSSFFSSFSSSSFSFFFSFSSF
- 保留1/4 'F'层: FSFSFSSSSFSSSFSSFFSSFSSFSSSSFSSSFSSSSFSSSSSSSSS
- 保留1/8 'F'层: FSSSFSSSSSSSSFSSSFSSSSSFSSSSFSSSSSSSSFSSSSSSSSS

GLM-5的搜索模式：
- 保留1/2 'F'层: FFSFSSSFSSFFFSSSFFFSFSSSSSSFFSFFSFFSSFFFFFFSFFFFFSFFSSSSSSFSFFFSFSSSFSFFSFFSSS
- 保留1/4 'F'层: FFSFSSSFSSFSFSSSSSSSFSSSSSSFSSSFSFSSSSFFFFFSSSFFSSSFSSSSSSSSFSSSFSSSSSSFSFSSSS

C 基于相似度的模式搜索

报告失败尝试的动机。我们认为论文不仅应报告正面结果，也应报告负面（或不成功）的结果。在得出第3.1.2节中描述的基于损失的贪心搜索之前，我们探索了一个看似自然的选择：通过直接测量当一个索引器被跨层重用时注意力输出的相似程度来选择共享模式。尽管这种方法在理论上是合理的，并且计算上比贪心搜索更便宜，但最终证明它不足以作为下游质量的代理。我们在此描述它，以求完整性并提供为什么基于损失的搜索是必要的见解。

相似度矩阵的构建。给定一个有$N$层的DSA模型，我们在一个校准集上执行$N$次单次前向传播，对于每个层对$(i, j)$（其中$i>j$），计算以下两者之间的余弦相似度：
1. 第$i$层使用其自身索引器时的核心注意力输出（即原始模型）。
2. 第$i$层重用第$j$层索引器时的核心注意力输出（即如同第$i$层是一个从第$j$层共享的S层）。
这产生一个$N \times N$的下三角相似度矩阵$S$，其中$S_{i,j}$量化了第$j$层的索引可以多好地替代第$i$层自身的索引。直观上，如果$S_{i,j}$接近1，那么第$i$层可以安全地跳过自己的索引器计算，并以最小的注意力输出失真重用第$j$层的索引。

动态规划求解。给定相似度矩阵$S$和目标F层数$M$（等价于$N-M$个S层），我们寻求最大化总相似度的模式$c^*$：

其中$src(ℓ)$表示第$ℓ$层继承其索引的最近的前一个F层（即$src(ℓ) = \max\{j < ℓ : c_j = F\}$）。
这可以通过动态规划精确求解。令$dp[i][k]$表示使用恰好$k$个F层，对层$1, \ldots, i$所能达到的最大累积相似度，且第$i$层本身是一个F层。转移方程考虑所有可能的先前F层：

其中求和项考虑了$j$和$i$之间的所有S层，每个都重用第$j$层的索引。我们通过回溯DP表来恢复最优模式。

基于相似度的搜索结果不佳。与均匀交错类似，DP搜索出的模式在下游任务上表现出与原始DSA模型相比同样显著的质量下降，如表5所示。换句话说，尽管明确优化了最大跨层相似度，但得到的模式与朴素的均匀基线相比没有提供任何有意义的优势。相比之下，基于损失的贪心搜索（第3.1.2节）产生的模式显著优于均匀和相似度最优模式，尤其是在激进的保留率下。

Table 5: 训练无关的基于相似度搜索模式的评估结果。

相似度作为代理指标失败的原因。根本问题在于，逐层输出相似度是一个局部指标：它衡量单个层的注意力输出在孤立情况下被保留得如何，而没有考虑微小的扰动如何跨越剩余的层传播。两个层的注意力输出可能几乎相同（$S_{i,j} \approx 1$），但在对下游质量至关重要的细微方面有所不同：例如，重用的索引可能错过少数关键token，其重要性仅在后续层的推理步骤中才变得明显。这些细微的不匹配通过各层累积，导致最终质量出现不可忽视的下降，而层级局部的相似度得分无法预测这一点。基于损失的贪心搜索通过直接优化一个全局指标，即LM损失，来避免这个陷阱，该损失捕捉了每个共享决策对模型输出分布的端到端影响。这使其能够识别出“关键”层（那些必须保留其索引器以避免级联错误的层），而基于相似度的方法则将这些层与其邻居视为可互换的。

D 评测设置

所有基准测试均使用温度1.0、top-p=0.95和top-k=40进行评估。对于长上下文任务，我们设置总上下文窗口为200K token，并为输出保留32K。对于通用与推理任务，我们允许最大输出长度为64K token。在MRCR v2上，我们报告2、4和8针设置的平均分。在GraphWalks上，我们报告Parent-type和BFS-type问题的平均分。在RULER上，我们报告上下文长度从4K到128K的所有实例的分数。对于MRCR v2和GraphWalks，我们仅包括输入长度在有效输入预算（200K-32K=168K token）内的实例。对于LongBench v2、RULER和AA-LCR，我们包括所有实例，并对超过168K token的实例应用中间截断。