NOSA: Native and Offloadable Sparse Attention

• Training-free offloading reduces GPU memory use and enables larger batches, but unseen sparse patterns cause performance degradation, particularly during long generation.The cited passages identify accumulated selection errors as especially problematic in long-generation scenarios.Source: §1 · p. 1
• Trainable sparse attention can alleviate the training-inference gap, but unrestricted KV access leaves cache size unchanged and cannot increase maximum batch size under GPU memory constraints.Source: §1 · pp. 1–2
• InfLLMv2 analysis finds strong block-selection locality, but PCIe bandwidth limitations can still make decoding communication-bound.Source: §1 · p. 2
• NOSA maintains good task performance and consistently outperforms KV-offloading baselines, especially on long-generation tasks.It limits KV transmission volume using query-aware and query-agnostic selection, with eviction over query-agnostic selections bounding fetched blocks.Source: §1 · p. 2
• NOSI improves decoding throughput by up to 5.04× over FullAttn, 1.92× over InfLLMv2, and 1.83× over ShadowKV.The system is tailored for NOSA and targets PCIe communication bottlenecks in vanilla offloading implementations.Source: §1 · p. 2

2.1 Preliminaries

• Query-aware selection scores KV positions using query-key similarity, while query-agnostic selection scores cache units independently and maintains a fixed eviction pattern.The passages describe query-aware top-k selection and query-agnostic eviction as having different recall and locality properties.Source: §2.1 · p. 3

2.2 Observations

• InfLLMv2 selection shows inherent locality: with 4K of 16K tokens selected, most layers have γ(t) ≥ 0.8, so fewer than 20% of KV entries change per step.Locality is the fraction of tokens selected at step t that were also selected at step t −1; this behavior emerges without explicit constraints.Source: §2.2 · p. 3
• Even ∼80% locality is insufficient to mitigate PCIe overhead: more than 80% of decoding time is spent in attention, and query-aware offloading remains communication-bound.Source: §2.2 · pp. 3–4
• Increasing cache-hit rate reduces attention's share of inference time and improves theoretical decoding throughput, motivating training-time constraints that promote higher hit rates.Source: §2.2 · p. 4

3.1 Offloading-aware Trainable Sparse Attention

• NOSA imposes a minimum overlap γ0 between selected tokens at adjacent decoding steps and guarantees a lower bound on locality.The selection divides its budget between query-aware and query-agnostic subsets.Source: §3.1 · pp. 4–5
• NOSA uses two-stage block-wise selection, applying mean pooling to sub-blocks and max pooling at the block level.It combines query-aware selection, which supports retrieval, with query-agnostic selection, which has an eviction constraint.Source: §3.1 · pp. 4–5
• NOSA training starts with dense short-context pretraining, then uses long-context continual pretraining and SFT with NOSA.Source: §3.1 · p. 6

3.2 Ablations on Algorithm Design

• ED-DMA outperforms the other eviction-head implementations on 1B-model RULER evaluation with 16K inputs.The compared implementations include Locret, vanilla DMA, and simple-DMA.Source: §3.2 · p. 6, Table 1 · p. 6
• NOSA achieves approximately 5.5% higher locality than InfLLMv2, equivalent to nearly a 2× reduction in KV-cache miss rate.Source: §3.2 · p. 6
• At batch size 64 and sequence length 4096, block-wise communication sustains high throughput while element-wise communication degrades sharply under high locality.The comparison uses d_kv = 128 and 2 KV heads.Source: §3.2 · p. 6
• Query-aware selection improves recall at test time, and NOSA further outperforms DMA-R on the evaluated recall tasks.The results support using query-aware selection in both training and inference.Source: §3.2 · p. 7

4 NOSI: Inference System for NOSA

• NOSI mitigates inefficient kernels and excessive launches that obscure communication overhead in vanilla Hugging Face inference.It is designed to expose the true communication overhead of NOSA inference.Source: §4 · p. 7
• Fused Triton kernels combine eviction-head computation with QKV splitting and merge query-aware and query-agnostic pooling; CUDA Graphs reduce launch overhead.Source: §4 · p. 7
• A custom CUDA kernel computes the set difference between old and newly selected KV blocks to create a replacement mapping.Only blocks unused in the current step are swapped out, since GPU-resident blocks need not be contiguous.Source: §4 · p. 7
• 83% of peak PCIe bandwidth is reached by NOSI’s custom transfer kernel, versus less than 2 GB/s for the PyTorch implementation.The kernel uses Unified Virtual Addressing for fine-grained, 16 KiB KV blocks.Source: §4 · p. 7

