IndexCache: Aceleración de la Atención Dispersa mediante Reutilización de Índices entre Capas

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

Resumen

Los flujos de trabajo agentes de contexto largo han surgido como un caso de uso definitorio para los modelos de lenguaje grandes, haciendo que la eficiencia de la atención sea crítica tanto para la velocidad de inferencia como para el costo de servicio. La atención dispersa aborda este desafío efectivamente, y DeepSeek Sparse Attention (DSA) es una solución representativa de nivel productivo: un indexador ligero selecciona los k tokens más relevantes por consulta, reduciendo la atención principal de O(L²) a O(Lk). Sin embargo, el indexador en sí mantiene una complejidad de O(L²) y debe ejecutarse independientemente en cada capa, a pesar de que las selecciones top-k resultantes son altamente similares entre capas consecutivas. Presentamos IndexCache, que explota esta redundancia entre capas dividiéndolas en un pequeño conjunto de capas Completas que ejecutan sus propios indexadores y una mayoría de capas Compartidas que simplemente reutilizan los índices top-k de la capa Completa más cercana. Proponemos dos enfoques complementarios para determinar y optimizar esta configuración. IndexCache sin entrenamiento aplica un algoritmo de búsqueda voraz que selecciona qué capas retener indexadores minimizando directamente la pérdida de modelado del lenguaje en un conjunto de calibración, sin requerir actualizaciones de pesos. IndexCache con entrenamiento introduce una pérdida de destilación multicapa que entrena cada indexador retenido contra las distribuciones de atención promediadas de todas las capas que sirve, permitiendo que incluso patrones intercalados simples igualen la precisión del indexador completo. Los resultados experimentales en un modelo DSA de 30B muestran que IndexCache puede eliminar el 75% de los cálculos del indexador con una degradación de calidad negligible, logrando hasta 1.82x de aceleración en prefilling y 1.48x de aceleración en decodificación en comparación con DSA estándar. Estos resultados positivos son confirmados además por nuestros experimentos preliminares en el modelo a escala productiva 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).