PaperlayerView exact revision on arXiv

Source-linked AI summary

MiniMax Sparse Attention

Xunhao Lai, Weiqi Xu, Yufeng Yang, Qiaorui Chen, Yang Xu, Lunbin Zeng, Xiaolong Li, Haohai Sun, Haichao Zhu, Vito Zhang, Jinkai Hu, Jiayao Li, Rui Gao, Zekun Li, Songquan Zhu, Jingkai Zhou, Pengyu Zhao

arXiv:2606.13392v2cs.AI

TL;DR

Ultra-long contexts create severe compute and memory bottlenecks because softmax attention scales quadratically. MSA uses group-specific block selection with exact sparse attention, matching GQA while achieving substantial compute and wall-clock speedups at 1M context.

  • Problem

    Ultra-long-context workloads face severe training and inference compute and memory bottlenecks because softmax attention has quadratic cost.

  • Method

    MSA adds a lightweight Index Branch to GQA that independently selects key-value blocks per group, followed by exact softmax attention over selected blocks.

  • Results

    MSA matches GQA on downstream benchmarks while delivering 14.2× prefill and 7.6× decoding speedups at 1M context length.

  • Takeaways & Limitations

    At 109B-MoE scale, MSA preserves GQA Full-Attention capability across most pretraining and agentic benchmarks while reducing per-token attention compute by 28.4× at 1M context.

  • Takeaways & Limitations

    MSA retains a residual long-context retrieval gap, motivating longer sparse training, larger inference-time selection budgets, or richer indexer scoring.

Abstract

from arXiv · show

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.

1. Introduction

Ultra-long-context LLM workloads create a severe efficiency challenge, motivating MiniMax Sparse Attention (MSA), a minimal blockwise sparse attention design that preserves existing infrastructure. MSA combines group-specific block selection with co-designed GPU kernels, matching GQA while delivering substantial 1M-context speedups.

  • Motivation: Ultra-long-context agentic workflows demand extensive reasoning, tool use, coding, web navigation, and document production, imposing severe compute and memory pressures.The supplied passage identifies longhorizon workflows spanning hundreds of interleaved reasoning and action steps, but its final compute statement is truncated.
  • Method: MSA uses blockwise token selection with a smaller top-k to reuse sparse softmax attention infrastructure, support broader GPU architectures, and relax prior head-dimension constraints.Its design follows Occam’s razor, retaining only essential components after extensive ablation.
  • Systems co-design: MSA co-designs an exp-free TopK kernel and KV-outer sparse attention path to convert theoretical sparsity into practical GPU speedups.The execution path bypasses unnecessary softmax computation before selection and gathers selected KV blocks with associated queries to fill tensor-core MMAs.
  • Results: 14.2× prefill and 7.6× decoding speedups at 1M context length are achieved while MSA matches GQA on downstream benchmarks in a 109B-parameter MoE model.The model was trained from scratch with a 3T-token budget to assess textual and multimodal capabilities.
  • Scope and validation: MSA supports training from scratch and near-lossless conversion from pretrained GQA checkpoints, with extensive ablations across model scales and modalities.The ablations scale to a 109B-parameter MoE model with native multimodal training.

2. Preliminary

The preliminary framework contrasts quadratic causal Softmax Attention with sparse attention, which separates index selection from exact attention over selected keys. MSA specializes this approach through GQA-shared, block-level selection to improve computational regularity while retaining causal attention within selected blocks.

  • Attention foundations: Θ(2𝐻𝑞𝑁2𝑑ℎ) FLOPs is the cost of causal Softmax Attention, growing quadratically with sequence length 𝑁.GQA reduces key-value heads to 𝐻𝑘𝑣 and ties 𝐺=𝐻𝑞/𝐻𝑘𝑣 adjacent query heads to each shared key-value head.
  • Sparse attention: Sparse attention factors causal attention into an Index Branch that selects keys and a Main Branch that applies scaled dot-product softmax attention to them.The selected index set may be fixed-rule or learned, and different query heads can select different key/value sets.
  • GQA-based sparsity: GQA-based sparse attention shares the index result across the 𝐺 query heads served by each key-value head, reducing fine-grained routing that maps poorly to GPU matrix operations.Per-head token-level selection is finest-grained but difficult to execute efficiently on GPUs; group sharing improves efficiency.
  • Block sparsity: Block selection reduces routing overhead and regularizes computation by selecting key/value blocks of size 𝐵𝑘 instead of individual tokens.For sequence length 𝑁, the number of blocks is 𝐵=⌈𝑁/𝐵𝑘⌉.
  • MSA formulation: MSA computes sparse attention for each query head over causally visible tokens in its group’s selected blocks, using that group’s shared key-value head.Its concrete indexer architecture and training objective are introduced in the next section.