5.2 Long Context Tasks

• NOSA scores highest in 9 of 12 LongBench and HELMET settings; baselines lose at least 10% on hard retrieval tasks, while NOSA’s degradation is modest.On 8B models, NOSA consistently scores highest among competitors, with less severe recall degradation than on smaller models.Source: Table 2 · p. 8, §5.2 · p. 9

5.1 Experimental Setup

• Experiments use 1B, 3B, and 8B models extended from 4K to 16K context and compare offloading, attention, and serving-system baselines across task types.Source: §5.1 · pp. 8–9

5.3 General and Reasoning Tasks

• NOSA and DMA show no significant degradation on general tasks, while training-free offloading performs worse than learned sparse methods on long-generation reasoning.ShadowKV perplexity rises sharply during generation, whereas NOSA’s remains low.Source: Table 3 · p. 9, Table 4 · p. 9, §5.3 · p. 10

5.4 Efficiency Benchmarks

• Up to 5.04× vs FullAttn, 1.92× vs InfLLMv2, and 1.83× vs ShadowKV, NOSA leads offloading throughput; larger feasible batches and locality accompany the gains.Source: Table 5 · p. 10, §5.4 · pp. 10–11

5.5.2 Hyperparameter Analysis

• Average HELMET scores rise with budget k; at sparsity 0.375, NOSA has no degradation versus FullAttn, while Recall and Rerank are more budget-sensitive.NOSA’s average HELMET score exceeds other baselines under full- and sparse-prefill settings, and it remains more robust as budgets decrease.Source: Table 6 · p. 11, §5.5.2 · pp. 12–13

5.5.1 Length Generalization.

• NOSA outperforms the other baselines at 32K and 64K inputs, which are 2× and 4× beyond its 16K training context.At 64K, evaluation increases RoPE’s θ and the sparsity budget k.Source: §5.5.1 · pp. 11–12, Table 7 · p. 12

5.5.2 Hyperparameter Analysis

• Decreasing kq/k raises throughput but weakens recall and task performance; increasing it improves performance at lower throughput, and 0.25 is the main-experiment setting.The ratio generalizes to other inference-time values; the accompanying figure reports HELMET scores and throughput.Source: Figure 8 · p. 12, §5.5.2 · p. 12

5.5.3 Wall-Time Breakdown Analysis; 5.5.4 Case Study

• NOSI makes KV-loading overhead measurable, reduces its time by about half versus InfLLMv2, and NOSA maintains generation quality while ShadowKV’s perplexity rises and its output becomes corrupted.The Hugging Face implementation’s compression and pooling costs obscure KV-loading; Figure 9 breaks down layer decoding time. The generation perplexity comparison uses FullAttn 8B as evaluator.Source: Figure 9 · p. 13, §5.5.3 · p. 13, §5.5.4 · p. 14

Sparse Attention

• Sparse-attention work ranges from fixed patterns to richer or dynamic structures, with query-aware selection widely used and query-agnostic selection associated with larger performance degradation.Source: §6 · p. 14
• Recent trainable sparse-attention methods integrate sparsity during training to address training-inference mismatch; some use switchable dense-sparse patterns for short and long contexts.Source: §6 · p. 14

LLM Offloading Systems

• LLM offloading research spans parameter offloading and block-wise KV offloading for long sequences and higher decoding throughput.Source: §6 · p. 14

7 Conclusion

• NOSA combines trainable sparse attention with an explicit locality constraint for KV-cache offloading, while NOSI is its inference system.Source: §7 · p. 14
• NOSA outperforms KV-cache offloading baselines on general, long-context input, and long-generation tasks across evaluations on 1B, 3B, and 8B models.Source: §7 · p. 14
• Up to 5.04×, 1.92×, and 1.83× higher decoding throughput is reported over FullAttn, InfLLMv2, and ShadowKV, respectively.Source: §7 · p. 14

Impact Statement

