通过复用跨层索引,加速稀疏注意力计算。
IndexCache: Accelerating Sparse Attention via Cross-Layer Index Reuse
- 跨层复用顶层令牌索引,减少重复计算。
- 可减少75%索引器计算量,推理速度提升1.82倍。
- 无需训练即可适配,适合大规模语言模型部署。
长上下文智能工作流已成为大语言模型的关键应用场景,注意力效率对推理速度和服务成本至关重要。稀疏注意力有效应对这一挑战,其中DeepSeek稀疏注意力(DSA)是代表性生产级方案:轻量级闪电索引器为每个查询选择最相关的前k个标记,将核心注意力从O(L²)降低至O(Lk)。然而,索引器本身仍具有O(L²)复杂度,且需在每一层独立运行,尽管连续层间的顶k选择高度相似。本文提出IndexCache,利用这种跨层冗余性,将层划分为少量运行独立索引器的Full层和多数复用最近Full层顶k索引的Shared层。提出两种互补方法优化配置:训练无关的IndexCache采用贪心搜索,在校准集上直接最小化语言建模损失以确定保留索引器的层,无需权重更新;训练感知的IndexCache引入多层蒸馏损失,训练每个保留的索引器以匹配其所服务所有层的平均注意力分布,使简单交错模式即可达到全索引器精度。在30B DSA模型上的实验表明,IndexCache可消除75%索引器计算,质量损失极小,预填充速度提升达1.82倍,解码速度提升1.48倍。这些结果在生产级GLM-5模型的初步实验中得到验证。
原文摘要 · Abstract (English)
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).
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。