Technical Summary: NOSA: Native and Offloadable Sparse Attention

Problem Statement

The decoding throughput of Large Language Models (LLMs) is increasingly constrained by GPU memory limitations, primarily due to the storage requirements of the Key-Value (KV) cache. While KV cache offloading (keeping redundant context on CPU and fetching subsets to GPU) alleviates memory pressure, existing approaches suffer from a critical trade-off:

1. Training-Free Offloading: Methods like ShadowKV and InfLLM select sparse subsets without training. This creates a training-inference mismatch, where sparse patterns unseen during training lead to significant performance degradation, particularly in long-generation scenarios where selection errors accumulate.
2. Trainable Sparse Attention: Methods that learn sparse patterns (e.g., InfLLMv2) avoid the mismatch but are incompatible with efficient offloading. Without constraints, these methods may fetch arbitrary KV blocks from the CPU, causing massive PCIe communication overhead that negates throughput gains.

Current solutions fail to simultaneously achieve high decoding efficiency (via offloading) and high generation quality (via training consistency).

Methodology

The authors propose NOSA (Native and Offloadable Sparse Attention), a trainable sparse attention mechanism explicitly designed for KV cache offloading, coupled with NOSI, a specialized inference system.

1. Algorithm: NOSA

NOSA introduces a locality constraint during training to ensure that the sparse attention patterns remain stable across decoding steps, thereby minimizing CPU-GPU data transfers.

• Dual-Selection Mechanism: NOSA decomposes KV selection into two components:
  • Query-Aware Selection ($\Gamma_q$): Uses the query to retrieve specific, potentially distant information (preserving recall).
  • Query-Agnostic Selection ($\Gamma_e$): Uses an eviction head (inspired by DMA) to assign importance scores independent of the query. This component enforces a strict eviction policy: if a token is not selected at step $t$, it cannot be selected at any future step $t' > t$.
• Locality Guarantee: By allocating a fixed budget $k_e$ to the query-agnostic selection, NOSA mathematically guarantees a minimum locality $\gamma(t) \geq k_e/k$. This ensures that the majority of selected KV blocks in step $t$ were also selected in step $t-1$, allowing them to remain cached on the GPU.
• Block-Wise Implementation: To maximize PCIe bandwidth utilization, NOSA operates on blocks of tokens rather than individual tokens. It employs a two-stage pooling strategy (mean pooling within sub-blocks, followed by max pooling at the block level) to compute selection scores.
• ED-DMA (Exp-Delayed DMA): To handle numerical precision issues in the eviction head, NOSA defers the exponential operation until the attention calculation stage, selecting blocks based on pre-exponential scores.

2. System: NOSI

The authors observe that vanilla implementations (e.g., Hugging Face Transformers) introduce excessive kernel launch overheads that mask the true communication bottlenecks. NOSI is a custom inference system built to expose and optimize these bottlenecks:

• Kernel Fusion: Fuses eviction head computations, QKV splitting, and pooling operations into single Triton kernels to reduce launch overhead.
• Optimized Memory Layout: Organizes KV cache into minimal block-sized units and uses custom CUDA kernels to efficiently compute set-differences for block swapping, avoiding contiguous memory requirements.
• High-Bandwidth Communication: Implements custom host-to-device transfer kernels using Unified Virtual Addressing (UVA) to achieve near-peak PCIe bandwidth, overcoming the fragmentation issues of standard PyTorch implementations.

Key Contributions

1. Feasibility Analysis: The paper identifies that while trainable sparse attention (like InfLLMv2) possesses inherent locality, it is insufficient to overcome PCIe bandwidth limitations without explicit constraints.
2. NOSA Algorithm: A novel trainable sparse attention mechanism that explicitly bounds the volume of CPU-GPU KV transfers by enforcing a locality constraint via a query-agnostic eviction policy.
3. NOSI System: An inference system tailored to NOSA that eliminates software overheads, revealing communication as the true bottleneck and enabling high-throughput offloading.
4. Comprehensive Evaluation: Extensive benchmarks on 1B, 3B, and 8B models demonstrating that NOSA outperforms both training-free offloading baselines and other trainable sparse attention methods.

Experimental Results

The authors evaluated NOSA on 1B, 3B, and 8B models across general tasks, long-context inputs, and long-generation reasoning tasks.

• Performance: NOSA consistently outperforms training-free offloading baselines (ShadowKV, InfLLM, ArkVale) and trainable baselines (InfLLMv2, DMA) on long-context benchmarks (LongBench, HELMET). It maintains high recall and reasoning capabilities where training-free methods degrade significantly due to the training-inference mismatch.
• Throughput:
  • NOSA achieves up to 5.04× higher decoding throughput compared to FullAttn (due to larger batch sizes enabled by offloading).
  • It achieves up to 1.92× higher throughput than InfLLMv2 (due to reduced communication from locality constraints).
  • It achieves up to 1.83× higher throughput than ShadowKV (due to better efficiency in handling sparse patterns and system optimizations).
• Long-Generation: In long-generation reasoning tasks, NOSA maintains low perplexity, whereas training-free offloading methods exhibit error accumulation and perplexity explosion.
• Ablation: The study confirms that the block-wise selection pattern is critical for PCIe efficiency and that the ED-DMA eviction head implementation yields the best performance.

Significance and Claims

The paper claims that NOSA represents a necessary shift from "training-free" or "unconstrained trainable" approaches to algorithm-system co-design. By explicitly constraining the volume of KV transfers during training, NOSA bridges the gap between the need for memory efficiency (offloading) and the need for generation quality (training consistency).

The authors state that their work enables the deployment of LLMs with long-context capabilities on memory-constrained hardware without sacrificing reasoning quality. They emphasize that their approach lowers compute and energy costs by improving decoding efficiency, potentially facilitating broader deployment of LLM applications. The code is released to encourage further research in native offloading mechanisms.