• The authors expect more efficient LLM decoding to lower inference compute and energy costs and potentially reduce environmental footprint through broader application deployment.The impact statement reports no anticipated negative ethical or societal impacts.Source: Impact Statement · p. 15

A Pseudocode of NOSA

• At each decoding step, NOSA selects top-k blocks, transfers cache-missed KV blocks from CPU to GPU, then computes attention with the KV cache.Selection combines query-aware and query-agnostic scores, selecting up to kq tokens to maintain a minimum locality of ke/k.Source: A Pseudocode of NOSA · p. 21

B.1 Environment

• Training uses 32 NVIDIA A800 GPUs across four nodes, while inference uses one GPU and 26 nearby NUMA-domain CPU cores.GPU–CPU communication uses PCIe 4.0×16, with a stated peak bandwidth of 31.5 GB/s.Source: §B.1 · p. 21

B.2.1 Model Training

• The training pipeline produces 12 models: 1B, 3B, and 8B versions of FullAttn, InfLLMv2, DMA, and a NOSA variant.The model architecture and training details are summarized in Table 8.Source: §B.2.1 · p. 21, Table 8 · p. 22
• The 1B and 3B models use the Llama-3 architecture, while the 8B models use MiniCPM4-base; all share the MiniCPM4 tokenizer.LongRoPE extends the 4K pretraining context to 16K during long-context continual pretraining.Source: §B.2.1 · p. 21
• Long-context continual pretraining uses InfLLMv2-data-5B, AdamW, and the WSD learning-rate scheduler.Training uses the Megatron framework.Source: §B.2.1 · p. 21

B.2.2 Hyperparameters of each method

• InfLLMv2 uses 4096 selected tokens, with one sink block, 16 sliding-window blocks, and 47 dynamically selected blocks for each query token.This configuration is used for both training and inference.Source: §B.2.2 · p. 22
• DMA follows the InfLLMv2 configuration, but selects its 47 dynamic blocks using importance scores predicted by the eviction head..Source: §B.2.2 · p. 22
• For NOSA, k = 4096, kq = 1024, and ke = 3072; 16 blocks are query-aware, while the remaining dynamic blocks are query-agnostic.Including the attention sink and sliding window, 48 blocks are query-agnostic, giving γ(t) ≥75%.Source: §B.2.2 · p. 22
• For long-context tasks, ShadowKV uses 4096 sparse-attention tokens and evaluates block sizes 8 and 64, with the same sliding-window and outlier allocation.The variants also include settings with MInference applied during prefill; the 8-block variant is used for long-generation reasoning tasks.Source: §B.2.2 · p. 22
• For 8192-token generation, ShadowKV is adapted to keep a 4K context budget; its efficiency configuration uses 2048 dynamic tokens, a 1280-token sliding window, and 768 outlier tokens.The authors state that this efficiency configuration favors ShadowKV and disadvantages NOSA, yet NOSA remains faster in that setting.Source: §B.2.2 · p. 23
• ArkVale uses 4096 sparse-attention tokens, block size 32, a 64-token sink, and a 1024-token sliding window; its implementation falls back to MHA for the 16:1 query:key GQA ratio.The stated configuration favors ArkVale’s performance, and its official implementation does not support bfloat16 inference.Source: §B.2.2 · p. 23
• InfLLM is evaluated with 4096 retained tokens in block-size-64 and block-size-128 variants, alongside a 1024-token sliding window and one sink block.The block-size-64 variant is included for comparison with NOSA; the original InfLLM design uses block size 128.Source: §B.2.2 · p. 23

B.3 Datasets

• Long-context evaluation uses LongBench English tasks averaging over 16K input tokens and HELMET at 16K; general, reasoning, and decoding efficiency use separate benchmark sets.LongBench and HELMET results distinguish full from sparse prefill; general-task abbreviations are mapped in the text, and sparsity and attention bias are not applied for short sequences.Source: §B.3 · p. 23

B.4 Benchmark Settings

