在NVIDIA Blackwell上优化MiniMax M3稀疏注意力

Fireworks AI团队在NVIDIA Blackwell（SM100）GPU上为MiniMax M3稀疏注意力开发了一款高效内核。长上下文推理中，注意力机制主导计算和内存开销。MiniMax M3采用的稀疏注意力让每个查询只关注最相关的16个128-token KV块，从而降低开销。然而，稀疏选择依赖于数据且随请求变化，导致不规则访存模式，常常抵消理论加速。

Fireworks AI的性能团队构建的Blackwell内核通过KV-固定（KV-stationary）执行路径弥补了这一差距。该思路源自MiniMax团队的稀疏注意力论文：在外循环中加载每个选中的KV块一次，然后在内循环中处理所有选择该块的查询。基于这一调度，优化聚焦于减少聚集加载/分散存储的内存流量，以及改善SM间的负载均衡。

在单块B200（fp8）上的测试结果令人瞩目：注意力内核效率达到约980 TFLOP/s的吞吐量，HBM带宽约4.1 TB/s。相比查询固定基线（FlashInfer）实现1.9–2.4倍加速，相比MiniMax开源MSA内核提升约1.6倍。完整模块（含索引映射和组合阶段）相比基线提升1.18–1.43倍，相比开源MSA提升1.32–1.41倍。

文章还深入分析了算法和内核设计空间。M3注意力在分组查询注意力（GQA）基础上实现块状稀疏：将KV序列分为128-token块，轻量级索引分支为每个GQA组评分并选择TopK（=16）块，主分支仅在选定块上计算精确注意力。内核结构分为查询固定（Q-outer）和KV固定（KV-outer）两种。查询固定模式下，每个查询仅读取其选定的KV块，但MMA矩阵尺寸小（16×128），张量核心利用率低，且KV块被多个查询重复读取。KV固定则反转映射：对每个KV头/块，收集选择它的查询，加载KV块一次，运行完整128×128矩阵乘法，每个查询写入最多16个部分输出，后续由组合内核进行LSE合并。这样计算高效，但存储开销增加。

I/O分析显示，两种结构受不同内存瓶颈限制：查询固定受L2带宽限制（重复读取命中L2），KV固定受HBM带宽限制（每个块只读一次，但每个查询写入多个部分）。基于B200的峰值带宽（L2约24 TB/s，HBM约7.4 TB/s）的简单罗彻斯特线模型表明，当nsb/N < 2.85时KV固定更快（其中N为查询数，nsb为KV序列长度）。该模型保守假设查询固定完美命中L2，实际中KV固定更早胜出。

总之，Fireworks AI的KV-固定内核成功优化了MiniMax M3稀疏注意力在Blackwell上的性能，为长上下文推理提供了显著加速。