MiniMax Opens MSA Sparse Attention with 2-branch Block Design

MSA Sparse Attention breaks the quadratic barrier with a two-branch architecture, keeping per-query cost fixed for long-context transformer models.

By Central
MiniMax MSA uses an Index Branch for block selection and a Main Branch for exact attention, achieving constant cost per query.
Highlights
  • MSA decomposes attention into an Index Branch and a Main Branch, with block-level selection limiting each query to 2,048 tokens.
  • The open-source MiniMax-M3 model uses 109 billion parameters and was trained on native multimodal data with support for FP4 precision.
  • Quality benchmarks show MSA remains competitive with full attention on MMLU, GSM8K, and long-context tests despite the fixed budget.

MiniMax has formally unveiled MSA (MiniMax Sparse Attention), a new attention mechanism engineered to break the quadratic complexity barrier that plagues long-context transformer models. Building directly on the foundation of Grouped Query Attention (GQA), MSA introduces a two-stage sparse architecture that drastically reduces the computational cost per query while maintaining output quality. The team validated the method within a 109-billion-parameter Mixture-of-Experts (MoE) model trained on native multimodal data, and they have simultaneously open-sourced both the production-grade inference kernel and the resulting model, MiniMax-M3.

What Is MiniMax Sparse Attention (MSA)?

MSA decomposes standard attention into two distinct branches: an Index Branch and a Main Branch. The Index Branch is responsible for determining which key-value blocks each query should attend to, while the Main Branch executes exact softmax attention strictly over those selected blocks. This selection operates at the block level rather than the token level, with a default block size Bk of 128 tokens. Each query and each GQA group retains k = 16 blocks, fixing the per-query budget at kBk = 2,048 key-value tokens — a constant that remains unchanged regardless of the full context length.

The cost distinction is fundamental: dense GQA attention scales as O(N) per query, where N is the full context length. MSA scales as O(kBk), a fixed value that grows linearly only with the budget, not the context. As context length increases, the compute gap between the two methods widens substantially. Importantly, selection is shared within each GQA group but independent across groups; one key-value head serves multiple query heads, and different groups can independently attend to different long-range regions of the sequence.

How the Two Branches Work

The Index Branch introduces only two additional projection matrices to a standard GQA layer. It defines one index query head per GQA group and a single shared index key head. It scores visible key tokens, then max-pools those scores at the block level. A Top-k operator then selects the highest-scoring blocks for each query and group. The local block containing the query is always included, preventing the selector from dropping the query’s immediate neighborhood — a critical guard for maintaining local coherence.

The Main Branch gathers causally visible tokens from the selected blocks and applies scaled dot-product softmax attention restricted exclusively to those tokens. Each query head retains its own query projection but shares the group’s block set. A visualization in the research report shows that the learned indexer concentrates attention density on the local diagonal and the first block, reserving the remainder of its budget for a small number of long-range stripes.

How MSA Is Trained

Top-k selection is inherently non-differentiable, so the language-modeling loss cannot directly train the index projections. MSA solves this through a KL alignment loss that matches the distribution produced by the Index Branch to the attention pattern of the Main Branch. The teacher is the group-averaged distribution over the selected tokens from the Main Branch. Three mechanisms stabilize this sparse training process:

  • Gradient Detach applies stop-gradient to the Index Branch input, confining the KL loss to the index projections themselves and preventing gradient spikes that would otherwise cause loss divergence with larger KL coefficients.
  • Indexer Warmup runs full attention in both branches for the first iterations, allowing the indexer to learn from the KL loss before it takes control of routing.
  • Forced Local Block reserves one block slot for nearby context, guaranteeing that the query always sees its immediate neighborhood.

Ablations informed the final recipe. An early variant that added an Index Branch value head with its own output proved unnecessary once warmup was used, and the final design dropped it for efficiency. MSA supports two training routes: MSA-PT (pre-training from scratch after a 40B-token indexer warmup) and MSA-CPT (conversion from a dense GQA checkpoint trained on 2.6 trillion tokens, followed by 400 billion tokens of continued training including 40 billion tokens of warmup).

The Kernel Co-Design

Theoretical sparsity does not translate into speed without a matching GPU path. MSA pairs its algorithm with two kernel innovations. The first is an exp-free Top-k selection: because softmax preserves order, ranking raw scores yields identical indices, and the kernel skips the max, exp, and sum steps entirely before selection. At 128K context with k=16, this kernel runs 5.1× faster than Torch’s topkcode and 3.7× faster than a TileLang radix-select kernel. The second is a KV-outer sparse attention with query gather: iterating over KV blocks rather than query positions raises arithmetic intensity, and the kernel packs ⌈128/G⌉ query positions into a single 128×128 score MMA. A two-phase forward pass splits attention and combination steps across CTAs for maximum throughput.

The open-source kernel, fmha_sm100code, targets NVIDIA SM100 GPUs. It ships dense FlashAttention alongside sparse Top-k kernels under an MIT license, supporting BF16, FP8, NVFP4, and FP4 precision.

How MSA Compares to Other Sparse Methods

The research team positions MSA against four native sparse designs in the literature: NSA (which uses compressed and selected blocks with a sliding window), InfLLM-V2 (a parameter-free block selection with sliding window), MoBA (which uses very large KV blocks with block-averaged keys), and DSA (token-level selection with a single Top-k shared across heads). MSA’s distinguishing feature is its per-GQA-group Top-k sharing combined with block-level selection — this keeps KV reads contiguous while giving each group its own independent retrieval. Quality benchmarks confirm that both MSA-PT and MSA-CPT remain broadly competitive with the Full-Attention baseline across MMLU, GSM8K, HumanEval, RULER, and VideoMME evaluations.

After long-context extension, MSA-CPT stays close to the dense baseline on HELMET-128K and RULER-128K, despite each query attending to only 2,048 key-value tokens.

Who Should Try This Now

Developers and research teams working with long-context models — particularly those operating at 128K tokens or beyond — should evaluate MSA as a drop-in replacement for dense GQA attention. The open-source kernel and model are available now, and the two-stage architecture requires no changes to the underlying transformer backbone. For teams building on NVIDIA SM100 hardware, the fmha_sm100code kernel provides a production-ready path to real-world speedups. The immediate next step is to clone the repository, replace your existing GQA attention with the MSA configuration, and run the indexer warmup on your own data — the 40B-token warmup regime is the only additional training overhead introduced.

Share This Article