NOSA (Native and Offloadable Sparse Attention) is a trainable sparse attention mechanism designed so that most of a language model's KV cache can live in CPU memory during decoding without the CPU-to-GPU transfer traffic swallowing the throughput gains. It was developed by the NLP group at Tsinghua University (THUNLP) with OpenBMB, and described in an arXiv preprint first posted on 15 October 2025 and substantially revised on 29 January 2026.[1][3] The paper pairs the attention mechanism with NOSI, a purpose-built inference system, and reports that on 1B, 3B and 8B models NOSA improves decoding throughput by up to 5.04x over full attention, 1.92x over InfLLM-V2 and 1.83x over ShadowKV while beating training-free KV cache offloading baselines on long-input, long-generation and general tasks.[1][2] Code is released under the MIT license on GitHub and NOSA-1B, NOSA-3B and NOSA-8B checkpoints are on Hugging Face under Apache 2.0.[4][5][6][7] In June 2026 NVIDIA's SparDA paper adopted NOSA-8B as one of its two evaluation models and built its efficiency benchmarks on the NOSI engine.[9][10]

Background

Decoding throughput scales with batch size, and batch size on a GPU is bounded by memory that the KV cache consumes. Two families of techniques attack this from different sides, and the NOSA paper's argument is that each fails a test the other passes.[2]

Training-free KV cache offloading systems such as ShadowKV, InfLLM and ArkVale keep most KV entries on the CPU and fetch a sparse subset to the GPU for each decoding step. They cut GPU memory and allow larger batches, but the sparse pattern was never seen during training, so a training-inference mismatch appears. The paper reports that this mismatch is worst on long generations, where selection errors accumulate; in a case study a ShadowKV-served 8B model's perplexity explodes partway through a MATH-500 solution and the output degenerates into whitespace and stray tokens, while the NOSA model finishes the solution.[2]

Trainable sparse attention methods such as Native Sparse Attention, MoBA and InfLLM-V2 learn the sparse pattern during training and so avoid the mismatch, but they do not shrink the KV cache: any entry may be selected by a future query, so everything stays on the GPU and the maximum batch size is unchanged.[2]

The paper then makes two empirical observations on InfLLM-V2, the trainable block-sparse mechanism used in OpenBMB's MiniCPM4 models.[11] First, block selection already has strong locality: with 16K-token inputs and 4K tokens selected, the fraction of selected tokens that were also selected on the previous step (the paper's "locality" gamma(t)) is at least 0.8 for most layers, so under 20 percent of the working set changes per step. Second, 80 percent locality is still not enough. On an NVIDIA A800-80GB with a PCIe 4.0 x16 link (31.5 GB/s peak), more than 80 percent of decoding time in a naive offloaded implementation sits inside attention, and a simulated throughput model shows throughput rising steeply as the cache hit rate climbs toward 1.0.[2] The design goal that follows is to force locality higher at training time rather than hoping for it.[2][8]

How NOSA works

NOSA is built on top of InfLLM-V2's two-stage block-wise selection and keeps its attention sink and sliding window. The change is in how the dynamically selected blocks are chosen. The per-step selection budget k is split into two parts, k = k_q + k_e:[2]

Component	Selection score	Constraint	Purpose
Query-aware selection (k_q blocks)	Dot product of the current query with mean-pooled block keys, as in InfLLM-V2	None	Preserves the model's ability to recall distant context, including blocks that were evicted earlier
Query-agnostic selection (k_e blocks)	An importance score computed from each token's value vector by a trainable "eviction head"	Once a block is not selected at step t it is never re-selected by this component	Behaves like a learned KV eviction policy, so its selections never require fetching from the CPU

