Flash-MSA 出手：把百万 Token 训练啃下来

Flash-MSA 出手：稀疏注意力 Kernel 开始瞄准百万 Token 训练

长上下文训练一直卡在同一个位置——注意力。序列长度线性拉长，计算和显存却是平方级涨，这事从 Transformer 诞生那天就没变过。过去两年大家的解法基本分两派：一派靠 FlashAttention 系列在 IO 层面榨性能，另一派靠稀疏化直接砍计算量。今年 7 月，一个叫 Flash-MSA（Flash Multi-Sparse Attention，也常被引用为 Flash Sparse Attention / FSA）的项目把这两条路合到一起，直接把矛头对准了百万 Token 训练场景。

这不是又一篇纸面论文。项目在 GitHub 上放出了完整的 Triton 实现，配套的性能数据也相当扎眼：在 H20 上相较 DeepSeek 的 NSA（Native Sparse Attention）最高 3.5 倍加速，H200 上 2.9 倍；相较全注意力最高 6.4 倍。放在 Llama3-8B、Qwen3-14B、Qwen2.5-32B 上做真实训练，端到端也能拿到平均 1.86 倍的加速。

NSA 的坑，Flash-MSA 是怎么填的

要理解 Flash-MSA 的意义，得先说清楚它在补谁的窟窿。

DeepSeek 今年年初推的 NSA（Native Sparse Attention）是稀疏注意力路线上一个里程碑——它把稀疏做进了训练阶段，而不是像早期方案那样只在推理时凑合用。NSA 的核心 idea 是把 Attention 拆成三条支路：压缩、选择、滑窗，其中「选择注意力」（selected attention）负责挑最相关的 KV block 去算，是稀疏收益的主要来源。

问题出在 Kernel 层面。NSA 的选择注意力用了一个**查询分组（query-grouping）**策略——外层循环遍历查询 token，内层循环遍历 KV block。这套写法的高效性完全依赖 GQA 的分组大小。GQA 组大，一次 matmul 塞满 Tensor Core，很爽；GQA 组一小，麻烦就来了。

现代 LLM 的 GQA 分组大多是 4、甚至更小（Llama3、Qwen3 都在这个量级）。这时候 NSA 为了满足 Hopper 架构 matmul 对操作数形状的硬件要求（每维至少 8），必须对 query head 做 padding。padding 出来的这些计算是纯粹的浪费——多读一次显存、多算一遍无效乘加。

结果就是一个挺尴尬的现象：在很多现代 LLM 的配置下，未经优化的 NSA 甚至跑不过全注意力。稀疏做了半天，反而更慢，这就是 Flash-MSA 团队盯上的空档。

三个 Kernel，一次解耦

Flash-MSA 的做法不修改 NSA 算法本身，而是把选择注意力的计算拆成三个独立 Kernel：

• 主 Kernel：把注意到同一个 KV block 的所有查询 token 批量化，一次算完，中间结果写到 buffer。这是核心——原本每个 query 单独处理导致的矩阵形状不匹配问题，被换到了「按 KV block 聚合 query」的维度上，matmul 形状自然凑齐了。
• Reduction Kernel：把 buffer 里的部分结果按 query token 累加回去。
• Online Softmax Kernel：单独处理 softmax 的在线统计量，避免和主计算路径纠缠。

这套结构最有意思的地方是它换了循环顺序的语义。原来的 NSA 是「query 找 KV」，Flash-MSA 变成「KV block 收 query」。听起来只是外内层调换，但在 GQA 组小的时候，一个 KV block 通常会被多个 query batch 引用，这一 batch 就把 Tensor Core 填满了，padding 需求消失。

代价当然有：引入了非连续内存访问和额外的 buffer 内存。团队做了「早期返回」和「紧凑 buffer 管理」两个优化——按论文的估算，在 64K token、T=16、d=128 的配置下，buffer 大约占 1GB HBM，可以接受。

值得一提的是 Flash-MSA 是条件切换的：GQA 分组 ≥ 8 的时候，原版 NSA 已经够高效，Flash-MSA 会直接 fallback 回去。这种「知道自己什么时候该退场」的设计在 Kernel 库里其实不多见，工程上的分寸感做得比较好。

数据说话：小 GQA 是它的主场

把关键数字拉出来对比一下：

| 配置 | 相较 NSA | 相较全注意力 | |------|---------|-------------| | H20 峰值 | 3.5x | 6.4x | | H20 平均 | 1.8x | 2.4x | | H200 峰值 | 2.9x | 4.9x | | H200 平均 | 1.4x | 2.3x |

分解到正反向：

• Forward：相较 NSA 最高 2.36 倍，平均 1.62 倍
• Backward：最高 4.32 倍，平均 2.59 倍

