MiniMax Sparse Attention (MSA)

一、论文概述

项目	内容
标题	MiniMax Sparse Attention (MSA)
作者	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 (主要), NVIDIA, 北京大学, 浙江大学, 华中科技大学, 南京大学, 杭州电子科技大学
论文	arXiv:2606.13392
代码	https://github.com/MiniMax-AI/MSA
模型	https://huggingface.co/MiniMaxAI/MiniMax-M3
发布	2026-06-11 (v1); 2026-07-15 (v2)
许可	arXiv.org perpetual non-exclusive license

核心贡献：

1. 提出 MSA——一种基于 GQA 的 blockwise sparse attention 机制，通过轻量 Index Branch 为每个 GQA group 独立选择 Top-k KV blocks，Main Branch 仅在选中 blocks 上执行精确 softmax attention
2. 设计 exp-free Top-k 选择和 KV-outer sparse attention GPU kernel，利用 pre-scheduled tile chunking + two-phase combine + query concatenation 解决 sink row 导致的负载不均衡问题
3. 提出 KL alignment loss 训练 Index Branch，配合 gradient detach、indexer warmup 和 forced local block 三项稳定训练的技术
4. 在 109B 参数 MoE 模型（3T tokens）上验证：MSA 在下游基准上与 GQA 表现相当，同时将 1M context 下的 per-token attention FLOPs 减少 28.4×
5. 配套 kernel 在 H800 上实现 14.2× prefill 和 7.6× decoding wall-clock speedup
6. 支持两种训练路线：从头稀疏预训练（MSA-PT）和从稠密 checkpoint 转换（MSA-CPT）

二、核心思想

问题定义

LLM 正从短对话转向长 horizon agentic workflows（数百个 interleaved reasoning 和 action steps），需要同时处理数十万到数百万 token。Softmax attention 的二次计算成本在部署规模下变得不可接受。

现有两条技术路线：

1. 混合架构：用 linear attention 或 sliding window 替换部分 softmax layers
2. 稀疏 softmax attention：直接对 softmax attention 本身做稀疏化

MSA 遵循 Occam’s razor——保留最必要组件，在 sparse softmax attention 范式下最大化复用现有软硬件基础设施。

解决方案概述

MSA 采用两阶段稀疏注意力架构：

Input X → Index Branch → Top-k block selection (per GQA group)
          ↓
Input X → Main Branch → Softmax attention over selected blocks only

Index Branch：引入一个 index query head（每个 GQA group 一个）和一个共享的 index key head，通过 block-level max-pooling 对 KV blocks 打分并选择 Top-k。local block 始终保留以保证训练稳定性。

Main Branch：在选中的 blocks 上执行标准 scaled dot-product attention。

训练策略：Top-k 选择不可微，通过 KL alignment loss 对齐 Index Branch 的分布与 Main Branch 在选中 token 上的分布。配合 stop-gradient 隔离辅助目标与 backbone。

三、技术架构

整体框架

MSA 建立在 Grouped Query Attention (GQA) 之上：

组件	参数量	功能
Main Branch Q/K/V	与 GQA 相同	标准注意力投影
Index Branch Q	$H_{kv} \times d_{\text{idx}}$	每个 GQA group 一个 index query head
Index Branch K	$1 \times d_{\text{idx}}$	全局共享 index key head
Output projection	与 GQA 相同	最终输出投影

核心公式

Causal Attention with GQA（Eq. 1）

$${\bm{o}}_{t}^{(h)} = \sum_{i \leq t} \alpha_{t,i}^{(h)} {\bm{v}}_{i}^{(h)}, \qquad \alpha_{t,i}^{(h)} = \frac{\exp\left(\langle{\bm{q}}_{t}^{(h)},{\bm{k}}_{i}^{(h)}\rangle/\sqrt{d_h}\right)}{\sum_{j \leq t}\exp\left(\langle{\bm{q}}_{t}^{(h)},{\bm{k}}_{j}^{(h)}\rangle/\sqrt{d_h}\right)} \tag{1}$$

