IndexCache: Accelerating Sparse Attention via Cross-Layer Index Reuse

TL;DR

IndexCache通过跨层索引重用加速稀疏注意力,减少75%计算量,提升1.82倍速度。

cs.CL 🔴 高级 2026-03-13 12 次浏览
Yushi Bai Qian Dong Ting Jiang Xin Lv Zhengxiao Du Aohan Zeng Jie Tang Juanzi Li
稀疏注意力 跨层索引 深度学习 高效计算 长上下文处理

核心发现

方法论

IndexCache通过将层划分为Full层和Shared层,前者保留索引器,后者重用最近Full层的top-k索引。提出两种优化方法:无训练的贪婪搜索算法直接最小化语言模型损失;训练感知的多层蒸馏损失训练保留的索引器,使其匹配所有服务层的平均注意力分布。

关键结果

  • 在30B DSA模型上,IndexCache移除75%的索引器计算,几乎不影响质量,相比标准DSA实现了高达1.82倍的预填速度提升和1.48倍的解码速度提升。
  • 在生产级GLM-5模型上的初步实验进一步证实了这些积极结果,IndexCache移除50%的索引器计算,同时在长上下文和推理任务中保持了可比的性能。
  • 实验显示,IndexCache在200K上下文长度下的解码吞吐量提高了22-51%,尤其在更长的上下文中效果显著。

研究意义

IndexCache在长上下文推理中显著减少了计算开销,尤其是在需要高效处理大规模数据的场景中。通过减少索引器的计算复杂度,它为大规模语言模型的推理效率提供了新的解决方案,解决了稀疏注意力在长上下文应用中的瓶颈问题。

技术贡献

IndexCache通过跨层索引重用,显著降低了稀疏注意力的计算复杂度。与现有方法不同,它不依赖于完整注意力层作为索引器的锚点,而是通过轻量级索引器实现高效的top-k选择。其创新的训练感知蒸馏损失为模型提供了新的训练和优化可能性。

新颖性

IndexCache首次在稀疏注意力中实现了跨层索引重用,显著减少了计算开销。与现有方法相比,它通过轻量级索引器实现了高效的top-k选择,而不依赖完整注意力层。

局限性

  • 在极端稀疏性下(仅保留1/8索引器层),长上下文性能下降显著,表明在某些情况下索引器的重用可能导致质量下降。
  • 虽然IndexCache在大多数情况下表现良好,但在某些特定任务中可能需要进一步的调优以确保性能不受影响。
  • 当前的实验主要集中在特定的模型和任务上,尚未在更广泛的应用场景中验证其通用性。

未来方向

未来的研究方向包括在更大规模的模型上应用训练感知的IndexCache,探索其在不同任务和数据集上的表现。此外,可以研究如何进一步优化索引器的选择和重用策略,以提高模型的适应性和性能。

AI 总览摘要

长上下文推理是现代大规模语言模型的一个重要应用场景,而稀疏注意力则是解决这一问题的有效方法。传统的稀疏注意力机制,如DeepSeek Sparse Attention (DSA),通过轻量级索引器选择每个查询最相关的top-k token,将核心注意力计算从O(L^2)减少到O(Lk)。然而,索引器本身仍然保留O(L^2)的复杂度,并且必须在每一层独立运行,尽管跨层的top-k选择高度相似。

IndexCache通过跨层索引重用解决了这一冗余问题。它将层划分为少量的Full层和大多数的Shared层,前者保留自己的索引器,后者则重用最近Full层的top-k索引。我们提出了两种互补的方法来确定和优化这一配置:无训练的IndexCache通过贪婪搜索算法直接最小化语言模型损失,而训练感知的IndexCache则引入多层蒸馏损失,训练每个保留的索引器,使其匹配所有服务层的平均注意力分布。

实验结果表明,在一个30B参数的DSA模型上,IndexCache可以移除75%的索引器计算,几乎不影响质量,相比标准DSA实现了高达1.82倍的预填速度提升和1.48倍的解码速度提升。这些积极结果在生产级GLM-5模型上的初步实验中得到了进一步证实,IndexCache移除50%的索引器计算,同时在长上下文和推理任务中保持了可比的性能。

IndexCache的意义在于它显著减少了长上下文推理中的计算开销,尤其是在需要高效处理大规模数据的场景中。通过减少索引器的计算复杂度,它为大规模语言模型的推理效率提供了新的解决方案,解决了稀疏注意力在长上下文应用中的瓶颈问题。

