NOSA: Нативная и Выгружаемая Разреженная Внимательность

NOSA: Native and Offloadable Sparse Attention

October 15, 2025

Авторы: Yuxiang Huang, Chaojun Xiao, Xu Han, Zhiyuan Liu

cs.AI

Аннотация

Обучаемое разреженное внимание стало перспективным решением для устранения узкого места в эффективности декодирования крупных языковых моделей (LLM) при обработке длинных контекстов, значительно сокращая количество обращений к памяти при минимальном влиянии на производительность задач. Однако существующие методы разреженного внимания не решают одну ключевую проблему: размер кэша ключей и значений (KV) остается неизменным, что ограничивает размеры пакетов на GPU и снижает пропускную способность декодирования, особенно при крупномасштабном пакетном выводе. В данной работе мы показываем, что обучаемое разреженное внимание естественным образом демонстрирует сильную локальность в выборе токенов на соседних шагах декодирования, что позволяет выгружать KV-кэш без изменения базовых вычислений внимания. Однако встроенной локальности недостаточно для эффективной выгрузки, так как передача выбранных KV-пар между CPU и GPU продолжает доминировать в общих затратах на декодирование. На основе этого наблюдения мы представляем NOSA — фреймворк обучаемого разреженного внимания, разработанный для нативной поддержки выгрузки KV-кэша. NOSA вводит явные ограничения локальности путем декомпозиции выбора токенов на компоненты, зависящие и не зависящие от запроса, что сокращает передачу KV-данных при сохранении тех же вычислений внимания, что и во время обучения. Мы предварительно обучаем модель с 1 миллиардом параметров с использованием NOSA и проводим обширные тесты, показывая, что она сохраняет почти без потерь производительность, достигая при этом улучшения пропускной способности декодирования до 2,3 раз по сравнению с базовым вариантом обучаемого разреженного внимания (InfLLM-V2).

English

Trainable sparse attention has emerged as a promising solution to address the decoding efficiency bottleneck of LLMs in long-context processing, significantly saving memory accesses while minimally impacting task performance. However, existing sparse attention methods leave a crucial limitation unresolved: the size of the key-value (KV) cache remains unreduced, which constrains on-GPU batch sizes and throttles decoding throughput, especially in large-scale batched inference. In this paper, we show that trainable sparse attention naturally exhibits strong locality in token selection across adjacent decoding steps, thereby enabling KV cache offloading without altering the underlying attention computation. However, the inherent locality remains insufficient to achieve efficient offloading, as the transfer of selected KV pairs between the CPU and GPU continues to dominate the overall decoding cost. Building on this insight, we present NOSA, a trainable sparse attention framework designed to natively support KV cache offloading. NOSA introduces explicit locality constraints by decomposing token selection into query-aware and query-agnostic components, thereby reducing KV transfers while preserving the same attention computation as used during training. We pretrain a 1B-parameter model with NOSA and conduct extensive benchmarks, showing that it preserves near-lossless performance while achieving up to a 2.3x improvement in decoding throughput compared with the vanilla trainable sparse attention baseline (InfLLM-V2).