3. MSA

MSA is a GQA-based sparse attention mechanism that independently selects key-value blocks for each GQA group before applying exact attention only within those blocks. Its fixed block budget reduces per-query attention complexity as sequence length grows, while auxiliary alignment and stabilization mechanisms train the non-differentiable indexer.

  • Architecture: MSA uses a lightweight Index Branch to select key-value blocks independently for each GQA group, followed by exact Main Branch attention over selected tokens.The Index Branch adds only two projection matrices to standard GQA and operates at block granularity.
  • Architecture: At most kB_k causally visible tokens are attended per query, reducing per-query attention cost from O(N) to O(kB_k), fixed as sequence length increases.The selected block index set is shared by all query heads within a GQA group.
  • Training: The non-differentiable Top-k selection is trained with a KL alignment loss that matches Index Branch scores to the Main Branch attention distribution.Gradient Detach restricts the auxiliary loss to the index projections rather than the backbone.
  • Training: Indexer Warmup begins with full attention in both branches, then switches to sparse attention while computing KL loss over selected positions.The same schedule is used when sparsifying a pretrained full-attention checkpoint.
  • Training: A local block containing each query is always selected, reserving one block slot and preventing selections from omitting the query’s immediate neighborhood.The remaining slots are chosen by the Index Branch during training and inference.
  • Efficiency: MSA’s FLOPs advantage over GQA grows with sequence length when kB_k≪N and H_kv d_idx≪H_q d_h, because its main path uses a fixed selection budget.GQA scales its main attention path with the full context length, whereas MSA adds only lightweight index computation.

4. Kernel Design

The kernel design combines exp-free Top-k selection with KV-outer sparse attention to improve blockwise GPU execution. It uses co-designed block and selection sizes, query gathering, pre-scheduled chunking, and a two-phase forward path to handle sparse workloads efficiently.

  • Top-k kernel: Exp-free Top-k selection bypasses softmax max/exp/sum operations because softmax preserves score ordering.The index module ranks raw index scores directly and passes them to selection.
  • Top-k kernel: 128 and 16 are the adopted KV block size B_k and selection size k, respectively, co-designed for the specialized register Top-k kernel.The design balances attention arithmetic intensity against the candidate-block and selection sizes handled by general-purpose Top-k methods.
  • Sparse attention forward: 3 B_k≫G motivates KV-outer iteration with query gathering to maximize arithmetic intensity.A persistent grid processes (KV block, KV head) tiles, while reverse sparse indices identify and gather the relevant query positions.
  • Sparse attention forward: Pre-scheduled tile chunking splits overloaded KV tiles into bounded query-dimension chunks and assigns each query-chunk pair an output slot, avoiding atomics during partial writes.The combine kernel uses per-query slot counts to determine how many partials to merge.
  • Sparse attention forward: Two-phase forward separates sparse attention and combination with HBM buffers for locally normalized partial outputs and per-partial logsumexp values.The split is required because KV-outer execution produces each query’s k partials in different CTAs; Programmatic Dependent Launch hides launch latency.

5. Experiment

The experiments validate MSA on a 109B-parameter native multimodal model through matched pretraining, sparse continued pretraining, long-context extension, and efficiency measurements. Across these settings, MSA maintains competitive training and evaluation behavior while substantially reducing attention computation.

  • Experimental setup: Two 109B-scale experiments evaluate MSA: MSA-PT trains natively from scratch, whereas MSA-CPT replaces dense attention in a Full-Attention checkpoint and continues pretraining.Both use the same architecture family as the Full-Attention baseline and are trained on native multimodal data.
  • Training dynamics: Over 3T training tokens, MSA-PT and Full Attention have nearly indistinguishable LM-loss curves, with gradient norms remaining in the same range.These dynamics indicate no noticeable optimization degradation, abnormal gradient fluctuations, or training instability relative to Full Attention.
  • Sparse continued pretraining: During MSA-CPT, indexer warmup rapidly reduces KL loss, while favorable block recall and higher score recall show that selected blocks recover important attention mass.The indexer remains well aligned after sparse attention is enabled.
  • Pretraining evaluations: Both sparse models remain broadly competitive with Full Attention across general language, reasoning, multimodal, and agent-oriented evaluation profiles.Table 2 compares Full, MSA-PT, and MSA-CPT under the same 3T-token training budget.
  • Long-context extension: After approximately 140B long-context training tokens, MSA-CPT remains close to the Full-Attention baseline on HELMET and RULER while each query and GQA group attends to 2,048 key-value tokens.The selected budget is kB_k = 16 × 128 = 2,048 key-value tokens.
  • Efficiency: 28.4×: at 1M tokens, MSA reduces per-token attention FLOPs relative to GQA under the same head configuration.Measured runtime speedups scale similarly but differ from FLOPs reduction because sparse execution adds indexing, selection, gathering, materialization, and load-balancing overheads.

