The attention mechanism is the backbone of every transformer model, but it carries a brutal cost: quadratic complexity with respect to sequence length. As LLMs push from 128K to 1M token contexts for agentic workflows, repository-scale code reasoning, and persistent memory, that quadratic scaling becomes the single biggest bottleneck in both training and inference. A recent paper from MiniMax introduces a practical approach to this problem that manages to be both remarkably simple and dramatically effective.
The Long-Context Bottleneck
Modern LLMs are shifting from short, single-turn interactions to long-horizon agentic workflows. Writing and deploying production code, navigating the open web, orchestrating diverse tools, and producing structured documents all require the model to jointly attend over hundreds of thousands to millions of tokens. The dominant approaches to handling this explosion have fallen into two camps.
The first replaces softmax attention entirely with cheaper alternatives like linear attention or state-space models (e.g., Mamba). Hybrid architectures interleave these efficient layers with a few full-attention layers, preserving some exact-softmax capacity while reducing the overall quadratic cost. The second camp keeps softmax attention but restricts its receptive field through fixed patterns like sliding windows and attention sinks, or through adaptive methods that learn which tokens to attend to.
Both approaches involve trade-offs. Replacing attention can lose the expressive power that makes transformers effective, while fixed sparse patterns are content-agnostic and can miss critical information. What’s needed is something that preserves the quality of full softmax attention while achieving real wall-clock speedups at production scale.
MiniMax Sparse Attention: The Architecture
MiniMax Sparse Attention (MSA) builds on Grouped Query Attention (GQA), the standard attention mechanism used in most current frontier models. The core idea is a two-branch design: a lightweight Index Branch that quickly decides which parts of the context matter, and a Main Branch that performs full softmax attention but only over the selected blocks.
The Index Branch introduces just two additional projection matrices to standard GQA. For each query token, it produces one index query head per GQA group and a single shared index key head. It scores all causally visible key tokens using a scaled dot-product, then aggregates those scores to the block level via max-pooling. From these block-level scores, it selects the top-k blocks for each GQA group independently. The local block (the one containing the current query position) is always included regardless of its score, which is critical for training stability.
The Main Branch then performs standard scaled dot-product softmax attention, but only over the tokens in those selected blocks. Since each query attends to at most k × Bk tokens (in the paper’s configuration, 16 blocks of 128 tokens = 2,048 tokens), the per-query attention cost drops from O(N) to O(k × Bk), which stays fixed as the sequence length grows.
The key design decision is per-GQA-group independent selection. Each of the 4 key-value heads (serving 16 query heads each, for 64 total query heads) independently picks its own top-16 blocks. This means different attention groups can focus on different parts of the context simultaneously — one group might retrieve code blocks while another retrieves documentation — all within the same layer. The block-level granularity (rather than per-token selection) keeps KV reads contiguous and maps efficiently to GPU memory operations.
Training the Index Branch
Since the top-k selection operation is non-differentiable, the standard language modeling loss can’t directly train the index projections. MSA solves this with a KL alignment loss: the Index Branch’s attention distribution over the selected tokens is trained to match the Main Branch’s attention distribution. The Main Branch acts as the teacher, and the Index Branch learns to predict where the Main Branch would want to attend.
Three additional mechanisms stabilize training:
- Gradient Detach — The Index Branch input is detached from the backbone with stop-gradient. The KL loss only updates the two index projection matrices, keeping the auxiliary objective from interfering with the main model.
- Indexer Warmup — During the first few billion tokens of training, both branches run full attention and the Index Branch learns from the KL loss before sparse selection is enabled. This also works for converting a pretrained dense checkpoint to sparse attention.
- Forced Local Block — The local block containing the current query position is always selected, preventing degenerate selections that omit the immediate context.
Making Sparsity Fast: The Kernel Design
Theoretical FLOPs reduction doesn’t automatically translate to wall-clock speedups. Sparse attention introduces index construction, top-k selection, reverse-index materialization, and irregular memory access patterns that can easily eat into the savings. MSA is co-designed with custom GPU kernels (available on GitHub) to address each of these challenges.
Exp-free TopK selection. Since softmax is order-preserving (si ≤ sj ⟺ softmax(s)i ≤ softmax(s)j), the index module skips the exp/max/sum steps entirely and passes raw dot-product scores directly to the top-k kernel. A specialized per-thread register top-k implementation with min-heaps and shuffle-merge outperforms both torch.topk and TileLang’s radix-select by 2–5x at the deployed k=16 setting.
KV-outer iteration. Rather than iterating over queries (Q-outer), the kernel iterates over KV blocks and gathers the queries that selected each block. This dramatically improves arithmetic intensity: with a block size of 128, the FLOPs/IO ratio becomes (2/3) × Bk ≈ 85, compared to just G = 16 for Q-outer. Each KV block’s associated queries are concatenated into 128×128 score MMAs to fill tensor cores efficiently.
Pre-scheduled tile chunking. A GPU scheduler kernel splits each KV tile along its query dimension into chunks of ~2k × Bk queries, distributing “hot” KV blocks (selected by many queries) across multiple CTAs that share the same K/V load. Each query-chunk pair is pre-assigned a slot in an output buffer, eliminating the need for atomic updates.
Two-phase forward. Because the KV-outer split means each query’s k partials are produced by k different CTAs, the forward pass is split into an attention kernel (writing locally-normalized partials) and a combine kernel (merging partials with numerically stable log-sum-exp reduction). Programmatic Dependent Launch hides the inter-kernel launch latency.
The Results: Quality Preserved, Speed Dramatically Improved
MSA was validated on a 109B-parameter Mixture-of-Experts model (6B activated parameters per token) with native multimodal training on a 3T-token budget. The model uses 41 layers (3 dense + 38 MoE), 64 query heads, 4 KV heads, 128 routed experts with top-4 routing, and a hidden dimension of 3,072.
Two training routes were tested: MSA-PT (from-scratch sparse pretraining) and MSA-CPT (converting a 2.6T-token dense checkpoint by replacing attention and continuing for 400B tokens). Both remain broadly competitive with the full-attention GQA baseline across a comprehensive evaluation suite spanning general reasoning (MMLU, MMLU-Pro, GPQA Hard), math (GSM8K, OlymMATH), code (HumanEval, EvalPlus, BigCodeBench), multimodal benchmarks (image: MMMU, ChartQA, CharXiv; video: VideoMME, MLVU, TemporalBench), and long-context retrieval (RULER, HELMET). MSA-PT actually outperforms full attention on many math, image, video, and long-context benchmarks, suggesting that native sparse pretraining can adapt model representations to the sparse attention pattern.
The efficiency numbers are where MSA truly stands out:
- 28.4x reduction in per-token attention compute at 1M context length
- 14.2x prefill wall-clock speedup on H800 GPUs
- 7.6x decoding wall-clock speedup on H800 GPUs
The gap between FLOPs reduction (28.4x) and runtime speedup (14.2x prefill, 7.6x decode) reflects the overhead of index construction, top-k selection, reverse-index materialization, query gathering, and the less regular memory access pattern of sparse attention. But the key insight is that as context length grows, the dense baseline continues to scale quadratically while MSA’s main attention budget stays fixed at 2,048 tokens per query — so the speedup advantage only increases with longer contexts.
A production-grade model powered by MSA, MiniMax-M3, is publicly available on HuggingFace, along with the inference kernels for anyone to reproduce or build upon.
Why This Matters
MSA follows a principle of deliberate simplicity. Unlike some approaches that add multiple parallel branches with different attention patterns, MSA is a single streamlined mechanism: score blocks, select top-k, attend to selected blocks. It adds just two projection matrices to GQA and uses a single KL loss to train them. This simplicity is what makes it practical to deploy across a range of GPU hardware.
The per-GQA-group independent selection is particularly elegant. Most current open-weight models already use GQA, which means MSA’s recipe transfers with minimal modification. The block-level granularity keeps KV reads contiguous and avoids the fragmentation problems of per-token selection, while still allowing each attention group to develop its own retrieval strategy — some groups focusing on recent context, others on semantically distant but relevant blocks.
For teams deploying LLMs in production, the implications are concrete. A 14x prefill speedup means that processing a 1M-token context — the scale needed for serious agentic workflows — goes from impractically expensive to genuinely deployable. The fact that this comes with no meaningful quality degradation, and in some cases actual improvements, makes MSA one of the most practical long-context solutions to emerge this year.
The attention mechanism in transformers has been due for an efficiency overhaul, and approaches like MSA show that you don’t need to abandon softmax attention to make it work at scale — you just need to be smarter about where you apply it.