然而,IndexCache在极端稀疏性下(仅保留1/8索引器层),长上下文性能下降显著,表明在某些情况下索引器的重用可能导致质量下降。未来的研究方向包括在更大规模的模型上应用训练感知的IndexCache,探索其在不同任务和数据集上的表现。此外,可以研究如何进一步优化索引器的选择和重用策略,以提高模型的适应性和性能。

深度分析

研究背景

自注意力机制是现代大规模语言模型的基石,但其在序列长度上的二次复杂度成为长上下文推理的瓶颈。随着大规模语言模型在需要扩展上下文的场景中被广泛应用,如长链推理、多步骤代理工作流和基于检索的生成,如何在不牺牲模型质量的情况下降低注意力成本成为一个关键研究问题。稀疏注意力提供了一种有效的解决方案,即每个查询只选择最相关的子集进行注意力计算。近年来,DeepSeek Sparse Attention (DSA)作为一种生产级的可训练稀疏注意力机制脱颖而出。DSA通过轻量级索引器对所有前序token进行评分,并选择top-k用于后续的核心注意力计算,从而将每层的核心注意力计算从O(L^2)减少到O(Lk)。然而,索引器本身仍然以O(L^2)复杂度运行,且必须在每一层独立运行,这使得在长上下文中索引器的计算成本成为总注意力预算中不可忽视的一部分。

核心问题

在长上下文推理中,稀疏注意力的索引器计算成本成为瓶颈。尽管跨层的top-k选择高度相似,但每层索引器仍需独立运行,导致大量冗余计算。如何有效利用跨层的索引稳定性,减少不必要的索引器计算,同时保持模型质量,是一个亟待解决的问题。

核心创新

IndexCache通过跨层索引重用显著减少了稀疏注意力的计算复杂度。具体创新包括:

1. 层划分策略:将层划分为Full层和Shared层,前者保留索引器,后者重用最近Full层的top-k索引。

2. 无训练优化:提出贪婪搜索算法,通过直接最小化语言模型损失来选择保留索引器的层。

3. 训练感知优化:引入多层蒸馏损失,训练保留的索引器,使其匹配所有服务层的平均注意力分布。

方法详解

IndexCache的方法论包括以下步骤:

  • �� 层划分:将模型层划分为Full层和Shared层,Full层保留索引器,Shared层重用最近Full层的top-k索引。
  • �� 无训练优化:应用贪婪搜索算法,通过直接最小化语言模型损失选择保留索引器的层。
  • �� 训练感知优化:引入多层蒸馏损失,训练保留的索引器,使其匹配所有服务层的平均注意力分布。
  • �� 实验验证:在30B参数的DSA模型和生产级GLM-5模型上进行实验,验证IndexCache的有效性。

实验设计

实验设计包括在30B参数的DSA模型和生产级GLM-5模型上进行验证。使用的基准数据集包括OpenAI的GraphWalks、LongBench v2、RULER和AA-LCR,以及四个通用和推理基准:AIME 2025、GPQA-Diamond、LiveCodeBench v6和IFBench。实验中比较了原始DSA基线与IndexCache在两种保留比例下的性能:1/2(保留一半的索引器层)和1/4(保留四分之一)。

结果分析

实验结果显示,IndexCache在30B DSA模型上移除75%的索引器计算,几乎不影响质量,相比标准DSA实现了高达1.82倍的预填速度提升和1.48倍的解码速度提升。在生产级GLM-5模型上的初步实验进一步证实了这些积极结果,IndexCache移除50%的索引器计算,同时在长上下文和推理任务中保持了可比的性能。此外,IndexCache在200K上下文长度下的解码吞吐量提高了22-51%,尤其在更长的上下文中效果显著。

应用场景

IndexCache的应用场景包括:

1. 大规模语言模型的长上下文推理,尤其是在需要高效处理大规模数据的场景中。

2. 提高推理效率,降低计算成本,适用于需要实时响应的在线服务。

3. 在资源受限的环境中部署大规模模型,减少计算资源的消耗。

局限与展望

尽管IndexCache在大多数情况下表现良好,但在极端稀疏性下(仅保留1/8索引器层),长上下文性能下降显著。此外,当前的实验主要集中在特定的模型和任务上,尚未在更广泛的应用场景中验证其通用性。未来的研究方向包括在更大规模的模型上应用训练感知的IndexCache,探索其在不同任务和数据集上的表现。

通俗解读 非专业人士也能看懂

