Flash-MSA：スパースアテンションの訓練を世界で初めて高速化したオープンソースカーネル

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

MiniMax Sparse Attention (MSA) を効率的に訓練するための、世界初の高性能オープンソースカーネル「Flash-MSA」が登場。HopperおよびBlackwell GPU向けにCuTeDSLで実装され、ブロック単位のスパーシティ、GQAベースのメインアテンション、プロキシヘッドのグループ特化など、MSAの特徴を活かした設計。プロキシとメインのアテンションを融合したバックワードや、KL損失を明示的に計算しない勾配計算など、独自の工夫でメモリと計算を効率化。Flash-Attentionと比較して、単一ステップで最大2倍以上の高速化を達成。