At each step the k_q query-aware blocks are picked first, their scores are set to positive infinity, and the remaining k_e blocks are filled from the top of the query-agnostic scores. Because the query-agnostic subset can only shrink or extend forward in time, the paper proves (Theorem 2, Appendix E) that the overlap between consecutive steps satisfies gamma(t) >= k_e / k for every step. With the main-experiment setting of k = 4096 and k_q = 1024, the guaranteed lower bound is 75 percent; measured locality on a 16K PG19 sample was a mean of 0.944 across layers for NOSA versus 0.889 for InfLLM-V2, which the authors describe as roughly halving the cache miss rate.[2]

The eviction head is adapted from DMA (Trainable Dynamic Mask Sparse Attention, Shi et al., August 2025), which computes a per-token score s = tau(v W1) W2 from the value vector, with tau a softplus, and injects it as an additive attention bias so the head receives gradients.[2][12] NOSA's variant, which the authors call ED-DMA (Exp-Delayed DMA), ranks tokens on the pre-exponential score rather than exp(b) because the selection turned out to be highly sensitive to numerical precision. In an ablation on a 1B model evaluated on RULER at 16K, ED-DMA scored 61.9 percent against 60.3 for plain InfLLM-V2, 55.9 for a Locret-style MLP head, 56.2 for vanilla DMA and 60.9 for a straight-through-estimator variant.[2]

Selection is block-wise rather than token-wise for a systems reason: element-wise gathers from host memory produce a large number of irregular, fragmented transfers with poor PCIe utilisation, and the paper's microbenchmark shows element-wise host-to-device throughput collapsing at high locality while block-wise transfers hold up.[2]

NOSA is applied from the long-context stage of training onward. Models are pretrained with ordinary dense attention at a 4K context, switched to NOSA for long-context continual pretraining at 16K (with LongRoPE for the extension), and kept on NOSA through supervised fine-tuning, so the released checkpoints have offloading behaviour "built in". For short sequences on general benchmarks no sparsity or attention bias is applied.[2]

Main-experiment configuration

Setting	Value
Block size n_b	64 tokens
Mean-pool stride / kernel	32 / 16 as printed in the paper (the SparDA paper and InfLLM-V2 give the same model a kernel of 32 and stride of 16)
Selected tokens per query, k	4096 (1 sink block, 16 sliding-window blocks, 47 dynamic blocks)
Query-aware budget k_q	1024 tokens (16 blocks)
Query-agnostic budget k_e	3072 tokens (31 dynamic blocks; with sink and window counted, 48 query-agnostic blocks)
Guaranteed locality	gamma(t) >= 0.75
Per-block transfer unit in NOSI	16 KiB

Source: paper Appendix B.2.2 and Section 4.[2] The k_q / k ratio is a tunable inference-time knob: the paper's Figure 8 shows HELMET score rising and throughput falling as k_q / k grows from 0.125 to 0.5 on the 8B model, and 0.25 was chosen as the balance point.[2]

NOSI: the inference system

The authors found that a straightforward Hugging Face Transformers implementation was so slow (block compression and pooling dominated, with many small kernel launches) that PCIe traffic was not even the bottleneck; NOSA under plain Transformers decoded at 14.69 tokens per second at 16K, slower than InfLLM-V2's 47.10. NOSI ("NOSA-based System for Inference") is the engine they wrote to expose the real communication cost.[2] Its main pieces:[2]

FlashInfer kernels for LayerNorm, RoPE and the feed-forward network.
Two hand-written Triton kernels: one fuses the eviction heads with the QKV split, the other fuses the query-aware and query-agnostic max-pooling stages. CUDA Graphs cover the pipeline from pooling to index generation.
A custom CUDA kernel that computes the set difference between the previous and current block selections, so only blocks that changed are swapped, with O(k) work per thread.
A Triton host-to-device kernel using Unified Virtual Addressing to read CPU memory directly. In the paper's microbenchmark it reaches up to 83 percent of peak PCIe bandwidth where a PyTorch copy of the same fragmented blocks stays under 2 GB/s.

