Command Palette
Search for a command to run...
MiniMax 稀疏注意力
MiniMax 稀疏注意力
摘要
超长上下文能力正成为前沿大语言模型不可或缺的一部分：agent 工作流、仓库级代码推理与持久化记忆均要求模型对数十万至数百万个 token 进行联合注意力计算，然而 softmax 注意力的二次方计算成本使得这在部署规模下难以实现。我们提出了 MiniMax 稀疏注意力（MSA），这是一种基于分组查询注意力（GQA）的块级稀疏注意力机制。一个轻量级的索引分支对键值块进行评分，并为每个 GQA 组独立选择一个 Top-k 子集，在保持高效块级执行的同时实现组特定的稀疏检索；主分支随后仅对选定的块执行精确的块稀疏注意力计算。MSA 围绕简洁性与可扩展性原则设计，经过刻意精简，使其能够轻松高效地部署在广泛的 GPU 硬件上。为了将稀疏性转化为实际的加速效果，我们将 MSA 与一条 GPU 执行路径协同设计，该路径采用免指数运算的 Top-k 选择和 KV-outer 稀疏注意力，以提升块粒度访问下的张量核心利用率。在具有原生多模态训练的 109B 参数模型上，MSA 的性能与 GQA 相当，同时在 1M 上下文长度下将每个 token 的注意力计算量降低了 28.4 倍。结合我们协同设计的内核，MSA 在 H800 上实现了 14.2 倍的预填充加速和 7.6 倍的解码实际运行时间加速。我们的推理内核已开源，地址为：https://github.com/MiniMax-AI/MSA。一款由 MSA 驱动的生产级原生多模态模型已在以下地址公开发布：https://huggingface.co/MiniMaxAI/MiniMax-M3。
一句话总结
MiniMax Sparse Attention (MSA) 是 Grouped Query Attention 的一种变体，采用轻量级 Index Branch 独立为每个组选择 Top-k key-value blocks。该方法实现了精确的 block-sparse 计算，在保持与标准方法相当性能的同时，将 100 万 tokens 上下文下的 per-token 计算量降低 28.4 倍。此外，通过结合 exp-free Top-k 选择与 KV-outer sparse attention 的协同设计 kernel，在 H800 GPU 上实现了 14.2 倍的 prefill 加速与 7.6 倍的 decoding 加速。
核心贡献
- MiniMax Sparse Attention (MSA) 是一种基于 Grouped Query Attention 的块级稀疏注意力机制。该机制在计算精确的 block-sparse 注意力之前，利用轻量级 index branch 独立评分并为每个 GQA 组选择 top-k 个 key-value blocks。
- 通过协同设计的 GPU 执行路径实现 exp-free top-k 选择与 KV-outer sparse attention，从而在块粒度内存访问模式下最大化 tensor-core 利用率。
- 在 1090 亿参数多模态模型上的评估表明，该方法在 100 万 token 上下文下将 per-token 注意力计算量降低 28.4 倍，同时性能与标准 Grouped Query Attention 保持一致。在 H800 硬件上，实现了 14.2 倍的 prefill 与 7.6 倍的 decoding 实际运行时间加速。
引言
基于 Transformer 的语言模型在进行长上下文建模时，需要高效的注意力机制来缓解密集 softmax 注意力带来的二次方计算与内存开销。现有方法通常采用线性或循环替代方案，应用固定的与内容无关的稀疏模式，或实现自适应稀疏化。然而，自适应稀疏化往往继承全注意力训练的开销，或面临内存访问碎片化与推理 kernel 未优化的问题。该方法采用基于每个 GQA 组的 Top-k 共享策略，并结合块级选择，在保持自适应上下文感知能力的同时，确保 KV cache 的连续读取。此外，通过调整 FlashAttention 算法骨架中的循环顺序以适配该访问模式，进一步加速了整体框架，成功将理论上的 FLOP 降低转化为可测量的实际运行时间加速。
实验
两项 109B 规模的实验验证了通过从头训练或从全注意力 checkpoint 继续预训练，将密集注意力替换为稀疏机制的可行性。原生稀疏方法表明，模型能够在无需硬编码约束的情况下稳定调整表示，以学习关键的注意力结构；而继续预训练路线则验证了从密集 checkpoint 进行转换的实用且稳定的路径。尽管 key-value token 预算严格，两种方法在语言、多模态及长上下文基准测试中均保持了具有竞争力的性能。最终，该稀疏架构提供了显著的计算效率与持久的长上下文能力，确立了其作为密集注意力可扩展替代方案的地位。