OC

Knowledge OS
百万 Token 不该让每个词都互相看:Flash-MSA 补上了稀疏注意力的训练内核
科技 · 2026-07-13 · AI 基础设施 · 阅读 1

百万 Token 不该让每个词都互相看:Flash-MSA 补上了稀疏注意力的训练内核

作者:林岚

作者:林岚

据开发者 Ganesh Nanduru 发布的 Flash-MSA,他用 CuTeDSL 为 Hopper 和 Blackwell GPU 编写了一套 MiniMax Sparse Attention 训练内核,试图解决一个很具体的问题:稀疏注意力已经能加快推理,但公开、可用的高性能训练实现仍然很少。

一句话结论: Flash-MSA 不是让普通显卡突然训练百万 Token 模型,而是把“只看重要上下文”的算法真正翻译成 GPU 能高效执行的训练内核,距离成熟通用方案仍有工程路要走。

标准注意力的麻烦是,文本长度翻倍,词与词之间需要计算的关系接近四倍。到了百万 Token,不可能让每个查询都完整扫描全部历史。MiniMax 的 MSA 先用轻量索引分支给上下文块打分,再只让主注意力处理最相关的一小部分块。它像在巨型档案馆里先查目录,而不是每问一句都把所有文件重新翻一遍。

难点是“稀疏”不自动等于“快”。GPU 擅长规则、连续的大矩阵运算;随机挑选小块可能把节省的计算又浪费在搬数据和调度上。Flash-MSA 采用 128 Token 的块级选择,并缓存前向阶段选出的块索引,让主注意力和大部分反向传播不再随上下文长度平方增长。

稠密注意力与块稀疏注意力示意图

作者还把代理注意力与主注意力的反向过程融合起来。原本 KL 损失似乎需要把两套完整概率分布都放进显存,他推导出代理分数的梯度可以化简为“代理概率减主概率”,从而避免完整展开。这类细节才是内核工作的价值:论文里一句公式,落到 GPU 上可能决定它是能跑,还是只在示意图里漂亮。

不过,别把“百万 Token”理解成已经完成百万长度的端到端训练验证。当前公开页面展示了若干配置下与 PyTorch eager 实现的余弦相似度,最低约 0.9983;作者同时坦言,融合反向内核受寄存器和共享内存压力限制,H100/B200 上只能做到每个 SM 一个 CTA,理论占用率约 12.5%。上下文并行也还在计划中,而长上下文训练没有它很快就会耗尽单卡显存。

本站编辑林岚认为,这个项目的新闻价值不在“又一个 10 倍加速”,而在训练基础设施开始追上稀疏注意力论文。MiniMax 官方论文给出了 MSA 在百万上下文推理上的收益;Flash-MSA 是独立开发者对训练侧的补课,并非 MiniMax 官方实现,两者的数字不能混着引用。

关键事实

  • Flash-MSA 面向 NVIDIA Hopper 与 Blackwell,使用 CuTeDSL。
  • 它实现的是 MiniMax Sparse Attention 的非官方训练内核。
  • 正确性测试使用 BF16,多组前向与梯度余弦相似度约为 0.9983 以上。
  • 反向内核占用率、上下文并行和更多模型验证仍未完成。

OC 判断

长上下文竞争正在从“窗口标多大”进入“每个 Token 花多少钱”。真正有意义的进展,是算法、内核和分布式训练三层同时成立;Flash-MSA 已经走通了关键一层,但还不是最终答案。

为什么重要

  • 对开发者:稀疏注意力的性能取决于具体内核,不能只看理论复杂度。
  • 对企业:百万上下文能否落地,成本和吞吐往往比模型页面上的窗口数字更关键。
  • 对用户:更长上下文未来可能更便宜,但不代表模型会自动记住或正确使用全部内容。

参考来源

相关阅读

基于标题、摘要和正文内容自动匹配。

更多科技

评论

围绕这篇文章补充信息、提出问题或分享观察。

0
暂无评论。

发表评论