MiniMax Sparse Attention

Abstract: Ultra-long-context capability is becoming indispensable for frontier LLMs: agentic workflows, repository-scale code reasoning, and persistent memory all require the model to jointly attend over hundreds of thousands to millions of tokens, yet the quadratic cost of softmax attention makes this untenable at deployment scale. We introduce MiniMax Sparse Attention (MSA), a blockwise sparse attention built upon Grouped Query Attention (GQA). A lightweight Index Branch scores key-value blocks and independently selects a Top-k subset for each GQA group, enabling group-specific sparse retrieval while maintaining efficient block-level execution; the Main Branch then performs exact block-sparse attention over only the selected blocks. Designed around a principle of simplicity and scalability, MSA is deliberately streamlined, making it straightforward to deploy efficiently across a broad range of GPUs. To translate sparsity into practical speedups, we co-design MSA with a GPU execution path that uses exp-free Top-k selection and KV-outer sparse attention to improve tensor-core utilization under block-granular access. On a 109B-parameter model with native multimodal training, MSA performs on par with GQA while reducing per-token attention compute by 28.4x at 1M context. Paired with our co-designed kernel, MSA achieves 14.2x prefill and 7.6x decoding wall-clock speedups on H800. Our inference kernel is available at: https://github.com/MiniMax-AI/MSA. A production-grade natively multimodal model powered by MSA has been publicly released at: https://huggingface.co/MiniMaxAI/MiniMax-M3.

• The paper introduces a minimal sparse attention mechanism leveraging a dual-branch design that selects top-k key-value blocks, reducing computation drastically.
• It employs Grouped Query Attention with an auxiliary KL loss to align sparse selection with dense attention, ensuring stable and effective training.
• Empirical results demonstrate near-lossless performance on benchmarks with orders-of-magnitude speedups in computation compared to full softmax attention.

MiniMax Sparse Attention: A Minimal, Scalable Sparse Attention Mechanism for Ultra-Long Context LLMs

Introduction and Motivation

Transformer-based LLMs face quadratic time and memory complexity in the softmax attention mechanism, especially problematic for ultra-long contexts (hundreds of thousands to millions of tokens) needed in agentic workflows, code understanding, and persistent memory. Existing efforts to address this bottleneck typically fall into two categories: hybridization (interleaving linear or sliding-window attention with dense softmax) or sparsification (directly imposing a sparse support on softmax attention). MiniMax Sparse Attention (MSA) introduces a principled, minimal sparse attention architecture, leveraging Grouped Query Attention (GQA) and a co-designed GPU execution path to deliver orders-of-magnitude computational savings with minimal loss in performance.

Architectural Overview

MSA’s architecture couples a lightweight Index Branch to a standard GQA layer, producing a dual-branch attention mechanism (Figure 1). For each query token, the Index Branch independently selects the top-$k$ key-value (KV) blocks for each GQA group at block granularity, using only two small projection matrices. The Main Branch executes exact softmax attention, restricted to the tokens in these selected blocks. The selection and computation are performed independently for each GQA group, ensuring block-level regularity and highly efficient mapping to GPU tensor-core operations.

Figure 1: Overview of MSA: the Index Branch (left) computes blockwise scores over the causal context with a lightweight projection and selects the top-$k$ blocks for each GQA group; the Main Branch (right) restricts attention to the selected blocks and outputs the attended values.

MSA is intentionally streamlined: it introduces minimal architectural modifications to typical GQA layers, maximally reusing existing infrastructure, and avoids per-head or per-token granularity, which has prohibitive implementation and efficiency drawbacks on commodity hardware.

Training Procedure and Index Branch Optimization

The selection by the Index Branch is non-differentiable, so a direct gradient signal from the LM objective is unavailable. MSA employs an auxiliary KL alignment loss: the Index Branch is trained to match the Main Branch’s distribution on the selected tokens. To isolate this auxiliary supervision, the KL gradient is strictly confined to the Index Branch parameters via stop-gradient, preventing degenerate self-distillation and gradient spikes in the backbone. Three stabilization strategies are combined: detached KL gradients, a warm-up phase (where the Index Branch is aligned to dense attention before sparsification is enabled), and forced selection of the local block to maintain local connectivity.

Figure 2: KL loss during continued pretraining, demonstrating rapid reduction during indexer warmup and subsequent stability post-sparsification.

Efficient GPU Implementation

MSA’s practical advantages hinge on efficient kernel design for both index selection and sparse attention:

