NOSA: Native und auslagerbare spärliche Aufmerksamkeit

NOSA: Native and Offloadable Sparse Attention

October 15, 2025

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

cs.AI

papers.abstract

Trainable Sparse Attention hat sich als vielversprechende Lösung erwiesen, um den Engpass bei der Dekodierungseffizienz von LLMs (Large Language Models) in der Langzeitkontextverarbeitung zu adressieren. Dabei wird der Speicherzugriff erheblich reduziert, während die Aufgabenleistung nur minimal beeinträchtigt wird. Allerdings bleibt eine entscheidende Einschränkung bestehender Sparse-Attention-Methoden ungelöst: Die Größe des Key-Value (KV)-Caches bleibt unverändert, was die Batch-Größen auf der GPU begrenzt und den Dekodierungsdurchsatz verringert, insbesondere bei großskaliger Batch-Inferenz. In diesem Artikel zeigen wir, dass trainable Sparse Attention natürlicherweise eine starke Lokalität bei der Token-Auswahl über benachbarte Dekodierungsschritte aufweist, wodurch ein KV-Cache-Offloading ermöglicht wird, ohne die zugrunde liegende Attention-Berechnung zu verändern. Die inhärente Lokalität reicht jedoch nicht aus, um ein effizientes Offloading zu erreichen, da der Transfer ausgewählter KV-Paare zwischen CPU und GPU weiterhin den Gesamtdekodierungsaufwand dominiert. Aufbauend auf dieser Erkenntnis präsentieren wir NOSA, ein trainable Sparse-Attention-Framework, das KV-Cache-Offloading nativ unterstützt. NOSA führt explizite Lokalitätsbeschränkungen ein, indem die Token-Auswahl in abfrageabhängige und abfrageunabhängige Komponenten zerlegt wird, wodurch KV-Transfers reduziert werden, während die gleiche Attention-Berechnung wie während des Trainings beibehalten wird. Wir trainieren ein 1B-Parameter-Modell mit NOSA vor und führen umfangreiche Benchmarks durch, die zeigen, dass es eine nahezu verlustfreie Leistung bewahrt und gleichzeitig eine bis zu 2,3-fache Verbesserung des Dekodierungsdurchsatzes im Vergleich zur Baseline der herkömmlichen trainable Sparse Attention (InfLLM-V2) erreicht.

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).