TLDR: NOSA is a new trainable sparse attention framework that significantly improves the decoding efficiency of Large Language Models (LLMs) for long contexts. It addresses the KV cache bottleneck by leveraging and enhancing the natural locality in token selection, enabling efficient offloading of the KV cache to the CPU. Through a combination of query-aware and query-agnostic selection components and an optimized inference system, NOSA achieves up to a 2.3x increase in decoding throughput with near-lossless task performance compared to existing methods.
Large Language Models (LLMs) have revolutionized artificial intelligence, enabling advanced agent systems and complex reasoning. However, as these models handle increasingly longer contexts, a significant bottleneck emerges during the decoding process: memory input/output (I/O). This memory constraint primarily stems from the Key-Value (KV) cache, which stores information from previous tokens to accelerate future computations. While existing sparse attention methods have helped reduce computational load and memory access, they haven’t addressed the fundamental issue of KV cache size, limiting the maximum batch size and overall decoding speed, especially in large-scale inference scenarios.
A new research paper, “NOSA: Native and Offloadable Sparse Attention,” introduces a novel framework designed to tackle this challenge head-on. The authors, Yuxiang Huang, Chaojun Xiao, Xu Han, and Zhiyuan Liu from Tsinghua University, observed that trainable sparse attention naturally exhibits a strong “locality” in how it selects tokens across consecutive decoding steps. This means that the set of tokens an LLM pays attention to in the current step largely overlaps with those from the previous step. This inherent locality opens the door for offloading parts of the KV cache to the CPU, freeing up valuable GPU memory.
However, the researchers found that this natural locality alone isn’t enough. Even with an 80% overlap in selected tokens, the transfer of KV pairs between the CPU and GPU via the PCIe bus still dominates the decoding cost. This insight led to the development of NOSA, which explicitly enhances this locality during the model’s training phase. NOSA achieves this by decomposing token selection into two components: query-aware and query-agnostic. Query-aware selection allows the model to maintain its ability to retrieve important, previously evicted tokens, crucial for performance. Query-agnostic selection, on the other hand, introduces explicit locality constraints, ensuring a higher cache hit rate and significantly reducing the need for frequent data transfers between the CPU and GPU.
The NOSA framework is not just a theoretical concept; it includes a carefully optimized inference system. This system incorporates an efficient memory layout for block-sparse patterns, custom-designed communication kernels that leverage unified virtual addressing (UVA) for parallel CPU-to-GPU transfers, and a high-performance C++ memory manager for block tables. These engineering optimizations ensure that the benefits of NOSA’s algorithmic design are fully realized in practice, achieving impressive data transfer speeds and management efficiency.
Extensive benchmarks were conducted using a 1-billion parameter model trained with NOSA. The results are compelling: NOSA maintains nearly lossless task performance across various short-context and long-context evaluations, including MMLU, LongBench, and RULER. More importantly, it delivers a substantial improvement in decoding efficiency, achieving up to a 2.3 times increase in decoding throughput compared to the vanilla trainable sparse attention baseline (InfLLM-V2). This efficiency gain is particularly pronounced with larger batch sizes and longer input sequences, making NOSA a promising solution for deploying LLMs in memory-constrained environments or for applications requiring high throughput.
Also Read:
- LouisKV: A New Approach to Efficient KV Cache Management for Long Language Model Sequences
- Efficient LLM Inference: Unpacking Direct Multi-Token Decoding
In essence, NOSA represents a significant step forward in optimizing LLM inference for long contexts. By intelligently combining inherent token selection locality with explicit locality constraints and a highly optimized offloading system, it allows LLMs to process longer sequences more efficiently without compromising performance. For more details, you can refer to the research paper: NOSA: Native and Offloadable Sparse Attention.