A wall-time breakdown of one 8B layer at 16K input and equivalent batch 16 shows that under NOSI the KV-loading segment for NOSA is about half of InfLLM-V2's, and the fused QKV-plus-scoring kernel is also faster than InfLLM-V2's PyTorch split.[2] The repository ships NOSI as an installable package (nosi/) alongside a Transformers-compatible modeling_llama_nosa.py, benchmark suites, and a Docker image distributed through ModelScope.[4] The paper lists integration into production serving engines and reinforcement-learning rollout as future work, and vLLM and SGLang appear only as full-attention baselines in its efficiency tables.[2]

Models and training

The paper trains twelve models: 1B, 3B and 8B sizes, each in FullAttn, InfLLM-V2, DMA and NOSA variants, all through the same pipeline. Only the NOSA checkpoints are public, gathered in an OpenBMB collection on Hugging Face; the README invites requests for the baselines.[2][4][13]

	1B	3B	8B
Architecture	Llama 3 style	Llama 3 style	MiniCPM4-base (uses muP)
Layers / hidden size	28 / 2048	32 / 2560	32 / 4096
Attention heads / KV heads	16 / 2	32 / 2	32 / 2
Vocabulary	73,448	73,448	73,448
Pretraining tokens (4K context)	1.5T	3.8T	8.0T
Long-context continual pretraining tokens (16K)	2.0B	3.4B	5.0B
SFT tokens	0.6B	1.3B	1.9B
Long-context CPT GPU hours	336	883	1,496
SFT GPU hours	78	190	377
Hugging Face repo	openbmb/NOSA-1B	openbmb/NOSA-3B (3.61B parameters)	openbmb/NOSA-8B (8.19B parameters)

Sources: paper Table 8; Hugging Face model metadata.[2][5][6][7] All three share the MiniCPM4 tokenizer, use AdamW with a WSD schedule, and were trained with Megatron-LM on 32 NVIDIA A800 GPUs across four nodes; long-context continual pretraining used OpenBMB's InfLLM-V2-data-5B dataset.[2] The 8B model's 16:1 ratio of query to KV heads (grouped-query attention) matters for one baseline: ArkVale's official implementation only supports 1:1, 4:1 or 8:1 and falls back to multi-head attention, which limits its batch size in the efficiency comparison.[2]

The Hugging Face cards list the checkpoints as Apache 2.0, English and Chinese, with custom_code remote modelling files; the NOSA-8B config declares 32,768 maximum position embeddings with LongRoPE scaling and a RoPE base of 10,000, even though the paper's long-context training ran at 16,384 tokens.[5][14] The NOSA-8B repository was created on 4 January 2026 and the 1B on 21 January, before the v2 paper appeared (the 8B repository was last modified on 12 February 2026); the 8B card's overview paragraph still says the repository "contains the weights for the NOSA-3B model", a copy-editing slip.[5][7]

Results

Every comparison in the paper is against models trained through the same pipeline, and the authors treat FullAttn and InfLLM-V2 as upper bounds rather than competitors; the "baselines" are the training-free offloading systems ShadowKV (ShadKV, at block size 8 or 64, with MInference-style sparse prefill in the "M" variants), InfLLM (block size 64 or 128) and ArkVale, plus the DMA models. NOSA_F uses full attention for prefill and sparse attention for decode; NOSA_S is sparse in both phases.[2]

Long-context input (LongBench and HELMET at 16K)

Model	Method	LongBench avg	HELMET avg	HELMET Recall
1B, full prefill	FullAttn	30.2	28.3	61.3
	ShadKV8	29.4	25.5	48.9
	ArkVale	20.6	19.8	26.7
	NOSA_F	29.8	24.0	34.3
1B, sparse prefill	InfLLM-V2	29.8	24.4	37.2
	ShadKV8 (M)	29.3	23.9	41.9
	NOSA_S	30.3	22.9	29.2
3B, full prefill	FullAttn	38.0	29.6	48.4
	ShadKV8	36.9	26.5	39.4
	ArkVale	37.9	26.1	28.4
	NOSA_F	37.4	28.2	45.1
