NOSA : Attention Sparse Native et Déchargeable

NOSA: Native and Offloadable Sparse Attention

October 15, 2025

papers.authors: Yuxiang Huang, Chaojun Xiao, Xu Han, Zhiyuan Liu

cs.AI

papers.abstract

L'attention creuse entraînable est apparue comme une solution prometteuse pour résoudre le goulot d'étranglement de l'efficacité de décodage des modèles de langage de grande taille (LLMs) dans le traitement de contextes longs, permettant d'économiser de manière significative les accès mémoire tout en impactant minimalement les performances des tâches. Cependant, les méthodes d'attention creuse existantes laissent une limitation cruciale non résolue : la taille du cache clé-valeur (KV) reste inchangée, ce qui limite les tailles de lots sur GPU et réduit le débit de décodage, en particulier dans le cadre d'inférences massives par lots. Dans cet article, nous montrons que l'attention creuse entraînable présente naturellement une forte localité dans la sélection de tokens à travers les étapes de décodage adjacentes, permettant ainsi le déchargement du cache KV sans altérer le calcul sous-jacent de l'attention. Cependant, la localité inhérente reste insuffisante pour parvenir à un déchargement efficace, car le transfert des paires KV sélectionnées entre le CPU et le GPU continue de dominer le coût global du décodage. Sur la base de cette observation, nous présentons NOSA, un cadre d'attention creuse entraînable conçu pour supporter nativement le déchargement du cache KV. NOSA introduit des contraintes explicites de localité en décomposant la sélection de tokens en composantes dépendantes et indépendantes de la requête, réduisant ainsi les transferts KV tout en préservant le même calcul d'attention utilisé lors de l'entraînement. Nous pré-entraînons un modèle de 1 milliard de paramètres avec NOSA et menons des benchmarks approfondis, montrant qu'il préserve des performances quasi sans perte tout en atteignant une amélioration allant jusqu'à 2,3 fois du débit de décodage par rapport à la base de référence d'attention creuse entraînable standard (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).