MiniMax Sparse Attention Mechanism: An Efficient Inference Solution for Million-Scale Long Contexts

The Sparse Attention Mechanism (MSA) proposed by the MiniMax team addresses the quadratic complexity issue of traditional softmax attention. Based on block-level sparse design of Grouped Query Attention (GQA), it achieves a 28.4x reduction in computation for million-scale contexts on a 109B-parameter model. Combined with GPU kernel optimizations, it delivers 14.2x faster pre-filling and 7.6x faster decoding speeds while maintaining model performance comparable to the original GQA, providing a practical solution for deploying large models with ultra-long contexts.

Current application scenarios of large language models (agent workflows, codebase-level reasoning, persistent memory systems, etc.) have seen a surge in demand for context length. However, the complexity of traditional Transformer attention mechanisms is proportional to the square of the sequence length. When expanding from 4K to 1M, the computation increases by over 60,000 times, leading to resource and latency challenges for ultra-long context deployment.

MSA adopts a dual-branch block-level sparse design based on GQA:

1. Index Branch: A lightweight component scores KV blocks and independently selects the Top-k subset for each GQA group, enabling fine-grained attention allocation. It is compatible with the GQA architecture and aligns with GPU memory access patterns.
2. Main Branch: Performs precise sparse attention computation only on the selected blocks, focusing resources on relevant contexts and avoiding information loss. The collaboration between the two branches significantly reduces computational overhead while maintaining performance.

MSA improves practical deployment performance through algorithm-hardware co-optimization:

1. Exponent-free Top-k Selection: Uses efficient numerical methods to replace softmax exponentiation, reducing the overhead of the index branch.
2. KV Outer Product Sparse Attention: Leverages GPU tensor cores and organizes computations via KV outer products to better utilize hardware parallelism and improve throughput.

Test results on a 109B-parameter multimodal model:

• Computational Efficiency: Reduces per-token attention computation by 28.4x compared to standard GQA for million-scale contexts.
• End-to-End Speed: Achieves 14.2x faster pre-filling and 7.6x faster decoding on the H800 platform.
• Performance Preservation: Maintains performance roughly equivalent to the original GQA in multiple benchmark tests, without sacrificing output quality due to sparsification.

MSA marks the transition of long-context technology from the lab to production:

• Application Value: Enables existing hardware to deploy longer context capabilities (agent long-term memory, codebase understanding, multimodal long-sequence processing).
• Research Reference: Demonstrates the importance of algorithm-hardware co-design and provides end-to-end optimization ideas for efficient Transformer architectures.

The MiniMax team has open-sourced the MSA inference kernel (https://github.com/MiniMax-AI/MSA) to lower the threshold for community reproduction and adaptation. The MSA-based MiniMax-M3 model has been released on Hugging Face (https://huggingface.co/MiniMaxAI/MiniMax-M3), providing an off-the-shelf foundation for downstream applications.