All posts

Native Sparse Attention: Why Trainable Sparsity Beats Post-Hoc Pruning

7 min read llmattentiongpulong-contextinference

Sparse attention has a credibility problem. Papers routinely report 90%+ sparsity with "negligible quality loss," yet almost none of it ships in production serving stacks. DeepSeek's Native Sparse Attention (NSA, arXiv:2502.11089, ACL 2025 Best Paper) is interesting because it diagnoses why — and the answer is mostly about hardware and training, not about which tokens you keep.

The two failure modes of prior sparse attention

Failure mode 1: theoretical FLOPs ≠ wall-clock time. Methods like H2O or Quest drop tokens at inference, but the surviving tokens are scattered across the KV cache. Attention kernels (FlashAttention and friends) get their throughput from coalesced, block-contiguous memory access. Gathering scattered tokens destroys that: you save FLOPs but pay in random HBM reads, and on an A100/H100 attention is memory-bound during decode anyway. Many published methods are slower than dense FlashAttention at the sparsity levels where quality holds.

There's a subtler version of this with GQA. Modern models share one KV head across a group of query heads. If each query head selects its own token subset, the kernel must load the union of all selections for the group — with 8 heads per group selecting independently, your effective sparsity can quietly collapse toward dense.

Failure mode 2: post-hoc sparsity fights a dense-trained model. If you pretrain with full attention and prune at inference, the model never learned to route information through a sparse topology. Retrieval heads that softly spread attention across many tokens break when you hard-threshold them. The pruning decision is also non-differentiable, so you can't fine-tune your way out.

NSA's thesis: make the sparsity pattern trainable end-to-end and block-aligned by construction, so the model learns to concentrate information where the kernel can read it cheaply.

The architecture: three branches, one gate

For each query, NSA replaces the single attention over all history with a learned mixture of three parallel branches, each producing a full attention output:

out = g_cmp · Attn(q, KV_compressed)     # coarse global view
    + g_slc · Attn(q, KV_selected)       # fine-grained retrieval
    + g_win · Attn(q, KV_window)         # local context

The gates g_* are per-token sigmoid outputs of a small MLP on the query — the model decides, per position, how much to trust each view.

Branch 1 — compression. Sequential blocks of 32 tokens (stride 16, so overlapping) are squashed into single summary tokens by a learned MLP with intra-block position encodings. A 64k context becomes ~4k compressed tokens: cheap global awareness, but too lossy for exact retrieval on its own.

Branch 2 — selection. This is the core trick. Instead of running a separate scorer to decide which tokens matter, NSA reuses the attention scores from the compression branch. If the query attended strongly to compressed block i, the underlying raw blocks are probably important. Scores are accumulated into selection blocks of 64 tokens, and the top-16 blocks (always including the first block and the two most recent) are attended at full resolution. Selection is per-block, not per-token — so the kernel reads contiguous 64-token spans, which is exactly what tensor cores and HBM want.

Crucially, selection is shared across all query heads in a GQA group (scores are summed over the group before top-k). That kills the union-blowup problem: one block list per group, one coalesced load.

Branch 3 — sliding window. The last 512 tokens, attended densely. This branch exists for a training reason as much as an inference one: local patterns are easy gradients, and without an explicit local branch the compression and selection branches lazily specialize into local attention instead of learning long-range routing. Isolating "local" into its own branch forces the other two to earn their keep. Each branch gets its own K/V projections to prevent gradient shortcut leakage between them.

Why it trains: differentiability where it counts

The top-k block choice is still discrete — NSA doesn't relax it with Gumbel tricks. What makes it natively trainable is that the scores feeding the top-k come from the compression branch, which receives gradients through its own attention output. So the scorer improves as a side effect of the compression branch doing its job, and the hard selection rides along. It's the same move EAGLE-style speculative decoding made for drafting (covered in an earlier post here): stop bolting a frozen heuristic onto a trained model and let training shape the mechanism itself.

The payoff shows up in pretraining economics too: at 64k sequence length the backward pass is ~6× faster than FlashAttention, which matters because long-context training — not just serving — is where quadratic attention actually burns money.

The kernel: group-centric loading

NSA ships a Triton kernel whose inner loop is organized around GQA groups rather than individual heads:

  1. Load all query heads of one group for one position into SRAM at once.
  2. Load the group's shared sparse KV blocks (contiguous, thanks to block-aligned selection).
  3. Compute attention for the whole group against those blocks; inner loop over selected blocks only.

Because every head in the group shares the block list, KV blocks are fetched from HBM exactly once per group. Arithmetic intensity stays balanced — you're not issuing a gather per head — and the schedule eliminates redundant KV transfers entirely. This is the part most sparse attention papers skip, and it's why their speedups evaporate outside of microbenchmarks.

Numbers

Experiments use a 27B-parameter MoE backbone (~3B active) pretrained from scratch with NSA on 260B tokens, versus an identical model with full attention:

  • Quality: the NSA model matches or beats full attention on the majority of general benchmarks (MMLU, GSM8K, etc.) and outperforms it on long-context suites — perfect needle-in-a-haystack retrieval across all depths at 64k, and higher LongBench averages than both full attention and inference-time sparse baselines like Quest and H2O.
  • Speed at 64k context: roughly 9× forward, 6× backward, and 11.6× decoding versus FlashAttention-2, with the advantage growing with sequence length (decode is memory-bound, and NSA reads a near-constant number of KV tokens per step regardless of context length).
  • Reasoning: after identical long-CoT distillation, the NSA model scores higher on AIME than its full-attention twin — sparsity learned during pretraining doesn't cap the ceiling of chain-of-thought fine-tuning; it seems to help.

That last point is the one that should update priors. Sparsity isn't just "acceptable loss for speed" — a model trained to route information through compression + selection can generalize better on long-range tasks, plausibly because the bottleneck acts as a structural prior against attention noise.

What to take away

If you're building or serving long-context models, NSA's transferable lessons are:

  1. Sparsity must be block-shaped. Any token-granular scheme will lose to HBM physics. Pick a block size that matches your kernel's tile size and select at that granularity.
  2. Share selection across GQA groups. Per-head selection silently destroys sparsity at the memory system level.
  3. Derive importance from computation you already do. A separate scorer network adds latency and a training target; reusing compression attention scores costs nothing.
  4. If you control pretraining, train sparse. Post-hoc pruning is a compromise; native sparsity is a design. The gap between Quest-style inference hacks and NSA on long-context quality is the empirical proof.

The pattern here rhymes with a broader shift in systems-ML: the winning designs (FlashAttention, MLA, EAGLE-3, NSA) are co-designed across algorithm, training procedure, and memory hierarchy simultaneously. Attention sparsity spent years stuck as an algorithms-only problem. It took a hardware-first reframing to make it real.