{"id":98054,"date":"2026-06-17T18:06:02","date_gmt":"2026-06-17T18:06:02","guid":{"rendered":"https:\/\/youzum.net\/minimax-sparse-attention-msa-a-two-branch-block-sparse-attention-trained-on-a-109b-parameter-moe-with-a-3t-token-budget\/"},"modified":"2026-06-17T18:06:02","modified_gmt":"2026-06-17T18:06:02","slug":"minimax-sparse-attention-msa-a-two-branch-block-sparse-attention-trained-on-a-109b-parameter-moe-with-a-3t-token-budget","status":"publish","type":"post","link":"https:\/\/youzum.net\/zh\/minimax-sparse-attention-msa-a-two-branch-block-sparse-attention-trained-on-a-109b-parameter-moe-with-a-3t-token-budget\/","title":{"rendered":"MiniMax Sparse Attention (MSA): a Two-Branch Block-Sparse Attention Trained on a 109B-Parameter MoE With a 3T-Token Budget"},"content":{"rendered":"

MiniMax released MSA (MiniMax Sparse Attention), a sparse attention method built directly on Grouped Query Attention (GQA). It targets one bottleneck: the quadratic cost of softmax attention at long context. The MiniMax research team tested it inside a 109B-parameter Mixture-of-Experts model trained with native multimodal data. They also open-sourced an inference kernel and shipped a production model, MiniMax-M3.<\/p>\n

What is MSA (MiniMax Sparse Attention)<\/strong><\/h2>\n

MSA (MiniMax Sparse Attention) factors attention into two stages: an Index Branch and a Main Branch. The Index Branch decides which key-value blocks each query should read. The Main Branch then runs exact softmax attention over only those blocks.<\/p>\n

Selection happens at block granularity, not per token. The default block size is Bk<\/sub> = 128<\/code> tokens. Each query and GQA group keeps k = 16<\/code> blocks. That fixes the per-query budget at kBk<\/sub> = 2,048<\/code> key-value tokens.<\/p>\n

The two cost structures differ. Dense GQA attention scales per query as O(N), the full context. MSA scales as O(kBk<\/sub>), which stays fixed as N grows. The compute gap therefore widens as context length increases.<\/p>\n

Selection is shared inside each GQA group but independent across groups. One key-value head serves several query heads, and they share one block set. Different groups can attend to different long-range regions.<\/p>\n

How the Two Branches Work<\/strong><\/h2>\n

The Index Branch adds only two projection matrices to a standard GQA layer. It defines one index query head per GQA group and one shared index key head. It scores visible key tokens, then max-pools those scores to the block level.<\/p>\n

A Top-k operator then selects the highest-scoring blocks per query and group. The local block containing the query is always included. This prevents the selector from dropping the query\u2019s immediate neighborhood.<\/p>\n

The Main Branch gathers causally visible tokens from the selected blocks. It applies scaled dot-product softmax attention restricted to those tokens. Each query head keeps its own query projection but shares the group\u2019s block set.<\/p>\n

A visualization in the report shows what the learned indexer selects. Heads concentrate on the local diagonal and the first block. They reserve the rest of the budget for a few long-range stripes.<\/p>\n

\n
\"\"
https:\/\/arxiv.org\/pdf\/2606.13392v1<\/figcaption><\/figure>\n<\/div>\n
\n
\"\"
https:\/\/arxiv.org\/pdf\/2606.13392v1<\/figcaption><\/figure>\n<\/div>\n

How MSA is Trained<\/strong><\/h2>\n

Top-k selection is non-differentiable, so the language-modeling loss cannot train the index projections. MSA solves this with a KL alignment loss. The loss matches the Index Branch distribution to the Main Branch attention pattern. The teacher is the group-averaged Main Branch distribution over the selected tokens.<\/p>\n

Three mechanisms stabilize sparse training. Gradient Detach applies stop-gradient to the Index Branch input. This confines the KL loss to the index projections, not the backbone. Without it, larger KL coefficients caused gradient spikes and loss divergence.<\/p>\n

Indexer Warmup runs full attention in both branches for the first iterations. The indexer learns from the KL loss before it controls routing. The forced Local Block reserves one slot for nearby context.<\/p>\n

Ablations shaped the final recipe. An early variant added an Index Branch value head with its own output. Once warmup is used, that value head is no longer necessary. The final design drops it on efficiency grounds.<\/p>\n

MSA supports two training routes. MSA-PT trains from scratch after a 40B-token indexer warmup. MSA-CPT converts a dense GQA checkpoint trained on 2.6T tokens. It then continues for 400B tokens, including 40B tokens of warmup.<\/p>\n

The Kernel Co-Design<\/strong><\/h2>\n

Theoretical sparsity does not become speed without a matching GPU path. MSA pairs the algorithm with two kernel ideas.<\/p>\n

The first<\/strong> is exp-free Top-k selection. Softmax preserves order, so ranking raw scores yields identical indices. The kernel skips the max, exp, and sum steps before selection. At 128K context with k = 16<\/code>, it ran 5.1\u00d7 faster than torch.topk<\/code>. It also beat the TileLang radix-select kernel by 3.7\u00d7.<\/p>\n

The second <\/strong>is KV-outer sparse attention with query gather. Iterating over KV blocks raises arithmetic intensity versus iterating over queries. The kernel packs \u2308128\/G\u2309 query positions into one 128\u00d7128 score MMA. A two-phase forward splits the attention and combine steps across CTAs.<\/p>\n

The open-source kernel, fmha_sm100<\/code>, targets NVIDIA SM100 GPUs. It ships dense FlashAttention plus sparse Top-k kernels under an MIT license. It supports BF16, FP8, NVFP4, and FP4 precision.<\/p>\n

How MSA Compares To Other Sparse Methods<\/strong><\/h2>\n

The research team positions MSA against four natively trained sparse designs. <\/p>\n

The table below summarizes the differences it describes.<\/strong><\/h4>\n
\n\n\n\n\n\n\n\n\n\n
Method<\/th>\nBackbone<\/th>\nSelection granularity<\/th>\nIndexer \/ selection signal<\/th>\n<\/tr>\n<\/thead>\n
MSA<\/strong><\/td>\nGQA<\/td>\nBlock-level (B_k = 128<\/code>), per-GQA-group Top-k<\/td>\nKL alignment loss<\/td>\n<\/tr>\n
NSA<\/strong><\/td>\nMQA \/ MHA<\/td>\nCompressed + selected blocks + sliding window<\/td>\nNative (end-to-end) training<\/td>\n<\/tr>\n
InfLLM-V2<\/strong><\/td>\nDense\"\u2194\"sparse switchable<\/td>\nParameter-free block selection + sliding window<\/td>\nParameter-free (no trained indexer)<\/td>\n<\/tr>\n
MoBA<\/strong><\/td>\nGQA<\/td>\nVery large KV blocks (block-averaged keys)<\/td>\nLM gradient only<\/td>\n<\/tr>\n
DSA<\/strong><\/td>\nMLA (MQA mode)<\/td>\nToken-level; single Top-k shared across heads<\/td>\nReLU lightning indexer<\/td>\n<\/tr>\n<\/tbody>\n<\/table>\n<\/figure>\n

MSA\u2019s distinguishing pair is per-GQA-group Top-k sharing combined with block-level selection. This keeps KV reads contiguous while giving each group its own retrieval.<\/p>\n

The quality side holds up. Both sparse models stay broadly competitive with the Full-Attention baseline. <\/p>\n

The table below shows representative results under the 3T-token budget.<\/strong><\/h4>\n
\n\n\n\n\n\n\n\n\n\n\n
Benchmark<\/th>\nFull<\/th>\nMSA-PT<\/th>\nMSA-CPT<\/th>\n<\/tr>\n<\/thead>\n
MMLU<\/td>\n67.0<\/td>\n67.2<\/td>\n66.8<\/td>\n<\/tr>\n
GSM8K<\/td>\n76.2<\/td>\n77.7<\/td>\n73.7<\/td>\n<\/tr>\n
HumanEval<\/td>\n61.0<\/td>\n64.0<\/td>\n57.9<\/td>\n<\/tr>\n
RULER-8K<\/td>\n79.8<\/td>\n84.2<\/td>\n77.2<\/td>\n<\/tr>\n
RULER-32K<\/td>\n75.0<\/td>\n77.5<\/td>\n75.7<\/td>\n<\/tr>\n
VideoMME<\/td>\n41.11<\/td>\n45.48<\/td>\n39.65<\/td>\n<\/tr>\n<\/tbody>\n<\/table>\n<\/figure>\n

After long-context extension, MSA-CPT stayed close to Full on HELMET-128K and RULER-128K. Each query still attends to only 2,048 key-value tokens.<\/p>\n

Explainer Playground<\/strong><\/h2>\n

<\/p>\n

\n

Use Cases With Examples<\/strong><\/h2>\n<\/p>

MSA targets workloads where context length is the binding deployment constraint.<\/p>\n