想象你在一个大型超市购物。超市里有成千上万种商品,每次购物你都需要在货架上找到你需要的商品。传统的方法是每次都浏览所有货架,找到你需要的商品,这就像是完整注意力机制,需要处理所有可能的选项。然而,这样做效率很低,尤其是在超市很大的时候。

现在,假设超市提供了一种智能购物助手,它会根据你的购物清单和历史记录,提前为你挑选出最相关的商品,并在你到达超市时告诉你这些商品的位置。这就像是稀疏注意力机制,它只关注最相关的选项,节省了大量时间和精力。

但是,这个智能助手每次购物都需要重新计算所有商品的位置,即使这些商品的位置几乎没有变化。IndexCache就像是一个记忆功能,它会记住上次购物时的商品位置,并在下次购物时重用这些信息,只在必要时更新。这大大减少了助手的计算量,提高了购物效率。

通过这种方式,IndexCache帮助我们在处理大量数据时节省计算资源,尤其在需要快速响应的场景中,显著提高了效率。

简单解释 像给14岁少年讲一样

想象一下,你在玩一个大型多人在线游戏。游戏中有很多任务和挑战,每个任务都有很多步骤,你需要完成这些步骤才能获得奖励。传统的方法是每次都从头开始完成每个步骤,这就像是完整注意力机制,需要处理所有可能的选项。

但是,这样做效率很低,尤其是在任务很复杂的时候。于是,游戏开发者引入了一种智能助手,它会根据你的游戏历史和当前任务,为你挑选出最相关的步骤,并指导你完成任务。这就像是稀疏注意力机制,它只关注最相关的选项,节省了大量时间和精力。

然而,这个助手每次任务都需要重新计算所有步骤,即使这些步骤几乎没有变化。IndexCache就像是一个记忆功能,它会记住上次任务时的步骤,并在下次任务时重用这些信息,只在必要时更新。这大大减少了助手的计算量,提高了游戏效率。

通过这种方式,IndexCache帮助我们在处理大量任务时节省计算资源,尤其在需要快速响应的场景中,显著提高了效率。

术语表

稀疏注意力 (Sparse Attention)

稀疏注意力是一种只关注最相关子集的注意力机制,减少了计算复杂度。

在IndexCache中用于选择每个查询最相关的top-k token。

索引器 (Indexer)

索引器是用于评分和选择最相关token的组件,决定了稀疏注意力的计算效率。

在DSA中,索引器对所有前序token进行评分并选择top-k。

跨层索引重用 (Cross-Layer Index Reuse)

跨层索引重用是指在不同层之间共享索引结果,减少冗余计算。

IndexCache通过跨层索引重用减少了75%的索引器计算。

贪婪搜索算法 (Greedy Search Algorithm)

贪婪搜索算法是一种通过逐步选择最优解来最小化损失的优化方法。

用于IndexCache的无训练优化,选择保留索引器的层。

多层蒸馏损失 (Multi-Layer Distillation Loss)

多层蒸馏损失是一种训练策略,使索引器匹配所有服务层的平均注意力分布。

在IndexCache的训练感知优化中用于训练保留的索引器。

Full层 (Full Layer)

Full层是保留索引器的层,负责计算新的top-k索引。

在IndexCache中,Full层保留自己的索引器。

Shared层 (Shared Layer)

Shared层是重用最近Full层top-k索引的层,减少了索引器计算。

在IndexCache中,Shared层重用最近Full层的top-k索引。

长上下文推理 (Long-Context Inference)

长上下文推理是指在处理长序列数据时的推理过程,通常需要高效的计算方法。

IndexCache通过减少索引器计算提高了长上下文推理的效率。

DeepSeek Sparse Attention (DSA)

DSA是一种生产级的可训练稀疏注意力机制,通过轻量级索引器选择最相关的top-k token。

IndexCache在DSA模型上进行实验验证。

GLM-5模型 (GLM-5 Model)

GLM-5模型是一个生产级的大规模语言模型,IndexCache在其上进行了初步实验。

用于验证IndexCache在生产环境中的有效性。

开放问题 这项研究留下的未解疑问

  • 1 如何在更大规模的模型上应用IndexCache并验证其在不同任务和数据集上的表现?当前的实验主要集中在特定的模型和任务上,尚未在更广泛的应用场景中验证其通用性。
  • 2 在极端稀疏性下,如何进一步优化索引器的选择和重用策略,以提高模型的适应性和性能?
  • 3 是否可以在其他类型的注意力机制中应用IndexCache的跨层索引重用方法,以提高计算效率?
  • 4 如何在不影响模型质量的情况下,进一步减少索引器计算的复杂度?
  • 5 在实际应用中,IndexCache的性能和效率如何受到硬件和软件环境的影响?

