FlashAttention Pt 1: 個人メモの理解

0. はじめに IOを意識したアテンションは、近似なしにTransformerを高速化する そのメカニズムは、3つの言葉で表せる：タイリング + オンラインソフトマックス + 再計算。 このハンドブックのすべては、その要約の詳述です。 正確なタイリングアテンションに関する技術ハンドブック：GPUメモリトラフィック、オンラインソフトマックス、フォワードパスとバックワードパス、IO複雑性、FlashAttention-1からFlashAttention-4までの進化、および現在のフレームワークの動作。 0.1 このハンドブックの読み方 このハンドブックは、このツイートに触発されました。 修正前、標準的なアテンション実装は次のようになります。 Q, K, V ∈ ℝ^(N × d) をHBMにロードし、その後： Q, K をHBMから読み込み、S を計算し、S をHBMに書き込む。 S をHBMから読み込み、P を計算し、P をHBMに書き込む。 P, V をブロックごとにHBMから読み込み、O を計算し、O をHBMに書き込む。 O を返す。 私にとって際立っているのは、HBMへの往復回数です。 すべての途中値 — S, P, O — は書き出され、読み戻される必要があります。 それがFlashAttentionが解決している問題です。 ハンドブック自体は、3つの異なる質問を分離して保つことで、主題を最も理解しやすいものとして提示しています： 計算されている数学的関数は何ですか？ 密なアテンションの場合、ターゲットは通常のスケール付きドット積アテンションのままです。 その関数にはどれくらいの計算量が必要ですか？ 密な全ペアクエリキーのスコアリングは、シーケンス長に対して依然として二次的です。 実装はGPUメモリ階層を通じてデータをどのように移動しますか？ これがFlashAttentionがアルゴリズムの実行を劇的に変える場所です。 私がこのフレームワークから得る中心的な教訓は、壁時計速度はFLOP数だけで決まるわけではないということです。 アルゴリズムは実質的に同じ数学的工作を実行したり、中間値を再計算したりすることさえできますが、高帯域幅メモリとの間で移動するデータ量がはるかに少ないため、より高速に実行できます。 中心的な区別。 密なFlashAttentionは正確なアテンションアルゴリズムです：低ランク、スパース、カーネル化、または近似的な公式でソフトマックスアテンションを置き換えるものではありません。 「正確」とは、数学的アテンション計算を指します。 浮動小数点カーネルは、操作の順序が変更されるため、小さな丸め誤差で異なる場合があります。 「正確」という言葉はここで重要な意味を持っています。 正確さは、計算される関数に関する声明であり、ビット単位の再現性に関するものではありません。 カーネルは浮動小数点操作の順序を変更する自由があります。 計算される関数を変更する自由はありません。 0.2 記法 1つのアテンションヘッドについて、次のようにします。 Q ∈ ℝ^(N_q × d), K ∈ ℝ^(N_k × d), V ∈ ℝ^(N_k × d_v)。 自己アテンションの場合、通常 N_q = N_k = N です。 スケール付きスコア行列は S = (QK^T / √d) + B ここで、B はオプションの加算マスクまたはバイアスを表し、 P = softmax_row(S), O = PV。 全体を通して、HBM はオフチップの大容量高帯域幅GPUメモリを指します。 オンチップメモリは、レジスタや共有メモリ/SRAMのような、はるかに小さく高速なストレージの広範な教育用語です。 正確なハードウェア詳細はGPU世代によって異なります。 私が常に立ち返る実用的な違い： HBM SRAM 遅い 速い 大きい 小さい オフチップ オンチップ システムの課題：GPUが計算している間、すべての中間値（S, P, O）はどこに格納されるのでしょうか？ その質問 — 計算ではなく — がFlashAttentionが構築された理由です。 目次 パート1：根本的な問題 1.1 FlashAttentionが実際に最適化するもの 1.2 アテンション方程式は実装ではない 1.3 GPUメモリ階層とIOが重要な理由 1.4 SとPを具体化することのコスト 1.5 密な計算量は依然として二次的 1.6 メモリ効率の良い正確なアテンションはFlashAttentionより前から存在する パート2：数学的なトリック 2.1 クエリ、キー、バリューのタイリング 2.2 ソフトマックスはストリーミングの難しい部分 2.3 オンラインソフトマックスをゼロから構築 2.4 ブロックごとのマージ再帰 2.5 完全な数値例 パート3：数学をGPUに実装する 3.1 FlashAttentionフォワードパス 3.2 なぜ密なFlashAttentionは正確なのか 3.3 IO複雑性：定理が実際に言っていること 3.4 メモリ複雑性：線形補助状態、線形計算ではない 3.5 因果マスキングとタイルスキップ 3.6 バックワードパス：保存ではなく再計算 3.7 なぜより多くのFLOPがより速くなるのか パート4：アーキテクチャ互換性 4.1 MHA, MQA, GQA互換性 4.2 可変長、ローカルアテンション、ドロップアウト パート5：FlashAttentionの進化 5.1 FlashAttention-2：何が変わったか 5.2 シーケンシャルタイルのFA2並列処理 5.3 FA2ワークパーティショニングと非GEMM FLOPs 5.4 FlashAttention-3：Hopper世代 5.5 FA3非同期処理：データ移動、GEMM、ソフトマックスのオーバーラップ 5.6 FA3 FP8：精度が無料ではないという前提でのパフォーマンス 5.7 FlashAttention-4：Blackwell世代 5.8 FA4と非対称ハードウェアスケーリング 5.9 FA4実装と現在のステータス 5.10 FlashAttention-1から-4の比較 パート6：フレームワークでのFlashAttentionの使用 6.1 現在のPyTorchスケール付きドット積アテンション 6.2 正確さはビット単位の同一性ではない パート7：FlashAttention vs. その他の技術 7.1 FlashAttention vs. PagedAttention 7.2 FlashAttention vs. スパースおよび線形アテンション パート8：トレーニング vs. 推論 8.1 トレーニング、プリフィル、デコードは異なるレジーム パート9：実践的なエンジニアリング 9.1 一般的な実装ミス 9.2 一般的な誤解と実践的なメンタルモデル パート1：根本的な問題 1.1 FlashAttentionが実際に最適化するもの 通常の注意関数から始めます。 O = softmax((QK^T / √d) + B) V。 FlashAttentionはIOを意識しています。 私の作業定義： 単に数学的操作（FLOPs）の数を減らすだけでなく、GPUメモリの異なるレベル間のデータ移動を最小限に抑えます。 現代のAIハードウェアにおける速度のボトルネックは、GPUがどれだけ速く数学を計算できるかではなく、データをどれだけ速く読み書きできるかであることがよくあります。 これがメモリと計算のトレードオフです。 教科書的な実装は、しばしばこれを3つの大きな操作のように見せます： スコア行列Sを形成し、行ごとのソフトマックスを適用してPを得て、次にVを掛けます。 数学的にはそれは問題ありません。 しかし、GPU上では、巨大な中間行列をHBMに書き込み、それを読み戻すことは、方程式が示唆するよりもはるかにコストがかかる可能性があります。 FlashAttentionの中心的な貢献は、アルゴリズムをIOを意識したものにすることです。 それは計算を高速なオンチップメモリに収まるタイルに分割し、KとVのブロックをストリーミングし、完全なN×Nアテンション行列をHBMに具体化することなく正確な出力を生成するために十分な行ごとのソフトマックス状態を維持します。 変わるもの：実行スケジュール、メモリトラフィック、保存される中間値。 変わらないもの：評価されている密なスケール付きドット積アテンション関数。 この区別が、「FlashAttentionはより高速なアテンションの一種である」という表現が誤解を招く可能性がある理由です。 それは、アクセラレータ上でアテンションを効率的に評価するためのアルゴリズムとカーネルファミリーとして考える方が良いでしょう。 モデルは、因果マスキング、RoPE、MQA/GQA、またはその他のアテンション機能を使用しながら、その下にFlashAttention実装を使用できます。 元の論文は、このアプローチを、数学的問題を変更することによって計算量を削減する近似アテンション手法と比較しています。 密なFlashAttentionは、そのようなトレードオフを行いません。 同じ論文ではブロックスパース拡張も導入されましたが、ブロックを省略すると計算される相互作用が変わるため、それは別のケースです。 1.2 アテンション方程式は実装ではない 方程式はテンソルがどこにあるかを教えてくれません。 標準的なナイーブアテンション実装： S ← QK^T / √d, P ← softmax(S), O ← PV。 問題はHBMの往復回数です。 単純な具体化実装は次のように行います。 S ← QK^T / √d, P ← softmax(S), O ← PV。 最初の行列乗算の後にSがHBMに書き込まれ、ソフトマックスのために読み込まれ、Pが書き戻され、その後PV乗算のためにPが再度読み込まれる場合、GPUはN^2オブジェクトをメモリに移動させるのにかなりの時間を費やします。