IndexCache：通过跨层索引复用加速稀疏注意力 IndexCache: Accelerating Sparse Attention via Cross-Layer Index Reuse
通过跨层复用索引，移除75%计算，实现1.82倍加速
前置知识
稀疏注意力（Sparse Attention）
稀疏注意力是一种优化技术，通过让每个查询只选择最相关的 token 子集进行计算，而不是关注所有前序 token。传统自注意力的计算复杂度为 O(L^2)，其中 L 是序列长度。稀疏注意力将其降低到 O(Lk)，其中 k 是选择的 token 数量，通常 k 远小于 L。这种技术对于长上下文推理（如 100K+ token）至关重要，可以显著减少计算量和推理延迟。
本文的 IndexCache 建立在稀疏注意力之上，理解稀疏注意力的基本原理是理解本文贡献的前提
DeepSeek Sparse Attention (DSA)
DSA 是 DeepSeek 提出的生产级稀疏注意力机制，采用两阶段架构：首先使用轻量级 lightning indexer 对所有前序 token 进行评分，然后通过 top-k 操作选择最高分的 token 子集，最后仅在该子集上计算主注意力。Lightning indexer 使用少头数、低秩投影和 FP8 算术，使其比主注意力便宜一个数量级。DSA 通过两阶段持续预训练引入：先用 KL 散度蒸馏训练 indexer，再联合优化整个模型。
本文的核心创新直接针对 DSA 的 indexer 进行优化，理解 DSA 的架构是理解本文方法的必要条件
跨层 token 选择稳定性
这是 Transformer 模型中观察到的一个经验现象：相邻层选择的重要 token 集合高度相似。已有研究发现相邻层共享绝大部分 top-k 注意力质量，这种稳定性允许在某些锚点层计算完整注意力，而中间层复用这些索引。在本文中，作者首次验证了 DSA 的 lightning indexer 输出也表现出这种跨层稳定性，相邻层共享 70-100% 的选择 token。
这是本文的核心洞察和出发点，理解这个现象才能理解为什么跨层索引复用是可行的
KV 缓存（KV Cache）
在自回归解码中，为了避免重复计算，模型会缓存每一层之前 token 的 Key 和 Value 向量。KV 缓存的大小与序列长度和层数成正比，在长上下文场景下会占用大量 GPU 显存。本文提到的跨层 KV 缓存共享是另一种优化方向，但与本文的索引复用是正交的。
KV 缓存是理解长上下文推理瓶颈的重要概念
语言建模损失（Language Modeling Loss）
语言建模损失是训练语言模型的核心目标，通常使用交叉熵损失函数，即最大化给定前序 token 预测下一个 token 的概率。在本文中，训练无关 IndexCache 使用 LM 损失作为评估指标来指导贪心搜索最优层选择模式，因为更低的 LM 损失通常与更好的下游任务性能正相关。
本文的贪心搜索算法使用 LM 损失作为代理指标来评估不同层选择模式的质量
研究动机
长上下文智能体工作流已成为大型语言模型的核心应用场景，这使得注意力效率对推理速度和服务成本至关重要。虽然 DSA 通过引入 lightning indexer 成功将主注意力从 O(L^2) 降低到 O(Lk)，但 indexer 本身仍保持 O(L^2) 复杂度，且必须在每一层独立运行。对 30B DSA 模型的性能分析表明，indexer 在总延迟中的占比随上下文长度急剧上升：在 200K token 的预填充阶段，indexer 占比达到 81%；在解码阶段占比也达到 38-41%。这意味着即使主注意力已经稀疏化，indexer 的累积成本跨 N 层总计 O(NL^2)，在长上下文下仍成为注意力预算的显著瓶颈。具体数据表明，当上下文长度从 10K 增加到 200K 时，预填充延迟从 0.57 秒增加到 19.5 秒，其中 indexer 的贡献增长最为显著。
本文的目标是本文的核心目标是大幅减少 DSA 中 indexer 的计算成本，同时保持模型质量。具体而言，作者希望验证：能否移除大部分 indexer 计算并让多数层复用少数保留层的 top-k 索引；在不降低质量的前提下，最大可达到的索引复用比例是多少；如何系统地确定最优的层选择模式；是否可以通过训练进一步提升复用效果。实验目标是在 30B DSA 模型上实现至少 75% 的 indexer 计算移除，同时在 9 个长上下文和推理基准测试上保持与原始 DSA 可比的性能。
与已有工作不同的是，本文的独特切入角度在于首次将跨层 token 选择稳定性原理从依赖完整注意力的锚点层方法迁移到完全不保留完整注意力的 DSA 架构。此前所有相关方法都假设锚点层使用完整注意力 O(L^2) 作为识别重要 token 的神谕，但 DSA 已经完全移除了完整注意力。本文的关键创新是认识到：DSA 的 lightweight indexer 输出本身也表现出足够的跨层稳定性（相邻层共享 70-100% 的选择 token），因此可以直接复用 indexer 的输出，而不需要完整注意力。此外，本文提出了两种互补的优化方法，即训练无关的贪心搜索和训练感知的多层蒸馏，为不同场景提供了系统化的解决方案。
核心方法
IndexCache 的整体思路基于一个关键观察：在 DSA 的多层架构中，不同层的 lightning indexer 选择的 top-k token 集合高度相似。基于这一洞察，作者提出将 N 个 Transformer 层划分为两种角色：F（Full）层保留其 indexer 并正常计算 top-k 索引，S（Shared）层则完全跳过 indexer 计算，直接复用最近的 F 层产生的索引。这种划分用二进制模式字符串编码，其中每层标记为 F 或 S。在推理时，S 层仅需一个条件分支：要么运行 indexer（F 层），要么复制缓存的索引（S 层）。技术路线包括两种互补方法：训练无关方法使用贪心搜索算法在预训练模型上直接确定最优层选择模式；训练感知方法通过多层蒸馏损失训练模型适应索引复用，使简单的均匀交错模式也能达到最优效果。
本文的核心创新在于两个方面。第一，发现 DSA 的 lightweight indexer 输出具有跨层稳定性，这与此前依赖完整注意力作为锚点的方法本质不同。传统方法需要在锚点层运行 O(L^2) 的完整注意力来获取真实的 top-k 索引，而 IndexCache 直接复用 indexer 的输出，一个已经设计为比完整注意力便宜一个数量级的模块。第二，提出多层蒸馏损失，将每个保留的 indexer 训练为同时服务多个层，而不仅仅是其原始对应层。数学上，这等价于将 indexer 蒸馏到其服务的所有层注意力分布的质心。形式化地，对于 F 层 l 及其服务的后续 S 层 l+1 到 l+m，多层蒸馏损失为 L_I^{multi} 等于 1/(m+1) 乘以对 j 从 0 到 m 求和，再对 t 求和 D_KL(p_t^{(l+j)} || q_t^{(l)})，其中 p_t^{(l)} 是聚合注意力分布，q_t^{(l)} = Softmax(I_t^{(l)}) 是 indexer 输出分布。
方法步骤详情
IndexCache 方法包含以下完整步骤。第一，初始化：将所有 N 层设为 F 模式，得到基线模式全 F，其中第一层始终为 F 以提供初始索引。第二，层选择模式确定：对于训练无关方法，使用贪心搜索算法。从所有层为 F 的基线出发，迭代 K 步（K 为目标 S 层数）。每步中，遍历所有当前为 F 的层（除第一层），逐一尝试将其翻转为 S，计算在小型校准集上的 LM 损失（校准集为 768 个样本，200K 上下文长度），选择损失最低的层完成翻转。对于训练感知方法，使用简单的均匀交错模式。第三，训练（仅训练感知方法）：采用 DSA 标准训练流程的两阶段。第一阶段 dense warm-up：仅用多层蒸馏损失训练 F 层的 indexer，其他参数冻结。第二阶段 sparse training：继续用多层蒸馏损失训练 indexer，同时加入 LM 损失训练其余参数。第四，推理：在推理循环中，对每层，若为 F 层则运行 indexer 获取索引并更新缓存；若为 S 层则直接使用缓存的索引。然后在选择的 token 子集上计算稀疏注意力和前馈网络。
技术新颖性
IndexCache 的技术新颖性体现在以下几个方面。首先，这是首个在不保留完整注意力的情况下实现跨层索引复用的方法。此前所有方法（如 TidalDecode、LessIsMore、OmniKV、DELTA、Kascade、HySparse）都依赖完整注意力层作为识别重要 token 的神谕，而 DSA 已经完全移除了完整注意力。IndexCache 证明了 lightweight indexer 的输出本身就足以作为跨层复用的基础。其次，贪心搜索算法使用端到端 LM 损失而非局部相似度指标来确定最优模式，这是一个关键设计选择。附录中的对比实验表明，基于相似度的动态规划搜索虽然计算成本更低，但产生的模式与均匀交错效果相当，而贪心 LM 损失搜索显著优于两者。第三，多层蒸馏损失的梯度等价性提供了理论保证：多层 KL 散度之和的梯度等价于蒸馏到平均分布的梯度。这使得 indexer 学习到跨层的共识 top-k 集合。
实验结果
实验结果验证了 IndexCache 在多个维度的有效性。在端到端推理性能方面，30B DSA 模型在 H100 节点上（dp attention, dp size=8, SGLang）的测试表明：保留 1/4 indexer 的 IndexCache 在 200K token 上下文下将预填充延迟从 19.5 秒降至 10.7 秒，实现 1.82 倍加速；单请求解码吞吐量从 58 tok/s 提升到 86 tok/s（1.48 倍）；满 KV 缓存时总吞吐量从 197 tok/s 提升到 297 tok/s（1.51 倍）。在模型质量方面，训练无关 IndexCache（贪心搜索模式）在保留 1/4 indexer 时，长上下文平均分 49.9 与原始 DSA 的 50.2 相当，推理平均分 74.9 甚至略优于基线 74.6。训练感知 IndexCache 在均匀交错 1/2 保留时超越基线（长上下文 51.6 vs 51.0）。关键发现包括：贪心搜索模式显著优于均匀交错，1/4 保留时均匀交错长上下文仅 43.0 分而搜索模式达 49.9 分；训练消除了模式敏感性，训练感知方法中均匀交错与搜索模式表现相当；跨层蒸馏提供实质性收益，移除跨层损失后长上下文从 51.6 降至 49.8 分。GLM-5（744B 参数）的初步实验进一步确认了可扩展性，1/4 保留的搜索模式在长上下文平均分上与基线仅差 0.4 点（78.0 vs 78.4）。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 预填充延迟（200K context） | 延迟（秒） | 10.7s（1/4保留） | 19.5s（原始DSA） | 1.82倍加速 |
| 单请求解码吞吐量（200K context） | tokens/s | 86 tok/s（1/4保留） | 58 tok/s（原始DSA） | 1.48倍提升 |
| 满KV缓存解码吞吐量（200K context） | tokens/s | 297 tok/s（1/4保留） | 197 tok/s（原始DSA） | 1.51倍提升 |
| 长上下文基准测试平均分 | 平均分 | 49.9（训练无关1/4搜索模式） | 50.2（原始DSA） | 仅差0.3点，保持可比性能 |
| 通用与推理基准测试平均分 | 平均分 | 74.9（训练无关1/4搜索模式） | 74.6（原始DSA） | 略优0.3点 |
| AIME 2025 | 准确率 | 92.6（训练无关1/4搜索模式） | 91.0（原始DSA） | 提升1.6点 |
局限与改进
尽管 IndexCache 取得了显著成果，仍存在以下局限性。首先，训练无关方法的贪心搜索成本较高：完整搜索（从全 F 到目标模式）需要 N(N-1)/2 次前向传播。虽然作者提出了基于流水线并行的加速策略（将层分成 P 个块并行搜索），但对于超大模型（如 744B 参数的 GLM-5），搜索成本仍然可观。其次，极端稀疏情况下（如仅保留 1/8 indexer）性能下降明显：均匀交错下长上下文平均分从 50.2 降至 35.3，即使使用搜索模式也仅恢复到 46.1，表明 1/8 可能是这种方法的极限。第三，训练感知方法需要额外的训练流程（1000 步 dense warm-up + 4000 步 sparse training），对于已经部署的模型需要额外的计算资源。第四，论文仅在 DSA 架构上验证了 IndexCache，虽然作者声称核心原则可推广到其他动态 token 选择的稀疏注意力方法（如 MoBA、NSA），但尚未进行实验验证。此外，论文未详细讨论 IndexCache 对不同任务类型（如代码生成、多轮对话）的影响，以及在不同硬件平台上的加速效果。
独立分析的弱点
基于对论文的深入分析，我认为存在以下几个弱点和改进方向。第一，贪心搜索的计算效率：当前贪心搜索需要多次完整前向传播来评估每个候选模式，虽然有流水线并行优化，但对于生产环境仍然较慢。改进方向可以探索基于梯度的快速近似方法，或者训练一个小的代理模型来预测层重要性。第二，模式泛化性：论文发现的贪心搜索模式是针对特定模型和数据集的，当模型更新或数据分布变化时可能需要重新搜索。改进方向是研究模式的迁移性，或者开发自适应的动态层选择机制。第三，均匀交错的局限性：训练感知方法虽然消除了模式敏感性，但仍需要预定义均匀间隔。改进方向是让模型自动学习最优的 F/S 层分配，类似于神经架构搜索。第四，与其他优化的协同：论文仅单独评估 IndexCache，未与 KV 缓存压缩、注意力核优化等其他技术组合。改进方向是研究多种优化的联合应用和相互影响。
未来方向
作者在论文中提出了几个未来研究方向：将训练感知 IndexCache 应用于 GLM-5 等生产级大模型，预期可进一步巩固性能增益并转化为实际部署效率收益；将 IndexCache 的核心原则推广到其他稀疏注意力方法，如 MoBA 和 NSA 的 block-level 选择。基于当前成果，还可以延伸以下方向：动态 IndexCache，根据输入内容自适应调整 F/S 层比例，在简单查询时使用更多 S 层，在复杂推理时使用更多 F 层；联合优化，将 IndexCache 与 KV 缓存共享、注意力核融合等技术结合，实现端到端的推理加速；硬件感知优化，针对不同 GPU 架构优化 IndexCache 的实现，利用硬件特性进一步降低开销；多模态扩展，将跨层索引复用的思想应用到视觉-语言模型的注意力优化中。
复现评估
从复现角度来看，本文提供了较好的信息但存在一定局限。作者详细描述了实验设置：使用 GLM-4.7-Flash（30B-A3B MoE，47 层）作为基模型，通过两阶段训练（1000 步 warm-up + 4000 步 sparse training）得到 DSA 模型。训练数据为 SFT 数据，上下文长度 200K，batch size 768。评估涵盖 5 个长上下文基准和 4 个推理基准，使用温度 1.0、top-p 0.95、top-k 40 的标准设置。附录给出了具体的搜索模式字符串。然而，复现面临几个挑战：基模型 GLM-4.7-Flash 未开源（仅提供 HuggingFace 链接但权重访问可能受限）；DSA 的两阶段训练流程需要大量计算资源；贪心搜索需要多次前向传播，对计算资源要求较高；评估数据集（如 MRCR v2、GraphWalks）的具体获取方式未完全说明。总体而言，完全复现需要中等到大量的计算资源和对模型权重的访问权限。
论文图表
该图以伪代码形式并排展示了标准 DSA 推理循环（左）和 IndexCache 推理循环（右）。标准 DSA 在每一层都运行 indexer，而 IndexCache 添加了一个条件分支：F 层计算并缓存新索引，S 层复用缓存索引。关键观察是缓存是临时缓冲区，仅存储当前索引张量，在每个 F 层被覆盖，不需要额外的 GPU 显存。
这张图是理解 IndexCache 推理机制的关键，清晰展示了方法的简洁性，仅需一个条件分支即可实现跨层索引复用，同时强调了显存开销几乎为零。