NOSA: Native en Offloadbare Sparse Attention

NOSA: Native and Offloadable Sparse Attention

October 15, 2025

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

cs.AI

Samenvatting

Trainbare sparse attention is naar voren gekomen als een veelbelovende oplossing om het decodeerefficiëntieprobleem van LLM's bij lange-contextverwerking aan te pakken, waarbij het geheugentoegang aanzienlijk bespaart terwijl het de taakprestatie minimaal beïnvloedt. Bestaande sparse attention-methoden laten echter een cruciaal probleem onopgelost: de grootte van de key-value (KV)-cache blijft onverminderd, wat de batchgroottes op de GPU beperkt en de decodeersnelheid vertraagt, vooral bij grootschalige batchinferentie. In dit artikel tonen we aan dat trainbare sparse attention van nature een sterke localiteit vertoont in tokenselectie over aangrenzende decodestappen, waardoor KV-cache-offloading mogelijk wordt zonder de onderliggende attention-berekening te wijzigen. De inherente localiteit blijft echter onvoldoende om efficiënte offloading te bereiken, aangezien de overdracht van geselecteerde KV-paren tussen de CPU en GPU het grootste deel van de totale decodeerkosten blijft uitmaken. Op basis van dit inzicht presenteren we NOSA, een trainbaar sparse attention-framework dat KV-cache-offloading van nature ondersteunt. NOSA introduceert expliciete localiteitsbeperkingen door tokenselectie op te splitsen in query-aware en query-agnostische componenten, waardoor KV-overdrachten worden verminderd terwijl dezelfde attention-berekening behouden blijft als tijdens de training. We pretrainen een 1B-parameter model met NOSA en voeren uitgebreide benchmarks uit, waaruit blijkt dat het near-lossless prestaties behoudt terwijl het een verbetering van tot 2,3x in decodeersnelheid bereikt vergeleken met de baseline van trainbare sparse attention (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).