Sparsity is Cool: Reverse-Engineering MoBA and NSA

TL;DR

A Tilde Research mechanistic study of two natively-trainable sparse-attention variants — Moonshot AI’s Mixture of Block Attention (MoBA) and DeepSeek’s Native Sparse Attention (NSA) — pretrained at 340M/1.3B scale on H100 clusters with a new Triton NSA kernel. The headline findings are mechanistic: (a) MoBA introduces periodic transient attention sinks, one per block, to compensate for not always reading the first block — at a rate proportional to its block size; (b) NSA’s three branches (compression / selection / sliding-window) collapse into one or two active branches per head, with sliding-window dominating later layers and compression dying off after the first few layers; (c) under aggressive GQA/MQA the key cloud becomes isotropic and the query manifold “snakes around” it, which the authors hypothesize is why sparse attention especially benefits high-KV-sharing regimes. Motivated by the gating analysis, the authors prune unused NSA branches at inference and recover most of base NSA’s perplexity while boosting per-block throughput ~10× on later layers.

Method

The training side uses torchtitan + FSDP across four H100 nodes from Nebius AI; models are transformer++ (RoPE, SwiGLU, pre-LN, QKV bias, QK-Norm), trained at 340M/15B-tokens and 1.3B/100B-tokens with 8192-token sequence length and tokenizer/vocab matching LLaMA/Mistral. MoBA uses block size 512 / top-k 3 (81% sparsity); NSA uses block size 64 / top-k 16 / sliding-window 512 (matching sparsity), with mean-pooling compression instead of NSA’s MLP for parameter parity. They sweep MHA / GQA-4 / MQA across all three attention classes (except MHA NSA, which is intentionally hardware-misaligned).

The kernel side reimplements NSA in PyTorch + Triton + FlexAttention, since the original NSA paper provides no reference code and the FLA implementation has performance bottlenecks. The selection-branch backward kernel is the main architectural challenge: each query attends to a unique selected subset of keys, so the standard Flash-Attention-style key/value-stationary backward doesn’t apply. They evaluate (i) a single-pass query-stationary kernel with Triton atomics and (ii) a two-pass query-and-KV-stationary kernel that reconstructs the reverse key→query mapping; empirically the two-pass variant wins despite redundant compute, because it avoids atomic contention from key reuse [§ Method, NSA Implementation Details]. Their NSA implementation is released as tilde-research/nsa-impl.

The mechanistic analysis uses long-context attention-map visualizations at 8192 tokens on GovReport, t-SNE projections of query/key manifolds per head/layer, gating-distribution aggregates across 30 long PG-19 examples, and a per-token sink heuristic computed only from immediately-following queries (rather than the all-subsequent-keys average) to capture transient sinks.

Key claims

MoBA introduces more attention sinks than dense attention, and they are periodic and transient — one sink per block, at a rate proportional to MoBA’s block size — because MoBA does not always read the first block and needs a sink token within each block it might retrieve [§ Attention Sinks].
A MoBA variant that always selects the first 16 keys/values for attention (mimicking NSA’s “always include the first block” trick) alleviates periodic sinks but has negligible effect on downstream performance — the model has multiple ways to satisfy the sink requirement [§ Attention Sinks].
Sparse attention has much stronger length generalization than dense — NSA in particular sees nearly no degradation up to 64k tokens without YaRN; within an attention class, more KV heads helps generalization, but the gain from extra KV heads is much smaller for MoBA/NSA than for MHA [§ Evaluation].
At sequences ≤1.5k, MoBA and NSA degenerate to fully dense attention; the dense transformer occasionally wins at short contexts, which the authors attribute to learned key geometry, not the attention mechanism itself [§ Evaluation].
MoBA heads specialize into four operational types — Extended-SWA (retrieved blocks contiguous with the sliding window), Prefix-attention (retrieve early blocks for semantic anchoring), Dilated-SWA (checkerboard pattern, rare), and Mixture (incompressible, dense-like) [§ Demystifying Sparse Attention Patterns].
NSA’s routing gate collapses heavily across heads in later layers: sliding-window averages ~0.5 weight in later layers, while the compression branch becomes inactive after the early layers (with strange last-layer behavior where all three branches reactivate) [§ Not all branches are created equal].
Under aggressive GQA/MQA, the NSA key cloud becomes isotropic and head-specific queries “snake around” it; the authors hypothesize that sparse attention thrives under high KV-sharing pressure because it reduces query-side interference from unrelated parts of the cloud [§ Sparse Attention Key Geometry].
Selectively replacing later NSA layers with gated-sliding-window-only attention (motivated by the gating analysis) recovers most of base NSA’s long-context perplexity on CodeParrot/GovReport and boosts those blocks’ throughput by ~10× at long contexts; uniformly using sliding-window-only across all layers degrades performance [§ Not all branches are created equal].
The two-pass query-and-KV-stationary backward kernel for NSA selection is empirically faster than the single-pass atomic variant, because atomic contention from cross-query key reuse outweighs the cost of an extra pass [§ NSA Implementation Details].
The compression branch can use a parametric pooling MLP (as in the NSA paper) because its scores enter the gated combination and receive direct downstream gradient flow; Tilde uses mean-pooling for parameter parity [§ B. Motivating Sparse Attention].

