MSA: Breaking the Million-Token Barrier with Co-Designed Sparse Attention

TL;DR

MiniMax has introduced MiniMax Sparse Attention (MSA), a system that allows Large Language Models to handle ultra-long contexts (1M+ tokens) with massive speedups. By replacing dense Grouped Query Attention (GQA) with a hardware-aligned block-sparse approach, MSA achieves a 14.2x prefill speedup on H800 GPUs while maintaining the accuracy of a 109B-parameter model.

Context: The Quadratic Wall

As AI agents and repository-scale code reasoning become standard, LLMs must process millions of tokens. However, the standard Transformer's Softmax Attention is governed by quadratic complexity: doubling the context quadruples the cost. This "Quadratic Wall" makes serving long-context models economically and technically prohibitive.

Methodology: The MSA Scalpel

MSA operates on a simple principle: Not all tokens are created equal. For any given query, only a small subset of the past context is actually relevant.

1. Dual-Branch Architecture

MSA splits attention into two stages:

• The Index Branch: A lightweight "scout" that scores and selects the Top-k blocks of Key-Value (KV) pairs for each GQA group. It uses single heads and max-pooling to keep overhead minimal.
• The Main Branch: The "heavy lifter" that performs standard attention—but only on the blocks the Index Branch selected.

2. The Training Secret: KL Alignment

How does the scout know what to pick? MSA uses a KL Alignment Loss. It takes the attention distribution from the Main Branch as a "teacher" and trains the Index Branch to mimic it. Critically, these gradients are detached from the rest of the model to prevent the auxiliary task from interfering with the primary language modeling objective.

3. Hardware-Software Co-Design

Theoretical FLOP reductions often fail to translate to real-world speed. MiniMax solved this by building custom GPU kernels:

• Exp-free Top-k Selection: Bypasses expensive softmax math during the indexing phase.
• KV-outer Iteration: Instead of checking every query against every KV (Q-outer), the kernel iterates over KV blocks and "gathers" relevant queries. This maximizes Tensor Core utilization and memory bandwidth efficiency.

Experimental Results: Near-Lossless Sparsity

The team tested MSA on a 109B Mixture-of-Experts (MoE) model trained on 3 Trillion tokens.

• Zero Degradation: In benchmarks like MMLU, GSM8K, and HumanEval, the sparse MSA model performed on par with the dense GQA version.
• Massive Speedups: At a 1M context length, MSA reduced attention compute by 28.4x. The wall-clock speedup for prefilling (the stage where the model "reads" the input) was 14.2x on H800 GPUs.

Critical Insight: The "Sink" and "Diagonal" Emergence

Interestingly, even without being forced, MSA's Index Branch naturally learned to focus on "Attention Sinks" (the first few tokens) and the local diagonal (immediate context). This confirms the inherent biases of successful LLMs and shows that MSA's learnable indexer is robust enough to discover these patterns autonomously.

Conclusion and Future Outlook

MiniMax Sparse Attention (MSA) proves that we don't need "dense" attention to achieve "smart" long-context models. By sharing index selections across GQA groups and optimizing for modern GPU architectures, MSA provides a production-ready pathway to million-token intelligence. For researchers and engineers, this represents a shift away from "more compute" toward "smarter retrieval" within the attention mechanism itself.

Explore the code at: GitHub - MiniMax-AI/MSA