SparDA: Разреженное разделенное внимание для эффективного вывода LLM с длинным контекстом

Аннотация

Разреженное внимание снижает вычислительные затраты и пропускную способность памяти при инференции LLM с длинным контекстом. Однако остаются две ключевые проблемы: (1) объём KV-кэша по-прежнему растёт с длиной последовательности, а выгрузка в память CPU создаёт узкое место передачи по PCIe; (2) сам этап разреженного выбора сохраняет сложность O(T²) и может доминировать в стоимости внимания при длинных контекстах. Мы предлагаем SparDA — развязанную архитектуру разреженного внимания, которая вводит четвёртую послойную проекцию, Прогноз (Forecast), наряду с Query, Key и Value. Прогноз предсказывает блоки KV, необходимые для следующего слоя, что позволяет выполнять опережающий выбор, совмещающий предварительную загрузку с CPU на GPU с выполнением текущего слоя. Поскольку Прогноз не зависит от запроса внимания, наша реализация GQA использует одну голову Прогноза на группу GQA, снижая накладные расходы на выбор по сравнению с исходным мультиголовочным селектором. SparDA добавляет менее 0,5% параметров и обучает только проекции Прогноза путём согласования с распределением внимания исходного селектора. На двух разреженно предобученных моделях с 8B параметров SparDA соответствует или незначительно улучшает точность и обеспечивает ускорение префилла до 1,25× и ускорение декодирования до 1,7× по сравнению с базовым методом с разреженным вниманием и выгрузкой. За счёт увеличения допустимых размеров батча на одной GPU SparDA дополнительно достигает пропускной способности декодирования до 5,3× выше, чем базовый метод с разреженным вниманием без выгрузки. Наш исходный код доступен по адресу https://github.com/NVlabs/SparDA.

English

Sparse attention reduces compute and memory bandwidth for long-context LLM inference. However, two key challenges remain: (1) KV cache capacity still grows with sequence length, and offloading to CPU memory introduces a PCIe transfer bottleneck; (2) the sparse selection step itself retains O(T^2) complexity and can dominate attention cost at long contexts. We propose SparDA, a decoupled sparse attention architecture that introduces a fourth per-layer projection, the Forecast, alongside Query, Key, and Value. The Forecast predicts the KV blocks needed by the next layer, enabling lookahead selection that overlaps CPU-to-GPU prefetch with current-layer execution. Because Forecast is decoupled from the attention query, our GQA implementation uses one Forecast head per GQA group, reducing selection overhead versus the original multi-head selector. SparDA adds <0.5% parameters and trains only the Forecast projections by matching the original selector's attention distribution. On two sparse-pretrained 8B models, SparDA matches or slightly improves accuracy and delivers up to 1.25times prefill speedup and 1.7times decode speedup over the sparse-attention offload baseline. By enabling larger feasible batch sizes on a single GPU, SparDA further reaches up to 5.3times higher decode throughput than the non-offload sparse baseline. Our source code is available at https://github.com/NVlabs/SparDA.