3B, sparse prefill	InfLLM-V2	38.1	27.5	39.7
	ShadKV8 (M)	36.1	26.3	43.7
	NOSA_S	38.0	27.8	36.1
8B, full prefill	FullAttn	43.9	42.6	89.5
	ShadKV8	43.3	38.8	78.3
	ArkVale	39.5	36.2	56.2
	NOSA_F	43.7	40.4	86.3
8B, sparse prefill	InfLLM-V2	44.3	40.3	72.7
	ShadKV8 (M)	41.1	36.5	74.6
	NOSA_S	43.3	37.6	67.2

Source: paper Table 2, which also reports ShadKV64, InfLLM64, InfLLM128 and DMA rows and per-task scores.[2] Counting the 12 combinations of model size, prefill mode and benchmark, NOSA has the best offloading-method score in 9, ShadowKV in 2 (both 1B HELMET settings) and ArkVale in 1 (3B LongBench with full prefill).[2] The paper is candid that sparse attention of any kind is unstable at 1B scale, where InfLLM-V2 alone loses more than 20 points of HELMET recall, and that NOSA's advantage grows with model size.[2]

LongBench was restricted to its English tasks with average input above 16K; HELMET was run at 16K with its recall subset consisting of the MK2, MK3 and MV tasks from RULER plus JsonKV, and its LongQA and summarisation tasks judged by GPT-4o.[2]

General and long-generation tasks

On eight short-context benchmarks (MMLU, MMLU-Pro, BBH, GSM8K, MATH, DROP, MBPP, HumanEval) evaluated with full attention, the NOSA models average 33.4, 43.9 and 53.6 at 1B, 3B and 8B against 33.4, 44.0 and 52.9 for the FullAttn models trained the same way, so the sparse training did not cost general ability.[2]

The long-generation test used the 8B model with a roughly 22K-token two-shot reasoning prompt (the base model was never trained on reasoning traces) and up to 8K generated tokens, scored by gpt-4o-2024-11-20:[2]

Dataset	FullAttn	InfLLM-V2	ShadKV	InfLLM	ArkVale	DMA	NOSA
MATH-500	52.8	50.6	19.0	38.8	47.0	50.4	50.4
Gaokao-MS	55.6	49.5	7.0	36.9	43.9	57.0	51.9
Gaokao-MH	61.5	57.8	15.6	42.2	47.7	58.7	62.4
Gaokao-Phy	36.3	34.8	8.2	29.7	30.5	31.6	33.6
Average	51.6	48.2	12.5	36.9	42.3	49.4	49.6

Source: paper Table 4.[2] The ShadowKV collapse is the training-inference mismatch made visible; the paper notes that ShadowKV's official implementation caps generation at 1,024 tokens and keeps generated tokens on the GPU as a growing window, so the authors added value-cache offloading and prefill-time SVD factorisation of the key cache to hold it to the same 4K budget for the 8K-token generations; in the throughput benchmark ShadowKV instead ran in its simulation mode with 2,048 dynamic tokens and a 1,280-token window, a setting the paper says favours ShadowKV.[2]

Decoding throughput

Throughput was measured on the 8B model on a single A800-80GB with 26 CPU cores, on PG19 text, averaged over four runs. Because methods keep different amounts of KV cache on the GPU, the paper fixes the on-GPU KV cache size through an "equivalent batch size" (EB): at 16K input NOSA keeps 4K of 16K tokens resident, so EB 128 means a real batch of 128 for NOSA but 32 for full attention. The table below takes the largest EB column at each input length; real batch sizes are in parentheses.[2]

Input length (EB)	Best FullAttn engine	InfLLM-V2 on NOSI, offloaded	ShadowKV	NOSA on NOSI, offloaded
16K (128)	1,267.20 SGLang (32)	1,441.32 (128)	1,071.36 (128)	1,961.17 (128)
32K (128)	818.53 vLLM (16)	652.60 (128)	995.81 (128)	1,536.19 (128)
64K (64)	240.28 vLLM (4)	717.65 (64)	866.01 (64)	1,378.82 (64)
96K (64)	214.42 SGLang (4)	725.26 (64)	730.48 (64)	1,080.45 (64)

