Native Sparse Attention: making sparsity trainable rather than bolted on

Why sparse attention kept disappointing

DeepSeek posted the Native Sparse Attention paper on February 16, with Wenfeng Liang among the authors. The starting observation is one anyone who has tried to ship sparse attention will recognise. The theoretical savings are large, since softmax attention is empirically sparse and most of the key-value cache is irrelevant to any given query, and the realised savings have been small. The paper gives two reasons and we think both are right.

The first is that most prior methods were phase-restricted. H2O prunes the KV cache during decoding and does nothing for prefilling. MInference sparsifies prefilling and does nothing for decoding. Training is left at full cost in either case. The second is that many methods were designed around multi-head attention and fall apart on grouped-query or multi-query attention, where the KV cache is shared across heads. In Quest, each head selects its own blocks independently, so under GQA the memory actually loaded is the union of every head's selection in the group, and the sparsity you paid for at the algorithm level evaporates at the memory level.

There is a third problem the paper spends less time on but which we think matters most. Methods that pick tokens after training use non-differentiable operations like top-k on a model that never learned to be sparse. The model's attention patterns were shaped under full attention, and pruning them is a lossy approximation of a computation that was never meant to be pruned.

Three branches and a gate

NSA replaces the full key-value set with three views of it and lets a learned gate combine them. The first branch compresses. Keys and values are cut into blocks of length 32 with a stride of 16, and a small learnable MLP with intra-block position encoding maps each block to a single compressed key and value. This branch is cheap and gives every query a coarse picture of the whole context.

The second branch selects. The attention scores that the query produces against the compressed keys are reused as importance scores for selection blocks of 64 tokens, so the ranking costs almost nothing extra. The top 16 blocks are kept, including one fixed initial block and two local blocks, and the query attends to every token inside them at full resolution. For GQA, the importance scores are summed across the heads in a group before ranking, so all heads in a group select the same blocks and the kernel loads each block once.

The third branch is a sliding window of 512 recent tokens. The justification is a training dynamics point rather than a modelling one. Local patterns are learned fast and would otherwise dominate, starving the compression and selection branches of gradient. Giving locality its own branch with its own keys and values, and gating the three outputs, keeps the long-range branches learning. The authors describe this as preventing shortcut learning, and the extra parameter cost is marginal.

The kernel is the paper

The reason this method trains where others did not is the Triton kernel. Compression and sliding window attention already fit FlashAttention-2 style kernels. Selection does not, because if you load a temporally contiguous block of queries into SRAM the way FlashAttention does, those queries want disjoint KV blocks and the memory access pattern falls apart. NSA flips the grouping. For each query position it loads all the query heads in a GQA group at once, since they share the same selected blocks, then streams the selected KV blocks through SRAM in an inner loop.

Because the number of selected blocks is fixed, the inner loop length is the same for every query, which lets the outer loop sit in Triton's grid scheduler and balance work across streaming multiprocessors. The paper's claim is near-optimal arithmetic intensity from eliminating redundant KV transfers and balancing compute. The reported result on A100s for 64k sequences is a 9.0 times forward speedup, a 6.0 times backward speedup and an 11.6 times decoding speedup over full attention, with the ratio growing as sequences get longer.

What the 27B run showed

The test model is a 27B parameter, 3B active mixture-of-experts transformer with 30 layers, GQA with 4 groups over 64 heads, and a DeepSeekMoE layout of 72 routed and 2 shared experts. It was pretrained on 260B tokens, once with full attention and once with NSA, and the loss curves show NSA slightly lower throughout. On the general benchmarks NSA averaged 0.456 against 0.443 for full attention, with wins on BBH, GSM8K and DROP and small losses on MMLU and MBPP.

On long context the paper reports perfect needle-in-a-haystack retrieval across all depths at 64k. On LongBench, with all sparse baselines held to 2,560 activated tokens per query to match NSA's average, NSA scored 0.469 against 0.437 for full attention and 0.423 for the exact top-k oracle baseline, with the biggest gains on multi-hop QA. After distilling 10B tokens of 32k reasoning traces from DeepSeek-R1, the NSA variant scored 0.121 and 0.146 on AIME 24 at 8k and 16k generation limits against 0.046 and 0.092 for the full attention variant. Those AIME numbers are small in absolute terms and we would not read much into them, but the direction is the one you would want.

What we take from it

The claim we believe most is the structural one. Sparse attention becomes usable when the sparsity pattern is decided by a cheap differentiable proxy, the block selection is shared at the granularity the hardware caches at, and the model is trained under that pattern from the first step. Each of those three had been tried separately. Doing all three at once is what let the authors report a sparse model that is faster in every phase and not worse on the benchmarks.

The results are from one 27B model and one training run, which is the usual caveat and a real one. What we would want to see is the same recipe at a larger scale with a public checkpoint, and an ablation that removes the sliding window branch to test whether the shortcut-learning story holds. If it does, this is the shape that long-context attention in trained-from-scratch models is going to take, and the interesting question becomes how far the selection budget can be cut before the multi-hop results go away.