GQA 将 $H_q$ 个 query heads 分为 $H_{kv}$ 个 groups，$G = H_q/H_{kv}$ 个相邻 query heads 共享一个 KV head。

Sparse Attention Two-Stage Formulation（Eq. 2）

$${\mathcal{I}}_i = \mathrm{Index}_\phi({\bm{q}}_i, {\bm{K}}_{\leq i}), \qquad {\bm{o}}_i = \mathrm{Attn}({\bm{q}}_i, {\bm{K}}[{\mathcal{I}}_i], {\bm{V}}[{\mathcal{I}}_i]) \tag{2}$$

Block Partition（Eq. 4）

$${\mathcal{B}}_b = \{(b{-}1)B_k+1,\dots,\min(bB_k,N)\}, \qquad b=1,\dots,B, \quad B=\lceil N/B_k\rceil \tag{4}$$

Index Branch Scoring（Eq. 5-7）

Index query/key projections：

$${\bm{Q}}^{\rm idx} = {\bm{X}}{\bm{W}}_q^{\rm idx} \in \mathbb{R}^{N\times H_{kv}\times d_{\rm idx}}, \qquad {\bm{K}}^{\rm idx} = {\bm{X}}{\bm{W}}_k^{\rm idx} \in \mathbb{R}^{N\times 1\times d_{\rm idx}} \tag{5}$$

Token-level scoring + block-level max pooling：

$$S^{\rm idx,(r)}_{i,j} = \frac{({\bm{Q}}^{\rm idx})^{(r)}_i ({\bm{K}}^{\rm idx})_j^\top}{\sqrt{d_{\rm idx}}}, \qquad M^{\rm idx,(r)}_{i,b} = \max_{\substack{j \in {\mathcal{B}}_b \\ j \leq i}} S^{\rm idx,(r)}_{i,j} \tag{6}$$

Top-k block selection（local block always included）：

$${\mathcal{I}}_i^{(r)} = \mathrm{TopK}_{b \in \{1,\dots,B\}}\!\big(M^{\rm idx,(r)}_{i,\cdot},\, k\big) \tag{7}$$

Main Branch Attention（Eq. 8）

$${\bm{O}}_i^{(h)} = \mathrm{softmax}\!\left(\frac{{\bm{Q}}_i^{(h)}\,\big({\bm{K}}^{(r)}[{\mathcal{I}}_i^{(r)}]\big)^\top}{\sqrt{d_h}}\right){\bm{V}}^{(r)}\!\big[{\mathcal{I}}_i^{(r)}\big] \tag{8}$$

每 query 注意力成本从 $O(N)$ 降至 $O(kB_k)$，随序列长度增加保持固定。

KL Alignment Loss（Eq. 9-10）

Index 和 Main Branch 在选中 token 集 ${\mathcal{I}}_{i,\mathrm{tok}}^{(r)}$ 上的分布：

$$P^{\rm idx,(r)}_{i,j} = \frac{\exp(S^{\rm idx,(r)}_{i,j})}{\sum_{u \in {\mathcal{I}}_{i,\mathrm{tok}}^{(r)}}\exp(S^{\rm idx,(r)}_{i,u})}, \qquad P^{(r)}_{i,j} = \frac{1}{G}\sum_{\ell \in \mathcal{H}_r}\frac{\exp(S^{(\ell)}_{i,j})}{\sum_{u \in {\mathcal{I}}_{i,\mathrm{tok}}^{(r)}}\exp(S^{(\ell)}_{i,u})} \tag{9}$$

KL 损失（teacher distribution detached）：

$$\mathcal{L}_{\rm KL} = \frac{1}{NH_{kv}}\sum_{i=1}^N\sum_{r=1}^{H_{kv}} D_{\mathrm{KL}}\big(\mathrm{stopgrad}(P^{(r)}_{i,\cdot}) \,\|\, P^{\rm idx,(r)}_{i,\cdot}\big) \tag{10}$$

Gradient Detach（Eq. 11）

$${\bm{Q}}^{\rm idx} = \mathrm{stopgrad}({\bm{X}}){\bm{W}}_q^{\rm idx}, \qquad {\bm{K}}^{\rm idx} = \mathrm{stopgrad}({\bm{X}}){\bm{W}}_k^{\rm idx} \tag{11}$$

