NOSA: Attenzione Sparse Nativa e Trasferibile

NOSA: Native and Offloadable Sparse Attention

October 15, 2025

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

cs.AI

Abstract

L'attenzione sparsa addestrabile è emersa come una soluzione promettente per affrontare il collo di bottiglia dell'efficienza di decodifica nei modelli linguistici di grandi dimensioni (LLM) durante l'elaborazione di contesti lunghi, consentendo un significativo risparmio negli accessi alla memoria con un impatto minimo sulle prestazioni delle attività. Tuttavia, i metodi esistenti di attenzione sparsa lasciano irrisolta una limitazione cruciale: la dimensione della cache chiave-valore (KV) rimane invariata, il che limita le dimensioni dei batch sulla GPU e riduce la velocità di decodifica, specialmente nell'inferenza su larga scala con batch multipli. In questo articolo, dimostriamo che l'attenzione sparsa addestrabile presenta naturalmente una forte località nella selezione dei token tra passaggi di decodifica adiacenti, consentendo così lo scaricamento della cache KV senza alterare il calcolo sottostante dell'attenzione. Tuttavia, la località intrinseca non è sufficiente per ottenere uno scaricamento efficiente, poiché il trasferimento delle coppie KV selezionate tra la CPU e la GPU continua a dominare il costo complessivo della decodifica. Basandoci su questa intuizione, presentiamo NOSA, un framework di attenzione sparsa addestrabile progettato per supportare nativamente lo scaricamento della cache KV. NOSA introduce vincoli espliciti di località scomponendo la selezione dei token in componenti dipendenti e indipendenti dalla query, riducendo così i trasferimenti KV pur preservando lo stesso calcolo dell'attenzione utilizzato durante l'addestramento. Abbiamo preaddestrato un modello da 1 miliardo di parametri con NOSA e condotto ampie valutazioni, dimostrando che preserva prestazioni quasi senza perdite mentre raggiunge un miglioramento fino a 2,3x nella velocità di decodifica rispetto alla baseline di attenzione sparsa addestrabile 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).