IndexCache: Ускорение разреженного внимания за счет межслойного повторного использования индексов

IndexCache: Accelerating Sparse Attention via Cross-Layer Index Reuse

March 12, 2026

Авторы: Yushi Bai, Qian Dong, Ting Jiang, Xin Lv, Zhengxiao Du, Aohan Zeng, Jie Tang, Juanzi Li

cs.AI

Аннотация

Агентные рабочие процессы с длинным контекстом стали определяющим сценарием использования больших языковых моделей, что делает эффективность механизма внимания критически важной как для скорости вывода, так и для стоимости обслуживания. Разреженное внимание эффективно решает эту задачу, и DeepSeek Sparse Attention (DSA) является представительным промышленным решением: легковесный индексер Lightning выбирает k наиболее релевантных токенов для каждого запроса, сокращая сложность базового механизма внимания с O(L²) до O(Lk). Однако сам индексер сохраняет сложность O(L²) и должен запускаться независимо на каждом слое, несмотря на то, что результирующие выборки top-k оказываются высоко схожими между последовательными слоями. Мы представляем IndexCache, который использует эту межслойную избыточность, разделяя слои на небольшое множество Полных слоев, запускающих собственные индексеры, и большинство Общих слоев, которые просто повторно используют индексы top-k от ближайшего Полного слоя. Мы предлагаем два взаимодополняющих подхода для определения и оптимизации этой конфигурации. IndexCache без дообучения применяет жадный алгоритм поиска, который выбирает, на каких слоях сохранить индексеры, напрямую минимизируя потери языкового моделирования на калибровочном наборе, не требуя обновления весов. IndexCache с дообучением вводит мультислойную дистилляционную функцию потерь, которая обучает каждый сохраненный индексер на усредненных распределениях внимания всех обслуживаемых им слоев, позволяя даже простым чередующимся шаблонам соответствовать точности полного индексера. Экспериментальные результаты на 30B модели DSA показывают, что IndexCache может удалить 75% вычислений индексера с незначительной деградацией качества, достигая до 1.82× ускорения фазы предзаполнения и 1.48× ускорения фазы декодирования по сравнению со стандартной DSA. Эти положительные результаты дополнительно подтверждаются нашими предварительными экспериментами на моделях производственного масштаба, таких как GLM-5 (Рисунок 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).