隔离辅助目标与 backbone，$\mathcal{L}_{\rm KL}$ 仅更新 ${\bm{W}}_q^{\rm idx}$ 和 ${\bm{W}}_k^{\rm idx}$。

Computational Complexity（Eq. 12）

$$F_{\rm GQA}(N) = 2H_qd_hN^2, \qquad F_{\rm MSA}(N) = \underbrace{H_{kv}d_{\rm idx}N^2}_{\text{Index Branch}} + \underbrace{4H_qd_hNkB_k}_{\text{Main Branch}} \tag{12}$$

当 $kB_k \ll N$ 时，FLOPs 差距随 $N$ 增长而增大。

训练流程（Algorithm 1）

Algorithm 1: One MSA layer training forward
Input: hidden states X ∈ R^(N×d_model), block size B_k, selected blocks k

1: Q, K, V ← XW_q, XW_k, XW_v                    // standard GQA projections
2: Q_idx, K_idx ← stopgrad(X)W_q_idx, stopgrad(X)W_k_idx  // detached index projections
3: M_idx ← BlockMaxPool(Q_idx, K_idx, B_k)         // per-group causal block scores
4: I ← TopK(M_idx, k)                              // selected block indices (local always included)
5: O ← TopKAttn(Q, K, V, I)                        // attends to selected blocks
6: output ← OW_o                                    // final projection
7: L_KL ← KLdiv(Q_idx, K_idx, stopgrad(Q), stopgrad(K), I)  // over selected tokens
8: return output, L_KL

完整训练损失：$\mathcal{L} = \mathcal{L}_{\rm LM} + \lambda \sum_{\text{layers}} \mathcal{L}_{\rm KL}$

训练技术

技术	说明
KL Loss	对齐 Index Branch 与 Main Branch 在选中 token 上的分布
Gradient Detach	stop-gradient 隔离 Index Branch 输入，防止 KL 梯度影响 backbone
Indexer Warmup	两阶段训练：前 40B tokens 全注意力+KL loss 初始化 indexer，之后切换稀疏注意力
Local Block	每个 query 位置的 local block 始终被选中，防止退化选择

四、核心创新

创新点	说明	依据
Per-GQA-group independent block selection	每个 GQA group 独立选择 Top-k blocks，兼顾多组特异性和块级执行效率	Eq. (5)-(7): 单 index key head 共享，每 group 一个 index query head
Exp-free Top-k selection	跳过 softmax 的 max/exp/sum 步骤，raw scores 直接用于排序	Section 4.1: 在 $k=16$ 设置下比 torch.topk 快 5.1×
KV-outer sparse attention	选择 KV-outer 而非 Q-outer 迭代以最大化 arithmetic intensity	FLOPs/IO: Q-outer ≈ G vs KV-outer ≈ 2/3 $B_k$
Pre-scheduled tile chunking	解决 sink row 导致的热点 CTA 问题，将热门 tile 拆分给多个 CTA	Section 4.2: 避免 atomic updates
Two-phase forward with PDL	Split-K softmax normalization + Programmatic Dependent Launch 隐藏 inter-kernel 延迟	Section 4.2
Sparse KL loss fusion	跳过 KL loss forward pass，在主 pass 中直接 emit LSE scalars	Section 4.3
Dynamic load balancing	Persistent grid + global atomic counter 处理 variable-length sequences 下的 per-tile work 差异	Section 4.3

五、GPU Kernel 设计

4.1 Index & TopK Kernel

Exp-free selection：由于 softmax 是 order-preserving，直接用 raw scores 进行 Top-k 选择，跳过 softmax 计算。

Per-thread register top-k：

• Block size $B_k = 128$，selection $k = 16$
• Warp 的 32 lanes 各流式处理 1/32 stride 的 input row
• 每个 lane 在 shared memory 中维护 $k$-element min-heap
• Heap root 缓存在 register 中，insertions 使用 deferred writes
• 最后 $k$-round shuffle merge 合并 32 个局部 Top-k 结果

