MiniMax Sparse Attention(MSA)块状稀疏注意力
MiniMax这个稀疏注意力把长上下文推理计算砍掉28倍,而且直接开源了高效CUDA kernel和模型,做agent和代码仓库级推理的团队可以赶紧试试。
MiniMax 提出块状稀疏注意力 MSA,基于 GQA 构建。轻量级 Index Branch 为每个 GQA 组独立选择 Top‑k KV 块,Main Branch 仅对选中块执行精确块稀疏注意力。在 109B 参数多模态模型上,MSA 与 GQA 性能持平,1M 上下文下每 token 注意力计算减少 28.4 倍。配合协同设计的 GPU 内核,H800 上实现 14.2 倍 prefill 和 7.6 倍 decoding 端到端加速。推理内核与基于 MSA 的多模态模型已公开发布。
超长上下文能力正成为前沿大语言模型不可或缺的能力:智能体工作流、仓库级代码推理以及持久记忆都要求模型能够联合关注数十万到数百万个模型token,然而softmax注意力的二次方开销使得这在规模化部署中难以为继。我们提出MiniMax稀疏注意力(MSA),这是一种基于分组查询注意力(GQA)构建的块级稀疏注意力机制。一个轻量级的索引分支对键值块进行评分,并为每个GQA分组独立选择Top-k子集,从而在保持高效块级执行的同时实现分组特定的稀疏检索;主分支则仅对所选块执行精确的块级稀疏注意力。MSA的设计遵循简洁与可扩展性原则,有意精简结构,使其能够在广泛的GPU上高效部署。为了将稀疏性转化为实际的速度提升,我们与MSA协同设计了GPU执行路径,该路径使用无指数Top-k选择和外层键值稀疏注意力,以在块粒度访问下提升张量核心利用率。在具有原生多模态训练的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。