• Index Top-K Selection: The top-$k$ kernel is optimized for the small-$k$ regime. It eliminates softmax computation before block selection by exploiting softmax’s order-preserving property, instead selecting blocks using raw scores. Register-level min-heaps coupled with shared memory reduce selection latency by over $5\times$ versus torch.topk.
• Sparse Attention Execution: The kernel iterates over KV blocks (KV-outer) instead of queries, enabling tensor-core utilization by packing queries that attend to the same KV tile. To balance load and minimize atomic operations, indexer outputs are chunked and pre-scheduled. The attention is split into two kernels to handle normalization, preventing launch latency bottlenecks.

  Figure 3: Efficiency comparison—MSA achieves a $28.4\times$ reduction in attention FLOPs at 1M context and $14.2\times$/ $7.6\times$ wall-clock speedups for prefill/decoding against dense GQA.

A dedicated backward kernel for the KL divergence emits normalization factors directly to memory, eliminating redundant forward passes.

Empirical Results

A series of large-scale experiments at the 109B-parameter multimodal MoE scale demonstrate the capabilities and efficiency of MSA. Both from-scratch (MSA-PT) and dense-to-sparse continued pretraining (MSA-CPT) variants match or marginally improve on full-attention baselines across a wide suite of benchmarks: general QA, math, code, multimodal (image, video), and agentic downstream perplexity. MSA-PT is often superior on math and multimodal retrieval, reflecting better adaptation to sparsity throughout all training, while MSA-CPT offers a near-lossless, practically relevant conversion pathway when a dense checkpoint is available.

Notably, even with $k=16$ and $B_k=128$ (i.e., 2,048 KV tokens per query out of up to $k$0 available), MSA nearly closes the long-context performance gap with dense attention, as evidenced by HELMET and RULER evaluations with minimal degradation.

Figure 4: LM loss trajectories for full attention vs MSA, demonstrating indistinguishable training dynamics and convergence.

Figure 5: The auxiliary KL alignment loss is essential: ablation studies show that relying solely on the language modeling objective produces suboptimal long-context retrieval—even at parity on standard QA.

Ablation and Analysis

A comprehensive ablation study confirms:

• KL alignment (with detached gradients) is critical—gradient leakage into the backbone from the auxiliary loss causes instability and regression in short-context benchmarks.
• Short indexer warmup is essential for stabilizing early sparse training and avoiding degenerate selections.
• Dynamic, learned sparse selection significantly outperforms fixed-pattern (sliding-window) baselines at FLOP-matched budgets, especially in agentic and retrieval-oriented evaluations.
• Attention “sink” phenomena are robust: heads naturally allocate focus to the sequence prefix even without explicit sink tokens or hard-coded rules.

  Figure 6: Heads in all layers naturally allocate significant attention to the sequence prefix, underscoring emergent sink behavior.

Theoretical and Practical Implications

From a systems perspective, MSA efficiently translates its theoretical $k$1 scaling into practical speedups via blockwise regularity and direct hardware alignment. Its indexer-only auxiliary loss and warmup schedule prevent the negative transfer and instability commonly observed in alternative sparse training regimens. Because MSA attaches to standard GQA layers, it is compatible with most open-source large models and can be applied post hoc to pretrained dense models, preserving most capabilities while enabling multi-million context inference on commodity hardware.

The results indicate that aggressively sparse selection budgets are tenable without significant performance loss and that dynamic, input-driven selection grants a distinct advantage in practical deployments requiring long-horizon reasoning or memory.

Outlook and Future Directions

Beyond pretraining, the MSA design can be directly applied to post-training RL or agentic LLMs, where context efficiency and dynamic memory retrieval are critical. Possible extensions include adaptive selection budgets, richer indexer architectures for higher selectivity, or learnable context allocation policies. There is also an opportunity to further close the residual long-context retrieval gap by combining longer sparse training, large inference budgets, or integrating content- and position-aware scoring functions.

Conclusion

MiniMax Sparse Attention establishes a compelling paradigm for scalable, hardware-aligned, and minimal sparse attention in the age of ultra-long context LLMs. Its methodical co-design of algorithm, kernel, and training procedure enables LLMs to sustain competitive performance while shrinking attention-based computation by over an order of magnitude. This architecture is poised for broad adoption in both pretraining and agentic inference scenarios as the frontier of context length continues to expand.

References: See (2606.13392) for full details and experimental results.