Benchmark（Table 1）：

Seq Len N	Blocks B	k	torch.topk	TileLang	Ours	vs. torch	vs. TileLang
128K	1024	16	3970 μs	2864 μs	779 μs	5.1×	3.7×
128K	1024	32	5378 μs	3630 μs	1991 μs	2.7×	1.8×
512K	4096	16	33810 μs	17779 μs	7880 μs	4.3×	2.3×
512K	8192	32	57659 μs	26100 μs	21326 μs	2.7×	1.2×

4.2 Sparse Attention Kernel

Q-outer vs KV-outer 分析：

• Q-outer: FLOPs/IO ≈ $G$（GQA ratio）
• KV-outer: FLOPs/IO ≈ $\frac{2}{3}B_k$

由于 $\frac{2}{3}B_k \gg G$ in practice，选择 KV-outer iteration with Q gather。

Two-phase forward：

1. Phase 1: Attention kernel 写入 per-partial outputs 到 $\mathbf{O}_{\text{buf}} \in \mathbb{R}^{k \times n \times H_q \times d}$
2. Phase 2: Combine kernel 计算 split-K softmax normalization

4.3 Sparse KL Loss Fusion

• 跳过 KL loss forward pass，直接在 main pass 中 emit LSE scalars
• Backward kernel 直接从 global memory 加载标量到 softmax
• 每个 block 保存 LSE，对 top-k blocks 做 reduction 得到 $\mathrm{LSE}_{\rm idx}$

六、实验结果

实验设置

配置项	值
Model	109B MoE (41 layers: 3 dense + 38 MoE)
Parameters	109B total, 6B active per token
Vocabulary	200K tokens
Hidden size	3072
Attention	64 query heads, 4 KV heads ($G=16$), $d_h=128$, RoPE=64
MoE	128 routed experts + 1 shared expert, top-4 routing
Training budget	3T tokens
Block size	$B_k = 128$
Selected blocks	$k = 16$ (2,048 KV tokens per query)
Hardware	H800 for kernel benchmarks; multi-GPU for training
Baselines	FlashAttention (Full GQA), SageAttention, SpargeAttention, SVG2, Jenga

主结果（Table 2）

General Knowledge & Reasoning：

Benchmark	Full	MSA-PT	MSA-CPT
MMLU	67.0	67.2	66.8
MMLU-Pro	38.5	38.8	39.1
BBH	67.7	66.6	66.1
GPQA Hard	25.9	26.3	26.3
ARC Challenge	82.7	82.5	82.9
TriviaQA	66.0	65.5	67.7
WinoGrande	58.3	60.9	62.0

Math：

Benchmark	Full	MSA-PT	MSA-CPT
GSM8K	76.2	77.7	73.7
MGSM	44.1	46.0	44.2
MathVista	43.8	46.8	44.5
OlymMATH Easy	23.0	26.0	22.0

Code：

Benchmark	Full	MSA-PT	MSA-CPT
HumanEval	61.0	64.0	57.9
EvalPlus	59.4	61.8	60.0
BigCodeBench	44.8	44.0	45.7
MBPP P@10	82.1	81.6	81.1

Long-context Retrieval：

Benchmark	Full	MSA-PT	MSA-CPT
RULER-8K	79.8	84.2	77.2
RULER-32K	75.0	77.5	75.7

MSA-PT 在多数数学、图像、视频和长上下文检索基准上表现最强，表明原生稀疏预训练可以让模型表示适应稀疏注意力模式。MSA-CPT 在文本、代码和 PPL 评估上更接近 Full Attention，是已有稠密 checkpoint 的实用转换路线。

长上下文扩展（Table 3）

从 MSA-CPT checkpoint 继续约 140B tokens 的长上下文训练后：

Benchmark	Full	MSA-CPT (extended)	Δ
HELMET-128K Overall	46.53	45.93	-0.60
HELMET ICL	70.40	72.80	+2.40
RULER-128K Overall	72.00	72.12	+0.12
RULER MK/MQ/MV	96.63	98.87	+2.24