Tokens per second, from paper Table 5.[2] The headline ratios come from this table: 1,080.45 / 214.42 = 5.04x over full attention at 96K, where NOSA's real batch is 16 times larger; 1,378.82 / 717.65 = 1.92x over InfLLM-V2 at 64K; and 1,961.17 / 1,071.36 = 1.83x over ShadowKV at 16K.[2] Two caveats follow from the paper's own protocol: the comparison fixes the on-GPU KV cache size, so the full-attention engines run at a real batch of 4 where NOSA runs 64 (only the Hugging Face Transformers rows actually run out of memory), and the paper itself attributes the 5.04x to the real batch being 16 times larger; and InfLLM-V2 was itself run on NOSI with offloading, so the 1.92x isolates the locality constraint rather than the engine. Each baseline was run in its authors' own implementation, which the paper defends as unavoidable given how tightly offloading systems are coupled to their algorithms.[2]

Ablations

The budget study on the 8B model uses HELMET with 20 sampled instances per task, as does the k_q / k sweep above. Varying the budget k, NOSA_F averaged 37.8 / 41.3 / 43.3 and NOSA_S 31.3 / 37.9 / 42.1 at k = 2048 / 4096 / 6144, against 42.1 for full attention; at k = 6144 (37.5 percent of a 16K context) neither NOSA variant trails full attention, and the recall and rerank tasks are the most sensitive to k.[2] For length generalisation beyond the 16K training window (32K evaluated under the main-experiment protocol), at 32K NOSA_F scored 36.0 (FullAttn 38.2, ShadKV8 34.6, ArkVale 32.5) and at 64K, with RoPE theta raised from 10,000 to 40,000 and k = 6144, NOSA_F scored 27.7 (FullAttn 30.2, ArkVale 27.6, ShadKV8 26.5).[2]

Version history

Version	Date	Scope
v1	15 October 2025	Marked "Working in Progress"; four authors (Yuxiang Huang, Chaojun Xiao, Xu Han, Zhiyuan Liu); a single 1B model; headline result 2.3x decoding throughput over InfLLM-V2 on an A800
v2	29 January 2026	Twelve authors; 1B, 3B and 8B models; NOSI system; ShadowKV, InfLLM, ArkVale, vLLM and SGLang comparisons; 5.04x / 1.92x / 1.83x headline

Sources: arXiv version pages.[1][3] The v2 author list is Yuxiang Huang and Pengjie Wang (equal contribution), Jicheng Han, Weilin Zhao, Zhou Su, Ao Sun, Hongya Lyu, Hengyu Zhao, Yudong Wang, Chaojun Xiao, Xu Han and Zhiyuan Liu, with Xiao, Han and Liu as corresponding authors; affiliations as printed are Tsinghua's NLP group (DCST, IAI, BNRIST), OpenBMB, BUPT and the Beijing Institute of Technology.[2] The GitHub repository was created on 28 January 2026, the day before v2, and as of 16 September 2026 the arXiv listing carried the comment "Preprint" and no venue.[1][4] A January 2026 Chinese-language blog post by the first author places NOSA in the lineage of DeepSeek's NSA, Kimi's MoBA and the group's own InfLLM-V2 (which the post says was accepted at ICLR 2026), and frames the contribution as adding a locality lower bound so that trainable sparse attention can finally enlarge batch size, the one lever it previously could not pull.[8]

Adoption in SparDA

