Blackwell向けFlashAttention-4、エンドツーエンドMXFP8対応へ

導入

FP8を注意機構で効果的に使うには、データ型を置き換えるだけでは不十分です。ブロック単位でスケーリングする行列乗算では、スケール係数を生成し、適切な次元に沿って配置し、メモリと同期の制約を満たしながら演算へ渡す必要があります。PyTorchチームは、NVIDIA Blackwell向けFlashAttention-4を拡張し、MXFP8による前向き計算と後ろ向き計算を実装しました。

主なポイント

• 前向きと後ろ向きの両方を低精度化。 LLM向け形状では前向き2.85 PFLOPS、後ろ向き2 PFLOPSを報告しています。内部形状では前向き2.54 PFLOPS、後ろ向き1.58 PFLOPSに達し、BF16比で最大1.6倍、1.52倍の向上とされています。
• 限られたTMEMを再利用。 BlackwellのTMEMは512列で、既存のFlashAttention-4でもほぼ使い切られます。実装では、ある計算段階で不要になった領域をスケール係数の格納に重ねて使い、非同期MMAとの競合を避けるために必要なバリアを追加しました。
• 量子化を前段の演算へ融合。 RMSNormと量子化、GEMMと量子化を一つのカーネルにまとめ、後続処理に必要なFP8データとスケール係数を同時に生成します。独立した変換カーネルによるオーバーヘッドを抑える狙いです。
• softmaxの結果をオンラインで変換。 softmaxの出力Pは、次のP・V演算に向けてMXFP8へ変換されます。softmaxですでに求めている行ごとの最大値を再利用し、Blackwell向け命令の組み合わせで追加処理を減らしています。
• 可変長入力で大規模なGatherを回避。 jaggedモジュールでは、FP8の活性値をパディング前の位置に保持し、より小さいスケール係数だけを散在・パディング・並べ替えします。これによりTMA転送に適した配置を作ります。

意義と限界

BlackwellはMXFP8、MXFP6、MXFP4、NVFP4などのマイクロスケーリング形式に対応するMMA命令を備えています。しかし、量子化やスケール係数の移動を主計算から切り離すと、理論上の演算性能が変換や同期のコストで失われる可能性があります。今回の実装は、低精度化を行列乗算だけの変更ではなく、メモリ配置、非同期実行、softmax、前後段のカーネルを含む全体設計として扱っています。

また、MMAが高速になるほど、softmaxや特殊関数演算のコストが目立ちます。チームはKVのアンロール方法とバリアの待機位置を調整し、タイル境界で発生するパイプラインの空白を抑えました。実装はMeta内部のGEMトレーニングで利用され、関連コードも公開されています。ただし、性能値は特定の形状と実装に基づくため、すべてのモデルで同じ向上が得られることを意味しません。

出典：PyTorch Blog