反向传播的加速比正向还大，这一点挺关键。稀疏注意力过去在训练场景没大规模铺开，一个很大的原因就是反向路径涉及索引张量的重复计算，很容易变成瓶颈。Flash-MSA 在反向阶段跳过了这部分开销，等于把稀疏注意力从「推理友好」拉到了「训练友好」。

端到端训练测试也做了——Llama3-8B、Qwen3-14B、Qwen2.5-32B 三个尺度：

• 相较 NSA：最高 1.25 倍，平均 1.09 倍
• 相较全注意力：最高 2.47 倍，平均 1.86 倍

1.86 倍是什么概念？如果你原本要花 30 天预训练一个长上下文模型，现在 16 天就能跑完。这在算力紧张的今天，比参数量拨零点几个百分点实在得多。

为什么这事值得开发者盯一眼

把镜头拉远一点看。这一年 AI 圈两条主线特别清楚：一是上下文军备竞赛——Gemini 2 百万 Token、Claude 的记忆机制、GPT-5 的长文档处理，都在把窗口往外顶；二是Kernel 层的效率革命——FlashAttention-3、FlashInfer、SGLang 的自定义 kernel、vLLM 的 PagedAttention，每一波都在证明「同样的算法、换个写法就能快一倍」。

Flash-MSA 站在这两条线的交叉口。它不是提出新的稀疏算法，也不是发明新的硬件，而是把已有的 NSA 算法「翻译」成一份 GPU 能吃得下的代码。这种工作看起来「不性感」，但对于真正在训练侧堆算力的团队，价值远大于又一个 SOTA 分数。

几个具体的场景：

• 百万 Token 预训练：过去在这个长度做完整训练几乎是奢侈品，反向传播的显存和时间都劝退。Flash-MSA 反向 2.59 倍的平均加速，直接把这道门槛降下来。
• 长上下文 SFT / RLHF：法律、代码库、科研文献这类场景，本身就需要模型看得下几十万 token 的输入。这些团队用不起全注意力，也不敢用没优化好的稀疏方案，Flash-MSA 提供了一个可落地的中间选项。
• 推理加速的溢出效应：训练用什么 Kernel，推理往往顺手用同一套。虽然 Flash-MSA 主要打训练场景，但选择注意力的 Kernel 优化对长上下文推理同样受益。

一些没解决的问题

泼点冷水。Flash-MSA 不是银弹：

1. 只针对 NSA 这一族。如果你用的是 MoBA、Longformer、BigBird 那种不同思路的稀疏方案，Flash-MSA 帮不上忙。稀疏注意力现在门派太多，Kernel 层的碎片化只会更严重。
2. 依赖 Triton。Triton 生态在成熟，但对 CUDA 深度调优选手来说仍然是黑盒。极致场景下自己写 CUDA 可能还是更快。
3. buffer 开销随 batch 和序列长度增长。1GB 在 64K 场景可以接受，128K、256K 呢？团队没给出更极端配置的数据。
4. 和 FlashAttention-3 的组合关系还没完全铺开。理论上 Flash-MSA 的思路可以和 FA3 的 IO 优化叠加，但目前实现里两者不是同一套代码路径。

从项目的迭代节奏看，团队更新挺勤，Issue 里能看到已经有人在跑 Qwen3-32B 长上下文微调的实际反馈，性能数据大体对得上论文。

更长远的信号

把 Flash-MSA 放到更大的背景里看，它其实印证了一个正在发生的转向：注意力机制的创新，主战场正在从「算法」滑到「Kernel」。

FlashAttention 之前，大家比谁的注意力算法更聪明；FlashAttention 之后，大家比谁的 Kernel 更能榨硬件。NSA 提出稀疏算法半年就被 Flash-MSA 追着优化 Kernel——这个反馈周期比过去缩短了一个量级。算法-硬件-系统三层的耦合在变紧，训练框架也在被迫向 Kernel 库的方向长。

对个人开发者的启示更直接：如果你在做长上下文相关的工作，光跟算法论文已经不够了，得学会看 Triton 代码、读 Kernel profile。这不是危言耸听，是接下来两年拿到工程 offer 的硬门槛。

OpenAI Hub 目前接入的模型里，Qwen3、DeepSeek 系列都属于 GQA 组较小的架构，如果你在这些模型上做长上下文测试，可以直接调 API 感受长窗口输入的实际表现，Flash-MSA 这类底层优化最终会以更快的推理响应和更长的上下文支持体现出来。

参考来源

• Relaxed-System-Lab/Flash-Sparse-Attention (GitHub) — Flash-MSA / FSA 官方仓库，含 Triton Kernel 实现与复现脚本
• 深度解析 DeepSeek NSA 原生稀疏注意力 idea 和 kernel (知乎) — NSA 算法与 Kernel 细节的中文长文解析，理解 Flash-MSA 前置知识