IndexCache：跨层索引复用，让稀疏注意力推理再快一倍

IndexCache跨层索引复用让稀疏注意力推理再快一倍论文标题IndexCache: Accelerating Sparse Attention via Cross-Layer Index Reuse作者Yushi Bai, Qian Dong, Ting Jiang, Xin Lv, Zhengxiao Du, Aohan Zeng, Jie Tang, Juanzi Li机构清华大学 Z.ai智谱AI发表日期2026年3月论文链接arXiv:2603.12201一句话总结当大模型处理超长文本时稀疏注意力Sparse Attention是加速推理的利器但其内部的索引器——负责挑选最相关token的模块——本身就很慢。IndexCache发现相邻层的索引选择结果高度重叠相似度70%~100%因此可以让多个层共享同一份索引从而跳过75%的索引器计算在200K上下文长度下实现1.82倍预填充加速和1.48倍解码加速且几乎不损失模型质量。一、问题背景稀疏注意力的索引器瓶颈1.1 长文本推理的核心痛点大语言模型LLM的自注意力机制复杂度为O ( L 2 ) O(L^2)O(L2)其中L LL是上下文长度。当L LL从4K增长到200K计算量增长了2500倍。这使得长文本场景下的推理延迟和显存占用成为严重瓶颈。稀疏注意力是一条主流的解决路径不让每个query关注所有token而是只选择最相关的top-k个token进行注意力计算。这样核心注意力的复杂度从O ( L 2 ) O(L^2)O(L2)降到O ( L k ) O(Lk)O(Lk)其中k ≪ L k \ll Lk≪L。1.2 DeepSeek Sparse AttentionDSADSA是目前生产级稀疏注意力方案的代表。它在DeepSeek-V3/R1等模型中被训练集成其核心设计包括闪电索引器Lightning Indexer一个轻量级的注意力模块为每个query快速计算与所有key的粗粒度相关性分数然后选出top-k个最相关的token索引。核心注意力只在选出的top-k个token上执行完整的注意力计算。看起来很完美问题在于索引器本身仍需要对所有L LL个token做一次完整的注意力计算。虽然索引器的参数量远小于核心注意力head维度更小但其计算复杂度仍然是O ( L 2 ) O(L^2)O(L2)。1.3 索引器的时间占比有多大论文用一张图清晰地揭示了这个问题图1在30B参数的DSA模型上索引器Lightning Indexer的时间占比随上下文长度的变化。数据非常惊人上下文长度Prefill阶段索引器占比Decode阶段索引器占比10K27%27%60K50%31%120K68%38%200K81%41%在200K上下文长度下预填充阶段有高达81%的时间花在索引器上这意味着尽管核心注意力已经被稀疏化了但索引器的O ( L 2 ) O(L^2)O(L2)计算反而成了新的瓶颈。上下文越长索引器的瓶颈越突出。这就是IndexCache要解决的问题能否在不显著损失质量的前提下大幅减少索引器的计算量二、核心发现跨层索引的惊人冗余2.1 相邻层选出的token几乎一样论文的核心洞察来自一个简单但关键的实验比较不同层的索引器选出的top-k索引集合的重叠程度。图247层DSA模型的跨层Top-k索引相似度热力图。横纵轴分别代表层编号颜色越亮黄色表示两层选出的top-k索引集合越相似。红色方框标注了贪婪搜索算法选择的分组。这张热力图传递了几个重要信息对角线附近一片亮黄相邻层的索引相似度非常高通常在70%以上很多甚至接近100%。这意味着第i ii层和第i 1 i1i1层的索引器费了同样的计算量却选出了几乎一模一样的token集合。存在明显的块状结构相似度不是均匀衰减的而是形成了清晰的块——同一个块内的层高度相似不同块之间的差异较大。这暗示了一种自然的分组结构。浅层和深层的差异模型的前几层和后几层内部相似度更高块更大中间层的块相对更小。这个发现的直觉解释是在深度神经网络中相邻层的表示representation变化通常是平滑的类似残差网络中的小步更新。因此相邻层的query和key的分布差异很小导致它们选出的最相关token高度一致。2.2 一个自然的想法既然相邻层选出的索引几乎一样那何不让一个层计算索引周围几个层直接复用这份索引这就是IndexCache的核心思路。三、IndexCache方法详解3.1 Full层与Shared层IndexCache将模型的所有层划分为两类Full层保留索引器的层正常运行索引器计算top-k索引并将索引缓存起来。Shared层共享索引的层跳过索引器计算直接使用最近的Full层缓存的top-k索引。假设模型有48层如果采用1/4的保留率keep ratio则只有12层是Full层剩下36层是Shared层。这意味着索引器的计算量直接减少到原来的1/4。这个设计的精妙之处在于只修改了索引的来源核心注意力的计算完全不变。每一层仍然在自己的query、key、value上做注意力只是关注哪些token这个决定由邻近的Full层代劳了。3.2 如何选择哪些层保留索引器一个关键问题是48层中选哪12层作为Full层最简单的方案是均匀间隔——每4层保留一个。但论文发现这并非最优选择。贪婪搜索算法论文提出了一种基于校准集的贪婪搜索算法初始状态所有层都是Full层。每一轮尝试移除每一个Full层变为Shared层评估移除后的语言建模损失LM Loss。选择移除后损失增加最少的那一层将其变为Shared层。重复步骤2-3直到剩余的Full层数量达到目标。这个算法的复杂度是O ( N 2 ) O(N^2)O(N2)N NN为层数但由于只需要在一个小的校准集上跑一次前向传播计算成本很低。图3随着移除的索引器层数增加横轴贪婪搜索蓝色实线与均匀交错橙色虚线的LM Loss变化。水平虚线标注了1/2、1/4、1/8三个保留率。从这张图可以看到当移除层数较少时保留率1/2左右两种策略差异不大因为相邻层本来就很相似。当移除层数增多时保留率1/4及以下贪婪搜索的优势开始显现——它总能找到损失增加最小的移除顺序。即使移除了3/4的索引器层约35层LM Loss的增加也非常微小从0.569增加到约0.571。回看图2的热力图红色方框标注的就是贪婪搜索选出的分组结构。可以看到算法自适应地为不同区域分配了不同大小的组相似度高的区域用更大的组更少的Full层相似度变化快的区域用更小的组。3.3 Training-free IndexCache将上述贪婪搜索得到的Full/Shared层划分直接应用到推理中就是Training-free IndexCache。它不需要任何额外训练即插即用。具体流程在校准集上运行贪婪搜索确定哪些层是Full层。推理时Full层正常运行索引器缓存top-k索引。Shared层跳过索引器使用最近Full层的缓存索引。3.4 Training-aware IndexCacheTraining-free版本在1/2保留率下效果几乎无损但在更激进的1/4保留率下会有轻微的质量下降。为了进一步提升效果论文提出了训练感知的方案。核心思路是既然Full层的索引器要服务多个层自己和若干Shared层那就让它在训练时意识到这一点学会选出对所有被服务层都好的索引。多层蒸馏损失具体做法是引入一个蒸馏损失distillation lossL distill ∑ l ∈ G ( f ) D KL ( A l full ∥ A l shared ) \mathcal{L}_{\text{distill}} \sum_{l \in \mathcal{G}(f)} D_{\text{KL}} \left( A_l^{\text{full}} \,\|\, A_l^{\text{shared}} \right)Ldistill​l∈G(f)∑​DKL​(Alfull​∥Alshared​)其中G ( f ) \mathcal{G}(f)G(f)是Full层f ff服务的所有Shared层的集合A l full A_l^{\text{full}}Alfull​是第l ll层使用自己索引器时的注意力分布教师信号A l shared A_l^{\text{shared}}Alshared​是第l ll层使用Full层f ff的缓存索引时的注意力分布学生信号D KL D_{\text{KL}}DKL​是KL散度总训练损失为L L LM α ⋅ L distill \mathcal{L} \mathcal{L}_{\text{LM}} \alpha \cdot \mathcal{L}_{\text{distill}}LLLM​α⋅Ldistill​这里有两个关键设计决策为什么是多层蒸馏而非单层如果只最小化Full层自身的注意力分布变化Full层的索引器会倾向于自私地优化自己的表现而忽略Shared层的需求。多层蒸馏让索引器学会兼顾所有被服务层。为什么不直接用相似度搜索替代论文也尝试了一种看似更直觉的替代方案不用索引器而是直接计算当前层的query与最近Full层的key之间的相似度来选top-k。但实验表明这种方法效果很差因为不同层的key空间分布不同跨层的相似度比较没有意义。四、实验结果4.1 实验设置论文在两个规模的模型上进行了实验30B参数的DSA模型47层主要实验平台GLM-5744B参数智谱AI的旗舰模型用于验证方法的扩展性评测包含两大类benchmarkLong-Context长上下文MRCR v2、Graph Walks、LongBench v2、RULER、AA-LCRGeneral Reasoning通用与推理HLE、HLE(w/ tools)、SciCode、AIME25、IFBench4.2 质量评估30B模型结果方法保留率Long AvgGeneral AvgDSA基线150.250.9Training-free1/250.150.2Training-free1/449.949.1Training-aware1/251.650.2Training-aware1/450.349.5几个值得注意的发现Training-free 1/2保留率几乎无损Long Avg从50.2只降到50.1减少了一半的索引器计算却几乎感知不到质量变化。Training-aware 1/2保留率甚至超越基线Long Avg达到51.6比原始DSA的50.2还高出1.4个百分点。这看似反直觉——减少计算反而更好论文的解释是多层蒸馏损失起到了正则化的作用帮助索引器学到了更具泛化性的索引选择策略。1/4保留率仍然具有竞争力Training-aware版本在1/4保留率下Long Avg为50.3与基线50.2几乎持平而索引器计算量仅为原来的1/4。GLM-5744B模型结果图4GLM-5原始版本与IndexCache1/2保留率版本在10个benchmark上的对比。蓝色标注的是IndexCache的分数。在这个744B的超大模型上10个benchmark中8个基本持平或有所提升只有2个Graph Walks和IFBench有轻微下降。Long-Context平均分从77.2到78.0General Reasoning平均分从59.8到60.1。端到端推理速度提升1.2倍。这验证了IndexCache的方法在超大规模模型上同样适用。4.3 速度评估图530B模型在不同上下文长度下的端到端加速效果。从左到右分别是Prefill时间、单请求Decode吞吐量、满载Decode吞吐量。速度提升的趋势非常清晰上下文越长加速比越大。这完全符合预期——上下文越长索引器的O ( L 2 ) O(L^2)O(L2)计算在总计算中的占比越大削减索引器带来的收益也越大。在200K上下文长度下1/4保留率指标加速比Prefill时间1.82x单请求Decode吞吐量1.48x满载Decode吞吐量1.51x即使在较短的10K上下文下也有1.21x~1.24x的加速说明方法在各种场景下都有收益。4.4 消融实验论文还做了几组有价值的消融实验贪婪搜索 vs 均匀交错层选择策略保留率 1/4 Long Avg保留率 1/8 Long Avg均匀交错49.345.9贪婪搜索49.947.3贪婪搜索在激进保留率下优势更明显1/8保留率下比均匀交错高出1.4个百分点。单层蒸馏 vs 多层蒸馏蒸馏策略保留率 1/4 Long Avg单层蒸馏仅Full层自身49.2多层蒸馏Full层 Shared层50.3多层蒸馏比单层蒸馏高出1.1个百分点验证了让索引器学会服务所有层的设计动机。相似度搜索的失败论文尝试了用跨层key相似度搜索替代索引器复用在Shared层中不使用Full层的缓存索引而是用当前层的query直接与Full层的key做内积来选top-k。结果表明这种方法效果很差Long Avg只有31.3 vs 基线50.2因为不同层的key处于不同的表示空间跨层的相似度计算没有意义。这个负面结果反过来说明了IndexCache直接复用索引策略的合理性复用的是位置信息哪些token重要而不是表示信息key向量本身。五、技术细节与实现5.1 与KV Cache的关系在标准的自回归解码中每一层都会维护一个KV Cache缓存所有已生成token的key和value。IndexCache在此基础上额外缓存一份索引Cache——即Full层选出的top-k token位置索引。两种Cache的生命周期不同KV Cache随着每个新token的生成不断增长。Index Cache在Full层计算后写入在后续的Shared层读取使用直到遇到下一个Full层才更新。在Prefill阶段所有层按顺序处理Index Cache自然地被逐层传递。在Decode阶段由于每次只处理一个新tokenFull层需要为这个新token重新计算索引因为新token与历史token的相关性分布可能不同但Shared层仍可复用。5.2 预填充与解码的不同加速特征从图5可以看到Prefill阶段的加速比1.82x显著大于Decode阶段1.48x。原因在于Prefill阶段需要处理所有输入token索引器对每个token都要做全局O ( L ) O(L)O(L)的相关性计算总复杂度O ( L 2 ) O(L^2)O(L2)。移除索引器直接将这部分从O ( L 2 ) O(L^2)O(L2)降为O ( 1 ) O(1)O(1)直接使用缓存。Decode阶段每次只处理一个新token索引器的单次计算为O ( L ) O(L)O(L)在总计算中的占比相对较小。5.3 训练成本Training-aware IndexCache的训练成本非常低。论文提到只需要在较小的数据集上进行微调训练数据校准集具体规模论文未详细说明但强调了效率训练对象仅更新Full层的索引器参数核心注意力参数冻结收敛速度由于蒸馏目标明确收敛很快相比于从头训练DSA模型的巨大成本需要在全部预训练数据上训练索引器IndexCache的后训练微调成本可以忽略不计。六、与相关工作的对比6.1 稀疏注意力方法谱系稀疏注意力领域的方法可以按是否需要训练和索引选择方式两个维度来分类方法类型代表工作特点固定模式Longformer, BigBird手工设计的局部全局注意力模式动态选择训练时集成DSA, Native Sparse Attention训练索引器推理时动态选择后训练剪枝Quest, RetrievalAttention基于重要性评分剪枝KV Cache跨层复用IndexCache利用跨层冗余复用索引计算IndexCache与上述方法是正交的——它不改变稀疏注意力的选择策略本身而是减少选择策略的执行次数。因此它可以与任何使用索引器的稀疏注意力方法结合。6.2 跨层共享的思想渊源跨层参数共享在Transformer研究中并不新鲜。ALBERT通过跨层共享全部参数来压缩模型Universal Transformer让所有层共享同一组参数。但这些方法共享的是参数而IndexCache共享的是计算结果索引。更接近的工作是跨层KV Cache共享如YOCO、CLA它们让多个层共享相同的KV Cache来减少显存占用。IndexCache的思路类似但共享的是索引而非KV目标也不同——前者减少显存后者减少计算。七、局限性与展望7.1 当前局限依赖DSA框架IndexCache目前仅在DSA及类似的索引器核心注意力架构上验证。对于其他稀疏注意力方案如基于hash的LSH Attention其跨层冗余假设是否成立还需验证。保留率的选择最优保留率因模型和任务而异。虽然贪婪搜索提供了自动化的层选择但保留率本身仍需人工指定。对极端长度的表现实验最长测试到200K token更长的上下文如1M token下的表现还有待探索。不过从趋势看上下文越长加速比越大前景应该是乐观的。7.2 未来方向动态保留率不同的输入可能需要不同的保留率。简单的输入可以用更低的保留率复杂的输入需要更高的保留率。设计一种根据输入动态调整的机制将是有价值的扩展。与其他加速技术的组合IndexCache可以与量化、知识蒸馏、投机解码等技术正交组合进一步提升推理效率。推广到其他架构Mamba等状态空间模型也面临长序列计算效率问题探索类似的跨层复用思想可能带来新的突破。八、总结IndexCache是一篇简洁而有效的工作。它从一个朴素的观察出发——相邻层选出的重要token几乎一样——提出了跨层索引复用的方案用极低的实现成本换取了显著的推理加速。这篇论文的价值不仅在于具体的技术方案更在于它揭示了一个重要的现象在深层Transformer中大量的计算是冗余的只是我们还没有找到安全地跳过它们的方法。IndexCache为此提供了一个优雅的解决方案而这个思路——“先发现冗余再设计复用”——对其他计算瓶颈同样具有启发意义。对于工业界来说IndexCache的实用价值尤为突出它不需要修改模型架构不需要重新训练Training-free版本可以直接应用于任何已部署的DSA模型且加速效果随上下文长度增加而更加显著——恰好是最需要加速的场景。觉得有启发的话欢迎点赞、在看、转发。跟进最新AI前沿关注我的微信公众号机器懂语言。