MiniMax Sparse Attention (MSA) Explained: Turning Sparse Attention Theory into Real GPU Speedups

> Long context is not the future—it is the present. Agent workflows, repository-scale code reasoning, and persistent memory all require models that can attend to millions of tokens simultaneously. But the quadratic cost of softmax attention makes this untenable at deployment scale. MiniMax's answer: don't reinvent attention—sparsify it, and use GPU kernels to reach over 90% of the theoretical speedup limit.

Key points

• Why long context is the bottleneck: attention FLOPs scale with N². At 1M-token context, attention costs are 62,500x higher than at 4K tokens, making real-time interaction infeasible even on H100/H800 clusters.
• Design philosophy: block-level sparsification on top of GQA rather than a new architecture—fully compatible with existing training frameworks, inference engines, and quantization tools, and supports near-lossless conversion from pretrained GQA checkpoints.

Architecture: two-stage block-sparse attention

• Index branch: one lightweight index query head per GQA group, a shared index key head, block-level scoring, and TopK selection of k=16 blocks (block size B_k=128, so each query attends to only 2,048 tokens). A local block containing the current position is always retained.
• Main branch: standard softmax attention computed only over the selected blocks—no approximation error inside the branch.
• Complexity: Main Branch cost drops from O(N²) to O(N) when k×B_k ≪ N, reducing per-token attention compute by 28.4x at 1M tokens (equivalent to full attention at ~35K context).

Training stability

TopK selection is discrete and non-differentiable, so the index projections cannot be trained directly by the language modeling loss. MSA uses:

• KL alignment loss: the index learns to match the consensus softmax distribution averaged across query heads in a GQA group, computed only over selected tokens.
• Gradient detaching: the teacher distribution and index branch inputs are stop-gradient; the KL loss updates only the index projections.
• Indexer warmup: the first 40B tokens run full attention while only training the indexer.
• Forced local block: guarantees each query at least attends to its 128 most recent tokens.

During 3T-token training, MSA's LM loss curve is nearly indistinguishable from the full-attention baseline, with stable gradient norms and high block/score recall after switching to sparse attention.

GPU kernel co-design: theory into wall-clock gains

• Exp-free TopK kernel: softmax is order-preserving, so raw scores go straight to TopK without max/exp/sum. A per-warp streaming design with a k-element min-heap in shared memory and shuffle merges outperforms torch.topk and TileLang radix-select, especially at the deployment setting (B=128, k=16).
• KV-outer iteration: arithmetic intensity is roughly (2/3)B_k vs. ≈2kB_k/(2kB_k/G+1) for Q-outer—decisively better at B_k=128, G=16. Queries are gathered via TMA and packed into full 128×128 MMA tiles to saturate Tensor Cores.
• Scheduling and load balancing: a persistent grid on (kv_block, kv_head) tiles, reverse sparse indexing, pre-scheduled query-dimension chunking to defuse hot "sink" tiles, atomic-free output buffers, and a two-stage forward (partial normalized outputs + combine kernel with global softmax weights).
• Training kernels: LSE fusion skips the KL forward pass; dynamic load balancing via global atomic counters handles variable-length, data-dependent sparsity.

Experimental results (109B MoE, 3T tokens)

• Model: 41-layer MoE, 109B total parameters, 6B activated per token; 64 query heads, 4 KV heads, head dim 128, RoPE dim 64.
• Two training modes: MSA-PT (from scratch, 40B warmup) and MSA-CPT (continued pretraining from a 2.6T full-attention checkpoint, 400B sparse tokens).
• Benchmarks: MSA matches or slightly exceeds full-attention GQA across text reasoning (MMLU, MMLU-Pro, BBH, GPQA), math (GSM8K, OlymMATH), code (HumanEval, BigCodeBench), multimodal image (MMMU, OCRBench, CharXiv), video (VideoMME, MLVU), and long context (RULER, HELMET).
• Speed on H800 at 1M context: 14.2x prefill speedup, 7.6x decoding speedup, 28.4x per-token attention compute reduction.

Comparison with related work

Compared with DeepSeek NSA, Quest, and Infini-attention, MSA's distinguishing features are a dedicated exp-free TopK kernel, KV-outer iteration with query concatenation, validated CPT conversion, native multimodal training at 109B, and open-sourced kernels.

Limitations and future directions

• Fixed block size (128) may miss cross-boundary fine-grained patterns; a fixed k=16 budget may be suboptimal per layer; multimodal index learning needs further study.
• Future work: adaptive per-layer k, hierarchical block+token sparsity, cross-modal indexing, and hybrid designs with linear attention.

Conclusion

MSA's deepest contribution is an engineering philosophy: in the long-context race, the simplest deployable solution wins—not the most complex one. It doesn't reinvent attention (incremental on GQA), doesn't trade accuracy for speed (algorithm-hardware co-design achieves both), and doesn't raise the adoption barrier (from-scratch and CPT training, full ecosystem compatibility). The result: a 109B MoE model trained on 3T tokens achieving 14.2x prefill and 7.6x decoding speedups at 1M context with capabilities on par with full attention. This is not a replacement for attention—it is attention's evolution.

Reference

Lai, X., Xu, W., Yang, Y., Chen, Q., Xu, Y., Zeng, L., Li, X., Sun, H., Zhu, H., Zhang, V., & Zhao, P. (2026). MiniMax Sparse Attention. *arXiv preprint arXiv:2606.13392*.

Open-source kernels: https://github.com/MiniMax-AI/MSA

A production-grade multimodal model based on MSA has been publicly released.