MiniMax稀疏注意力

超长上下文能力对于前沿大语言模型来说变得不可或缺：代理工作流程、存储库规模的代码推理和持久内存都需要模型共同关注数十万到数百万个词元，但 softmax 注意力的二次成本使得这在部署规模上难以维持。我们引入了 MiniMax Sparse Attention (MSA)，这是一种基于分组查询注意力 (GQA) 的块式稀疏注意力。轻量级索引分支对键值块进行评分，并为每个 GQA 组独立选择 Top-k 子集，从而实现组特定的稀疏检索，同时保持高效的块级执行；然后，主分支仅对选定的块执行精确的块稀疏注意力。 MSA 围绕简单性和可扩展性原则进行设计，经过精心简化，使其能够在各种 GPU 上轻松高效地部署。为了将稀疏性转化为实际的加速，我们共同设计了 MSA 和 GPU 执行路径，该路径使用无 exp 的 Top-k 选择和 KV 外部稀疏注意力来提高块粒度访问下的张量核心利用率。在具有本机多模态训练的 109B 参数模型上，MSA 的性能与 GQA 相当，同时在 1M 上下文中将每个标记的注意力计算量减少了 28.4 倍。与我们共同设计的内核配合使用，MSA 在 H800 上实现了 14.2 倍预填充和 7.6 倍解码挂钟加速。我们的推理内核位于：https://github.com/MiniMax-AI/MSA。由 MSA 支持的生产级本机多模态模型已公开发布于：https://huggingface.co/MiniMaxAI/MiniMax-M3。