NVIDIA's SparDA paper (Sparse Decoupled Attention, arXiv 2606.04511, June 2026) is a published external use of the NOSA checkpoints and code. It characterises NOSA as pairing "a query-agnostic eviction head with a query-aware selector" and issuing UVA-based CPU-to-GPU transfers synchronously per layer, "reducing volume rather than hiding latency"; SparDA's own contribution is to forecast the next layer's block selection one layer ahead so the fetch can overlap compute.[9] SparDA evaluates on MiniCPM4.1-8B and NOSA-8B, implements all of its throughput measurements on the NOSI engine, and configures NOSA-8B with a 96-block budget of which 24 blocks (25 percent) are query-aware and a 1,024-token window, "following Huang et al.".[9][10] On NOSA-8B at its 32K native maximum, SparDA reports aggregate averages of 52.4 for dense attention, 49.4 for the NOSA sparse path, 45.6 for InfiniGen and 51.7 for SparDA, and decode speedups over the NOSA path of up to 1.40x at 128K, lower than on MiniCPM4.1-8B because, in SparDA's words, NOSA's "query-agnostic eviction head already reduces KV fetch traffic, leaving less room for overlap".[9]

Limitations

The paper's own results mark out where the method is weakest. NOSA_S trails ShadowKV on HELMET at 1B scale, and its recall at 1B with sparse prefill (29.2) is well below InfLLM-V2's (37.2), so the locality constraint is not free at small scale.[2] Throughput comparisons run each baseline in its original implementation rather than in a common engine, and the largest speedup comes from NOSA running a 16 times larger real batch under the matched on-GPU cache protocol.[2] The released models were long-context trained at 16K, so their 32K and 64K results rely on RoPE scaling, and NOSI is a research engine rather than a serving system; production integration is listed as future work.[2] The authors also state that the query-agnostic selection is sensitive to numerical precision, which is why the exp-delayed variant was needed.[2]

References

"NOSA: Native and Offloadable Sparse Attention" (arXiv abstract page, v2, 29 January 2026). https://arxiv.org/abs/2510.13602
Yuxiang Huang, Pengjie Wang, Jicheng Han, Weilin Zhao, Zhou Su, Ao Sun, Hongya Lyu, Hengyu Zhao, Yudong Wang, Chaojun Xiao, Xu Han, Zhiyuan Liu. "NOSA: Native and Offloadable Sparse Attention" (full paper, v2 PDF). https://arxiv.org/pdf/2510.13602v2
"NOSA: Native and Offloadable Sparse Attention" (arXiv v1, 15 October 2025). https://arxiv.org/abs/2510.13602v1
thunlp/NOSA, "The official implementation of NOSA" (GitHub repository, MIT license). https://github.com/thunlp/NOSA
openbmb/NOSA-8B (Hugging Face model card). https://huggingface.co/openbmb/NOSA-8B
openbmb/NOSA-3B (Hugging Face model card). https://huggingface.co/openbmb/NOSA-3B
openbmb/NOSA-1B (Hugging Face model card). https://huggingface.co/openbmb/NOSA-1B
Yuxiang Huang. "NOSA: 原生可卸载稀疏注意力" (author blog post, January 2026). https://huangyuxiang03.github.io/blogs_nosa
Yaosheng Fu, Guangxuan Xiao, Xin Dong, Song Han, Oreste Villa. "SparDA: Sparse Decoupled Attention for Efficient Long-Context LLM Inference" (arXiv 2606.04511, June 2026). https://arxiv.org/abs/2606.04511
NVlabs/SparDA (GitHub repository, "Related Projects" section). https://github.com/NVlabs/SparDA
Weilin Zhao et al. "InfLLM-V2: Dense-Sparse Switchable Attention for Seamless Short-to-Long Adaptation" (arXiv 2509.24663, September 2025). https://arxiv.org/abs/2509.24663
Jingze Shi et al. "Trainable Dynamic Mask Sparse Attention" (arXiv 2508.02124, August 2025). https://arxiv.org/abs/2508.02124
NOSA collection (Hugging Face, OpenBMB). https://huggingface.co/collections/openbmb/nosa
openbmb/NOSA-8B config.json (Hugging Face). https://huggingface.co/openbmb/NOSA-8B/blob/main/config.json

Editing: NOSA (Native and Offloadable Sparse Attention)