IndexCache: Beschleunigung von Sparse Attention durch Wiederverwendung von Indizes über Schichten hinweg

IndexCache: Accelerating Sparse Attention via Cross-Layer Index Reuse

March 12, 2026

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

cs.AI

Zusammenfassung

Langfristige agentenbasierte Workflows haben sich als wegweisende Anwendungsfälle für große Sprachmodelle etabliert, wodurch die Aufmerksamkeitseffizienz für sowohl Inferenzgeschwindigkeit als auch Betriebskosten entscheidend wird. Sparse Attention adressiert diese Herausforderung effektiv, und DeepSeek Sparse Attention (DSA) ist eine repräsentative produktionsreife Lösung: Ein leichtgewichtiger Lightning-Indexer wählt die k relevantesten Tokens pro Query aus und reduziert die Kern-Aufmerksamkeitsberechnung von O(L²) auf O(Lk). Der Indexer selbst behält jedoch eine O(L²)-Komplexität bei und muss unabhängig in jeder Schicht ausgeführt werden, obwohl die resultierenden Top-k-Auswahlen über aufeinanderfolgende Schichten hinweg sehr ähnlich sind. Wir stellen IndexCache vor, das diese schichtübergreifende Redundanz ausnutzt, indem Schichten in eine kleine Menge von Voll-Schichten, die ihre eigenen Indexer ausführen, und eine Mehrheit von Shared-Schichten unterteilt werden, die einfach die Top-k-Indizes der nächstgelegenen Voll-Schicht wiederverwenden. Wir schlagen zwei komplementäre Ansätze zur Bestimmung und Optimierung dieser Konfiguration vor. Training-freies IndexCache wendet einen gierigen Suchalgorithmus an, der auswählt, in welchen Schichten Indexer beibehalten werden, indem direkt der Sprachmodellierungsverlust auf einem Kalibrierungsdatensatz minimiert wird, ohne Gewichtsaktualisierungen zu benötigen. Training-bewusstes IndexCache führt einen Multi-Layer-Distillationsverlust ein, der jeden beibehaltenen Indexer gegen die gemittelten Aufmerksamkeitsverteilungen aller von ihm bedienten Schichten trainiert, was es sogar einfachen, verschachtelten Mustern ermöglicht, die Genauigkeit eines Voll-Indexers zu erreichen. Experimentelle Ergebnisse mit einem 30B DSA-Modell zeigen, dass IndexCache 75 % der Indexer-Berechnungen entfernen kann bei vernachlässigbarem Qualitätsverlust und dabei eine bis zu 1,82-fache Prefill-Beschleunigung und 1,48-fache Decode-Beschleunigung im Vergleich zu Standard-DSA erreicht. Diese positiven Ergebnisse werden durch unsere vorläufigen Experimente mit dem produktionsskalierten GLM-5-Modell weiter bestätigt (Abbildung 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).