learn.ifs.name

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

Yuan et al., 2025, Native Sparse Attention (Proceedings of the 63rd Annual Meeting of ACL (Volume 1: Long Papers))

Subscriber deck — free during early access

12 cards

1 Introduction

  1. Why is attention the bottleneck for long contexts, and why hadn’t sparse attention solved it?
    Attention cost grows with context length; when decoding 64k-token contexts it is estimated at 70–80% of latency. Earlier sparse methods often failed to turn theoretical savings into real speedups, and most were applied only at inference, without training support.
    • Yuan et al., 2025, Native Sparse Attention (Proceedings of the 63rd Annual Meeting of ACL (Volume 1: Long Papers)), §1, p. 23078
  2. What two requirements does NSA set for practical sparse attention?
    Hardware-aligned speedups in both prefilling and decoding (memory access and GPU scheduling must actually benefit), and training-aware design: sparse operators that can be trained end to end, so pretraining itself becomes cheaper and the model learns to use the sparsity.
    • Yuan et al., 2025, Native Sparse Attention (Proceedings of the 63rd Annual Meeting of ACL (Volume 1: Long Papers)), §1, p. 23079

2.1 Background

  1. Why does sparse attention need different optimizations for training and for decoding?
    Training and prefilling process many queries at once, so they are compute-bound and the goal is fewer FLOPs. Decoding produces one token per step but must load the whole KV cache, so it is memory-bound and the goal is to read less memory.
    • Yuan et al., 2025, Native Sparse Attention (Proceedings of the 63rd Annual Meeting of ACL (Volume 1: Long Papers)), §2.1, p. 23079

2.2 Overall Framework

  1. What is the overall idea of NSA?
    For each query, replace all previous keys and values by a much smaller set built from them in three branches (compressed, selected, sliding window). Each branch runs its own attention, and the outputs are combined with gates from a small MLP with sigmoid outputs.
    compressed× gselected× gwindow× g← earlier keysquery →output = gated sum of the three branches
    Every block is seen coarsely, a few blocks in full, recent tokens always (illustrative)
    • Yuan et al., 2025, Native Sparse Attention (Proceedings of the 63rd Annual Meeting of ACL (Volume 1: Long Papers)), §2.2, p. 23080

2.3.1 Token Compression

  1. What does NSA’s compression branch do?
    It slides over the keys (and values) in blocks of length \(l\) with stride \(d\) (32 and 16 in the experiments) and maps each block, with a learnable MLP and intra-block position encoding, to one compressed key or value. Attending to these gives cheap, coarse access to the whole context.
    • Yuan et al., 2025, Native Sparse Attention (Proceedings of the 63rd Annual Meeting of ACL (Volume 1: Long Papers)), §2.3.1, p. 23080

2.3.2 Token Selection

  1. How does NSA choose which blocks to attend to in full, and why blocks?
    It reuses the attention scores of the compressed keys as block importance scores (summed over the query heads of a GQA group so they share blocks) and keeps the top-\(n\) blocks (16 blocks of 64 tokens). Contiguous blocks suit GPUs and Tensor Cores, and attention scores tend to be spatially continuous anyway.
    • Yuan et al., 2025, Native Sparse Attention (Proceedings of the 63rd Annual Meeting of ACL (Volume 1: Long Papers)), §2.3.2, p. 23081

2.3.3 Sliding Window

  1. Why does NSA have a separate sliding-window branch?
    Local patterns are learned fast and can dominate, keeping the model from learning to use compressed and selected tokens. A dedicated window of recent tokens (512) handles local context, with separate keys and values per branch to prevent shortcut learning.
    • Yuan et al., 2025, Native Sparse Attention (Proceedings of the 63rd Annual Meeting of ACL (Volume 1: Long Papers)), §2.3.3, p. 23082

2.4 Kernel Design

  1. How does NSA’s kernel differ from FlashAttention’s?
    FlashAttention loads blocks of consecutive queries, but neighboring queries may select different KV blocks. NSA’s kernel instead loads, for one position, all query heads of a GQA group, which share the same selected blocks, then streams those contiguous KV blocks into SRAM, avoiding redundant KV transfers.
    • Yuan et al., 2025, Native Sparse Attention (Proceedings of the 63rd Annual Meeting of ACL (Volume 1: Long Papers)), §2.4, p. 23082

3.1 Pretraining Setup

  1. How was NSA evaluated against full attention?
    A 27B-parameter MoE model (3B active, with GQA) was pretrained on 270B tokens of 8k-length text with either full attention or NSA, then extended to 32k with YaRN, both trained to convergence, and compared on general, long-context and reasoning benchmarks.
    • Yuan et al., 2025, Native Sparse Attention (Proceedings of the 63rd Annual Meeting of ACL (Volume 1: Long Papers)), §3.1, p. 23083

3.3 Performance Comparison

  1. How did NSA’s quality compare with full attention?
    Despite its sparsity it was better on average: 7 of 9 general benchmarks (notably DROP and GSM8K), perfect needle-in-a-haystack retrieval at 64k, the best LongBench average (beating full attention and inference-time sparse methods), and higher AIME accuracy after distillation from DeepSeek-R1.
    • Yuan et al., 2025, Native Sparse Attention (Proceedings of the 63rd Annual Meeting of ACL (Volume 1: Long Papers)), §3.3, p. 23084

4 Efficiency Analysis

  1. What speedups did NSA achieve at 64k context?
    Compared with Triton FlashAttention-2 on A100s: up to 9.0× in the forward pass and 6.0× in the backward pass during training, and up to 11.6× in decoding, where speed tracks how much KV cache must be read. Speedups grow with context length.
    • Yuan et al., 2025, Native Sparse Attention (Proceedings of the 63rd Annual Meeting of ACL (Volume 1: Long Papers)), §4, p. 23085

Limitations

  1. What do NSA’s authors leave for future work?
    Sparsity across layers (NSA only exploits sparsity within one layer’s attention map) and a native CUDA implementation in place of the Triton kernels, which add some abstraction overhead.
    • Yuan et al., 2025, Native Sparse Attention (Proceedings of the 63rd Annual Meeting of ACL (Volume 1: Long Papers)), Limitations, p. 23086

Source

  • 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 and Wangding Zeng. 2025. Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 23078–23097. License: CC BY 4.0. Flashcards adapted from this work. doi:10.18653/v1/2025.acl-long.1126 https://aclanthology.org/2025.acl-long.1126/