Flash-MSA：百万Token训练加速神器

Flash-MSA: Accelerating Million-Token Training with Sparse Attention Kernels

我是开发者，今天开源了全球首个针对 MiniMax Sparse Attention 的高性能训练内核。虽然 DeepSeek、GLM 等前沿模型早已利用稀疏注意力加速推理，但高效的训练代码一直是个空白。我在 Spheron 的 H100 和 B200 上，基于 CuTeDSL 为 Hopper 和 Blackwell GPU 打造了这套方案。核心突破在于将 MLA 替换为 GQA，并引入 Blockwise sparsity，让代理头能独立选择 KV 块。通过巧妙的 Kernel 设计，我们实现了线性时间的反向传播，甚至用数学技巧避免了 KL 散度损失的全量计算。实测显示，该内核在 bf16 精度下与 PyTorch 实现的高度一致，为百万 Token 级训练带来了质的飞跃。

"虽然许多前沿模型利用稀疏注意力大幅加速了推理，但至今无人发布能高效训练它的代码。"