6. Related Works

Related work on long-context efficiency spans cheaper linear or recurrent alternatives, restricted-receptive-field softmax attention, adaptive support selection, and kernels that convert sparse computation into wall-clock gains.

  • Efficient Attention: Long-context efficiency broadly follows two directions: replacing dense softmax attention with linear or recurrent alternatives, or restricting softmax attention’s receptive field.Linear attention uses a linear-complexity surrogate, while Mamba replaces attention with a state-space model.
  • Adaptive Sparse Attention: Adaptive sparse attention makes attended support input-dependent, with methods differing in support-construction timing and whether selection is trained jointly.Inference-time methods sparsify a pretrained Full-Attention backbone during serving; H2O and SnapKV prune the KV cache during decoding using accumulated attention statistics.
  • Efficient Kernels: Efficient kernels are essential for translating sparse attention’s theoretical FLOP reduction into wall-clock speedups.FlashAttention and FlashAttention-2 introduced IO-aware tiled softmax attention, while FlashDecoding extended this approach to memorybound decoding.

7. Conclusion

The conclusion presents MSA as a lightweight, GQA-compatible block-sparse mechanism that preserves capability while substantially reducing long-context attention compute. It also identifies transferability and long-context retrieval as key outlook priorities.

  • Core mechanism: MSA combines GQA with a lightweight Index Branch that independently selects key-value blocks for each GQA group, followed by Main Branch attention over those blocks.The Main Branch applies softmax attention only to the selected blocks.
  • Training: The Index Branch is trained as a pure selector using KL alignment against the Main Branch, with two-stage warmup and stop-gradient on the index input.These choices confine the auxiliary loss to the index projections.
  • Results: 28.4× lower per-token attention compute at 1M context accompanies capability preservation versus a GQA Full-Attention baseline across most pretraining and agentic benchmarks.This result is reported at the 109B-MoE scale.
  • Outlook: MSA’s per-GQA-group selection, block-level granularity, and KL-trained indexer are expected to transfer with little modification to GQA backbones used by current open-source frontier models.The outlook identifies closing the residual long-context retrieval gap as a natural next step.

A. Visualization

MSA’s learned Index Branch recovers dense-attention structures while producing group-specific long-range sparse patterns. It also naturally identifies the initial key-value block as an attention sink across layers and heads.

  • Selection patterns: All GQA groups consistently select the local diagonal and sink column, while tracing different long-range stripes.The visualizations cover heads from early and later layers and show visibly different stripe patterns across groups.
  • Attention sink: The Index Branch naturally assigns high selection probability to the initial key-value block without explicit enforcement.This behavior appears across all layers and heads in the examined MSA models.

B. Preliminary Experiments · B.1. Setup

The preliminary experiments use small-scale ablations to identify training choices for stable optimization and downstream performance. The setup evaluates these choices on a 10B-parameter pilot Transformer matching the main model’s architecture family.

  • B. Preliminary Experiments: Small-scale ablations identify training-design choices essential for stable optimization and strong downstream performance.These results provide the empirical basis for the final recipe described in Section 3.
  • B. Preliminary Experiments: Figure 6 shows a significant fraction of attention allocated to the first token across heads in Layers 4 and 24.The figure reports a pervasive attention sink effect across heads and layers.
  • B.1. Setup: 10B-parameter pilot Transformer ablations use the same architecture family as the main paper MSA model, with 16 layers.The pilot model has 10.53B total parameters and 1.47B active parameters per token.
  • B.1. Setup: 200K-token vocabulary and hidden size 𝑑model = 2048 define the pilot Transformer’s representation setup.These dimensions are used consistently across the ablations in this section.
  • B.1. Setup: 32 query heads, 4 KV heads, head dimension 128, and RoPE dimension 64 configure each GQA attention module.The pilot Transformer uses grouped-query attention throughout these ablations.
  • B.1. Setup: 64 experts with top-4 expert routing and expert inner dimension 1536 configure the pilot model’s MoE component.This routing and expert-width configuration is part of the fixed ablation setup.

B.2. Gradient Sources for the Index Branch … C. Additional Ablation Study

