FLAMIA、Blackwell向けに因果安全な低精度Attentionを高速化

Transformerの学習では、コンテキストが長くなるほどAttentionの計算コストが大きな課題になります。NVIDIA BlackwellはFP8やFP4による低精度行列演算を高いスループットで実行できますが、その性能をAttentionで利用するには、単純な低ビット化だけでは不十分です。量子化処理そのもののコストに加え、共有スケールが因果性を損なう可能性や、順方向と逆方向で数値処理がずれる問題があります。

Yandex Researchが提案するFLAMIA（FLAsh MIcroscaling Attention）は、FlashAttentionの考え方をBlackwell向けに再設計したAttentionカーネル群です。主な学習経路はMXFP8と、Query-Keyスコア計算およびその逆方向の再計算にNVFP4、その他の積にFP8を使う混合方式です。すべてをFP4にする方式は、量子化の追加コストが増え、損失もわずかに悪化するため採用していません。

主な工夫

• 量子化表現の再利用。 32トークン単位で2のべき乗スケールを共有し、スコア勾配では32×32ブロックを利用します。これにより、同じ低精度オペランドを異なる行列方向で再利用でき、重複量子化やメモリアクセスを減らします。
• 因果性を考慮したスケーリング。 因果Attentionでは未来のトークンを参照してはいけません。しかし、未来位置を含むブロックから共有スケールを計算すると、マスク適用後でもスケールを通じて情報が伝わる可能性があります。FLAMIAは文書単位のスケール、飽和回避、keyごとのスケール処理でこのリスクを抑え、対応するバッチ呼び出しでは厳密な因果オプションも提供します。
• 順方向と逆方向の整合。 逆方向で確率を再計算する際、量子化、タイル分割、スケーリング、指数計算の規約を順方向と合わせます。1ステップの勾配が似ているだけでは、長時間の学習安定性を保証できないためです。
• ハードウェア向けの融合とスケジューリング。 量子化をカーネルに融合し、softmaxと行列積を重ねて実行します。低精度の逆方向処理では協調スレッド配列を使い、GPU利用率を高めます。長さの異なる文書をpacked layoutで扱う場合も、再パッキングせず、文書単位の量子化と負荷を考慮したスケジューリングを行います。

結果と意義

複数の因果packed-documentワークロードで、量子化を含む順方向・逆方向Attention全体を評価したところ、MXFP8パスは各ワークロードで最速だった基準実装に対し、B200で約1.10～1.40倍、B300で約1.06～1.42倍の高速化を示しました。Llama系モデルの事前学習でも、学習の安定性と品質維持が検証されています。

FLAMIAが示す重要な点は、低精度化の成果がデータ形式だけで決まらないことです。量子化粒度、因果性、順逆方向の数値処理、GPUスケジューリングを一体で設計することで、ハードウェアの理論性能を実際の学習速度に近づけられます。ただし、報告された性能は評価されたワークロードに基づくものであり、すべてのモデルや系列長分布に同じ倍率が適用されるわけではありません。

出典：Hugging Face Daily Papers