Flash-MSA: Sparse Attention Kernelsによるミリオン・トークン・トレーニングの高速化

[Github] [MiniMax Paper] [Trainer] Flash-MSA vs Flash-Attention 隔離されたトレーニングステップ。1 いくつかの最先端モデル[1, 2, 3, 4, 5]は、推論を大幅に高速化するために疎なアテンションを使用していますが、効率的にトレーニングするためのコードを投稿した人はいません。本日、HopperおよびBlackwell GPU向けのCuTeDSLで記述された、Minimax Sparse Attentionの、世界初の高性能オープンソーストレーニングカーネルであるFlash-MSAを紹介します。私はSpheron H100およびB200のレンタルで全ての開発作業を行い、FA4、MSA推論、およびCodexを参考にしました。免責事項: これは公式の実装ではなく、MiniMaxとは提携していません。 MSAについて MSAはDeepseek Sparse Attentionに似ていますが、いくつかのコアな変更があります（MSA PaperのFig. 1より）。 1. ブロックワイズ疎性 プロキシアテンションがメインアテンションのために個々のKVを選択する代わりに、最大プーリングを使用してプロキシスコアをブロック単位（128）で選択します。これにより、カーネルにいくつかの優れたキャッシングプロパティが導入されます。 2. メインアテンションのためのMLAではなくGQA これは特に重要です。なぜなら、私の知る限り、西側のラボはMLAをトレーニングに採用しておらず、GLM-5.2、DSv4のような最先端モデルで普及している疎なアテンションの定式化は、MLAに適しているため、ここにあるモデルではアクセスできません。 3. プロキシヘッドのグループごとの特殊化 MLAをGQAに置き換えることで、各レイヤー内にクエリの独立したグループが導入され、DSAが行うようなアテンションレイヤー全体を合計してスコアリングするのではなく、各プロキシヘッドで異なるKVのサブセットを選択するオプションが得られます。アテンションヘッドが有機的に異なるトークンに注意を向けるという証拠がいくつか示唆されているため、この変更はメインアテンションの表現力を高めるはずです。 カーネル設計 カーネルシーケンスの概要 MSAを効率的に実行するには、可能な限り作業を繰り返し、レジスタ/共有メモリを過負荷にしないようにする必要があります。通常のフラッシュレジスタ（Qタイル、KVタイル、Oアキュムレータ、LSEアキュムレータ）に加えて、フォワードパスでのストリーミングトップkアキュムレータを考慮する必要があります。バックワードパスでは、メインアテンションの勾配とプロキシアテンションの勾配を計算するために、ダブルアテンションの結合パスのためのスペースを確保する必要があります。なぜなら、プロキシ勾配はプロキシとメインのアテンション確率の両方にアクセスする必要があるからです。ブロック疎性の1つの良い利点は、DSAのように個々のトークンではなくブロックのインデックスのみをキャッシュする必要があるため、バックワードパスまでブロックインデックスを保持することが可能であるということです。つまり、トレーニングステップ全体で、プロキシフォワードパスのみがコンテキスト長に対して二次的であり、他のすべてはプロキシフォワードパスからキャッシュされた疎ブロックを使用します。 フォワードパス 操作順序は、プロキシアテンション -> 疎なメインアテンション -> メインアテンション出力を次のレイヤーに送信、バックワードパスのためにメインLSEを保存です。 プロキシアテンション プロキシドット積は、通常のFlash Attnとは少し異なります。出力の累積は不要ですが、キーをストリーミングする際にトップkアテンションスコアとその対応するインデックスを追跡する必要があります。Flashとは異なり、バックワードパスのためにLSEを累積せず、代わりにバックワードパス中にLSEを取得するために、疎なアクティベーションに対するプロキシドット積の非常に安価な再計算を実行します。これは、フォワードパスでのLSE+topkの統合よりも実用的には高速でした。QK^Tの各タイルが計算されると、各チャンクのローカル最大スコアを取得し、レジスタに保持されている各クエリ行の現在のトップk値に挿入ソートを行います。これらのトップkレジスタのためのスペースを確保するために、キーブロックを半分にスライスする必要がありました。また、MSAは各トークンのローカルブロックがマスクされていないスライディングウィンドウスタイルでなければならないと規定しているため、各クエリに対するローカルKVブロックのアテンションスコアをinfに設定しました。 メインアテンション メインアテンションは単なるブロック疎性フラッシュアテンションのフォワードパスです。これは以前MoBAで実行されており、ブロック疎性アテンションを変長フラッシュに再パラメータ化する巧妙なトリックをコピーしました。 バックワードパス プロキシヘッドの勾配を計算するには、プロキシトレーニング信号がプロキシのアテンション確率とメインのアテンション確率の両方に同時にアクセスする必要があるため、プロキシとメインのアテンションのバックワードパスを統合する必要があります。フォワードパスからブロックインデックスを保存し、疎なKVアクティベーションのみで両方のアテンションをトレーニングするため、バックワードパスは線形時間で実行できます。まず、キャッシュされたblock_indicesを取得し、$B$[batch, proxy head, query, top_k_slot] -> [key block]のマッピングを、$B^*$[batch, proxy head, key block] -> [this blockを使用するqueries]に反転させます。私たちは、$B^*$を使用してクエリチャンクをスケジュールし、共有疎KVブロックの再利用を最適化します。次に、選択されたブロックに対して迅速な疎プロキシアテンションパスを実行し（ここでもMoBAの変長トリックを使用）、プロキシLSEを取得し、次に融合プロキシ+メインバックワードタスクをストリーミングし、QKV、Q_proxy、K_proxy、およびmain_lseのチャンクをロードします。レジスタにロードする多くのヘッドを考慮するために、一度に使用するQチャンクとKVチャンクのサイズを減らす必要があります。各ストリームで、メインとプロキシのアテンション確率を計算し、dQ、dK、dVを計算し、次にKLトレーニング項からプロキシのdQ、dKを計算します。 KLダイバージェンス損失 DSAの元のKL損失項を思い出してください: $L^ abla = abla_t{D_{KL}(p_t,s_t Vert Softmax(I_t, s_t))}$ インデクサーとメインアテンション確率分布の両方を具体化してKLダイバージェンスを累積するには、共有メモリへの多くの読み書きと追加のレジスタの使用が必要になり、トレーニングが大幅に遅くなる可能性があります。幸いなことに、アトミックにバックプロパゲートし、完全なKL損失と同等の数学的効果を得るために使用できるトリックがあります。プロキシアテンション確率を$p_{px}$、メインアテンション確率を$p$として、KL項を展開します: $L^ abla = abla_t{D_{KL}(p_t Vert p_{px,t})} = abla_t{p_t * log(rac{p_t}{p_{px,t}})}$ 対数規則を使用して再度展開します: $L^ abla = abla_t{(p_t*log(p_t) - p_t*log(p_{px,t}))} = abla_t{(p_t*log(p_t))} - abla_t{(p_t*log(p_{px,t}))}$ 次の潜在変数、つまり位置i（tではない）のソフトマックス前のスコア、$z_{px,i}$に勾配を計算したいとします。 $rac{ abla L^ abla}{ abla z_{px,i}} = rac{ abla}{ abla z_{px,i}} abla_t{(p_t*log(p_t) - p_t*log(p_{px,t}))}$ メイン確率$p_t$はプロキシ/KL損失グラフから切り離されているため、この偏微分では定数です。 $rac{ abla L^ abla}{ abla z_{px,i}} = -rac{ abla}{ abla z_{px,t}} abla_t{p_t*log(p_{px,t})}$ ソフトマックス出力$p_{px,t}$のソフトマックスログ確率の偏微分（ソフトマックス前のロジット$z_{px,i}$に関する）は、$rac{ abla log(p_{px,t})}{ abla z_{px,i}} = abla_{it}-p_{px,i}$であることがわかっています。 $rac{ abla L^ abla}{ abla z_{px,i}} = - abla_t{ p_t*( abla_{it}-p_{px,i})} = - abla_t{ p_t abla_{it}} + abla_t{p_tp_{px,i}}$ クロネッカーのデルタ関数$ abla_{it}$はi==tのときにのみ非ゼロなので、$ abla_t{p_t abla_{it}}=p_i$です。 $rac{ abla L^ abla}{ abla z_{px,i}} = -p_i + abla_t{p_tp_{px,i}} = -p_i + p_{px,i} abla_t{p_t}$ $p_t$は確率分布なので、$ abla_t{p_t}=1$です。 $rac{ abla L^ abla}{ abla z_{px,i}} = -p_i+p_{px,i}$ 言い換えれば、KL損失からプロキシスコアへの勾配は、プロキシ確率 - メイン確率です。これは、KLを完全に具体化することなくプロキシ勾配を計算するためにカーネルで使用する項です。 ウォームアップカーネル ウォームアップモードでは、メインアテンションのフォワードパスは密であり、ブロックインデックスを使用しないため、プロキシフォワードパスは完全にスキップでき、バックワードパスで完全にトレーニングできます。メインアテンションのウォームアップフォワードカーネルでは、フラッシュを呼び出し、返された出力とlseを保存し、プレースホルダーKLを返します。バックワードパスでは、インデクサーに対して密なフラッシュを呼び出してLSEを取得し、疎なMSAカーネルからの融合プロキシ+メインアテンションバックワードを再利用します。 正しさ カーネルのフォワードパスとバックワードパスの正しさを検証するために、PyTorchのEagerモードでMSAを実装し、フォワード出力とバックワード勾配のコサイン類似度をいくつかの設定でスイープしました。スイープはbf16精度で実行され、バックワードパスにはターゲット出力損失と内部KL損失の両方が含まれます。通常、bf16での精度の許容範囲は0.01です。 Eager vs. kernel コサイン類似度 Bat