NVIDIA BlackwellでのMiniMax M3スパースアテンションの最適化

Fireworks AIチームは、NVIDIA Blackwell（SM100）GPU上でMiniMax M3スパースアテンション向けの高効率カーネルを開発しました。長コンテキスト推論において、アテンションは計算とメモリの大部分を占めます。MiniMax M3が採用するスパースアテンションは、各クエリが最も関連性の高い16個の128トークンKVブロックのみに注目することでコストを削減します。しかし、スパース選択はデータ依存でリクエストごとに異なるため、不規則なメモリアクセスパターンが発生し、理論上の高速化が打ち消されることがよくあります。

Fireworks AIのパフォーマンスチームが構築したBlackwellカーネルは、KV固定（KV-stationary）実行パスを採用することでこのギャップを埋めました。このアイデアはMiniMaxチームのスパースアテンション論文に基づいています。外側ループで選択された各KVブロックを1回ロードし、内側ループでそれを選択したすべてのクエリを処理します。このスケジュールを固定した上で、最適化は主に2つの重要な領域に焦点を当てています。ギャザードロード/スキャッターストアのメモリトラフィックの削減と、SM間の負荷分散の改善です。

単一B200（fp8）でのテスト結果は印象的です。アテンションカーネル効率は約980 TFLOP/sのスループット、約4.1 TB/sのHBM帯域幅を達成しました。これはクエリ固定ベースライン（FlashInfer）比1.9–2.4倍、MiniMaxのオープンソースMSAカーネル比約1.6倍の高速化です。完全モジュール（インデックスマッピングと組み合わせ段階を含む）はベースライン比1.18–1.43倍、オープンソースMSA比1.32–1.41倍の向上を示しました。

記事ではさらに、アルゴリズムとカーネル設計空間について深く分析しています。M3アテンションはグループ化クエリアテンション（GQA）の上にブロック単位のスパース性を実装しています。KVシーケンスは128トークンの固定ブロックに分割され、軽量なインデックスブランチが各GQAグループのブロックをスコアリングしてTopK（=16）を選択し、メインブランチが選択されたブロックのみで正確なアテンションを計算します。カーネル構造にはクエリ固定（Q-outer）とKV固定（KV-outer）の2つがあります。クエリ固定では、各クエリは選択されたKVブロックのみを読み取りますが、MMAタイルは16×128と小さく、テンソルコアの利用率が低く、KVブロックは複数のクエリによって繰り返し読み取られます。KV固定はマッピングを逆転させ、各（KVヘッド、KVブロック）について、それを選択したクエリを収集し、KVブロックを1回ロードして完全な128×128タイルを実行し、各（クエリ、KVブロック）に対して部分出力（O_partial, m, l）を生成します。各クエリは最大TopK（=16）個の部分出力を書き込み、後続の組み合わせカーネルがLSEマージして最終出力を生成します。計算は効率的ですが、ストアオーバーヘッドが増加します。

I/O分析によれば、2つの構造は異なるメモリに制約されます。クエリ固定はL2帯域幅に制約され（再読み取りがL2にヒット）、KV固定はHBM帯域幅に制約されます（各ブロックは1回のみ読み取り、各クエリは複数の部分出力を書き込む）。B200のピーク帯域幅（L2約24 TB/s、HBM約7.4 TB/s）に基づく単純なルーフラインモデルは、nsb/N < 2.85の場合にKV固定が高速であることを示しています（Nはクエリ数、nsbはKVシーケンス長）。このモデルは保守的で、クエリ固定が完全にL2にヒットすると仮定していますが、実際にはKV固定の方が早い段階で優位に立つとされています。

結論として、Fireworks AIのKV固定カーネルは、Blackwell上でのMiniMax M3スパースアテンションの最適化に成功し、長コンテキスト推論に顕著な高速化をもたらしました。