IndexCache: Accelerating Sparse Attention via Cross-Layer Index Reuse

📄 arXiv: 2603.12201v1 📥 PDF

作者: Yushi Bai, Qian Dong, Ting Jiang, Xin Lv, Zhengxiao Du, Aohan Zeng, Jie Tang, Juanzi Li

分类: cs.CL, cs.LG

发布日期: 2026-03-12


💡 一句话要点

提出IndexCache以加速稀疏注意力计算

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 稀疏注意力 索引重用 长上下文 推理加速 模型优化 深度学习

📋 核心要点

  1. 现有的稀疏注意力方法在每层都需要独立运行索引器,导致计算复杂度高且效率低下。
  2. 论文提出IndexCache,通过重用相邻层的top-k索引,显著降低索引器的计算需求。
  3. 实验表明,IndexCache在保持模型性能的同时,能够实现高达1.82倍的预填充加速和1.48倍的解码加速。

📝 摘要(中文)

长上下文的智能工作流已成为大型语言模型的重要应用场景,因此注意力效率对推理速度和服务成本至关重要。稀疏注意力有效应对这一挑战,而DeepSeek稀疏注意力(DSA)是一个典型的解决方案。然而,现有的索引器在每层都需独立运行,复杂度为O(L^2),且相邻层的选择高度相似。为此,本文提出IndexCache,通过将层划分为少量的全层和大多数共享层,重用最近全层的top-k索引,从而减少计算复杂度。实验结果表明,IndexCache能够在质量损失极小的情况下,去除75%的索引器计算,达到1.82倍的预填充加速和1.48倍的解码加速。

🔬 方法详解

问题定义:本文旨在解决现有稀疏注意力方法中索引器在每层独立运行导致的高计算复杂度问题,复杂度为O(L^2),影响了推理速度和服务成本。

核心思路:论文的核心思路是通过将层划分为少量的全层和大多数共享层,重用全层的top-k索引,从而降低计算复杂度。这样设计的原因在于相邻层的top-k选择高度相似,可以有效减少冗余计算。

技术框架:整体架构包括两个主要部分:全层负责独立运行索引器,而共享层则重用最近全层的top-k索引。为优化这一配置,提出了两种方法:训练无关的贪婪搜索算法和训练相关的多层蒸馏损失。

关键创新:最重要的技术创新在于通过跨层索引重用,显著降低了索引器的计算需求,解决了现有方法的冗余计算问题。与DSA相比,IndexCache在保持性能的同时,减少了计算复杂度。

关键设计:在训练无关的IndexCache中,采用贪婪搜索算法直接最小化语言建模损失;而在训练相关的IndexCache中,引入多层蒸馏损失,使得保留的索引器能够与所有服务层的平均注意力分布进行训练。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示,IndexCache能够去除75%的索引器计算,且质量损失极小。与标准DSA相比,IndexCache在预填充速度上实现了1.82倍的加速,在解码速度上实现了1.48倍的加速,验证了其有效性。

🎯 应用场景

该研究的潜在应用领域包括大型语言模型的推理加速和服务成本优化,尤其适用于需要处理长上下文的智能工作流。通过提高注意力计算的效率,IndexCache能够在实际应用中显著提升模型的响应速度和降低资源消耗,对未来的AI系统设计具有重要影响。

📄 摘要(原文)

Long-context agentic workflows have emerged as a defining use case for large language models, making attention efficiency critical for both inference speed and serving cost. Sparse attention addresses this challenge effectively, and DeepSeek Sparse Attention (DSA) is a representative production-grade solution: a lightweight lightning indexer selects the top-k most relevant tokens per query, reducing core attention from $O(L^2)$ to $O(Lk)$. However, the indexer itself retains $O(L^2)$ complexity and must run independently at every layer, despite the fact that the resulting top-k selections are highly similar across consecutive layers. We present IndexCache, which exploits this cross-layer redundancy by partitioning layers into a small set of Full layers that run their own indexers and a majority of Shared layers that simply reuse the nearest Full layer's top-k indices. We propose two complementary approaches to determine and optimize this configuration. Training-free IndexCache applies a greedy search algorithm that selects which layers to retain indexers by directly minimizing language modeling loss on a calibration set, requiring no weight updates. Training-aware IndexCache introduces a multi-layer distillation loss that trains each retained indexer against the averaged attention distributions of all layers it serves, enabling even simple interleaved patterns to match full-indexer accuracy. Experimental results on a 30B DSA model show that IndexCache can remove 75% of indexer computations with negligible quality degradation, achieving up to 1.82$\times$ prefill speedup and 1.48$\times$ decode speedup compared to standard DSA. These positive results are further confirmed by our preliminary experiments on the production-scale GLM-5 model (Figure 1).