应用场景

近期应用

实时在线服务

通过减少计算开销,IndexCache可以在需要快速响应的在线服务中提高推理效率,降低计算成本。

大规模数据处理

在需要处理大规模数据的场景中,IndexCache通过减少索引器计算复杂度,提高了长上下文推理的效率。

资源受限环境中的模型部署

在计算资源有限的环境中,IndexCache可以减少资源消耗,使大规模模型的部署更加可行。

远期愿景

大规模语言模型的普及

通过提高推理效率,IndexCache有望促进大规模语言模型在更多应用场景中的普及。

智能助手的优化

IndexCache的跨层索引重用方法可以用于优化智能助手的计算效率,提高用户体验。

原文摘要

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).

cs.CL cs.LG

参考文献 (20)

DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models

DeepSeek-AI, A. Liu, Aoxue Mei 等

2025 222 引用 ⭐ 高影响力 查看解读 →

FlexPrefill: A Context-Aware Sparse Attention Mechanism for Efficient Long-Sequence Inference

Xunhao Lai, Jianqiao Lu, Yao Luo 等

2025 74 引用 ⭐ 高影响力 查看解读 →

Improving Model Representation and Reducing KV Cache via Skip Connections with First Value Heads

Zhoutong Wu, Yuan Zhang, Yiming Dong 等

2025 1 引用 ⭐ 高影响力 查看解读 →

DELTA: Dynamic Layer-Aware Token Attention for Efficient Long-Context Reasoning

Hossein Entezari Zarch, Lei Gao, Chaoyi Jiang 等

2025 3 引用 查看解读 →

Kimi Linear: An Expressive, Efficient Attention Architecture

Yu Zhang, Zongyu Lin, Xingcheng Yao 等

2025 33 引用 查看解读 →

LServe: Efficient Long-sequence LLM Serving with Unified Sparse Attention

Shang Yang, Junxian Guo, Haotian Tang 等

2025 28 引用 查看解读 →

SpargeAttention: Accurate and Training-free Sparse Attention Accelerating Any Model Inference

Jintao Zhang, Chendong Xiang, Haofeng Huang 等

2025 49 引用 查看解读 →

H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models

Zhenyu (Allen) Zhang, Ying Sheng, Tianyi Zhou 等

2023 561 引用 查看解读 →

GLM-5: from Vibe Coding to Agentic Engineering

GLM-4.5 Team Aohan Zeng, Xin Lv, Zhenyu Hou 等

2026 10 引用 查看解读 →

MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding

Z. M. K. Zuhri, Farid Adilazuarda, Ayu Purwarianti 等

2024 17 引用 查看解读 →

Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

Tri Dao, Albert Gu

2024 1245 引用 查看解读 →

SeerAttention: Learning Intrinsic Sparse Attention in Your LLMs

Yizhao Gao, Zhichen Zeng, Dayou Du 等

2024 90 引用 查看解读 →

InfLLM-V2: Dense-Sparse Switchable Attention for Seamless Short-to-Long Adaptation

Weilin Zhao, Zihan Zhou, Zhou Su 等

2025 12 引用 查看解读 →

OmniKV: Dynamic Context Selection for Efficient Long-Context LLMs

Jitai Hao, Yuke Zhu, Tianjian Wang 等

2025 23 引用

MiniMax-01: Scaling Foundation Models with Lightning Attention

MiniMax, Aonian Li, Bangwei Gong 等

2025 118 引用 查看解读 →

XAttention: Block Sparse Attention with Antidiagonal Scoring

Ruyi Xu, Guangxuan Xiao, Haofeng Huang 等

2025 81 引用 查看解读 →

LiveCodeBench: Holistic and Contamination Free Evaluation of Large Language Models for Code

Naman Jain, King Han, Alex Gu 等

2024 1187 引用 查看解读 →

Gemma 3 Technical Report

Gemma Team Aishwarya Kamath, Johan Ferret, Shreya Pathak 等

2025 1087 引用 查看解读 →

RULER: What's the Real Context Size of Your Long-Context Language Models?

Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman 等

2024 704 引用 查看解读 →

TidalDecode: Fast and Accurate LLM Decoding with Position Persistent Sparse Attention

Lijie Yang, Zhihao Zhang, Zhuofu Chen 等

2024 19 引用 查看解读 →