Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
Authors: Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, Yuxing Wei, Lean Wang, Zhiping Xiao, Yuqing Wang, Chong Ruan, Ming Zhang, Wenfeng Liang, Wangding Zeng
Accepted at: acl25

Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention

Jingyang Yuan1*, Huazuo Gao2, Damai Dai2, Junyu Luo1, Liang Zhao2, Zhengyan Zhang2, Zhenda Xie2, Y. X. Wei2, Lean Wang2, Zhiping Xiao3, Yuqing Wang2, Chong Ruan2, Ming Zhang1,†, Wenfeng Liang2,†, Wangding Zeng2,†
*Contribution during internship at DeepSeek-AI. †Corresponding authors.
1Key Laboratory for Multimedia Information Processing, School of Computer Science, Peking University, PKU-Anker LLM Lab. 2DeepSeek-AI. 3University of Washington.
Correspondence: yuanjy@pku.edu.cn, mzhang_cs@pku.edu.cn, wenfeng.liang@deepeek.com, zengwangding@deepeek.com

Abstract
Long-context modeling is crucial for next-generation language models, yet the high computational cost of standard attention mechanisms poses significant computational challenges. Sparse attention offers a promising direction for improving efficiency while maintaining model capabilities. We present NSA, a Natively trainable Sparse Attention mechanism that integrates algorithmic innovations with hardware-aligned optimizations to achieve efficient long-context modeling. NSA employs a dynamic hierarchical sparse strategy, combining coarse-grained token compression with fine-grained token selection to preserve both global context awareness and local precision. Our approach advances sparse attention design with two key innovations: (1) We achieve substantial speedups through arithmetic intensity-balanced algorithm design, with implementation optimizations for modern hardware. (2) We enable end-to-end training, reducing pretraining computation without sacrificing model performance. As shown in Figure 1, experiments show the model pretrained with NSA maintains or exceeds Full Attention models across general benchmarks, long-context tasks, and instruction-based reasoning. Meanwhile, NSA achieves substantial speedups over Full Attention on 64k-length sequences across decoding, forward propagation, and backward propagation, validating its efficiency throughout the model lifecycle.

1 Introduction
The research community increasingly recognizes long-context modeling as a crucial capability for next-generation large language models, driven by diverse real-world applications ranging from in-depth reasoning (DeepSeek-AI, 2025; Zelikman et al., 2022), repository-level code generation (Zhang et al., 2023a; Zhang et al.) and multi-turn autonomous agent systems (Park et al., 2023). Recent breakthroughs, including OpenAI’s o-series models, DeepSeek-R1 (DeepSeek-AI, 2025), and Gemini 1.5 Pro (Google et al., 2024), enabling models to process entire codebases, lengthy documents, maintain coherent multi-turn conversations over thousands of tokens, and perform complex reasoning across long-range dependencies. However, the high complexity (Zaheer et al., 2020) of Attention (Vaswani et al., 2017) mechanisms emerges as a critical latency bottleneck as sequence length increases. Theoretical estimates indicate that attention computation with softmax architectures accounts for 70–80% of total latency when decoding 64k-length contexts, underscoring the urgent need for more efficient attention mechanisms.

A natural approach to efficient long-context modeling is to take advantage of the inherent sparsity of softmax attention (Ge et al., 2023; Jiang et al., 2023), where selectively computing critical query-key pairs can significantly reduce computational overhead while preserving performance. Recent advances demonstrate this potential through diverse

![Comparison of performance and efficiency between Full Attention model and our NSA. Left: Despite being sparse, NSA surpasses Full Attention baseline on average across general benchmarks, long-context tasks, and reasoning evaluation. Right: For 64k-length sequence processing, NSA achieves substantial computational speedup compared to Full Attention in all stages: decoding, forward, and backward propagation.](page_1012_370_393_312.png)

Figure 1: Comparison of performance and efficiency between Full Attention model and our NSA. Left: Despite being sparse, NSA surpasses Full Attention baseline on average across general benchmarks, long-context tasks, and reasoning evaluation. Right: For 64k-length sequence processing, NSA achieves substantial computational speedup compared to Full Attention in all stages: decoding, forward, and backward propagation.
strategies: KV-cache eviction methods (Zhang et al., 2023b; Li et al., 2024; Zhou et al., 2024), blockwise KV-cache selection methods (Tang et al., 2024; Xiao et al., 2024a; Gao et al., 2024), and sampling, clustering or hashing-based selection methods (Chen et al., 2024b; Liu et al., 2024; Desai et al., 2024). More related sparse attention methods are introduced in Appendix A. Despite these promising strategies, existing sparse attention methods often fall short in practical deployments. Many approaches fail to achieve speedups comparable to their theoretical gains; moreover, most methods lack effective training-time support to fully exploit the sparsity patterns of attention.

To address these limitations, the deployment of effective sparse attention must tackle two key challenges: (1) Hardware-aligned inference speedup: Converting theoretical computation reductions into actual speed improvements requires hardware-friendly algorithm design during both prefilling and decoding stages to mitigate memory access and hardware scheduling bottlenecks; (2) Training-aware algorithm design: Enabling end-to-end computation with trainable operators to reduce training costs while maintaining model performance. These requirements are crucial for real-world applications to achieve fast long-context inference or training. When considering both aspects, existing methods still exhibit a noticeable gap.

To achieve more effective and efficient sparse attention, we present NSA, a Natively trainable Sparse Attention architecture that integrates hierarchical token modeling. As shown in Figure 2, NSA reduces per-query computation by organizing keys and values into temporal blocks and processing them through three attention paths: compressed coarse-grained tokens, selectively retained fine-grained tokens, and sliding windows for local contextual information. Then we implement specialized kernel to maximize its practical efficiency. NSA introduces two core innovations corresponding to the key requirements above: (1) Hardware-aligned system speedup: Optimize blockwise sparse attention for Tensor Core utilization and memory access, ensuring balanced arithmetic intensity. (2) Training-aware design: Enable stable end-to-end training through efficient algorithms and backward operators. This optimization enables NSA to support both efficient deployment and end-to-end training.

We evaluate NSA through comprehensive experiments on real-world language corpora. Pre-training on a 27B-parameter transformer backbone with 260B tokens, we assess NSA’s performance across general language evaluations, long-context evaluations, and chain-of-thought reasoning evaluation. We further compare the kernel speed on A100 GPUs with optimized Triton (Tillet et al., 2019) implementations. Experimental results demonstrate that NSA achieves comparable or superior performance to Full Fttention baseline, while outperforming existing sparse attention approaches. Additionally, NSA delivers substantial speedups across decoding, forward, and backward stages compared to vanilla Full Attention, with the speedup ratio increasing for longer sequences. These results validate that our hierarchical sparse attention design effectively balances model capability and computational efficiency.

2 Methodology

Our technical approach spans algorithm design and kernel optimization. In the following subsections, we first introduce the background of our methodology. Then we present the overall framework of NSA, followed by its key algorithmic components. Finally, we detail our hardware-optimized kernel design that maximizes practical efficiency.

2.1 Background

Attention Mechanism is widely used in language modeling where each query token \( q_t \) computes relevance scores against all preceding keys \( k_{<t} \) to generate a weighted sum of values \( v_{<t} \). Formally, for an input sequence of length \( t \), the attention operation is defined as:

\[
o_t = \text{Attn}(q_t, k_{<t}, v_{<t})
\]

where Attn denotes the attention function:

\[
\text{Attn}(q_t, k_{<t}, v_{<t}) = \frac{\sum_{i=1}^t \alpha_{t,i} v_i}{\sum_{j=1}^t \alpha_{t,j}}, \quad \alpha_{t,i} = e^{\frac{q_t \cdot k_i}{\sqrt{d_k}}}.
\]

Here, \( \alpha_{t,i} \) represents the attention weight between \( q_t \) and \( k_i \), and \( d_k \) is the feature dimension of keys. As sequence length increases, attention computation becomes increasingly dominant in the overall computational cost, presenting significant challenges for long-context processing.

Arithmetic Intensity is the ratio of compute operations to memory accesses. It intrinsically shapes algorithm optimization on hardware. Each GPU has a critical arithmetic intensity determined
Figure 2: Overview of NSA’s architecture. Left: The framework processes input sequences through three parallel attention branches: For a given query, preceding keys and values are processed into compressed attention for coarse-grained patterns, selected attention for important token blocks, and sliding attention for local context. Right: Visualization of different attention patterns produced by each branch. Green areas indicate regions where attention scores need to be computed, while white areas represent regions that can be skipped.

by its peak compute capability and memory bandwidth, calculated as the ratio of these two hardware limits. For computation tasks, arithmetic intensity above this critical threshold becomes compute-bound (limited by GPU FLOPS), while below it becomes memory-bound (limited by memory bandwidth).