Results

Their reimplementation reproduces the headline NSA finding — NSA outcompetes dense attention in the GQA-4 and MQA regimes on held-out SlimPajama PPL [§ Evaluation, SlimPajama PPL].
MoBA also beats the dense transformer when KV sharing is introduced (GQA/MQA) [§ Evaluation].
NSA’s reported FlashAttention-2 speedups (from the original DeepSeek paper) — 9× forward, 6× backward — are noted as motivation but are not the focus of the Tilde benchmarks, which instead profile per-component latency and throughput from 512 → 65k token contexts; selection-attention and the top-k op dominate latency [§ Hardware-Aware Implementation, Benchmarking Results].
Optimal sink-detection threshold uses a heuristic over the immediately following query window (not all subsequent keys), which the authors validate as necessary because the all-subsequent average erases early-transient sinks once they fall out of use [§ Attention Sinks].
Position-1 sinking behavior, sink ratio under different optimization regimes, and exact per-component latencies are described qualitatively in the post body but the underlying tables/figures are gated behind interactive visualizations that the static snapshot does not render — see the live blog for the quantitative receipts [§ Demystifying Sparse Attention Patterns, § Benchmarking Results].

Why it’s interesting

This is the third independent mechanistic look at attention sinks now living on the wiki, and it answers a question the other two don’t: do sinks survive when you change the attention mechanism itself? A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Training frames sinks as outlier-driven implicit rescaling in standard softmax + RMSNorm stacks and proposes Gated Attention / GatedNorm to remove the need for them; The Spike, the Sparse and the Sink: Anatomy of Massive Activations and Attention Sinks traces the layer-by-layer causal chain through SwiGLU step-up blocks. Tilde’s answer here is that MoBA creates more sinks, periodically, to compensate for its block-skipping behavior — which is consistent with both prior papers’ framing of sinks as a load-bearing functional pressure rather than an artifact you can engineer away by changing one component. The Gated Attention story would predict that adding learned gates to MoBA / NSA should similarly eliminate periodic sinks while preserving long-context PPL; no filed paper has run that experiment.

The branch-collapse-and-prune finding is the more immediately actionable result. It’s structurally similar to the Better MoE model inference with warp decode “MoE inference is bottlenecked by routing/expert decisions that collapse in practice” line, and to DynaMoE: Dynamic Token-Level Expert Activation with Layer-Wise Adaptive Capacity for Mixture-of-Experts Neural Networks‘s dynamic-capacity expert routing — sparse architectures empirically don’t use all their capacity uniformly, and inference-time pruning recovers performance while shifting the throughput Pareto frontier. For a Luma context where DiT attention blocks are increasingly the latency bottleneck, the methodology (gate-distribution analysis → branch-pruning → layer-replacement with cheaper attention) ports cleanly even if the specific architecture (NSA) does not. The GQA-amplifies-sparse-attention hypothesis is also a clean prediction for video-DiT design: the more shared-KV across heads, the more headroom sparse-attention recipes should have.

It’s worth noting the kernel companion artifact: Tilde’s open-source NSA implementation (tilde-research/nsa-impl) is currently one of the few production-quality NSA references outside DeepSeek’s internal stack, and the two-pass backward is a documented design point for anyone building selection-style sparse attention — a different IO discipline than FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling‘s asymmetric-hardware kernel pipelining but the same “kernel choices dominate the speedup envelope” lesson.