每个 query 仅 attends to $kB_k = 16 \times 128 = 2,048$ KV tokens，MSA-CPT 仍保持接近 Full Attention 的长上下文能力。

效率对比（Figure 4）

• Prefill speedup: 在 1M context 下达到 14.2×（H800）
• Decoding speedup: 在 1M context 下达到 7.6×（H800）

训练效率消融（Table 4）

方法	Token 数	Prefill Speedup	Decode Speedup
FlashAttention	1M	1.00×	1.00×
FlashAttention	32K	1.00×	1.00×
MSA	1M	14.2×	7.6×
MSA	32K	~2.5×	~1.8×

性能开销（Table 5）

MSA 引入的额外计算开销极小：

指标	变化
Index computation overhead	< 0.1% of total FLOPs
Memory overhead	两个额外的投影矩阵 $W_q^{\rm idx}, W_k^{\rm idx}$
KV cache reduction	从 $N$ 降至 $kB_k = 2048$ tokens per query

七、消融实验与可视化分析

训练动态可视化

Index selection pattern（Figure 5）：不同 GQA group 选择不同 long-range stripes，同时共享 local diagonal 和 sink column 模式，说明 learned indexer 捕获了 group-specific 的稀疏注意力模式而非坍缩为单一全局模式。

Attention sink（Figure 6）：即使没有显式强制 indexer 选择第一个 KV block，learned Index Branch 自然地跨所有层和 heads 对初始 block 分配高选择概率。

Indexer 训练信号消融（Appendix B.2, Figure 7）

配置	短上下文	长上下文	平衡
LM Loss only	好	差	❌
KL Loss only	差	好	❌
LM + KL Loss	好	好	✅

Gradient Detach 消融（Appendix B.3, Figures 8-9）

Detaching KL gradient 从 backbone 防止 auxiliary loss 污染主训练目标，在通用基准上保持一致的性能。

Indexer Warmup（Appendix B.4, Figure 11）

Index warmup 在训练早期显著改善收敛行为，使 indexer 在全注意力阶段学习合理的 block 选择模式后再切换到稀疏模式。

Block Size 消融（Appendix C.1, Table 4）

Block Size	TAU2 PPL	RULER-8K	RULER-32K
32	1.176	72.5	66.1
64	1.176	72.8	65.3
128	1.176	73.8	64.6

较大 block size 对模型质量影响有限，同时提升 kernel 效率。

Forced Sink & Local Selection（Appendix C.2, Table 5）

移除 forced first-block 和 fixed local window 后，模型仍然自然学习到这些模式，说明不需要硬编码选择规则。

Index Branch Value Head（Appendix C.3, Table 6）

一旦使用 indexer warmup，index value head 不再是必需的。移除后在大多数基准上没有系统性退化。

Learnable Sink（Appendix B.3, Figures 12-13）

引入 GPT-OSS 风格的 learnable sink 参数后，部分 heads 的 sink-like 注意力被 learnable sink 吸收，部分仍集中在第一个 token，但对下游 agent-oriented 评估的 perplexity 没有一致改善。

Sliding-Window Ablation（Appendix B.4, Figure 14）

与 FLOPs-matched sliding window baseline 相比，MSA 在多项 agent-oriented 评估上达到更低 perplexity，说明选择性稀疏优于固定窗口。

八、与相关工作对比

方法	稀疏时机	是否训练	Selector	粒度	GQA 支持
MSA (Ours)	Native	是	Per-GQA-group Top-k	Block	✅
NSA	Native	是	3 parallel branches	Block	❌ (MQA/MHA)
MoBA	Native	是	Block-averaged keys	Large block	✅
DSA	Native	是	Multi-head ReLU lightning	Token-level	❌ (MLA)
InfLLM-V2	Inference	否	Parameter-free	Block+SW	❓
H2O	Inference	否	Accumulated stats	Token	❓
SnapKV	Inference	否	Attention stats	Token	❓
Quest	Inference	否	Page-level importance	Page	❓
Sliding Window	Fixed	否	N/A	Fixed window	✅

MSA 与相近工作的两个区别轴：