Specifically for causal self-attention mechanism, during training and prefilling phases, batched matrix multiplications and attention computations exhibit high arithmetic intensity, making these stages compute-bound on modern accelerators. In contrast, auto-regressive decoding becomes memory-bandwidth constrained because it generates one token per forward pass while requiring loading the entire key-value cache, resulting in low arithmetic intensity. This leads to different optimization goals — reducing computation cost during training and prefilling, while reducing memory access during decoding.

2.2 Overall Framework

To leverage the potential of attention with natural sparse pattern, we propose replacing the original key-value pairs \( k_{t}, v_{t} \) in Equation 1 with a more compact and information-dense set of representation key-value pairs \( \tilde{K}_{t}, \tilde{V}_{t} \) given each query \( q_{t} \). Specifically, we formally define the optimized attention output as follows:

\[
\tilde{K}_{t} = f_{K}(q_{t}, k_{t}, v_{t}), \tilde{V}_{t} = f_{V}(q_{t}, k_{t}, v_{t}) \tag{3}
\]

\[
o_{t}^{*} = \text{Attn}\left(q_{t}, \tilde{K}_{t}, \tilde{V}_{t}\right) \tag{4}
\]

where \( \tilde{K}_{t}, \tilde{V}_{t} \) are dynamically constructed based on the current query \( q_{t} \) and the contextual memory \( k_{t}, v_{t} \). We can design various mapping strategies to get different categories of \( \tilde{K}_{t}^{c}, \tilde{V}_{t}^{c} \), and combine them as follows:

\[
o_{t}^{*} = \sum_{c \in \mathcal{C}} g_{t}^{c} \cdot \text{Attn}(q_{t}, \tilde{K}_{t}^{c}, \tilde{V}_{t}^{c}). \tag{5}
\]

As illustrated in Figure 2, NSA have three mapping strategies \( \mathcal{C} = \{\text{cmp, slc, win}\} \), representing compression, selection, and sliding window for keys and values. \( g_{t}^{c} \in [0, 1] \) is the gate score for corresponding strategy \( c \), derived from input features via an MLP and sigmoid activation. Let \( N_{t} \) denote the total number of remapped keys/values:

\[
N_{t} = \sum_{c \in \mathcal{C}} \text{size}[\tilde{K}_{t}^{c}].
\tag{6}
\]

We maintain a high sparsity by ensuring \( N_{t} \ll t \).

2.3 Algorithm Design

In this subsection, we introduce the design of our remapping strategies \( f_{K} \) and \( f_{V} \): token compression, token selection, and sliding window.

2.3.1 Token Compression

By aggregating sequential blocks of keys or values into block-level representations, we obtain compressed keys and values that capture the information of the entire block. Formally, the compressed key representation \( \tilde{K}_{t}^{\text{cmp}} \) is calculated by:

\[
f_{K}^{\text{cmp}}(k_{t}) = \left\{ \phi(k_{id+1:id+l}) \middle| 1 \leq i \leq \left\lfloor \frac{t-l}{d} \right\rfloor \right\}
\tag{7}
\]

