百万 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 已经走通了关键一层，但还不是最终答案。
为什么重要
- 对开发者：稀疏注意力的性能取决于具体内核，不能只看理论复杂度。
- 对企业：百万上下文能否落地，成本和吞吐往往比模型页面上的窗口数字更关键。
- 对用户：更长上下文未来可能更便宜，但不代表模型会自动记住或正确使用全部内容。
评论
围绕这篇文章补充信息、提出问题或分享观察。