MiniMax Sparse Attention

TL;DR: MiniMax's recipe for million-token attention in a production model is deliberately boring in the best way: keep standard GQA, bolt on a tiny learned "index" head that lets each group of query heads pick its own top-16 blocks of 128 tokens, train that indexer with a detached KL loss against the real attention distribution, and co-design the GPU kernels so the sparsity actually shows up on the wall clock. On a 109B-parameter multimodal MoE trained from scratch on 3T tokens, MSA matches full attention on their eval suite while cutting per-token attention FLOPs 28.4× at 1M context, with measured 14.2× prefill and 7.6× decode speedups on H800s. The kernels are open-sourced and the method powers the released MiniMax-M3. The interesting content is less the headline numbers than the subtraction story — what they ablated away — and the kernel design that makes block-sparse attention tensor-core-friendly.

The problem, and why the obvious answers fall short

Softmax attention costs \Theta(N^2) in sequence length N, and agentic workloads are pushing N toward a million tokens. Everyone knows attention at that length is mostly wasted: for any given query, the attention mass concentrates on a small fraction of the context. The question is not whether to sparsify but how, and the design space has awkward trade-offs:

Fixed patterns (sliding windows, sinks, dilated blocks) are fast but content-blind — they can't retrieve the one relevant function definition 800K tokens back.
Inference-time pruning (H2O, Quest, MInference) leaves the model trained on full attention, so at least one phase of inference stays near dense speed, and the model never learned to live inside a sparse budget.
Natively trained sparse attention (NSA, MoBA, DeepSeek's DSA) trains a selector alongside the model. But NSA carries three parallel branches (compression + selection + sliding window); DSA does token-level selection with one index shared across all query heads on top of MLA; MoBA uses huge blocks scored by averaged keys and hopes the LM gradient trains the router.

MSA's pitch is Occam's razor applied to this last family: one branch, block-level selection, GQA-native, and a training recipe stripped to the components that ablations show are actually load-bearing.

The mechanism: a cheap librarian per group of heads

Start with standard GQA: H_q = 64 query heads sharing H_{kv} = 4 key-value heads, so G = 16 query heads per group. Partition the KV cache into blocks of B_k = 128 consecutive tokens.

The Index Branch adds exactly two projection matrices: one index query head per GQA group (H_{kv} of them) and a single index key head shared by everyone, both at small dimension d_{\rm idx}. For query token i and group r, it computes token-level dot-product scores S^{{\rm idx},(r)}_{i,j} against every visible index key, max-pools them within each block to get a block score, and takes the top-k blocks (k = 16). The block containing the query itself is always included. Crucially, each of the 4 groups gets its own top-16 — the selection is shared across the 16 heads within a group (so KV reads stay contiguous and shared) but independent across groups (so the model retains diverse retrieval patterns; their Appendix A visualizations show different groups tracing visibly different long-range stripes).

The Main Branch is then just exact softmax attention restricted to the selected blocks: at most kB_k = 2{,}048 key-value tokens per query per group, independent of context length.

Each GQA group (g1–g4) selects its own top-k blocks; selection is shared by the 16 query heads within a group, keeping KV reads contiguous. The local block is the only hard-coded selection; the sink block gets selected because the indexer learns to pick it.

FLOPs-wise (their Eq. 12): dense GQA costs 2H_q d_h N^2; MSA costs H_{kv} d_{\rm idx} N^2 for the indexer plus 4 H_q d_h N k B_k for the main branch. Note what this implies: MSA is not linear. The indexer is still quadratic — it's just ~30× cheaper per token pair (one tiny shared key head instead of the full attention). At 1M tokens the main branch has become negligible and the quadratic indexer dominates, which is why the reduction lands at 28.4× and will plateau near there rather than keep growing.

Training a non-differentiable selector without wrecking the backbone

Top-k selection kills gradients, so the LM loss can't train the index projections. The fix is a KL alignment loss: within the selected token support, make the indexer's softmax distribution match the group-averaged Main-Branch attention distribution (teacher detached). Three supporting mechanisms, each motivated by a documented failure mode:

Stop-gradient at the indexer input. If the KL gradient is allowed to flow back through the hidden states into the backbone, two bad things happen: large KL coefficients cause gradient spikes and LM-loss divergence within a few hundred steps, and even at stable coefficients short-context benchmarks gradually regress — the backbone discovers it can lower the KL loss by simplifying its own attention distribution to be easier to imitate, a self-distillation pathology. Detaching so that the KL loss updates only W_q^{\rm idx}, W_k^{\rm idx} eliminates both (Appendix B.3). This is the kind of mechanistic "why it works" detail that's worth remembering for any auxiliary-loss design.

Indexer warmup. Attention entropy collapses rapidly in the first few hundred steps of pretraining. If you enable sparse routing from step zero, a random indexer routes attention to garbage while chasing a fast-moving target. So for the first 40B tokens both branches run full attention and the indexer just learns to imitate; then sparse mode switches on. The same schedule converts a pretrained dense checkpoint.

Forced local block, and nothing else. Early versions forced selection of the first block (attention sink) and a local window. Ablations (Appendix C.2) showed the model learns both patterns on its own — the trained indexer assigns high probability to the sink block and local blocks without being told to. The final recipe forces only the query's own (incomplete) block. They also tried a GPT-OSS-style learnable sink logit and an index-branch value head that adds output to the layer (NSA-style); both were dropped after ablations showed the warmup makes them unnecessary. The paper is unusually explicit about this pruning process, and it's arguably its main contribution: a minimal recipe where every remaining component has an ablation justifying it.

The kernel: loop over KV blocks, not queries

Sparsity on paper routinely fails to become sparsity on the clock, and Section 4 is where MSA earns its speedups. Two ideas stand out.

Exp-free top-k. Softmax is order-preserving, so you can rank raw scores and skip max/exp/sum entirely before selection. Combined with a bespoke small-k kernel (per-lane min-heaps in shared memory, register-cached root, shuffle merge), their top-k runs 5.1× faster than torch.topk and 3.7× faster than TileLang's radix select at the deployed setting.

KV-outer iteration. The natural FlashAttention-style loop puts queries on the outside: each query gathers its k blocks. But then every query re-reads its KV blocks from HBM, and arithmetic intensity works out to roughly G = 16 FLOPs per byte — memory-bound. Flip the loop: iterate over KV blocks, and for each block gather all the queries that selected it. Now one KV load is amortized across every query that wants that block, and intensity rises to about \tfrac{2}{3}B_k \approx 85. Bonus: since all gathered queries share the same KV operands, you can pack \lceil 128/G \rceil = 8 query positions × 16 heads into a full 128×128 tensor-core MMA, instead of running MMAs with an M-dimension of 16.

Flipping the loop order is the key kernel insight: with KV blocks on the outer loop, one KV load serves every query that selected the block, raising arithmetic intensity ~5×.

The cost of KV-outer is that a query's k partial results come from k different CTAs, so softmax can't be normalized inline. They split the forward into an attention kernel writing locally-normalized partials plus per-partial log-sum-exps to an HBM buffer, and a combine kernel that merges them — with a scheduler kernel that pre-assigns buffer slots (no atomics) and splits "hot" blocks like the sink, which nearly every query selects, across many CTAs. It's split-K attention adapted to data-dependent, heavily skewed sparsity.

The evidence

The headline experiment: same 41-layer, 109B-total/6B-active MoE, same 3T-token multimodal budget, three runs — full-attention GQA, MSA from scratch (MSA-PT), and MSA converted from a 2.6T dense checkpoint with 400B tokens of continued pretraining (MSA-CPT). LM-loss curves for MSA-PT and full attention are reported as nearly indistinguishable across all 3T tokens, which is itself a notable stability result for from-scratch sparse training at this scale.

Parity holds broadly, and MSA-PT actually beats full attention on most video and several image benchmarks — VideoMME 45.5 vs 41.1, VisualWebBench 68.4 vs 55.6, EgoSchema 37.6 vs 29.6. Long visual token streams are exactly where a forced 2,048-token budget might act as useful inductive bias, but these are single runs and the gaps are large enough to want replication before drawing conclusions.

After a 140B-token long-context extension, MSA-CPT vs full attention at 128K: HELMET overall 45.93 vs 46.53 (−0.6, with Rerank/RAG at −2.1), RULER overall 72.12 vs 72.00. Remember each query still sees only 2,048 tokens per group — recovering 128K-scale retrieval through a 1.6% attention budget is the substantive quality result. Two smaller ablations matter for interpretation: a FLOP-matched sliding-window baseline has consistently worse agent-task perplexity, so the dynamic selection is doing real work; and the block-size ablation shows RULER-32K slipping from 72.5 to 64.6 as blocks grow from 32 to 128 tokens (in a short-training regime), hinting the coarse granularity isn't free.

Efficiency at 1M context: 28.4× attention-FLOPs reduction, 14.2× prefill and 7.6× decode wall-clock speedups on H800. The gap between 28.4× and 14.2× is the honest cost of indexing, top-k, reverse-index materialization, and gather overheads.

What to be skeptical about, and what changes

The biggest gap: quality is measured at 128K, speed at 1M. There are no retrieval or agentic-task numbers at the million-token lengths the speedups are advertised for, and HELMET's Rerank/RAG regression at 128K is exactly the kind of signal that could widen at 8× the length while the per-query budget stays fixed at 2,048 tokens. Second, there is no head-to-head with NSA, MoBA, or DSA — the paper positions MSA against them architecturally but only compares empirically against dense GQA and a sliding window. Whether per-group block selection beats DSA's shared token-level selection at matched cost is the live question this paper doesn't answer. Third, all quality numbers come from the authors' own eval suite on their own model family; "on par" summarizes a table where wins and losses are scattered.

That said, the deployability case is strong. Almost every current open-source frontier model uses GQA, and MSA is a strict add-on: two projection matrices, a KL loss, a warmup schedule, and open-sourced kernels — with a demonstrated near-lossless conversion path from an existing dense checkpoint in 400B tokens. Where DSA is entangled with MLA and NSA carries three branches, MSA composes with what people already have. If the recipe transfers as claimed, this is the kind of method that quietly becomes the default way GQA models get long contexts. It also plants a flag in a broader argument: that trainable sparse softmax attention, not linear/hybrid attention, is the practical path to million-token contexts — notable coming from MiniMax, who shipped hybrid linear attention in a previous generation.

If you read two sections, make them Section 4 (the KV-outer kernel design, useful well beyond this paper) and Appendix B (the failure modes — KL gradient leakage, early-training entropy collapse — that explain why the recipe is shaped the way it is).