where \( l \) is the block length, \( d \) is the sliding stride between adjacent blocks, and \( \phi \) is a learnable MLP with intra-block position encoding to map keys
in a block to a single compressed key. \( \tilde{K}_t^{\text{cmp}} \in \mathbb{R}^{d_k \times \lfloor \frac{l'}{d} \rfloor} \) is tensor composed by compression keys. Usually, we adopt \( d < l \) to mitigate information fragmentation. An analogous formulation holds for the compressed value representation \( \tilde{V}_t^{\text{cmp}} \). Compressed representations capture coarser-grained higher-level semantic information and reduce computational burden of attention.

2.3.2 Token Selection
Using only compressed keys, values might lose important fine-grained information, motivating us to selectively preserve individual keys, values. Below we describe our efficient token selection mechanism that identifies and preserves the most relevant tokens with low computational overhead.

Blockwise Selection. Our selection strategy processes key and value sequences in spacial continuous blocks, motivated by two key factors: hardware efficiency considerations and inherent distribution patterns of attention scores. Blockwise selection is crucial to achieve efficient computation on modern GPUs. That is because modern GPU architectures exhibit significantly higher throughput for continuous block accesses compared to random index-based reads. Also, blockwise computation enables optimal utilization of Tensor Cores. This architectural characteristic has established blockwise memory access and computation as a fundamental principle in high-performance attention implementations, as exemplified by FlashAttention’s block-based design. Blockwise selection follows the inherent distribution patterns of attention scores. Prior works (Jiang et al., 2024) have shown that attention scores often exhibit spatial continuity, suggesting that neighboring keys tend to share similar importance levels. Our visualization in Appendix F also shows this spatial continuous pattern.

To implement blockwise selection, we first divide key, value sequences into selection blocks. To identify the most important blocks for attention computation, we need to assign importance scores to each block. Below we present our method for computing these block-level importance scores.

Importance Score Computation. Computing block importance scores could introduce significant overhead. Fortunately, the attention computation of compression tokens produces intermediate attention scores that we can leverage to induce selection block importance scores, formulated as:

\[
\mathbf{p}_t^{\text{cmp}} = \text{Softmax}\left( \mathbf{q}_t^T \tilde{K}_t^{\text{cmp}} \right)
\] (8)

where \( \mathbf{p}_t^{\text{cmp}} \in \mathbb{R}^{\lfloor \frac{l'}{d} \rfloor} \) is the attention scores between \( q_t \) and compression keys \( \tilde{K}_t^{\text{cmp}} \). Let \( l' \) denote the selection block size. When compression blocks and selection blocks share the same blocking scheme, i.e., \( l' = l = d \), we can directly obtain the selection block importance scores \( \mathbf{p}_t^{\text{slc}} \) by \( \mathbf{p}_t^{\text{slc}} = \mathbf{p}_t^{\text{cmp}} \) straightforwardly. For cases where the blocking schemes differ, we derive the importance scores for selection blocks according to their spatial relationship. Given \( d \mid l \) and \( d \mid l' \), we have:

\[
\mathbf{p}_t^{\text{slc}}[j] = \sum_{m=0}^{\frac{l'}{d}-1} \sum_{n=0}^{\frac{l}{d}-1} \mathbf{p}_t^{\text{cmp}}\left[ \frac{l'}{d} j - m - n \right],
\] (9)

where [ ] denotes the indexing operator for accessing vector element. For models employing GQA or MQA where key-value caches are shared across query heads, consistent block selection across these heads has to be ensured to minimize KV cache loading during decoding. The shared importance scores across heads in a group are formally defined as:

\[
\mathbf{p}_t^{\text{slc}'(h)} = \sum_{h=1}^H \mathbf{p}_t^{\text{slc}(h)}
\] (10)

where (h) in the superscript denotes the head index, and \( H \) is the number of query heads in each group. This aggregation ensures consistent block selection across heads within the same group.

Top-n Block Selection. After obtaining the selection block importance scores, We retain tokens within the top-n sparse blocks ranked by block importance scores, formulated as:

\[
\mathcal{I}_t = \{ i \mid \text{rank}(\mathbf{p}_t^{\text{slc}'[i]}) \leq n \}
\] (11)

\[
\tilde{K}_t^{\text{slc}} = \text{Cat} \left[ \{ \mathbf{k}_{il'+1:(i+1)l'} | i \in \mathcal{I}_t \} \right],
\] (12)

where \( \text{rank}(\cdot) \) denotes the ranking position in descending order, with \( \text{rank} = 1 \) corresponding to the highest score, \( \mathcal{I}_t \) is the set of selected blocks’ indices, Cat denotes the concatenation operation. \( \tilde{K}_t^{\text{slc}} \in \mathbb{R}^{d_k \times n l'} \) is tensor composed by compression keys. An analogous formulation applies to the fine-grained value \( V_t^{\text{slc}} \). The selected keys and values then participate in the attention computation with \( q_t \) as defined in Equation 5.

2.3.3 Sliding Window
In attention mechanisms, local patterns typically adapt faster and can dominate the learning process, potentially preventing the model from effectively learning from compression and selection to-
kens. To address this issue, we introduce a dedicated sliding window branch that explicitly handles local context, allowing other branches (compression and selection) to focus on learning their respective features without being shortcutted by local patterns. Specifically, we maintain recent tokens \( K_t^{\text{win}} = k_{t-w:t}, V_t^{\text{win}} = v_{t-w:t} \) in a window \( w \), and isolate attention computations of different information sources (compression tokens, and selected tokens, sliding window) into separate branches. These branch outputs are then aggregated through a learned gating mechanism. To further prevent shortcut learning across attention branches with marginal computational overhead, we provide independent keys and values for three branches. This architectural design enables stable learning by preventing gradient interference between local and long-range pattern recognition, while introducing minimal overhead.

After obtaining all three categories of keys and values (\( \tilde{K}_t^{\text{cmp}}, \tilde{V}_t^{\text{cmp}}; \tilde{K}_t^{\text{sc}}, \tilde{V}_t^{\text{sc}}; \) and \( \tilde{K}_t^{\text{win}}, \tilde{V}_t^{\text{win}} \)), we compute the final attention output following Equation 5. Together with the compression, selection, and sliding window mechanisms described above, this forms the complete algorithmic framework of NSA.

2.4 Kernel Design

To achieve FlashAttention-level speedup during the training and prefilling, we implement hardware-aligned sparse attention kernels upon Triton. Given MHA is memory-intensive and inefficient for decoding, we focus on architectures with shared KV caches like GQA and MQA following the current state-of-the-art LLMs. While compression and sliding window attention computations are readily compatible with existing FlashAttention-2 kernels, we introduce the specialized kernel design for sparse selection attention. Figure 3 illustrates the forward pass implementation. If we were to follow FlashAttention’s strategy of loading temporally continuous query blocks into SRAM, it would result in inefficient memory access since queries within a block may require disjoint KV blocks. To address this, our key optimization lies in a different query grouping strategy: for each position on the query sequence, we load all query heads within a GQA group (they share the same sparse KV blocks) into SRAM. The proposed kernel architecture is characterized by the following key features:

1. Group-Centric Data Loading. For each inner loop, load all heads’ queries \( Q \in \mathbb{R}[h,d_k] \) in the group at position \( t \) and their shared sparse key/value block indices \( \mathcal{I}_t \).

2. Shared KV Fetching. In the inner loop, Sequentially load continuous key/value blocks indexed by \( \mathcal{I}_t \) into SRAM as \( K \in \mathbb{R}[B_k,d_k], V \in \mathbb{R}[B_k,d_v] \) to minimize memory loading, where \( B_k \) is the kernel block size satisfying \( B_k|l' \).

3. Outer Loop on Grid. Since the inner-loop length (proportional to the selected block count \( n \)) remains nearly identical for different query blocks, we put query/output loops in Triton’s grid scheduler to simplify and optimize the kernel.

This design achieves near-optimal arithmetic intensity by (1) eliminating redundant KV transfers through group-wise sharing, and (2) balancing compute workloads of GPU streaming multiprocessors.

![Kernel design for NSA. The kernel loads queries by GQA groups (Grid Loop), fetches corresponding sparse KV blocks (Inner Loop), and performs attention computation on SRAM. Green blocks indicate data on SRAM, while blue indicates data on HBM.](page_768_180_627_393.png)

Figure 3: Kernel design for NSA. The kernel loads queries by GQA groups (Grid Loop), fetches corresponding sparse KV blocks (Inner Loop), and performs attention computation on SRAM. Green blocks indicate data on SRAM, while blue indicates data on HBM.

3 Experiments

We evaluate NSA through three lenses: (1) general benchmarks performance, (2) long-context benchmarks performance, and (3) chain-of-thought reasoning performance, comparing against Full Attention baseline and state-of-the-art sparse attention methods. We defer the efficiency analysis of our sparse computation paradigm to Section 4, where
<table>
  <tr>
    <th rowspan="2">Model</th>
    <th colspan="2">MMLU</th>
    <th colspan="2">MMLU-PRO</th>
    <th colspan="2">CMMLU</th>
    <th colspan="2">BBH</th>
    <th colspan="2">GSM8K</th>
    <th colspan="2">MATH</th>
    <th colspan="2">DROP</th>
    <th colspan="2">MBPP</th>
    <th colspan="2">HumanEval</th>
    <th rowspan="2">Avg.</th>
  </tr>
  <tr>
    <th>Acc. 5-shot</th><th>Acc. 5-shot</th>
    <th>Acc. 5-shot</th><th>Acc. 5-shot</th>
    <th>Acc. 3-shot</th><th>Acc. 3-shot</th>
    <th>Acc. 8-shot</th><th>Acc. 4-shot</th>
    <th>Fi 1-shot</th><th>Pass@1 3-shot</th>
    <th>Pass@1 3-shot</th><th>Pass@1 0-shot</th>
    <th>Pass@1 0-shot</th><th>Pass@1 0-shot</th>
    <th>Pass@1 0-shot</th><th>Pass@1 0-shot</th>
  </tr>
  <tr>
    <td>Full Attn</td>
    <td>0.567</td><td>0.279</td>
    <td>0.576</td><td>0.497</td>
    <td>0.486</td><td>0.263</td>
    <td>0.503</td><td>0.482</td>
    <td></td><td>0.335</td>
    <td></td><td>0.443</td>
  </tr>
  <tr>
    <td>NSA</td>
    <td>0.565</td><td>0.286</td>
    <td>0.587</td><td>0.521</td>
    <td>0.520</td><td>0.264</td>
    <td>0.545</td><td>0.466</td>
    <td></td><td>0.348</td>
    <td></td><td>0.456</td>
  </tr>
</table>

Table 1: Pretraining performance comparison between the full attention baseline and NSA on general benchmarks, across knowledge (MMLU, MMLU-PRO, CMMLU), reasoning (BBH, GSM8K, MATH, DROP), and coding (MBPP, HumanEval) tasks. NSA achieves superior average performance on most benchmarks despite high sparsity.

<table>
  <tr>
    <th rowspan="2">Model</th>
    <th colspan="3">SQA</th>
    <th colspan="4">MQA</th>
    <th colspan="2">Synthetic</th>
    <th colspan="2">Code</th>
    <th rowspan="2">Avg.</th>
  </tr>
  <tr>
    <th>MFQA-en</th><th>MFQA-zh</th><th>Qasper</th>
    <th>HPQ</th><th>2Wiki</th><th>GovRpt</th><th>Dur</th>
    <th>PassR-en</th><th>PassR-zh</th>
    <th>LCC</th>
  </tr>
  <tr>
    <td>H2O</td>
    <td>0.428</td><td>0.429</td><td>0.308</td>
    <td>0.112</td><td>0.101</td><td>0.231</td><td>0.208</td>
    <td>0.704</td><td>0.421</td><td>0.092</td><td>0.303</td>
  </tr>
  <tr>
    <td>InfLLM</td>
    <td>0.474</td><td>0.517</td><td>0.356</td>
    <td>0.306</td><td>0.250</td><td>0.277</td><td>0.257</td>
    <td>0.766</td><td>0.486</td><td>0.143</td><td>0.383</td>
  </tr>
  <tr>
    <td>Quest</td>
    <td>0.495</td><td>0.561</td><td>0.365</td>
    <td>0.295</td><td>0.245</td><td>0.293</td><td>0.257</td>
    <td>0.792</td><td>0.478</td><td>0.135</td><td>0.392</td>
  </tr>
  <tr>
    <td>Exact-Top</td>
    <td>0.502</td><td>0.605</td><td>0.397</td>
    <td>0.321</td><td>0.288</td><td>0.316</td><td>0.291</td>
    <td>0.810</td><td>0.548</td><td>0.156</td><td>0.423</td>
  </tr>
  <tr>
    <td>Full Attn</td>
    <td>0.512</td><td>0.623</td><td>0.409</td>
    <td>0.350</td><td>0.305</td><td>0.324</td><td>0.294</td>
    <td>0.830</td><td>0.560</td><td>0.163</td><td>0.437</td>
  </tr>
  <tr>
    <td>NSA</td>
    <td>0.503</td><td>0.624</td><td>0.432</td>
    <td>0.437</td><td>0.356</td><td>0.307</td><td>0.341</td>
    <td>0.905</td><td>0.550</td><td>0.232</td><td>0.469</td>
  </tr>
</table>

Table 2: Performance comparison between our NSA and baselines on LongBench, including subsets in single document QA, multi-document QA, synthetic and code task categories. NSA outperformed most of the baselines including Full Attention.

we provides detailed discussions on training and inference speed.

3.1 Pretraining Setup
Following the common practice in state-of-the-art LLMs, our experiments adopt a backbone combining Grouped-Query Attention (GQA) and Mixture-of-Experts (MoE), featuring 27B total parameters with 3B active parameters. We chose this architecture as it achieves an effective trade-off between computation cost and model performance. More details about our framework are provided in Appendix C. For NSA, we set compression block size \( l = 32 \), sliding stride \( d = 16 \), selected block size \( l' = 64 \), selected block count \( n = 16 \) (including fixed activating the 1 initial block and 2 local blocks), and sliding window size \( w = 512 \). Both Full Attention and sparse attention models are pre-trained on 270B tokens of 8k-length texts, followed by continued training and supervised fine-tuning on 32k-length texts with YaRN (Peng et al., 2024) to achieve long-context adaptation. Both models are trained to full convergence to ensure fair comparison. We provide the loss curve in Appendix D.

3.2 Baselines Methods
In addition to comparing with Full Attention baseline, we evaluate several state-of-the-art inference-stage sparse attention methods: H2O (Zhang et al., 2023b), infLLM (Xiao et al., 2024a), Quest (Tang et al., 2024), and Exact-Top, which first computes full attention score and select the top-\( n \) scores keys corresponding to each query and then calculates attention on these positions. These methods span diverse sparse attention paradigms, including KV-cache eviction, query-aware selection, and exact top-\( n \) sparse selection.

For general evaluation, where most samples have lengths within the local context window of sparse attention baselines, these methods are effectively equivalent to Full Attention. Therefore, we present only the comparison results between NSA and Full Attention baseline in this setting. In the long-context evaluation, we conduct comparisons across all baseline methods, with the sparsity of all sparse attention methods set to the same to ensure a fair comparison. For chain-of-thought reasoning evaluation, which requires long-text supervised fine-tuning, we also limit our comparison to Full Attention, as most sparse attention baselines do not support training.

3.3 Performance Comparison
General Evaluation. We evaluated the pretrained NSA and Full Attention baseline, on a comprehensive suite of benchmarks spanning knowledge, reasoning, and coding capabilities, including MMLU (Hendrycks et al., 2020), MMLU-PRO (Wang et al., 2024), CMMLU (Li et al., 2023), BBH (Suzgun et al., 2022), GSM8K (Cobbe
Figure 4: Needle-in-a-Haystack retrieval accuracy across context positions with 64k context length. NSA achieves perfect accuracy through its hierarchical sparse attention design.

<table>
  <tr>
    <th>Generation Token Limit</th>
    <th>8192</th>
    <th>16384</th>
  </tr>
  <tr>
    <td>Full Attention-R</td>
    <td>0.046</td>
    <td>0.092</td>
  </tr>
  <tr>
    <td>NSA-R</td>
    <td><b>0.121</b></td>
    <td><b>0.146</b></td>
  </tr>
</table>

Table 3: AIME Instruction-based Evaluating after supervised fine-tuning. Our NSA-R demonstrates better performance than Full Attention-R at both 8k and 16k sequence lengths

<table>
  <tr>
    <th>Context Length</th>
    <th>8192</th>
    <th>16384</th>
    <th>65536</th>
  </tr>
  <tr>
    <td>Full Attention</td>
    <td>8192</td>
    <td>16384</td>
    <td>65536</td>
  </tr>
  <tr>
    <td>NSA</td>
    <td>2048</td>
    <td>2560</td>
    <td>5632</td>
  </tr>
  <tr>
    <td>Expected Speedup</td>
    <td><b>4×</b></td>
    <td><b>6.4×</b></td>
    <td><b>11.6×</b></td>
  </tr>
</table>

Table 4: Memory access volume (in equivalent number of tokens) per attention operation during decoding. Due to the low arithmetic intensity and memory-bound nature of decoding, the expected speedup is approximately linear with the volume of memory access.

et al., 2021), MATH (Hendrycks et al., 2020), DROP (Dua et al., 2019), MBPP (Austin et al., 2021), and HumanEval (Chen et al., 2021). The results are shown in Table 1. Despite its sparsity, NSA achieves superior overall performance, outperforming all baselines including Full Attention on 7 out of 9 metrics. This indicates that although NSA may not fully leverage its efficiency advantages on shorter sequences, it shows strong performance. Notably, NSA demonstrates significant gains in reasoning-related benchmarks (DROP: +0.042, GSM8K: +0.034), suggesting that our pretraining helps models to develop specialized attention mechanisms. This sparse attention pretraining mechanism force model to focus on the most important information, potentially enhancing performance by filtering out noise from irrelevant attention pathways. The consistent performance across diverse evaluations also validates NSA’s robustness as a general-purpose architecture.

Long-Context Evaluation. As shown in Figure 4, NSA achieves perfect retrieval accuracy across all positions in 64k-context needle-in-a-haystack (Kamradt, 2023) test. This performance stems from our hierarchical sparse attention design, which combines compression tokens for efficient global context scanning, and selection tokens for precise local information retrieval. The coarse-grained compression identifies relevant context blocks at low computational cost, while the token-level attention on selected token ensures the preservation of critical fine-grained information. This design enables NSA to maintain both global awareness and local precision.

We also evaluate NSA on LongBench (Bai et al., 2023) against state-of-the-art sparse attention methods and Full Attention baseline. To ensure consistent sparsity, we set the token activated by each query in all sparse attention baselines to 2560 tokens, which corresponds to the average number of tokens activated in NSA when handling 32k sequence lengths. Following StreamLLM (Xiao et al., 2023), this token budget includes the leading 128 tokens and 512 local tokens. We exclude certain subsets from LongBench due to their low scores across all models, which may not provide meaningful comparisons. As shown in Table 2, NSA achieves the highest average score 0.469, outperforming all baselines (+0.032 over Full Attention and +0.046 over Exact-Top). This improvement arises from two key innovations: (1) our native sparse attention design, which enables end-to-end optimization of sparse patterns during pretraining, facilitates synchronized adaptation between the sparse attention module and other model components; and (2) the hierarchical sparse attention mechanism achieves an balance between local and global information processing.

Notably, NSA demonstrates exceptional performance on tasks requiring complex reasoning over long contexts, achieving +0.087 and +0.051 improvements over Full Attention on multi-hop QA tasks (HPQ and 2Wiki), exceeding the performance of baselines on code understanding (LCC: +0.069), and outperforming other methods on passage retrieval (PassR-en: +0.075). These results validate NSA’s capability to handle diverse long-context challenges, with its natively pretrained sparse at-
Figure 5: Comparison of Triton-based NSA kernel with Triton-based FlashAttention-2 kernel. Our implementation significantly reduces latency across all context lengths, with the improvement becoming more pronounced as input length increases.

tention providing additional benefits in learning task-optimal patterns.

Chain-of-Thought Reasoning Evaluation. To evaluate NSA’s compatibility with advanced downstream training paradigms, we investigate its capacity to acquire chain-of-thought mathematical reasoning abilities via post-training. Given the limited effectiveness of reinforcement learning on smaller-scale models, we employ knowledge distillation from DeepSeek-R1, conducting supervised fine-tuning (SFT) with 10B tokens of 32k-length mathematical reasoning traces. This produces two comparable models: Full Attention-R (Full Attention baseline) and NSA-R (our sparse variant). We assess both models on the challenging American Invitational Mathematics Examination (AIME 24) benchmark. We use a sampling temperature of 0.7 and a top-p value of 0.95 to generate 16 responses for each question and obtain the average score. To validate the impact of reasoning depth, we conduct experiments with two generation context limits: 8k and 16k, measuring whether extended reasoning chains improve accuracy. Example comparisons of model predictions are provided in Appendix G.

As shown in Table 3, NSA-R achieves significantly higher accuracy than Full Attention-R under the 8k context setting (+0.075), with this advantage persisting at 16k contexts (+0.054). These results validate the pretrained sparse attention enable efficient capture of long-range logical dependencies critical for complex mathematical derivations.

4 Efficiency Analysis

We evaluate the computational efficiency of NSA against Full Attention on an 8-GPU A100 system. For both methods, we configure the model with GQA group \( g = 4 \), head per group \( h = 16 \), query/key dimension \( d_k = 192 \), and value dimension \( d_v = 128 \). The hyperparameters for NSA are kept consistent with Section 3.

4.1 Training Speed

We compare the Triton-based implementations of our NSA attention and Full Attention with Triton-based FlashAttention-2 to ensure fair speed comparison across the same backend. As shown in Figure 5, our NSA achieve progressively greater speedups as context length increases, up to 9.0× forward and 6.0× backward speedup at 64k context-length. Notably, the speed advantage becomes more pronounced with longer sequences. This speedup stems from our hardware-aligned algorithm design to maximize the efficiency of sparse attention architecture: (1) The Blockwise memory access pattern maximizes Tensor Core utilization, (2) The delicate loop scheduling in the kernel eliminates redundant KV transfers.

4.2 Decoding Speed

The decoding speed of Attention is primarily determined by the memory access bottleneck, which is closely tied to the amount of KV cache loading. In each decoding step, Our NSA just needs to load at most \( \left\lceil \frac{s - l}{d} \right\rceil \) compression tokens, \( n l' \) selected tokens, and \( w \) neighbor tokens, where \( s \) is the cached sequence length. As shown in Table 4, our method exhibits a significant reduction in latency as the decoding length increases, achieving up to 11.6× speedup at 64k context-length. This advantage in memory access efficiency also amplifies with longer sequences.

5 Conclusion

We present NSA, a hardware-aware sparse attention architecture for efficient long-context modeling. By integrating hierarchical token compression with block selection within a trainable architecture, we achieve faster training and inference while matching or surpassing Full Attention performance. Our approach demonstrates that hardware-aligned sparse attention can rival Full Attention implementations in efficiency and performance. NSA advances the state-of-the-art through: superior base model performance, exceptional long-context performance, and improved reasoning ability, all while achieving significant reductions in computational latency. We present NSA, a hardware-
aligned sparse attention architecture for efficient long-context modeling. By integrating hierarchical token compression with blockwise token selection within a trainable architecture, our architecture achieves accelerated training and inference while maintaining Full Attention performance. NSA advances the state-of-the-art by demonstrating general benchmark performance matches full-attention baselines, exceeding modeling capability in long-context evaluations, and enhanced reasoning ability, all accompanied by measurable reductions in computational latency.

Limitations

In this work, we focus on sparsity within attention maps of a single layer. Exploring cross-layer sparsity remains an intriguing avenue for future research. Also, our Triton-based implementation, while introducing some abstraction overhead compared to native CUDA kernels, offers opportunities for further hardware optimization.

Acknowledgments

The authors from Peking University (Jingyang Yuan, Junyu Luo, and Ming Zhang) were partially supported by grants from the National Key Research and Development Program of China with Grant No. 2023YFC3341203 and the National Natural Science Foundation of China (NSFC Grant Number 62276002).

References

Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. 2023. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13245.

Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, et al. 2021. Program synthesis with large language models. arXiv preprint arXiv:2108.07732.

Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. 2023. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308.14508.

Iz Beltagy, Matthew E Peters, and Arman Cohan. 2020. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150.

Guoxuan Chen, Han Shi, Jiawei Li, Yihang Gao, Xiaozhe Ren, Yimeng Chen, Xin Jiang, Zhenguo Li, Weiyang Liu, and Chao Huang. 2024a. Sepilm: Accelerate large language models by compressing one segment into one separator. arXiv preprint arXiv:2412.12094.

Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. 2021. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374.

Zhuoming Chen, Ranajoy Sadhukhan, Zihao Ye, Yang Zhou, Jianyu Zhang, Niklas Nolte, Yuandong Tian, Matthijs Douze, Leon Bottou, Zhihao Jia, et al. 2024b. Magicpig: Lsh sampling for efficient llm generation. arXiv preprint arXiv:2410.16179.

Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. 2021. Training verifiers to solve math word problems, 2021. URL https://arxiv.org/abs/2110.14168.

Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y Wu, et al. 2024. Deepseek-moe: Towards ultimate expert specialization in mixture-of-experts language models. arXiv preprint arXiv:2401.06066.

DeepSeek-AI. 2024. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model.

DeepSeek-AI. 2025. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning. Preprint, arXiv:2501.12948.

Aditya Desai, Shuo Yang, Alejandro Cuadron, Ana Klimovic, Matei Zaharia, Joseph E Gonzalez, and Ion Stoica. 2024. Hashattention: Semantic sparsity for faster inference. arXiv preprint arXiv:2412.14468.

Dheeru Dua, Yizhong Wang, Pradeep Dasigi, Gabriel Stanovsky, Sameer Singh, and Matt Gardner. 2019. Drop: A reading comprehension benchmark requiring discrete reasoning over paragraphs. arXiv preprint arXiv:1903.00161.

Tianyu Fu, Haofeng Huang, Xuefei Ning, Genghan Zhang, Boju Chen, Tianqi Wu, Hongyi Wang, Zixiao Huang, Shiyao Li, Shengen Yan, et al. 2024a. Moa: Mixture of sparse attention for automatic large language model compression. arXiv preprint arXiv:2406.14909.

Yu Fu, Zefan Cai, Abedelkadir Asi, Wayne Xiong, Yue Dong, and Wen Xiao. 2024b. Not all heads matter: A head-level kv cache compression method with integrated retrieval and reasoning. arXiv preprint arXiv:2410.19258.
Yizhao Gao, Zhichen Zeng, Dayou Du, Shijie Cao, Hayden Kwok-Hay So, Ting Cao, Fan Yang, and Mao Yang. 2024. Seerattention: Learning intrinsic sparse attention in your llms. arXiv preprint arXiv:2410.13276.

Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. 2023. Model tells you what to discard: Adaptive kv cache compression for llms. arXiv preprint arXiv:2310.01801.

Gemini Team Google, Petko Georgiev, Ying Ian Lei, Ryan Burnell, Libin Bai, Anmol Gulati, Garrett Tanzer, Damien Vincent, Zhufeng Pan, Shibo Wang, et al. 2024. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv:2403.05530.

Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. 2020. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300.

Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H Abdi, Dongsheng Li, Chin-Yew Lin, et al. 2024. Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention. arXiv preprint arXiv:2407.02490.

Huiqiang Jiang, Qianhui Wu, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. 2023. Llmlingua: Compressing prompts for accelerated inference of large language models. arXiv preprint arXiv:2310.05736.

G. Kamradt. 2023. LLMTest NeedleInAHaystack. GitHub repository. Accessed: [Insert Access Date Here].

Haonan Li, Yixuan Zhang, Fajri Koto, Yifei Yang, Hai Zhao, Yeyun Gong, Nan Duan, and Timothy Baldwin. 2023. Cmmlu: Measuring massive multitask language understanding in chinese. arXiv preprint arXiv:2306.09212.

Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. 2024. Snapkv: Llm knows what you are looking for before generation. arXiv preprint arXiv:2404.14469.

Guangda Liu, Chengwei Li, Jieru Zhao, Chenqi Zhang, and Minyi Guo. 2024. Clusterkv: Manipulating llm kv cache in semantic space for recallable compression. arXiv preprint arXiv:2412.03213.

Joon Sung Park, Joseph C. O’Brien, Carrie Jun Cai, Meredith Ringel Morris, Percy Liang, and Michael S. Bernstein. 2023. Generative agents: Interactive simulacra of human behavior. In Proceedings of the 36th Annual ACM Symposium on User Interface Software and Technology. UIST 2023, San Francisco, CA, USA, 29 October 2023– 1 November 2023, pages 2:1–2:22. ACM.

Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. 2024. Yarn: Efficient context window extension of large language models. In ICLR. OpenReview.net.

Noam Shazeer. 2019. Fast transformer decoding: One write-head is all you need. CoRR, abs/1911.02150.

Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoc V Le, Ed H Chi, Denny Zhou, et al. 2022. Challenging big-bench tasks and whether chain-of-thought can solve them. arXiv preprint arXiv:2210.09261.

Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. 2024. Quest: Query-aware sparsity for efficient long-context llm inference. arXiv preprint arXiv:2406.10774.

Philippe Tillet, Hsiang-Tsung Kung, and David Cox. 2019. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pages 10–19.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. Advances in Neural Information Processing Systems.

Yubo Wang, Xueguang Ma, Ge Zhang, Yuansheng Ni, Abhranil Chandra, Shiguang Guo, Weiming Ren, Aaran Arulraj, Xuan He, Ziyan Jiang, et al. 2024. Mmlu-pro: A more robust and challenging multi-task language understanding benchmark. arXiv preprint arXiv:2406.01574.

Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. 2022. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824–24837.

Wei Wu, Zhuoshi Pan, Chao Wang, Liyi Chen, Yunchu Bai, Kun Fu, Zheng Wang, and Hui Xiong. 2024. Tokenselect: Efficient long-context inference and length extrapolation for llms via dynamic token-level kv cache selection. arXiv preprint arXiv:2411.02886.

Chaojun Xiao, Pengle Zhang, Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, and Maosong Sun. 2024a. Inflilm: Training-free long-context extrapolation for llms with an efficient context memory. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.

Guangxuan Xiao, Jiaming Tang, Jingwei Zuo, Junxian Guo, Shang Yang, Haotian Tang, Yao Fu, and Song Han. 2024b. Duoattention: Efficient long-context llm inference with retrieval and streaming heads. arXiv preprint arXiv:2410.10819.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. 2023. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453.

Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. 2020. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283–17297.

Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah D. Goodman. 2022. Star: Bootstrapping reasoning with reasoning. In Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 – December 9, 2022.

Fengji Zhang, Bei Chen, Yue Zhang, Jacky Keung, Jin Liu, Daoguang Zan, Yi Mao, Jian-Guang Lou, and Weizhu Chen. 2023a. Repocoder: Repository-level code completion through iterative retrieval and generation. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, EMNLP 2023, Singapore, December 6–10, 2023, pages 2471–2484. Association for Computational Linguistics.

Kechi Zhang, Jia Li, Ge Li, Xianjie Shi, and Zhi Jin. Codeagent: Enhancing code generation with tool-integrated agent systems for real-world repo-level coding challenges. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2024, Bangkok, Thailand, August 11–16, 2024, pages 13643–13658.

Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuan-dong Tian, Christopher Ré, Clark Barrett, et al. 2023b. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36:34661–34710.

Junqi Zhao, Zhijin Fang, Shu Li, Shaohui Yang, and Shichao He. 2024. Buzz: Beehive-structured sparse kv cache with segmented heavy hitters for efficient llm inference. arXiv preprint arXiv:2410.23079.

Zihan Zhou, Chong Li, Xinyi Chen, Shuo Wang, Yu Chao, Zhili Li, Haoyu Wang, Rongqiao An, Qi Shi, Zhixing Tan, et al. 2024. Lim × mapreduce: Simplified long-sequence processing using large language models. arXiv preprint arXiv:2410.09342.

A Related Works

We review existing approaches that improve the efficiency of attention computation through sparse attention. These methods can be broadly categorized into three groups based on their core strategies: (1) fixed sparse pattern, (2) dynamic token pruning, and (3) query-aware selection. We introduce several representative works from each category.

A.1 Fixed Sparse Pattern

SlidingWindow is a commonly used approach that allows the query to compute attention only within a fixed window. StreamingLLM (Xiao et al., 2023) combines attention sinks with local windows to process continuous text streams. MoA (Fu et al., 2024a) and DuoAttention (Xiao et al., 2024b) also adopt similar local and sink information for long sequence modeling. Longformer (Beltagy et al., 2020) interleaves local windowed attention with global tokens to process long sequences. Compared with them, our NSA does not rely on pre-defined sparse patterns, but learns the patterns automatically, unlocking the potential to utilize full context.

A.2 Dynamic Token Pruning

Some methods are designed to reduce memory and computation costs during inference through dynamic KV-cache pruning. H2O (Zhang et al., 2023b), BUZZ (Zhao et al., 2024), and Se-pLLM (Chen et al., 2024a) implement adaptive approaches to reduce KV-cache memory usage during decoding. These methods dynamically evict tokens deemed less important for future predictions. FastGen (Ge et al., 2023) and HeadKV (Fu et al., 2024b) optimize computation by assigning different strategies to individual attention heads. SnapKV (Li et al., 2024) introduces a token pruning strategy that reduces the KV-cache by selectively retaining only the most crucial features, enabling efficient memory usage. Unlike these inference-focused approaches, our NSA incorporates sparsity natively during the training phase.

A.3 Query-Aware Selection

Other works focus on query-dependent token selection methods to reduce attention computation while preserving attention quality. Quest (Tang et al., 2024) employs a blockwise selection strategy where each chunk’s importance is estimated by product between query and coordinate-wise min-max of the key chunks. InfLLM (Xiao et al., 2024a) combines fixed patterns with retrieval by maintaining attention sinks, local context, and retrievable chunks. This method selects representative keys from each chunk to estimate chunk importance. HashAttention (Desai et al., 2024) formulates pivotal token identification as a recommendation prob-
lem by mapping queries and keys to Hamming space using learned functions. ClusterKV (Liu et al., 2024) achieves sparsity by firstly clustering keys and then selecting the most relevant clusters for attention computation based on query-cluster similarity. MInference (Jiang et al., 2024) and TokenSelect (Wu et al., 2024) select KV pairs for computation based on token-level importance scoring. SeerAttention (Gao et al., 2024) separates queries and keys into spatial blocks and performs block-wise selection for efficient computation. Compared to these methods, our NSA achieves hardware-aligned sparse attention computation throughout the full model lifecycle, including training, prefilling, and decoding.

B Rethinking Sparse Attention Methods

Modern sparse attention methods have made significant strides in reducing the theoretical computational complexity of transformer models. However, most approaches predominantly apply sparsity during inference while retaining a pretrained Full Attention backbone, potentially introducing architectural bias that limits their ability to fully exploit sparse attention’s advantages. Before introducing our native sparse architecture, we systematically analyze these limitations through two critical lenses.

B.1 The Illusion of Efficient Inference

Despite achieving sparsity in attention computation, many methods fail to achieve corresponding reductions in inference latency, primarily due to two challenges:

Phase-Restricted Sparsity. Methods such as H2O (Zhang et al., 2023b) apply sparsity during autoregressive decoding while requiring computationally intensive pre-processing (e.g. attention map calculation, index building) during prefilling. In contrast, approaches like MInference (Jiang et al., 2024) focus solely on prefilling sparsity. These methods fail to achieve acceleration across all inference stages, as at least one phase remains computational costs comparable to Full Attention. The phase specialization reduces the speedup ability of these methods in prefilling-dominated workloads like book summarization and code completion, or decoding-dominated workloads like long chain-of-thought (Wei et al., 2022) reasoning.

Incompatibility with Advanced Attention Architecture. Some sparse attention methods fail to adapt to modern decoding efficient architectures like Multiple-Query Attention (MQA) (Shazeer, 2019) and Grouped-Query Attention (GQA) (Ainslie et al., 2023), which significantly reduced the memory access bottleneck during decoding by sharing KV across multiple query heads. For instance, in approaches like Quest (Tang et al., 2024), each attention head independently selects its KV-cache subset. Although it demonstrates consistent computation sparsity and memory access sparsity in Multi-Head Attention (MHA) models, it presents a different scenario in models based-on architectures like GQA, where the memory access volume of KV-cache corresponds to the union of selections from all query heads within the same GQA group. This architectural characteristic means that while these methods can reduce computation operations, the required KV-cache memory access remains relatively high. This limitation forces a critical choice: while some sparse attention methods reduces computation, its scattered memory access pattern conflicts with efficient memory access design from advanced architectures.

These limitations arise because many existing sparse attention methods focus on KV-cache reduction or theoretical computation reduction, but struggle to achieve significant latency reduction in advanced frameworks or backends. This motivates us to develop algorithms that combine both advanced architectural and hardware-efficient implementation to fully leverage sparsity for improving model efficiency.

B.2 The Myth of Trainable Sparsity

Our pursuit of native trainable sparse attention is motivated by two key insights from analyzing inference-only approaches: (1) Performance Degradation: Applying sparsity post-hoc forces models to deviate from their prtrained optimization trajectory. As demonstrated by Chen et al. (2024b), top 20% attention can only cover 70% of the total attention scores, rendering structures like retrieval heads in prtrained models vulnerable to pruning during inference. (2) Training Efficiency Demands: Efficient handling of long-sequence training is crucial for modern LLM development. This includes both pretraining on longer documents to enhance model capacity, and subsequent adaptation phases such as long-context fine-tuning and reinforcement learning. However, existing sparse attention methods primarily target inference, leaving the computational challenges in training largely
unaddressed. This limitation hinders the development of more capable long-context models through efficient training. Additionally, efforts to adapt existing sparse attentions for training also expose challenges:

Non-Trainable Components. Discrete operations in methods like ClusterKV (Liu et al., 2024) (includes k-means clustering) and MagicPIG (Chen et al., 2024b) (includes SimHash-based selecting) create discontinuities in the computational graph. These non-trainable components prevent gradient flow through the token selection process, limiting the model’s ability to learn optimal sparse patterns.

Inefficient Back-propagation. Some theoretically trainable sparse attention methods suffer from practical training inefficiencies. Token-granular selection strategy used in approaches like HashAttention (Desai et al., 2024) leads to the need to load a large number of individual tokens from the KV cache during attention computation. This non-contiguous memory access prevents efficient adaptation of fast attention techniques like FlashAttention, which rely on contiguous memory access and blockwise computation to achieve high throughput. As a result, implementations are forced to fall back to low hardware utilization, significantly degrading training efficiency.

B.3 Native Sparsity as an Imperative
These limitations in inference efficiency and training viability motivate our fundamental redesign of sparse attention mechanisms. In the main paper, we proposed NSA, a natively sparse attention framework that addresses both computational efficiency and training requirements.

C Details of Backbone
We adopt a transformer framework based on GQA and MoE. The model consists of 30 layers with a hidden dimension of 2560. For GQA, we set the number of groups to 4, with a total of 64 attention heads. For each head, the query, key, and value dimensions are configured as \( d_q = d_k = 192 \) and \( d_v = 128 \), respectively. For MoE, we utilize the DeepSeekMoE (Dai et al., 2024; DeepSeek-AI, 2024) structure, with 72 routed experts and 2 shared experts, and set the top-k experts to 6. To ensure training stability, the MoE in the first layer is replaced by an MLP in the form of SwiGLU.

D Training loss curve
As shown in Figure 6, the pretraining loss curve of our NSA and Full Attention baseline demonstrates stable and smooth decline, with NSA consistently outperforming the Full Attention model. This indicates that our sparse model effectively maintains training stability while achieving superior performance.

![Pretraining loss comparison between Full Attention and our NSA on 27B-parameter model. Both models exhibit stable convergence, with NSA achieving marginally lower loss values.](page_186_120_1077_482.png)

Figure 6: Pretraining loss comparison between Full Attention and our NSA on 27B-parameter model. Both models exhibit stable convergence, with NSA achieving marginally lower loss values.
E Challenges with Alternative Token Selection Strategies

Before designing NSA, we explored adapting existing sparse attention methods to the training stage. However, these attempts encountered various challenges, prompting us to design a different sparse attention architecture:

Key-Clustering Based Strategies. We examined clustering-based strategies like ClusterKV (Liu et al., 2024). These methods store Keys and Values from the same cluster in contiguous memory regions. While theoretically feasible for training and inference, they face three significant challenges: (1) Non-trivial computational overhead introduced by dynamic clustering mechanisms; (2) Operator optimization difficulties exacerbated by inter-cluster imbalances, especially in Mixture-of-Experts (MoE) systems, where skewed Expert Parallelism (EP) group execution times lead to persistent load imbalances; (3) Implementation constraints arising from the need for mandatory periodic reclustering and chunk-sequential training protocols. These combined factors create substantial bottlenecks, significantly limiting their effectiveness for real-world deployment.

Other Blockwise Selection Strategies. We also considered blockwise key, value selection strategies different from NSA, such as Quest (Tang et al., 2024) and InfLLM (Xiao et al., 2024a). These methods rely on computing an importance score for each KV block and selecting the top-\( n \) blocks based on their similarity with \( q_i \). However, existing methods face two critical issues: (1) Since the selection operation is non-differentiable, importance score computation based on neural networks relies on auxiliary loss, which increases operator overhead and often degrades model performance; (2) Heuristic parameter-free importance score computation strategy suffer from low recall rates, leading to sub-optimal performance. We evaluate both approaches on a 3B-parameter model with similar architecture and compare their loss curve with NSA and Full Attention. For the auxiliary loss-based selection method, we introduce additional queries for each token and representative keys for each block to estimate the block importance scores. We compute block-level supervision signals by mean-pooling attention scores within each key block, and use KL divergence to supervise block importance prediction. We maintain individual query granularity instead of block-averaged queries to accommodate efficient decoding. This auxiliary loss-based importance estimation shares conceptual similarity with SeerAttention (Gao et al., 2024). For the heuristic parameter-free selection method, following the strategy of Quest, we implement direct selection using the product between queries and coordinate-wise min-max of the key chunks, without introducing additional parameters. We also explore a cold-start training approach where Full Attention is applied for the initial 1000 steps before transitioning to the heuristic blockwise selection. As shown in Figure 7, both methods exhibited inferior loss.

![Training Loss Curve comparing Full Attention, Auxiliary Loss-based Learnable Selection, Heuristic Parameter-free Selection, and NSA](page_820_98_563_312.png)

Figure 7: Compare training loss on a 3B-parameter model with Full Attention and different token selection strategies and . Our NSA achieves better performance.

![Visualization of Attention Map on a Full Attention transformer. Light-colored regions indicate higher attention values. As shown in the figure, attention scores exhibit blockwise clustering distribution.](page_820_420_563_312.png)

Figure 8: Visualization of Attention Map on a Full Attention transformer. Light-colored regions indicate higher attention values. As shown in the figure, attention scores exhibit blockwise clustering distribution.

F Visualization

To explore potential patterns in transformer attention distributions and seek inspiration for our design, we visualize the attention map from our pre-trained 27B Full Attention model in Figure 8. The
visualization reveals interesting patterns where attention scores tend to exhibit blockwise clustering characteristics, with nearby keys often showing similar attention scores. This observation inspired our design of NSA, suggesting that selecting key blocks based on spatial continuity might be a promising approach. The blockwise clustering phenomenon indicates that tokens adjacent in the sequence may share certain semantic relationships with query tokens, though the exact nature of these relationships requires further investigation. This observation motivated us to explore a sparse attention mechanism that operates on continuous token blocks rather than individual tokens, aiming to enhance computational efficiency and preserve high-attention pattern.
G  Example of AIME Results

Prompt:
“!User!Let x, y and z be positive real numbers that satisfy the following system of equations:

\[
\log_2 \left( \frac{x}{yz} \right) = \frac{1}{2}
\]
\[
\log_2 \left( \frac{y}{xz} \right) = \frac{1}{3}
\]
\[
\log_2 \left( \frac{z}{xy} \right) = \frac{1}{4}
\]

Then the value of \(|\log_2(x^4 y^3 z^2)|\) is \( \frac{m}{n} \) where m and n are relatively prime positive integers. Find \( m + n \). Please reason step by step, and put your final answer within □!Assistant!”

Our Result:

(think).../think> (2275 tokens)

Given the system of equations:
1. \( \log_2 \left( \frac{x}{yz} \right) = \frac{1}{2} \)
2. \( \log_2 \left( \frac{y}{xz} \right) = \frac{1}{3} \)
3. \( \log_2 \left( \frac{z}{xy} \right) = \frac{1}{4} \)

We convert these logarithmic equations to their exponential forms:
1. \( \frac{x}{yz} = \sqrt{2} \)
2. \( \frac{y}{xz} = 2^{1/3} \)
3. \( \frac{z}{xy} = 2^{1/4} \)

Expressing each variable in terms of x, y, and z:
1. \( x = \sqrt{2} \cdot yz \)
2. \( y = 2^{1/3} \cdot xz \)
3. \( z = 2^{1/4} \cdot xy \)

Substituting these expressions into each other, we find:
- From \( x = \sqrt{2} \cdot yz \) and \( y = 2^{1/3} \cdot xz \), we get \( x = \sqrt{2} \cdot (2^{1/3} \cdot xz) \cdot z \), leading to \( z = 2^{-5/12} \).
- Substituting z into the equation for y, we get \( y = 2^{-3/8} \).
- Substituting z and y into the equation for x, we get \( x = 2^{-7/24} \).

Next, we compute \( x^4 y^3 z^2 \):
- \( x^4 = (2^{-7/24})^4 = 2^{-7/6} \)
- \( y^3 = (2^{-3/8})^3 = 2^{-9/8} \)
- \( z^2 = (2^{-5/12})^2 = 2^{-5/6} \)

Multiplying these together:

\[
x^4 y^3 z^2 = 2^{-7/6} \cdot 2^{-9/8} \cdot 2^{-5/6} = 2^{-75/24} = 2^{-25/8}
\]

Taking the logarithm base 2:

\[
\log_2(x^4 y^3 z^2) = -25/8
\]

The absolute value is:

\[
|\log_2(x^4 y^3 z^2)| = 25/8
\]

Thus, \( m = 25 \) and \( n = 8 \), giving \( m + n = 33 \).
Baseline Result:

⟨think⟩...⟨/think⟩ (9392 tokens)

Given the system of logarithmic equations:

\[
\log_2 \left( \frac{x}{yz} \right) = \frac{1}{2}
\]
\[
\log_2 \left( \frac{y}{xz} \right) = \frac{1}{3}
\]
\[
\log_2 \left( \frac{z}{xy} \right) = \frac{1}{4}
\]

We convert these equations to exponential form:

\[
\sqrt{2} = \frac{x}{yz} \implies x = \sqrt{2} \cdot yz
\]
\[
2^{\frac{1}{3}} = \frac{y}{xz} \implies y = 2^{\frac{1}{3}} \cdot xz
\]
\[
2^{\frac{1}{4}} = \frac{z}{xy} \implies z = 2^{\frac{1}{4}} \cdot xy
\]

Substituting these into each other, we find:
1. \( x = \sqrt{2} \cdot yz \)
2. \( y = 2^{\frac{1}{3}} \cdot xz \)
3. \( z = 2^{\frac{1}{4}} \cdot xy \)

By solving these equations step-by-step, we find:
- \( z = 2^{-5/12} \cdot xy = 2^{-2/3} \cdot y = 2^{-1/24} \cdot x = 2^{1/24} \)

However, upon verifying the third equation, we realize the system might not have a solution. Therefore, we use matrix inversion to solve the system:

\[
\begin{cases}
\log_2(x) - \log_2(y) - \log_2(z) = \frac{1}{2} \\
\log_2(y) - \log_2(x) - \log_2(z) = \frac{1}{3} \\
\log_2(z) - \log_2(x) - \log_2(y) = \frac{1}{4}
\end{cases}
\]

Let \( A = \log_2(x), B = \log_2(y), C = \log_2(z) \). The system becomes:

\[
\begin{cases}
A - B - C = \frac{1}{2} \\
B - A - C = \frac{1}{3} \\
C - A - B = \frac{1}{4}
\end{cases}
\]

We solve this system using Cramer’s rule and matrix inversion. The determinant of the matrix is -2, and the solutions are:

\[
A = \frac{61}{48}, \quad B = \frac{13}{16}, \quad C = -\frac{1}{6}
\]

Thus, we have:

\[
x = 2^{61/48}, \quad y = 2^{13/16}, \quad z = 2^{-1/6}
\]
We need to find \( |\log_2(x^4 y^3 z^2)| \):

\[
4 \log_2 x + 3 \log_2 y + 2 \log_2 z = 4 \cdot \frac{61}{48} + 3 \cdot \frac{13}{16} + 2 \cdot \left( -\frac{1}{6} \right)
\]

Converting to a common denominator (48):

\[
4 \cdot \frac{61}{48} = \frac{244}{48}, \quad 3 \cdot \frac{13}{16} = \frac{117}{48}, \quad 2 \cdot \left( -\frac{1}{6} \right) = -\frac{16}{48}
\]

Summing these:

\[
\frac{244}{48} + \frac{117}{48} - \frac{16}{48} = \frac{345}{48} = \frac{115}{16}
\]

The fraction \( \frac{115}{16} \) is in simplest form, so \( m = 115 \) and \( n = 16 \). Therefore, \( m + n = 115 + 16 = 131 \).

Prompt:

“|User |Real numbers x and y with x, y > 1 satisfy \( \log_x(y^x) = \log_y(x^{4y}) = 10 \). What is the value of xy?
Please reason step by step, and put your final answer within □. |Assistant|”

Our Result:

〈think〉...〈/think〉 (15147 tokens)

Given the equations:

1. \( \log_x(y^x) = 10 \)
2. \( \log_y(x^{4y}) = 10 \)

We start by applying the logarithm power rule to both equations.

For the first equation:

\[
\log_x(y^x) = x \cdot \log_x(y) = 10
\]
Let \( z = \log_y(x) \). Then, \( \log_x(y) = \frac{1}{z} \). Substituting this into the equation, we get:
\[
x \cdot \frac{1}{z} = 10 \implies x = 10z
\]

For the second equation:

\[
\log_y(x^{4y}) = y \cdot \log_y(x^4) = 10
\]
Using the power rule, we get:

\[
y \cdot 4 \log_y(x) = 10 \implies 4y \log_y(x) = 10
\]

Since \( \log_y(x) = z \), this simplifies to:

\[
4yz = 10 \implies yz = \frac{5}{2}
\]

We now have two expressions:
1. \( x = 10z \)
2. \( y = \frac{5/2}{z} \)

Substituting \( x = 10z \) and \( y = \frac{5/2}{z} \) into the product \( xy \), we get:

\[
xy = (10z) \cdot \left( \frac{5/2}{z} \right) = 10z \cdot \frac{5}{2z} = 10 \cdot \frac{5}{2} = 25
\]

Thus, the value of \( xy \) is \( 25 \).

Baseline Result:

(think).../(think) (16223 tokens)

Given the problem where real numbers \( x \) and \( y \) with \( x, y > 1 \) satisfy the equations \( \log_x(y^x) = 10 \) and \( \log_y(x^4y) = 10 \), we need to find the value of \( xy \).

First, we convert the logarithmic equations to their exponential forms using the change of base formula.

1. \( \log_x(y^x) = 10 \)

\[
x \log_x(y) = 10
\]

Using the change of base formula, \( \log_x(y) = \frac{\ln y}{\ln x} \). Since \( \log_y x = \frac{9}{4} \), we have \( \log_x y = \frac{4}{9} \). Thus,

\[
x \cdot \frac{4}{9} = 10 \implies x = \frac{10 \cdot 9}{4} = \frac{90}{4} = \frac{45}{2}
\]

2. \( \log_y(x^4y) = 10 \)

\[
4 \log_y x + 1 = 10 \implies 4 \log_y x = 9 \implies \log_y x = \frac{9}{4}
\]

Again, using the change of base formula, \( \log_x y = \frac{4}{9} \).

Next, we express \( x \) and \( y \) in terms of each other:

\[
x = \frac{45}{2} \quad \text{and} \quad y = \left( \frac{45}{2} \right)^{4/9}
\]
We then compute the product \( xy \):

\[
xy = \left( \frac{45}{2} \right) \cdot \left( \frac{45}{2} \right)^{4/9} = \left( \frac{45}{2} \right)^{1+4/9} = \left( \frac{45}{2} \right)^{13/9}
\]

However, upon re-evaluating and using another approach by setting \( x = 2^a \) and \( y = 2^b \), we find:

\[
\log_x (y^x) = 10 \implies \frac{x \log_x y}{\log_x y} = 10 \implies x \cdot \frac{4}{9} = 10 \implies x = \frac{45}{2}
\]
\[
\log_y (x^4 y) = 10 \implies \frac{4 \log_y x + 1}{\log_y x} = 10 \implies 4 \cdot \frac{9}{4} + 1 = 10 \implies 9 + 1 = 10 \quad \text{consistent}
\]

Finally, we find that the product \( xy \) simplifies to:

\[
xy = \left( \frac{45}{2} \right)
\]

Reviewer paper summaries: This paper introduces the NSA sparse attention mechanism, which, unlike previous sparse attention mechanisms, supports both forward and backward propagation. NSA combines three sparse strategies: compression, top-k selection, and local windowing, all performed at the block level for compression, retrieval, and computation. By designing a Triton kernel similar to FlashAttention, NSA achieves a tenfold decoding speedup for sequences with a 64K context length and demonstrates strong performance across various tasks.

Aiming to the sparse attention for LLM training and inference efficiency improvement, this work adopts three distinct attention architectures with a gate mechanism to achieve long-context modeling. Additionally, hardware-aligned sparse attention kernels upon Triton are designed for training and prefilling optimization. Comprehensive experiments are conducted to demonstrate the effectiveness of the proposed method on both performance and efficiency.
Summary of strengths: - Practical Value: The biggest strength lies in its practical applicability. Sparse attention mechanisms that supports both forward and backward propagation is highly sought after in the LLM community. However, due to engineering challenges, most existing works fail to significantly outperform FlashAttention 2 in practice. NSA overcomes these engineering hurdles, making sparse attention a viable solution for large-scale real-world applications.

- Thorough Experiments: The authors conducted extensive experiments across multiple tasks (e.g., language modeling, long-context processing) and scenarios (e.g., training, inference) to validate NSA's effectiveness. The results demonstrate that NSA performs on par with full attention, which is particularly impressive.

- While the paper does not propose fundamental algorithmic innovations in sparse attention, its outstanding engineering implementation and comprehensive experimental validation more than compensate for this limitation.

1. The paper is well-organized and easy to read.
2. Comprehensive experimental results on various benchmarks against five baselines verify the effectiveness of both performance and efficiency.
3. An innovative method with three different attention mechanisms is proposed for sparse attention, speeding up the training, and inference processes. Moreover, a tailored kernel is designed for blockwise sparse attention optimization.
4. The authors further provide an in-depth analysis regarding the sparse attention solutions from the efficient inference and trainable sparsity two perspectives.
Summary of weaknesses: - The effective data density of currently popular long-text benchmarks is relatively low, with only a small fragment of the text containing useful information. This makes it difficult to verify the performance of sparse attention methods in scenarios with extremely high data density.

- NSA (Neural Sparse Attention) demonstrates superior performance only with extremely long texts. However, extending the context length of language models currently requires only a few hundred to a thousand steps to converge, meaning NSA does not bring significant practical benefits to long-text model training. For short texts, MLPs dominate the computational overhead (a point the authors seem not to have discussed), making it even less necessary to adopt NSA, a method that is not widely applied and offers limited benefits.



===========post rebuttal==================

Thanks for the responses. I believe this paper is above ACL acceptance.

1. Why are the three different attention architectures required for sparse attention? More detailed explanations and theoretical analysis are needed to demonstrate the rationality of these modules. Besides, how about each attention module affect the final results on both performance and efficiency?
2.  As this work utilizes a gate mechanism for attention module selection (Eq. 5), I am curious about the distribution of the output of the gate. Why a hard selection is adopted in an adaptive manner. Are these three attentions always required for long-context modeling? If so, which mechanism should used when facing different scenarios? How about the inductive bias of each attention mechanism? A more in-depth investigation would provide valuable insights into the effectiveness and underlying principles of the proposed sparse attention mechanism.
3. Also, the motivation, innovation, or intuitive rationale behind the gate mechanism requires further elaboration to enhance the clarity of this work.
4. How about the computation cost of your method against the baseline models?
5. How about the performance under various block size settings?
Comments, suggestions and typos: N/A

See weaknesses for details.
Overall assessment scores: 4

4