IndexCache: Acelerando a Atenção Esparsa por meio da Reutilização de Índices entre Camadas
IndexCache: Accelerating Sparse Attention via Cross-Layer Index Reuse
March 12, 2026
Autores: Yushi Bai, Qian Dong, Ting Jiang, Xin Lv, Zhengxiao Du, Aohan Zeng, Jie Tang, Juanzi Li
cs.AI
Resumo
Os fluxos de trabalho agenticos de contexto longo emergiram como um caso de uso definidor para modelos de linguagem de grande escala, tornando a eficiência da atenção crítica tanto para a velocidade de inferência quanto para o custo de serviço. A atenção esparsa aborda esse desafio de forma eficaz, e a DeepSeek Sparse Attention (DSA) é uma solução representativa de nível de produção: um indexador leve e rápido seleciona os k tokens mais relevantes por consulta, reduzindo a atenção principal de O(L²) para O(Lk). No entanto, o próprio indexador mantém complexidade O(L²) e deve ser executado independentemente em cada camada, apesar do facto de as seleções top-k resultantes serem altamente semelhantes entre camadas consecutivas. Apresentamos o IndexCache, que explora essa redundância entre camadas ao particioná-las em um pequeno conjunto de camadas Completas (que executam seus próprios indexadores) e uma maioria de camadas Compartilhadas (que simplesmente reutilizam os índices top-k da camada Completa mais próxima). Propomos duas abordagens complementares para determinar e otimizar esta configuração. O IndexCache sem treino aplica um algoritmo de busca gulosa que seleciona quais camadas reter os indexadores minimizando diretamente a perda de modelagem de linguagem num conjunto de calibração, sem exigir atualizações de pesos. O IndexCache com treino introduz uma perda de destilação multicamada que treina cada indexador retido contra as distribuições médias de atenção de todas as camadas que ele serve, permitindo que até padrões intercalados simples atinjam a precisão do indexador completo. Resultados experimentais num modelo DSA de 30B mostram que o IndexCache pode remover 75% dos cálculos do indexador com degradação de qualidade negligenciável, alcançando até 1,82x de aceleração no preenchimento prévio (prefill) e 1,48x de aceleração na descodagem (decode) em comparação com o DSA padrão. Estes resultados positivos são ainda confirmados pelas nossas experiências preliminares no modelo de escala de produção GLM-5 (Figura 1).
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.82times prefill speedup and 1.48times 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).