• LongBench uses greedy decoding and per-task scores; HELMET uses greedy decoding and reports mean scores with standard deviations for sampled evaluation instances.Long-context results are reported separately for full and sparse prefill where baselines differ in their prefill sparsity.Source: §B.4 · pp. 23–24
• General-task evaluation uses LM-Evaluation-Harness with task-specific shot counts and scoring protocols across eight benchmarks.Chat templates are applied for BBH, GSM8K, DROP, and HumanEval, but not the other listed tasks.Source: §B.4 · p. 24
• Reasoning evaluation uses a two-shot, approximately 22K-token prompt and permits up to 8K generated tokens; the 1B and 3B models are not evaluated on these tasks.Gaokao-Phy is repeated four times because the dataset is small and high variance; GPT-4o-2024-11-20 judges the outputs.Source: §B.4 · p. 24
• The LLM judge checks whether a generation explicitly contains the answer, accepting semantic equivalence but ignoring reasoning correctness..Source: §B.4 · p. 25
• Equivalent batch size controls on-GPU KV-cache size, and the corresponding real batch size is set to match that budget across methods.For example, at 16K input length NOSA retains 4K GPU tokens, so its GPU KV-cache footprint is one quarter of FullAttn’s.Source: §B.4 · p. 25

C Baselines

• The baselines section analyzes how NOSA differs from reference and baseline methods and briefly discusses potentially related methods..Source: C Baselines · p. 25

C.1 Reference and Baseline Methods

• InfLLMv2 uses block-wise sparsity without an explicit locality constraint, so queries may fetch arbitrary numbers of CPU blocks; NOSA adds locality constraints and query-aware and query-agnostic selection.The authors follow InfLLMv2’s long-context continual-pretraining setting and report higher decoding efficiency with no substantial performance degradation.Source: §C.1 · p. 25
• DMA estimates token importance from value vectors and supports query-agnostic eviction, but discarded tokens cannot be recalled; NOSA combines this with query-aware selection.The section notes that DMA is ill-suited to recall-heavy tasks, while DMA-R may underperform because of training–inference mismatch.Source: §C.1 · p. 26
• ShadowKV introduces sparsity only at decoding and may accumulate approximation errors over long generations; ArkVale and InfLLM also face efficiency or long-context performance drawbacks.ShadowKV overlaps K reconstruction with V fetching; ArkVale’s implementation overhead makes it slower than other baselines, while InfLLM is not optimized for large batches.Source: §C.1 · p. 26

C.2 Other Potentially Related Methods

• KV-cache compression spans token-, head-, layer-, and dimension-level approaches, and is described as orthogonal to NOSA because it targets the cache itself.The authors leave combinations with compression techniques for future work; they note that MLA can require substantially more training effort than NOSA.Source: §C.2 · pp. 26–27
• KV-cache quantization reduces memory use, with cited methods reaching 4-bit, 2-bit, and 1-bit representations; NOSA is also described as orthogonal to quantization.A systematic study of combining the two is left for future work.Source: §C.2 · p. 27
• KV-cache eviction reduces memory burden by discarding units, but eviction is irreversible; the section identifies eviction methods as possible NOSA eviction heads.The text states that eviction is expected to underperform offloading because discarded tokens cannot be retrieved.Source: §C.2 · p. 27

D Ablation on Eviction Head Implementation

• The ablation pretrains a 1B model at 8K with dense attention, then extends to 16K for NOSA continual pretraining; selection budget k is 4096.The budget includes 64 attention sink tokens, 1024 sliding-window tokens, and kq = 1024; models are evaluated on RULER.Source: D Ablation on Eviction Head… · p. 27
• 5.7% accuracy improvement results from shifting the exponential operator from bias calculation to attention computation, motivating the ED-DMA implementation.ED-DMA uses the preexponential value from DMA for query-agnostic selection.Source: D Ablation on Eviction Head… · pp. 27–28
• ED-DMA performs best and remains nearly lossless against the original InfLLMv2 on RULER; other implementations degrade noticeably, especially on MK and MV.The evaluation uses a 1B model and a 16K sequence length for NOSA continual pretraining.Source: D Ablation on Eviction Head… · p. 27, Table 10 · p. 28

E Theoretical Guarantee of Locality

• With selection budget k = kq + ke, NOSA guarantees at least ke selected entries overlap between consecutive decoding steps.The proof states the overlap as |Γ(t −1) ∩Γ(t)| ≥ ke.Source: E Theoretical Guarantee of Locality · p. 28