【第666期】MiniMax Sparse Attention-算力缩减28倍的MSA架构Seventy3

【第666期】MiniMax Sparse Attention-算力缩减28倍的MSA架构

24分钟 ·
播放数12
·
评论数0

Seventy3 是一档借助 NotebookLM 解读前沿论文的播客——不是念摘要,也不是泛泛而谈,而是让 AI 帮你把论文掰开揉碎,说人话。我们蹲在人工智能、大模型、机器人算法、crypto 这几个领域,每期挑一篇值得关注的工作,用对话的方式聊给你听。你可以在通勤路上听,做实验的时候听,也可以当背景音放着,让新知识自然地长进脑子里。

如果你正在做有意思的研究,想让更多人看到——不管你是博士生、独立研究者还是实验室的博士后——把论文发给我们,我们帮你用 AI 做一期深度解读,让你的工作被更多同路人发现。

联系小助手微信:seventy3_podcast(加群/投稿论文)

今天的这篇的主题是:

MiniMax Sparse Attention

Summary

超长上下文能力正在成为前沿大语言模型(LLM)不可或缺的能力:智能体工作流(agentic workflows)、代码仓库级别的代码推理(repository-scale code reasoning)以及持久化记忆(persistent memory)都要求模型能够联合关注数十万到数百万级别的 token。然而,softmax 注意力机制的二次计算复杂度,使得这种能力在实际部署规模下难以实现。

我们提出了 MiniMax 稀疏注意力(MiniMax Sparse Attention, MSA),这是一种基于分组查询注意力(Grouped Query Attention, GQA)构建的块级稀疏注意力机制。MSA 包含一个轻量级的索引分支(Index Branch),用于对键值(key-value)块进行评分,并为每个 GQA 组独立选择一个 Top-k 子集,从而实现针对不同组的稀疏检索,同时保持高效的块级执行;随后,**主分支(Main Branch)**仅在被选中的块上执行精确的块稀疏注意力计算。

MSA 的设计遵循简单性与可扩展性原则,经过刻意精简,使其能够在广泛的 GPU 平台上直接高效部署。为了将稀疏性转化为实际的速度提升,我们进一步针对 MSA 设计了专用的 GPU 执行路径:该路径采用无指数运算(exp-free)的 Top-k 选择机制以及 KV 外置稀疏注意力(KV-outer sparse attention),在块粒度访问模式下提升张量核心(tensor core)的利用率。

在一个具有原生多模态训练能力的 109B 参数模型上,MSA 的性能与 GQA 基本相当,同时在 100 万 token 上下文长度下,将每个 token 的注意力计算量降低了 28.4 倍。结合我们共同设计的 GPU 内核(kernel),MSA 在 H800 GPU 上实现了 14.2 倍的预填充(prefill)加速以及 7.6 倍的解码(decoding)端到端速度提升

原文链接:arxiv.org