Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
Make sparsity part of training, not a mask applied after: three attention branches, a learned gate, and kernels that read memory the way the hardware wants.
arXiv:2502.11089Where this sits
NSA is a research satellite rather than a model release: it uses the DeepSeekMoE substrate and the author group around the V3 effort, but it does not replace V3. Its value in this curriculum is as the trainable-sparse-attention line that the production DSA mechanism in V3.2-Exp later adapts to serving.
- Trains sparse attention end to end from scratch instead of applying sparsity only at inference, so selection and attention co-adapt.
- Combines three complementary paths - compressed blocks, selected blocks, sliding window - fused by learned gates, avoiding the global-awareness loss of pure local or pure selection patterns.
- Designs selection at block granularity and loads each block once per GQA group, converting theoretical sparsity into measured memory-bandwidth savings.
- Reports up to 9.0x forward, 6.0x backward, and 11.6x decoding speedup at 64K context in Triton on A100, with speedups growing with sequence length.
- Shows a 27B/3B model pretrained with NSA matches or exceeds a full-attention baseline on general, long-context, and chain-of-thought evaluations.
What it is
NSA is the DeepSeek-AI research paper that asks what sparse attention would look like if it were designed into training and into the GPU from the start. Each query reads three paths: compressed coarse blocks that summarize global context, top-n selected fine-grained blocks, and a sliding window for local detail, combined by learned scalar gates. All components are trainable, so there is no auxiliary loss and no train/inference mismatch. The paper pretrains a 27B-parameter GQA (grouped-query attention) plus MoE model (3B activated) on 270B tokens at 8K length with NSA throughout, then continues at 32K with YaRN, and reports parity or better against a full-attention baseline on general, long-context, and reasoning evaluations, with up to 9.0x forward, 6.0x backward, and 11.6x decoding speedups at 64K context in Triton kernels on A100 GPUs. DSA in DeepSeek-V3.2 is the production descendant of this line.
1Theory from first principles
Attention is dense by construction, sparse in practice
Softmax attention computes a weight for every query-key pair, but the learned weights are usually concentrated: a few keys carry most of the output, and many contribute almost nothing. In principle you could skip the near-zero ones. In practice three things stop you. First, you have to know which keys matter before you compute the expensive product, and computing a proxy is itself work. Second, skipping keys produces irregular memory access, and GPUs reward contiguous block reads. Third, if you impose a mask the model never trained under, quality drops, so you end up paying for a cheap-looking method that loses accuracy. NSA is a careful answer to all three.
Three paths, learned gates
NSA splits each query's attention into paths with different jobs. The compression path groups past tokens into blocks and projects each block into a single summary token, giving a cheap coarse view of the whole history. The selection path keeps the top n blocks of original tokens by importance, preserving fine detail where it matters. The sliding window path reads the most recent w tokens directly, so local continuity and recency never need to survive compression. A small learned scalar gate per path, computed per query position, mixes the three outputs. Because every branch is inside the differentiable graph, the model learns how much to rely on global summary, selected detail, and local context at each position.
Gated combination
o(t) = g_cmp(t) * Attn_cmp(t) + g_sel(t) * Attn_sel(t) + g_win(t) * Attn_win(t)
Attn_sel(t) = softmax over the selected blocks onlyThree attention computations of different granularity, each cheap or local in its own way, summed with per-query weights. Compression gives the model a summary it can always afford to read; selection gives it the ability to zoom into important distant detail; the window guarantees the last w tokens are never summarized away. The gates are learned, so no one has to hand-tune the mixture - the model discovers, position by position, which kind of context it needs.
What 'natively trainable' buys
Inference-time sparse methods treat the model as fixed and hope the mask does not hurt. NSA trains with the mask and the compression, so gradients flow into every part of the mechanism: the compression projection learns what to preserve, the selection scoring learns what to pick, and the gates learn the mixture. The paper's practical claims follow from this: no auxiliary losses to tune against the main objective, stable pretraining loss curves, and quality that matches or beats full attention rather than degrading against it. The model also gets training-time speedups, not just inference-time ones, which matters because pretraining is where the largest compute bill sits.
Hardware alignment: why blocks and groups
The paper's second contribution is making the sparsity hardware-shaped. Selection happens at block granularity: the top-n set indexes contiguous blocks of keys, which the kernel can load into SRAM with coalesced access, instead of chasing random token indices. Data loading is organized per GQA group: all query heads in a group share the same selected blocks, so the kernel loads each block once for the whole group and keeps arithmetic intensity high. The compression path uses fixed-size segments processed with the same blockwise pattern as FlashAttention-2. These choices are why the measured speedups grow with context length instead of flattening out: the algorithm saves arithmetic, and the kernels convert the saved arithmetic into saved memory traffic.
What the experiments show
The backbone is a 27B-parameter model with 3B activated parameters combining GQA and DeepSeekMoE (30 layers, 72 routed experts, 2 shared experts, top-6). Both NSA and the full-attention baseline are pretrained on 270B tokens of 8K text, then continued and fine-tuned at 32K with YaRN. Reported configuration: compression block 32 with stride 16, selected block size 64, 16 selected blocks including one initial and two local fixed blocks, sliding window 512. NSA matches or exceeds full attention on general benchmarks and long-context tasks, and shows up to 9.0x forward and 6.0x backward speedup at 64K in Triton on A100, with decoding at up to 11.6x as memory traffic falls from 65536 equivalent tokens to 5632 at that length.
def select_blocks(scores, block_size, n, fixed):
"""scores: importance per token, length L. fixed: block indices always kept
(e.g. the first block and the local blocks). Returns block indices to read."""
num_blocks = len(scores) // block_size
block_scores = []
for b in range(num_blocks):
start = b * block_size
window = scores[start:start + block_size]
# one summary score per contiguous block keeps access coalesced
block_scores.append(max(window) if window else 0.0)
ranked = sorted(range(num_blocks), key=lambda b: block_scores[b], reverse=True)
chosen = set(fixed)
for b in ranked:
if len(chosen) >= n:
break
chosen.add(b)
return sorted(chosen)
# The kernel then loads exactly these contiguous blocks into SRAM, once per GQA
# group, which is what turns a theoretical sparsity win into bandwidth savings.- Block granularity is the hardware-alignment trick: random token indices would break coalesced access and erase the speedup.
- The fixed local blocks guarantee recent context survives selection even when the scoring proxy is wrong.
- Importance is scored with max-pooling over the block, so a block containing one crucial token cannot be dropped.
The honest limits
The experiments run at 27B/3B scale on 270B tokens - small by frontier standards - and the speedups are measured with Triton kernels on A100 hardware; other hardware and hand-written kernels can move the numbers. The evaluation suite is the authors' selection. Most importantly, the production path diverged from the purest version of the idea: DSA keeps the top-k selection but replaces end-to-end trained selection with a KL-supervised lightning indexer and applies sparsity to an already trained dense backbone through continued training, because that is what serving constraints and existing checkpoints allowed. NSA is best read as the existence proof that native sparsity can be quality-neutral, and the design space it opened is what later releases mined.
2Inside the paper
Problem the paper sets out to solve
Sparse attention promises large savings and usually fails to deliver either the speedup or the quality, because existing methods are inference-time heuristics bolted onto dense-trained models and because irregular sparse access does not match GPU memory behavior. Can sparsity be designed into training and into the kernel at the same time?
Key idea to look for while reading
Three parallel attention paths with learned gates (Section 2), and the kernel design that makes block selection and GQA-group loading efficient (Section 3 and the appendix). Read the pretraining setup closely: the model trains with NSA throughout rather than upcycling a dense checkpoint, which is the 'native' in the title.
Evidence the paper musters
A 27B/3B GQA+MoE model pretrained on 270B tokens of 8K text and continued at 32K with YaRN, compared against an identically trained full-attention baseline: parity or better on general benchmarks, long-context evaluations, and chain-of-thought reasoning. Kernel measurements in Triton on A100 at 64K context: up to 9.0x forward, 6.0x backward, and 11.6x decoding speedup, with decode memory access falling from 65536 to 5632 equivalent tokens. Published at ACL 2025.
Limits the paper and its sequel admit
The scale is a research model, not a frontier one; speedups are backend- and hardware-specific; benchmarks are author-selected. The production adaptation in V3.2-Exp changed two important things - supervised indexer training instead of fully learned selection, and continued training instead of from-scratch native training - which is itself evidence about where pure native sparsity was hard to deploy.
Practitioner takeaway
If you are going to skip work, make skipping part of the model and part of the kernel. Train the selector with the attention, keep a cheap always-on path for global and local context so nothing catastrophic is missed, and choose access patterns (contiguous blocks, shared groups) that the hardware can actually exploit. Quality-neutral sparsity is possible; it is an architecture and systems problem at once.
3Implementation check
0 of 5 questions answered
1.What does 'natively trainable' mean for NSA?
2.Why does NSA keep a sliding window branch in addition to compression and selection?
3.Why does NSA select contiguous blocks rather than individual tokens?
4.NSA reports a larger speedup for decoding (11.6x) than for forward passes (9.0x) at 64K. Why?
5.What design difference separates DSA in V3.2 from the original NSA approach?
4Build it yourself
Three-branch sparse attention with learned gates
Intermediate3-5 hoursBuild NSA's three paths on scalar toy sequences where the answer lives in a hot block far from the query: a local window, a selected block, and a compressed global summary, mixed by gates trained with finite differences. Compare held-out MSE and tokens read against dense and window-only attention.
Milestones
- Generate sequences where the last token queries a hot block four positions into the sequence, so local attention alone cannot see it.
- Reproduce the provided dense and window baselines and record MSE and reads per query.
- Implement branches: local read of the last W values, block selection by highest max query-value score, and a global read of block means.
- Mix the three reads with softmax gates and train the three gate logits by finite differences on the train split.
- Report held-out MSE, reads, and the learned gate vector, and compare with dense at 15 reads and window at 4.
- Repeat the training with a second seed and check the gate ordering is stable.
Success criteria
- Gated held-out MSE is within 1.2x of dense attention while reading at most 12 of the 15 available tokens.
- The learned gates put most of the mixture weight on the selected-block path, and the gate vector is stable across runs with the same seed.
- The script exits 0, prints all three MSE rows, and finishes in under 10 seconds.
Stretch goals
- Read the top-2 selected blocks and check whether MSE improves enough to justify the extra reads.
- Replace finite differences with an analytic gradient through the gate softmax.
- Score blocks with max pooling over their values, as the paper does, and compare with mean pooling.