The ablations establish a training recipe for MSA’s Index Branch: combine language-modeling and KL supervision, confine KL gradients, and warm up with full attention. Additional studies find no consistent benefit from a learnable sink, while dynamic selection outperforms a FLOP-matched sliding window.

  • B.2. Gradient Sources for the Index Branch: The non-differentiable top-k routing leaves the Index Branch without useful language-modeling gradients, motivating an auxiliary output path and direct KL supervision.The output path trains the indexer through next-token prediction, while KL loss matches its selection distribution to the Main Branch on selected support.
  • B.2. Gradient Sources for the Index Branch: LM Loss only preserves short-context ability but performs poorly on long-context retrieval, whereas KL Loss only improves retrieval but reduces short-context ability.The combined LM Loss + KL Loss configuration is therefore used for subsequent ablations; the final recipe later retains KL supervision while removing the value head after warmup.
  • B.3. Confining the KL Gradient to the Index Branch: KL gradients that update the backbone can cause gradient-norm spikes, LM-loss divergence, and gradual short-context regression.Stopping the KL gradient at the Index Branch input makes the loss local to each indexer and removes these failure modes under the same coefficients.
  • B.4. Indexer Warmup: Main Branch attention entropy drops rapidly early in training, making sparse selection fragile while the Index Branch is still nearly random.During indexer warmup, the Main Branch uses full attention and the Index Branch learns from the full-sequence distribution before top-k selection is enabled.
  • B.4. Indexer Warmup: Indexer warmup improves short-context performance and long-context retrieval, so it is adopted for both pretraining and continued conversion of Full-Attention checkpoints.The warmup provides a meaningful initialization before sparse routing controls token selection.
  • B.5. Learnable Attention Sink: The learnable attention sink absorbs substantial mass in some heads but does not eliminate the original first-token sink.Because it provides no clear or consistent downstream perplexity improvement while adding parameters and implementation complexity, it is excluded from the final recipe.
  • B.6. Dynamic Sparse Selection vs. Sliding Window: Under the same sparse selection budget, the FLOP-matched sliding-window model has higher perplexity than MSA throughout the training trajectory.The fixed pattern attends to the first key block and a local window, whereas MSA selects tokens dynamically.

C.1. Block Size · C.2. Forced Sink & Local Selection

MSA’s block size trades selection granularity against kernel efficiency, while ablations show larger blocks can preserve quality with limited loss. Forced sink and local selection stabilize early training but are not required for the trained model to retain these patterns.

  • C.1. Block Size: MSA’s Main Branch processes key-value pairs in consecutive B_k-token blocks, making block size relevant to both performance and efficiency.Larger blocks can improve kernel efficiency but may reduce retrieval quality through coarser selection granularity.
  • C.1. Block Size: The block-size ablation varies B_k while keeping the total number of selected tokens constant, using fewer training iterations and a subset of evaluations.
  • C.1. Block Size: PPL results are nearly unchanged across different B_k values, indicating limited model-quality sensitivity to block size in this setting.
  • C.1. Block Size: Table 4 reports perplexity and long-context retrieval scores for different key-value block sizes, with lower perplexity and higher RULER scores preferred.
  • C.2. Forced Sink & Local Selection: Early sparse-training experiments forced selection of the first sequence block and a fixed local window around each query position.The first block reflects attention-sink behavior, while the local window preserves nearby context and provides dense supervision for the indexer.
  • C.2. Forced Sink & Local Selection: Removing forced sink and fixed local selection leaves the trained model still concentrating on useful prefixes and frequently selecting nearby tokens.The priors were introduced mainly to stabilize training before the indexer became reliable.
  • C.2. Forced Sink & Local Selection: After removing the forced priors, reasoning, code, and PPL metrics remain nearly unchanged in standard model-quality evaluations.The passage also reports long-context retrieval results in Table 5, but the supplied text does not provide their conclusion.

C.3. Index Branch Value Head

With Index Branch warmup already improving sparse-training initialization, the index value head is not critical for downstream quality and is removed from the final design. This simplifies inference by eliminating value aggregation and exponential calculations.

  • Motivation: Index Branch warmup motivates testing whether the additional value head remains necessary for sparse training.The value head initially helps training begin from step zero but adds computation and complexity.
  • Ablation result: Removing the index value head causes no systematic degradation across the evaluation suite, with mixed benchmark-dependent differences.The no-value variant is slightly better on some reasoning benchmarks, while the with-value variant retains small advantages on some math and code tasks.
  • Design choice: The final design drops the index value head because neither variant consistently dominates and its downstream effect is small.The earlier role of O_idx was primarily to provide an additional early training signal rather than essential converged capacity.
  • Inference implementation: At inference, the top-k indexer uses only the block-wise maximum of Q_idxK_idx^T, avoiding value aggregation and exponential calculations.This removes the value aggregation path from the indexer.
Loading 2606.13392v2…