MiniMax Sparse Attention (MSA): Turning Theoretical Sparse Attention Gains into Real GPU Speedups

Forum topic · 小凯 · 2026-06-14

Summary

MiniMax Sparse Attention (MSA) is a two-stage block-sparse attention architecture built on top of Grouped Query Attention (GQA), designed to make million-token contexts practical. A lightweight learned indexer scores visible key tokens per GQA group, selects the top-k key blocks (k=16, block size 128), and the main branch runs exact softmax attention only over selected blocks, with the local block always retained. Training stability is preserved via a KL alignment loss with gradient detach, a 40B-token indexer warmup, and forced local blocks. On a 109B-parameter MoE model trained on 3T tokens, MSA matches full-attention baselines across text, math, code, multimodal, and long-context benchmarks while cutting per-token attention compute by 28.4x. Custom CUDA kernels—an exp-free TopK, KV-outer iteration with query concatenation, persistent grids, and two-phase forward—convert sparsity into 14.2x prefill and 7.6x decoding speedups on H800 GPUs at 1M context. MSA supports both from-scratch training and continued pretraining from full-attention checkpoints, with open-source kernels on GitHub.

Key points

MiniMax Sparse Attention (MSA) is MiniMax's approach to making million-token contexts deployable: instead of redesigning attention, it sparsifies standard softmax attention and optimizes GPU kernels to convert theoretical FLOP savings into real wall-clock speedups.

Why long context is the bottleneck

Attention FLOPs scale as O(N²): a 1M-token context requires ~62,500x more attention compute than a 4K context.
Prefill must process the entire context; decoding suffers linearly growing KV-cache memory.
Two solution paths exist: hybrid architectures (e.g., MiniMax-Text-01, Qwen3) that replace some layers, and sparse softmax (e.g., DeepSeek-V3, NSA) that sparsifies within attention. MSA takes the second path, aiming for maximal simplicity and compatibility.

Architecture: two-stage block-sparse attention

Built directly on GQA; only two extra projection matrices are added.
Index branch: one lightweight index query head per GQA group plus a shared index key head scores visible key tokens, aggregates into blocks, and selects top-k blocks (k=16, B_k=128 → ~2048 tokens per query). The local block containing the current position is always kept.
Main branch: exact softmax attention computed only over selected blocks — no approximation error inside selected regions.
With k × B_k ≪ N, main-branch complexity drops from O(N²) to O(N). At 1M context, per-token attention compute is reduced 28.4x.

Training stability

TopK is non-differentiable, so index projections are trained with a KL alignment loss: the teacher distribution averages each query head's softmax over selected tokens, with stop-gradient isolation so only index projections update.
Three stabilization mechanisms: gradient detach for index inputs, a 40B-token indexer warmup with full attention, and forced local blocks.
On 3T tokens, MSA-PT's LM loss is nearly indistinguishable from the full-attention baseline; gradient norms stay stable. CPT (continued pretraining from a 2.6T full-attention checkpoint) shows rapidly falling KL loss and high block/score recall.

GPU kernel co-design (the hard part)

Exp-free TopK: softmax is order-preserving, so TopK runs on raw scores. A warp-per-row streaming kernel with a k-element min-heap in shared memory, register-cached root, delayed writes, and shuffle merges beats torch.topk and TileLang radix-select, especially at the deployment setting (B=128, k=16).
KV-outer iteration: arithmetic intensity ≈ (2/3)·B_k ≫ G (85 vs 16 with B_k=128, G=16), so KV-outer with query gathering beats Q-outer for Tensor Core utilization.
Implementation details: persistent grid over (kv_block, kv_head) tiles, reverse sparse indexing of queries, TMA loads, pre-scheduled tile chunking to fan out hot "sink" tiles, no atomic updates via preallocated output slots, and a two-phase forward (partial attention → global softmax combine kernel).
Query concatenation packs ⌈128/G⌉ query positions and their G heads into 128×128 score MMAs to fill Tensor Cores.
Training kernels add LSE fusion (skip KL forward) and dynamic load balancing via atomic work claiming.

Experimental results (109B MoE, 3T tokens)

Model: 41 layers, 109B total / 6B active parameters, 64 query heads, 4 KV heads, head dim 128, RoPE dim 64.
MSA-PT and MSA-CPT match or slightly exceed full-attention GQA on MMLU, MMLU-Pro, BBH, GPQA, GSM8K, OlymMATH, HumanEval, BigCodeBench, MMMU, OCRBench, CharXiv, VideoMME, MLVU, RULER, and HELMET.
On H800 at 1M context: 14.2x prefill speedup, 7.6x decoding speedup, 28.4x per-token compute reduction.

Comparison with related work

Why accuracy holds

Limitations and future directions

Fixed 128-token block granularity may miss cross-boundary patterns; a fixed k=16 may be suboptimal across layers; multimodal indexing behavior needs more study.
Future work: adaptive k per layer/task, hierarchical (block + token) sparsity, cross-modal indexing, and hybrid designs with linear attention.

Conclusion

MSA's core lesson is engineering minimalism: don't reinvent attention (incremental change on GQA), don't trade accuracy for speed (algorithm-hardware co-design achieves both), and don't raise deployment barriers (compatible with from-scratch training and CPT). The 14.2x prefill speedup comes not from sparsity itself but from how sparsity is executed on hardware.

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.