NOSA: Atenção Esparsa Nativa e Descarregável

NOSA: Native and Offloadable Sparse Attention

October 15, 2025

Autores: Yuxiang Huang, Chaojun Xiao, Xu Han, Zhiyuan Liu

cs.AI

Resumo

A atenção esparsa treinável emergiu como uma solução promissora para abordar o gargalo de eficiência na decodificação de LLMs no processamento de contextos longos, economizando significativamente acessos à memória enquanto impacta minimamente o desempenho da tarefa. No entanto, os métodos existentes de atenção esparsa deixam uma limitação crucial não resolvida: o tamanho do cache de chave-valor (KV) permanece inalterado, o que restringe o tamanho dos lotes na GPU e limita a taxa de decodificação, especialmente em inferências em lote em grande escala. Neste artigo, mostramos que a atenção esparsa treinável naturalmente exibe forte localidade na seleção de tokens entre passos de decodificação adjacentes, permitindo assim o descarregamento do cache KV sem alterar o cálculo subjacente da atenção. No entanto, a localidade inerente ainda é insuficiente para alcançar um descarregamento eficiente, pois a transferência dos pares KV selecionados entre a CPU e a GPU continua a dominar o custo total da decodificação. Com base nessa percepção, apresentamos o NOSA, um framework de atenção esparsa treinável projetado para suportar nativamente o descarregamento do cache KV. O NOSA introduz restrições explícitas de localidade ao decompor a seleção de tokens em componentes dependentes e independentes da consulta, reduzindo assim as transferências de KV enquanto preserva o mesmo cálculo de atenção utilizado durante o treinamento. Pré-treinamos um modelo de 1 bilhão de parâmetros com o NOSA e realizamos extensos benchmarks, mostrando que ele preserva um desempenho quase sem perdas enquanto alcança uma melhoria de até 2,3x na taxa de decodificação em comparação com a linha de base de atenção esparsa treinável convencional (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).