1. Per-GQA-group Top-k sharing combined with block-level selection：多组块级稀疏检索，同时保持 KV reads 连续性
2. Indexed by a lightweight dot-product scorer：相比 MoBA 的 block-averaged keys 更精细

九、总结

核心贡献

1. MSA 机制：minimal、scalable、accelerated 的 blockwise sparse attention，支持从头训练和 near-lossless 转换
2. Kernel co-design：exp-free Top-k + KV-outer sparse attention + two-phase forward，将理论 FLOP 节省转化为实际 wall-clock 加速
3. 109B MoE 大规模验证：native multimodal 训练 3T tokens，MSA-PT 和 MSA-CPT 两种训练路线均在下游基准上与 GQA 持平
4. 28.4× FLOPs 减少：在 1M context 下，配合 kernel 实现 14.2× prefill 和 7.6× decoding 加速

局限性

1. 长上下文检索仍有残差 gap：与 Full Attention 相比，某些长上下文检索任务（如 HELMET Rerank/RAG -2.10）仍存在性能差距
2. Block size 选择依赖 workload：虽然消融显示 block size 影响有限，但最优值可能因具体任务而异
3. 当前仅验证 pretraining：RL post-training 和 agentic deployment 场景下的适用性尚待探索
4. Index Branch 增加少量 FLOPs：虽然占比 < 0.1%，但在极端资源受限场景下仍需权衡

未来方向

1. 通过更长稀疏训练、更大 selection budget 或更丰富的 indexer scoring function 缩小残差长上下文检索 gap
2. 将 selector-only 设计扩展到 RL post-training 和 agentic deployment
3. 探索不同 GQA ratio 和 block size 组合的自动化搜索

十、参考资源

• 论文: arXiv:2606.13392
• 代码: https://github.com/MiniMax-AI/MSA
• 模型: https://huggingface.co/MiniMaxAI/MiniMax-M3
• FlashAttention: Dao et al., 2022
• GQA: Ainslie et al., 2023
• NSA: Yuan et al., 2025
• MoBA: Lu et al., 2025
• DSA: DeepSeek-AI et al., 2025
• SnapKV: Li et al., 2024
• H2O: Zhang et al., 2023
• RULER: Hsieh et al., 2024
• HELMET: Yen et al., 2025
• VBench: Huang et al., 2024

附图索引

编号	文件名	说明
Figure 1	figure-1-msa-architecture.png	MSA 架构总览：Index Branch + Main Branch
Figure 2	figure-2-training-dynamics.png	Pretraining dynamics：LM loss 和 gradient norm
Figure 3	figure-3-cpt-dynamics.png	Sparse CPT dynamics：KL loss 和 selection recall
Figure 4	figure-4-efficiency.png	效率对比：FLOPs 减少和 runtime speedup
Figure 5	figure-5-index-visualization.png	Index selection probability across layers and GQA groups
Figure 6	figure-6-attention-sink.png	Attention sink effect across heads and layers
Figure 7	figure-7-indexer-training-signals.png	Indexer training signal ablation (LM/LM+KL)
Figure 8	figure-8-gradient-detach-loss.png	Gradient detach effect on LM loss and gradient norm
Figure 9	figure-9-gradient-detach-benchmarks.png	Gradient detach effect on benchmark scores
Figure 10	figure-10-per-layer-entropy.png	Per-layer entropy of Main Branch attention during training
Figure 11	figure-11-index-warmup.png	Index warmup effect on training curves
Figure 12	figure-12-learnable-sink-vis.png	Learnable sink vs first token attention distribution
Figure 13	figure-13-learnable-sink-results.png	Perplexity comparison with/without learnable sink
Figure 14	figure-14-ablation-swa-ppl.png	MSA vs FLOPs-matched sliding window ablation

附表格索引

编号	说明
Table 1	Top-k latency benchmark (torch.topk vs TileLang vs Ours)
Table 2	Main evaluation results: Full vs MSA-PT vs MSA-CPT
Table 3	Long-context extension on HELMET and RULER
Table 4	Block size ablation (32/64/128)
Table 5	Forced sink and local selection ablation
Table 6	Index Branch value head ablation (with-value vs no-value)