IndexCache: Accelerare l'Attenzione Sparsa tramite Riuso degli Indici tra Livelli

IndexCache: Accelerating Sparse Attention via Cross-Layer Index Reuse

March 12, 2026

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

cs.AI

Abstract

I flussi di lavoro agentici a contesto lungo sono emersi come un caso d'uso fondamentale per i grandi modelli linguistici, rendendo l'efficienza dell'attenzione critica sia per la velocità di inferenza che per i costi di servizio. L'attenzione sparsa affronta efficacemente questa sfida, e DeepSeek Sparse Attention (DSA) è una soluzione rappresentativa di livello production: un indicizzatore leggero e rapido seleziona i k token più rilevanti per query, riducendo l'attenzione principale da O(L²) a O(Lk). Tuttavia, l'indicizzatore stesso mantiene una complessità O(L²) e deve essere eseguito indipendentemente a ogni layer, nonostante le selezioni top-k risultanti siano altamente simili tra layer consecutivi. Presentiamo IndexCache, che sfrutta questa ridondanza cross-layer suddividendo i layer in un piccolo insieme di layer Completi che eseguono i propri indicizzatori e una maggioranza di layer Condivisi che semplicemente riutilizzano gli indici top-k del layer Completo più vicino. Proponiamo due approcci complementari per determinare e ottimizzare questa configurazione. IndexCache senza addestramento applica un algoritmo di ricerca greedy che seleziona quali layer mantenere gli indicizzatori minimizzando direttamente la loss di modellazione linguistica su un set di calibrazione, senza richiedere aggiornamenti dei pesi. IndexCache con addestramento introduce una loss di distillazione multi-layer che addestra ogni indicizzatore mantenuto rispetto alle distribuzioni medie di attenzione di tutti i layer che serve, permettendo anche a semplici pattern interlacciati di eguagliare l'accuratezza degli indicizzatori completi. I risultati sperimentali su un modello DSA da 30B mostrano che IndexCache può rimuovere il 75% dei calcoli degli indicizzatori con una degradazione della qualità trascurabile, raggiungendo fino a 1.82x di accelerazione nella fase di prefill e 1.48x nella fase di decodifica rispetto al DSA standard. Questi risultati positivi sono ulteriormente confermati dai nostri esperimenti preliminari sul modello GLM-5 di scala production (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).