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
MiniMax Sparse Attention (MSA) uses a lightweight, group-specific indexer to enable efficient, block-sparse attention.
How can we enable ultra-long-context LLMs by replacing dense attention with a learnable, sparse selection mechanism that maintains performance while reducing compute?
Ultra-long-context tasks like agentic workflows require models to attend over millions of tokens, but the quadratic cost of standard softmax attention makes this computationally prohibitive at scale. MSA introduces a two-branch architecture: a lightweight Index Branch independently selects top-k key-value blocks for each Grouped Query Attention (GQA) group, and a Main Branch performs exact attention only over those selected blocks. On a 109B-parameter model, MSA matches GQA performance while reducing per-token attention compute by 28.4× at 1M context length.
Paper Primer
MSA treats sparse attention as a two-stage process: an Index Branch scores key-value blocks using a block-level dot-product indexer, and a Main Branch executes the attention computation. By sharing the selected block indices across all query heads within a single GQA group, MSA maintains block-granular execution that maps efficiently to GPU tensor cores.
To translate theoretical sparsity into wall-clock speed, MSA uses a co-designed GPU kernel that employs exp-free top-k selection and a KV-outer iteration order. This kernel packs gathered query positions into 128x128 matrix-multiply-accumulate (MMA) tiles, maximizing arithmetic intensity while avoiding atomic operations through pre-scheduled chunking.
MSA achieves significant inference speedups over dense GQA at long context lengths.
Measured wall-clock performance on H800 GPUs at 1M context.
MSA preserves the performance profile of dense GQA models across diverse benchmarks.
Evaluation of a 109B-parameter MoE model on MMLU, GSM8K, HumanEval, and multimodal tasks.
Why does MSA use block-level selection instead of token-level selection?
Token-level selection is difficult to map efficiently to GPU matrix operations; block-level selection reduces routing overhead and allows for more regular memory access patterns that better utilize tensor cores.
How does the model learn to select relevant blocks without a differentiable top-k operation?
The Index Branch is trained using an auxiliary KL alignment loss that supervises the indexer to match the attention distribution of the Main Branch, combined with a two-stage warmup schedule and gradient detachment to ensure stable training.
The Long-Context Bottleneck
Long-context LLM tasks expose quadratic cost of dense attention, motivating sparse alternatives.
Agentic applications—code generation, web navigation, and multi‑step reasoning—need models to attend over ultra‑long token streams, but dense softmax attention scales quadratically in both compute and memory, quickly becoming a deployment bottleneck.
When every token compares to every other token, the number of pairwise interactions grows as the square of the sequence length, exploding both compute and memory as contexts grow.
The shift from dense to sparse attention is required for agentic workflows.
Foundations of Causal Attention
Formal definitions of causal attention, GQA, and the two‑stage sparse attention pipeline.
We formalize causal attention, introduce Grouped Query Attention, and describe the two‑stage sparse attention pipeline that underlies MSA.
Multiple query heads are tied together so that a small set of key/value heads serves a whole group of queries, cutting the number of KV heads without sacrificing the ability to attend to different positions.
The FLOP count of the above attention scales as $\Theta(2\,H_q\,N\,2d_h)$, i.e. quadratic in the sequence length $N$.
We call the indexer the **Index Branch** and the subsequent restricted attention the **Main Branch**.
These definitions let MSA apply GQA‑based block‑sparse attention: groups share index results, and attention is performed only over the selected blocks, dramatically cutting routing overhead while preserving causal constraints.
The MSA Mechanism
Describes the MSA architecture and its Index Branch selection mechanism.
MSA splits attention into a fast Index Branch that picks a few key blocks, then a Main Branch that runs full attention only on those blocks — like a mail sorter first deciding which bins to open before delivering letters.
Compute index queries: $\text{idx}_q = X W_q$ yields a scalar per token per group; assume values $[0.2, 0.5, 0.1, 0.4]$ for group 1 and $[0.3, 0.6, 0.2, 0.5]$ for group 2.
Compute index keys: $\text{idx}_k = X W_k$ yields a scalar per token; assume values $[0.1, 0.2, 0.3, 0.4]$.
Score block 1 (tokens 1‑2) for query 3 in group 1: max over $j\le3$ gives $\max(0.2\cdot0.1, 0.5\cdot0.2)=0.10$.
Score block 2 (tokens 3‑4) for query 3 in group 1: max over $j\le3$ gives $\max(0.1\cdot0.3)=0.03$.
Select top‑$k=1$ block: block 1 wins, so $I^{(1)}(3)=\{ \text{block 1} \}$; the local block (block 2) is also kept, giving two blocks total.
Main Branch attention now attends only to the two tokens in block 1 (tokens 1‑2) for query 3, reducing the cost from $4$ tokens to $2$.
This toy example shows how the Index Branch prunes the attention budget to a fixed number of blocks, independent of sequence length, while still preserving the immediate local context.
**Figure 1 | Overview of MSA.** The Index Branch (left) scores the full causal context with a single lightweight head and selects, for each query and GQA group, a set $\mathcal{I}$ of $k$ key blocks; the local block is always included regardless of its score. The Main Branch (right) attends only to the selected blocks and produces the layer output. During training, a KL loss aligns the index distribution with the group-averaged Main Branch distribution on the selected blocks, and the Index Branch gradient is detached from the Main Branch.
Why does the Index Branch use a max‑pool over tokens in a block instead of averaging?
Max‑pool captures the strongest relevance signal within a block, ensuring that a single highly compatible token can make the block attractive. Averaging would dilute that signal, potentially causing the block to be missed even though it contains a useful token.
Training Sparse Selection
Training MSA uses KL alignment, gradient detach, warmup, and a forced local block to enable sparse attention.
The index branch is taught to mimic the main‑branch attention distribution on the tokens it actually selects.
How does this KL loss differ from a regular KL‑regularization term on the attention weights?
Standard KL regularization penalises a single distribution against a fixed prior; here the “teacher” distribution is the full‑attention output itself, and the loss is computed only over the tokens that the indexer actually selects, making it a targeted alignment rather than a generic regularizer.
For each group we compute a softmax over the four tokens, yielding $P^{\text{idx}}_{(1)} = [0.12,\,0.53,\,0.09,\,0.26]$ and $P^{\text{idx}}_{(2)} = [0.31,\,0.19,\,0.38,\,0.12]$.
The main‑branch teacher distributions (averaged over the two heads) are $P_{(1)} = [0.25,\,0.45,\,0.15,\,0.15]$ and $P_{(2)} = [0.30,\,0.20,\,0.35,\,0.15]$.
Compute KL for each group: $D_{\text{KL}}(P_{(1)}\|P^{\text{idx}}_{(1)}) \approx 0.18$, $D_{\text{KL}}(P_{(2)}\|P^{\text{idx}}_{(2)}) \approx 0.12$.
Average over groups and positions: $L_{\text{KL}} = (0.18+0.12)/(N H_{kv}) = 0.0375$.
The KL loss forces the indexer to assign higher probability to exactly those tokens the full attention would attend to, even though it only sees a tiny subset of them.
We stop gradients from flowing from the KL loss into the main‑branch backbone, so the indexer learns independently.
Why not simply omit the KL loss instead of detaching its gradients?
Omitting the KL loss would leave the indexer without any learning signal; detaching ensures the loss still provides a signal but does not corrupt the already‑trained main‑branch parameters.
Before the indexer is trusted to select blocks, we let the model run full attention while still training the new index projections with the KL loss.
Why not train the indexer from scratch without a warmup phase?
Training from scratch would expose the model to arbitrary block selections, causing the KL loss to receive meaningless targets and leading to unstable gradients; warmup gives the indexer a sensible initialization before it influences routing.
Across layers many heads allocate a large fraction of their attention mass to the first token, creating a “sink” that wastes capacity.
Is the sink effect a bug of the architecture or a learned behavior?
It emerges from the softmax dynamics: the first token is always visible, so the model can cheaply satisfy the attention budget by repeatedly attending to it. It is not a design flaw but a tendency that sparse training aims to mitigate.
**Figure 6.** Mean attention score on the first token for each attention head in Layer 4 and Layer 24. All heads allocate a significant fraction of attention to the first token, confirming a pervasive attention sink effect across heads and layers.
Training forward pass for one MSA layer with auxiliary KL loss.
**Algorithm 1** One MSA layer: training forward and the auxiliary KL loss. The layer returns its output and per-layer $\mathcal{L}_{KL}$; the model loss $\mathcal{L} = \mathcal{L}_{LM} + \lambda \sum_{layers} \mathcal{L}_{KL}$ is assembled by the training loop. **Require:** hidden states $X \in \mathbb{R}^{N \times d_{model}}$; block size $B_k$, number of selected blocks $k$. 1: $Q, K, V \leftarrow XW_q, XW_k, XW_v$ // $(N, H_q, d_h), (N, H_{kv}, d_h), (N, H_{kv}, d_h)$ 2: $Q^{idx}, K^{idx} \leftarrow \text{stopgrad}(X)W_q^{idx}, \text{stopgrad}(X)W_k^{idx}$ // $(N, H_{kv}, d_{idx}), (N, 1, d_{idx})$; detached 3: $M^{idx} \leftarrow \text{BlockMaxPool}(Q^{idx}, K^{idx}, B_k)$ // $(N, H_{kv}, B)$; per-group, causal 4: $I \leftarrow \text{TopK}(M^{idx}, k)$ // selected block indices; local block included 5: $O \leftarrow \text{TopKAttn}(Q, K, V, I)$ // $(N, H_q, d_h)$; attends to selected blocks 6: output $\leftarrow OW_o$ // $(N, d_{model})$ 7: $\mathcal{L}_{KL} \leftarrow \text{KLdiv}(Q^{idx}, K^{idx}, \text{stopgrad}(Q), \text{stopgrad}(K), I)$ // over tokens induced by $I$ 8: **return** output, $\mathcal{L}_{KL}$
GPU Kernel Implementation
Efficient GPU kernels enable sparse prefill without softmax overhead.
Sparse prefill must locate the most relevant KV blocks quickly and keep the GPU busy on the selected data; naïve dense kernels would waste cycles on irrelevant tokens.
The kernel stitches together three pieces—raw‑score top‑k indexing, KV‑outer iteration, and a two‑phase softmax—to turn a sparse attention problem into a dense‑compute workload on only the chosen blocks.
How does this sparse prefill kernel differ from a standard dense attention kernel?
A dense kernel multiplies every query by every key, incurring $O(N^2)$ work. Our kernel first discards $N\!-\!kB_k$ keys via the index, then performs dense matrix‑multiply only on the $kB_k$ selected blocks, turning quadratic work into linear‑in‑$k$ work while still using the highly optimized GEMM path.
Each lane of a warp reads one element (stride = 1/32) and inserts it into a $k\!=\!2$ min‑heap kept in a register.
After processing the first four elements, the heap holds the two largest seen so far: $[0.85,\,0.33]$.
Continuing through the remaining values, the heap finally contains $[0.91,\,0.85]$, the top‑2 scores.
The corresponding block indices (e.g. blocks 4 and 2) are emitted as the selected KV blocks for this query row.
The min‑heap lets each thread keep only the best $k$ candidates without ever sorting the full $B_k$ row, yielding $O(B_k\log k)$ work per lane.
Per‑thread register top‑k (min‑heap) kernel – one warp processes a row.
**Table 1.** Top-k latency ($\mu s$) for fp32 inputs of shape ($N, B$), with rows processed independently. The deployed setting uses $B_k = 128$, $k = 16$, while for reference we also report $k = 32$ with $B_k = 64$. All implementations produce identical index sets.
Because each query’s $k$ partial results are produced by different CTAs, we cannot normalize them inline; we therefore split the forward pass into a compute kernel that writes raw partials to $O_{\text{buf}}$ and a combine kernel that performs the softmax reduction.
The KL alignment loss only contributes to gradients, so we emit the necessary log‑sum‑exp scalars during the main pass and skip a separate forward kernel.
Why does emitting $LSE$ values during the main pass avoid a separate KL‑loss kernel?
Because the KL term’s only role is to provide a scalar for the softmax denominator in the backward pass; writing it once with the main attention output eliminates a full‑precision reduction and a second kernel launch, cutting both latency and memory bandwidth.
Empirical Performance
Dense attention wastes compute; MSA’s Index Branch selects only the relevant key blocks.
Recall that dense attention spends compute on irrelevant tokens; MSA replaces it with a lightweight Index Branch that dynamically picks the most useful key blocks.
MSA‑CPT matches Full‑Attention on most tasks and outperforms it by +2.4 points on the HELMET‑128K Rerank benchmark.
Table 3 reports 72.80 for MSA‑CPT versus 70.40 for Full, a +2.40 improvement.
**Figure 2.** Pretraining dynamics for the experiment model. LM loss and gradient norm are shown for Full Attention and MSA-PT over 3T training tokens. The inset in (a) zooms in on the final 50B-token window, where the two LM-loss curves remain nearly overlapping.
**Figure 3.** Sparse continued-pretraining dynamics. (a) Average KL loss during MSA-CPT. The solid segment denotes indexer warmup, and the dashed segment denotes sparse continued pretraining; the vertical dashed line marks the switch between the two stages. (b) Average block recall and score recall of the MSA-CPT indexer during sparse continued pretraining.
Efficiency and Context
Efficiency numbers for MSA and a survey of long‑context attention methods.
We instantiate the complexity analysis from Section 3.3 on our experimental model configuration and report both theoretical FLOPs reduction and measured runtime speedups.
**Figure 4.** Efficiency comparison between GQA and MSA under the shared experimental model configuration. The left subfigure reports the theoretical per-token attention-FLOPs. The middle and right subfigures report the measured implementation speedups for prefill and decode, respectively. All tests are conducted with 64 query heads, 4 key-value heads, and a head dimension of 128. MSA uses $B_k = 128$ and $k = 16$, corresponding to a selected budget of 2,048 tokens per query.
Long‑context efficiency has motivated two broad families of work: (1) replacing dense softmax attention with cheaper linear or recurrent kernels, and (2) keeping softmax but restricting its receptive field.
Beyond fixed patterns, adaptive sparse attention methods construct input‑dependent supports (e.g., H2O, SnapKV, Quest, MInference, FlexPrefill, Inf LLM), while natively trained sparse‑attention designs such as NSA, Inf LLM‑V2, MoBA, and DSA train an indexer jointly with the language model.
Efficient kernels (FlashAttention, FlashAttention‑2, FlashDecoding, Flash‑Sparse‑Attention, FlashMoBA) are essential for turning theoretical FLOP savings into real‑world speedups; MSA reuses the FlashAttention skeleton with a loop order tuned to its block‑granular access pattern.
Ablations and Visualizations
Ablation studies dissect the contributions of indexer training signals, warmup, and design choices.
**Figure b.** Layer 18, four GQA groups. Long-range selection sharpens into a few stripes per group; the four groups pick visibly different stripes.
Visualizing the Index Branch reveals that, across layers, heads consistently allocate mass to the local diagonal and the first token, yet each GQA group selects a unique set of long‑range positions.
When the Index Branch is trained with only the LM loss, short‑context perplexity stays low but long‑context retrieval suffers because the indexer receives no direct pressure to pick distant blocks.
Training with only the KL loss improves long‑range retrieval but harms short‑context performance, as the Index Branch output path is removed and the model loses capacity for immediate token prediction.
Combining LM loss and KL loss yields the best trade‑off, preserving short‑context ability while still learning useful long‑range selections; this configuration is used for all subsequent ablations.
Allowing the KL gradient to flow into the backbone causes occasional spikes that destabilize training; detaching the KL gradient at the Index Branch input eliminates these spikes and stops short‑context regression.
A short warmup phase where the Main Branch runs full attention while the Index Branch learns via KL loss gives the indexer a stable initialization, leading to better short‑context perplexity and stronger long‑context retrieval after sparse selection starts.
Adding a learnable attention‑sink logit does not consistently improve perplexity and fails to fully suppress the implicit first‑token sink, so the extra parameters are omitted from the final design.
Replacing dynamic selection with a FLOP‑matched sliding‑window pattern raises perplexity throughout training, confirming that content‑dependent token selection is superior to a fixed positional window for the evaluated tasks.
Forcing the selector to always include the first block and a fixed local window does not materially affect perplexity, reasoning, or long‑context retrieval; the model learns to attend to these patterns organically.
Removing the Index Branch value head after applying the warmup yields comparable performance across benchmarks; the head only provided an early training signal and is unnecessary at convergence.