FlashAttention Pt 1: 個人メモの理解

FlashAttention: Part 1

---

これはFlashAttentionシリーズの2部構成のパート1です。次にパート2をご覧ください。パート2ではFlashAttentionの進化（FA-1/2/3/4）を扱います。

---

IOを意識したAttentionが近似なしでTransformerを高速化する方法

このブログ記事の参考資料は、このツイートに触発されたもので、タイトルは「Understanding FlashAttention: How IO-Aware Attention Makes Transformers Faster Without Approximating Attention」です。

3つのアイデアでメカニズムを説明します：タイリング + オンラインSoftmax + 再計算。このハンドブック1のすべては、その要約の詳述です。ハンドブックは、正確なタイルドAttentionに関する技術的な深掘りです：GPUメモリトラフィック、オンラインSoftmax、フォワードおよびバックワードパス、IO複雑性、FlashAttention-1からFlashAttention-4までの進化、および現在のフレームワークの動作。

この記事では、FlashAttention-1（FA1）が対処する根本的な問題、使用される数学的トリック、それらのトリックのGPU実装、およびFlashAttentionのアーキテクチャ互換性（MHA/MQA/GQA）を扱います。

---

目次

0. はじめに

• 0.1 このハンドブックの読み方
• 0.2 記法

1. 根本的な問題

• 1.1 FlashAttentionが実際に最適化するもの
• 1.2 Attentionの等式は実装ではない
• 1.3 GPUメモリ階層とIOが重要な理由
• 1.4 スコア行列と確率行列を具体化することが高コストな理由
• 1.5 密な算術演算は依然として二次的である
• 1.6 メモリ効率の良い正確なAttentionはFlashAttention以前から存在した

2. 数学的トリック

• 2.1 クエリ、キー、バリューのタイリング
• 2.2 Softmaxはストリーミングの難しい部分である
• 2.3 原理からのオンラインSoftmax
• 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互換性（共有K/Vヘッド）
• 4.2 可変長、ローカルAttention、ドロップアウト

5. まとめ

付録

• 参考マップ
• 参考文献
• 引用

---

0. はじめに

0.1 このハンドブックの読み方

標準的なAttentionにおけるHBMへのラウンドトリップの多さが際立っています。すべての途中値（$S$, $P$, $O$）は書き出され、読み戻される必要があります。これがFlashAttentionが解決する問題です。ハンドブックでは、この主題を3つの異なる質問を分離して考えるのが最も理解しやすいとフレーム化しています。

1. どのような数学的関数が計算されているのか？ 密なAttentionの場合、ターゲットは通常のスケールドドット積Attentionのままです。

2. その関数にはどれだけ多くの算術演算が必要か？ 密な全ペアクエリ-キーのスコアリングは、シーケンス長に対して依然として二次的です。

3. 実装はGPUメモリ階層をどのようにデータを移動させるか？ ここでFlashAttentionはアルゴリズムの実行を劇的に変えます。

私が繰り返し思い出す教訓：FLOP数だけでは壁時計速度は決まりません。アルゴリズムは実質的に同じ数学を実行したり、途中結果をゼロから再計算したりしても、高帯域幅メモリ（HBM）との間で移動するデータがはるかに少ないため、より速く完了することができます。

コアな区別

密なFlashAttentionは正確なAttentionアルゴリズムです。Softmax Attentionを低ランク、スパース、カーネル化、または近似的な式で置き換えるものではありません。「正確」とは、数学的なAttention計算を指します。浮動小数点カーネルは、演算が並べ替えられるため、小さな丸め誤差によって異なる場合があります。

「正確」という言葉はここで重要な意味を持っています。正確さは、ビット単位の再現性に関するものではなく、数学的関数に関する声明です。カーネルは浮動小数点演算の順序を自由に変更できますが、計算される関数を変更することはできません。2 $^,$ 3 $^,$ 4

0.2 記法

1つのAttentionヘッドについて、以下のように定義します。

([Q \in \mathbb{R}^{N ext{}q \times d}, \quad K \in \mathbb{R}^{N ext{}k \times d}, \quad V \in \mathbb{R}^{N ext{}k \times d ext{}v}.)

自己Attentionの場合、通常は $N ext{}q = N ext{}k = N$ です。スケールされたスコア行列は

([S = \frac{QK^T}{\sqrt{d}} + B,)

ここで、$B$ はオプションの加算マスクまたはバイアスを表し、

([P = \mathrm{softmax} ext{_}(\mathrm{row}(S), \quad O = PV.)

全体を通して、HBM は大きなオフチップ高帯域幅GPUメモリを指します。オンチップメモリは、レジスタや共有メモリ/SRAMのような、はるかに小さく高速なストレージの広範な教育用語です。実際のハードウェアの詳細はGPU世代によって異なります。

1つの表での対比：

システムの問題

GPUが計算中に、これらのすべての途中値（$S, P, O$）はどこに格納されるのでしょうか？

その質問、つまり算術演算ではなく、FlashAttentionが構築された目的です。

1. 根本的な問題

1.1 FlashAttentionが実際に最適化するもの

通常のAttention関数から始めます。

([O = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}} + B\right) V.)

教科書的な実装は、これを3つの大きなステップに要約します。

• スコア行列 $S$ を形成する。
• 行ごとのSoftmaxを適用して $P$ を得る。
• $V$ と乗算する。

数学的には、それは完全に問題ありません。しかし、GPU上では、巨大な中間値をHBMに書き出して読み戻すことは、等式が示唆するよりもはるかに高コストになることがよくあります。

FlashAttentionの真の貢献は、IOを意識した実装です。計算を高速なオンチップメモリに収まるタイルに分割し、$K$と$V$のブロックをストリーミングし、正確な出力を生成するために必要なだけの行ごとのSoftmax状態を保持します。完全な $N \times N$ Attention行列がHBMに格納される必要は決してありません。2

IO意識の私の作業定義：

数学的演算（FLOPs）の数を減らすだけでなく、GPUメモリの異なるレベル（HBMとSRAM）間のデータ移動を最小限に抑える。

その理由は、現代のAIハードウェアにおける速度のボトルネックは、GPUがどれだけ速く数学を計算できるかではなく、どれだけ速くデータの読み書きができるかであることが多いからです。これがメモリと計算のトレードオフです。

問題の核心（1行で）

計算 → HBM書き込み → HBM読み込み → 計算

算術演算がボトルネックになることはめったにありません。ラウンドトリップが問題です。タイリングがその解決策です。1つの大きな $QK^T$ の代わりに、ブロックごとに処理します。

([Q K^T \rightarrow Q ext{}i K ext{}j^T \qquad (\text{小さいブロック } Q ext{}i, K ext{}j \text{ の場合}))

各スコアブロックは以下を受け取ります。

1. 計算される。
2. Softmax処理される。
3. 対応する $V ext{_}j$ と乗算される。
4. 破棄される。

$N = 8192$ の場合、$8192 \times 8192$ の中間行列の代わりに、$128 \times 128$ のようなブロックを処理します。すべての $QK$ の相互作用は依然として発生しますが、管理しやすい断片で行われます。

重要な事実

変わるもの：実行スケジュール、メモリトラフィック、および格納される中間値。
変わらないもの：評価されている密なスケールドット積Attention関数。

このフレームワークは重要です。FlashAttentionを「より高速なAttentionの一種」と呼ぶのは、ポイントを外しています。それは、アクセラレータ上でAttentionを効率的に評価するためのアルゴリズムおよびカーネルファミリーとして理解する方が良いでしょう。モデルは、因果マスキング、RoPE、MQA/GQAなどの他のAttention機能を使用でき、その下でFlashAttentionカーネルを実行できます。

元の論文は、これを近似Attention手法と区別しています。近似Attentionは、数学自体を変更することで算術演算を削減します。密なFlashAttentionはそうしません。数学をそのまま維持します。同じ論文のブロックスパース拡張は、ブロックをスキップすることが文字通り計算されるクエリ-キー相互作用を変更するため、別の話です。2

1.2 Attentionの等式は実装ではない

等式はテンソルがどこに存在するかを教えてくれません。標準的なナイーブなAttention実装を見てみましょう。

1. $S$ を計算し、$S$ を格納する。
2. $S$ を読み込み、$P$ を計算し、$P$ を格納する。
3. $P$ を読み込み、$O$ を計算する。

問題はHBMへのラウンドトリップの数です。単純な具体化実装は概念的に以下を行います。

([S \leftarrow QK^T / \sqrt{d}, \quad P \leftarrow \mathrm{softmax}(S), \quad O \leftarrow PV.)

HBMから $Q, K, V \in \mathbb{R}^{N \times d}$ をロードした後、標準的なAttention実装はメモリ通信に関してより詳細には次のようになります。

1. $Q, K$ をHBMから読み込み、$S$ を計算し、$S$ をHBMに書き込む。
2. $S$ をHBMから読み込み、$P$ を計算し、$P$ をHBMに書き込む。
3. ブロックごとに $P, V$ をHBMから読み込み、$O$ を計算し、$O$ をHBMに書き込む。
4. $O$ を返す。

GPUは $N^2$ オブジェクトをメモリ内で移動させるのにかなりの時間を費やします。FlashAttentionは、巨大な中間値を繰り返し移動させることを回避します。

1. 小さな $S$ タイルを計算し、タイルにSoftmaxを適用し、そのタイルを $V ext{_}j$ と共に使用し、タイルを破棄する。
2. 次のタイルを計算し、すべてのタイルが完了するまでステップ1を繰り返す。

Attentionの等式を見る際に私が問う重要な質問：

1. 必要な操作はいくつですか？

2. それらの操作を実行するために、どのデータがメモリレベル間を移動する必要がありますか？

FlashAttentionの移動はシンプルです：作業を近くに保つ。$Q$, $K$, $V$ のブロックは計算ユニットの近くに配置され、スコアタイルはオンザフライで構築および消費され、行ごとの実行統計と出力アキュムレータのみがタイル間で保持されます。それ以外は保持する必要はありません。

アルゴリズムの教訓

計算グラフはメモリスケジュールではありません。紙の上で $P = \mathrm{softmax}(QK^T)$ と書くことは、実装が $QK^T$ または $P$ のすべてを一度にオフチップメモリに格納する必要があることを意味しません。

このパターンはAttentionに固有のものではありません。融合カーネル、タイリング、再計算、および演算子スケジューリングはすべて同じことを行います：巨大な中間値をメモリ経由で引きずるのを避けるために、少し余分な算術演算を行います。現代のハードウェアでは、それは通常、勝利の賭けです。行列乗算のスループットはメモリ階層の他の部分よりも進んでおり、数学は安価で、データ移動が問題となります。FlashAttention-2（FA2）、FlashAttention-3（FA3）、およびFlashAttention-4（FA4）はこれをさらに進め、それらはそれを前面に出しています：アルゴリズムはハードウェアに適合し、その逆ではありません。5 $^,$ 6 $^,$ 7

1.3 GPUメモリ階層とIOが重要な理由

GPUは、均一な高速メモリプールではなく、階層を公開します。名前と容量はアーキテクチャによって異なりますが、メンタルモデルは次のとおりです。

flowchart BT HBM[HBM/デバイスメモリ] SRAM[オンチップ共有メモリ/SRAM] Registers[レジスタ/非常にローカルな状態] HBM -->|タイルをロード| SRAM SRAM -->|計算に供給| Registers classDef memoryDevice stroke:#818cf8,fill:#eef2ff class HBM memoryDevice class SRAM memorySRAM class Registers memoryRegisters

図1: GPUメモリ階層。

FA1はHBM/SRAMの非対称性をモデル化し、それらの間の転送回数を明示的に最適化します。2 上記の図における各GPUメモリレベルの私の内訳：

1. HBM:

• 比較的大きい
• アクセスが比較的遅い
• 計算ユニットから物理的に遠いため、データはそれらに到達するために長く移動する必要がある
• モデルの重み、$Q/K/V$、アクティベーション、その他の大きなテンソルを格納する

2. SRAM / 共有メモリ:

• はるかに小さい
• 速い
• 再利用コストが低い

3. レジスタ:

• さらに小さい
• よりローカル

タイリングが役立つ理由： $Q ext{}i$（クエリタイル/ブロック）が、その出力状態を書き出す前に多くの $K/V$ タイル/ブロックと相互作用する必要があると仮定します。 $Q ext{}i$ を絶えず前後に移動させる代わりに、処理中に計算ユニットの近くに保持できます。

([K ext{}1, V ext{}1 \to K ext{}2, V ext{}2 \to K ext{}3, V ext{}3 \to \cdots)

したがって、1回ロードされた $Q ext{_}i$ は多くの計算に参加できます。これは再利用と呼ばれます。タイリングはテンソルを小さくするだけでなく、タイルがオンチップに常駐している間に再利用を増加させます。逆に、$K/V$ ブロックは、選択されたスケジュールに従ってクエリブロックをストリーミングできます。

IO意識

カーネルがIOを意識するようになるのは、データの配置と移動がアルゴリズムの一部であり、個別に起動されたテンソル演算のシーケンスに任された後付けではない場合です。

これはAttentionが「常にメモリバウンド」であることを意味するわけではありません。FA3とFA4の存在は、ハードウェアが変化するにつれて、主要なボトルネックも変化したためです。6 $^,$ 7 ボトルネックは以下に依存します。

• シーケンス長、$N$
• ヘッド次元、$d$
• dtype
• マスクパターン
• GPU世代（Ampere/Hopper/Blackwell）
• フォワードパスかバックワードパスか
• 実行されているカーネル

1.4 スコア行列と確率行列を具体化することが高コストな理由

二次的な中間値は非常に急速に具体的になります。バッチに1つのシーケンス、32個のAttentionヘッド、シーケンス長 $N = 8192$、およびFP16またはBF16のような2バイトdtypeが含まれていると仮定します。具体的に、私の計算：

([\begin{aligned} \text{batch} &= 1 \ \text{heads} &= 32 \ N &= 8192 \ \text{dtype } &= \text{FP16/BF16} \ \text{tensor shape } &= \text{ [1, 32, 8192, 8192]} \ 2 & \text{ バイト/要素} \ \end{aligned})

要素あたり2バイトとすると、総メモリは： $32 \times 8192^2 \times 2 = 4,294,967,296 \text{ バイト} = 4 \text{ GiB}$ です。

形状が $[1, 32, 8192, 8192]$ の1つの密なテンソルには、$32 \times 8192^2$ 個の要素が含まれています。要素あたり2バイトで、これは4 GiBです。これは、この例における1つの完全なスコアまたは確率のようなテンソルのサイズです。この量のデータを読み書きすることを想像してください。

ナイーブなパイプライン：

ナイーブな分解は、いくつかの段階でこのサイズの配列を生成する可能性があります。これは、すべてのフレームワークが $S$ と $P$ を同時に保持しているという意味でも、コンパイラ融合がすでに一部のトラフィックを削減しているという意味でもありません。この例は、問題を具体化するためです。密な $N^2$ 中間値は、それを書き換えて再読み込みすることがメモリと帯域幅を支配するのに十分な大きさです。

flowchart LR QKT["QKᵀ"] --> S["巨大なS"] S --> Softmax["softmax"] Softmax --> P["巨大なP"] P --> PV["PV"] classDef boxed fill:#e3f2fd,stroke:#1565c0,stroke-width:2px,font-size:22px classDef emphasized fill:none,stroke:none,font-size:26px,font-weight:bold class QKT,Softmax,PV boxed class S,P emphasized

図2: ナイーブなAttention実装パイプライン。

FlashAttentionパイプライン：

FlashAttentionはナイーブな分解を完全にスキップします。スコアタイルは形成され、オンチップでSoftmax化され、出力への寄与を蓄積するために使用され、その後破棄されます。2 完全な行列はHBMに触れることはありません。

flowchart LR QK["QᵢKⱼᵀ"] --> Softmax["softmaxタイル"] Softmax --> PV["タイル × Vⱼ"] PV --> Discard["タイルを破棄"] classDef boxed fill:#e3f2fd,stroke:#1565c0,stroke-width:2px,font-size:22px classDef plain fill:none,stroke:none,font-size:26px,font-weight:bold class QK,Softmax,PV boxed class Discard plain

図3: FlashAttention実装パイプライン。

重要な言葉遣い

FlashAttentionは、完全なAttention行列をオフチップ中間値として具体化する必要性を排除します。それは、密なAttentionに必要な論理的なペアワイズ相互作用を排除するものではありません。

これはトレーニング中に最も重要です。なぜなら、ナイーブな自動微分はバックワードパスのためにそれらの大きな中間値を保存したいからです。

重要な区別： FlashAttentionは、$N^2$ の相互作用が存在することを認識し、「各部分をすぐに消費できる場合は、それらの相互作用の全体の結果を巨大な中間値として格納しない」と言います。

1.5 密な算術演算は依然として二次的である

FlashAttentionはメモリスケジュールを変更しますが、数学は変更しません。関数自体は依然として密なAttentionです。これは、ヘッド次元 $d$ を持つ $N$ トークンに対するすべてのクエリ-キーのスコアを形成するには、依然として $N^2 d$ に比例するコストがかかることを意味します。それを回避する方法はありません。

確率を値で乗算することは、同じ広範な順序の別の密なペアワイズ行列乗算です。FlashAttentionはこれらの操作を並べ替えますが、密なクエリ-キー相互作用のセットをスキップするわけではありません。それらは依然として発生します。2 $^,$ 5

したがって、3つのことを混同してはなりません。

• 算術複雑性： 密なAttentionは $O(N^2 d)$ のままです。
• 大きな中間ストレージ： FlashAttentionはHBMに具体化された $O(N^2)$ スコア/確率テンソルを回避します。
• HBMトラフィック： 元の論文は、その2レベルメモリモデルの下で、標準的な具体化Attentionよりも低いIOコストを証明しています。2

これは、メソッドが長大なシーケンスを実用的にする方法です。長大なコンテキストを「無料」にすることなく。 $N$ を2倍にすると、密なクエリ-キーペアの数は約4倍になります。FlashAttentionは、そのデータ移動とメモリフットプリントの側面を攻撃するだけです。

重要な事実

クエリ-キーペアの数を実際に減らしたい場合は、異なる数学的構造が必要です：スパース性、ローカルパターン、または異なるAttention定式化。それらの選択はモデルの動作を変更し、密なFlashAttentionとは概念的に別個のものです。

スピードアップを報告するときは、常に漸近的な算術と測定された実行時間を分離してください。カーネルは、同じ $O(N^2 d)$ の複雑性で数倍速く実行できます。なぜなら、定数係数、占有率、融合、およびメモリトラフィックが劇的に変化するからです。

本質的に、主要なテイクアウェイは以下のように要約されます。

1.6 メモリ効率の良い正確なAttentionはFlashAttention以前から存在した

正確なAttentionが二次的なメモリを回避できることをFlashAttentionが最初に発見したと言うのは、歴史的に不正確でしょう。RabeとStaatsによる以前の研究は、二次的な計算でサブ二次的なメモリを持つ正確なメモリ効率の良いAttentionを実証しました。彼らは、Attentionが $N$ に対して $O(N^2)$ のメモリを必要としないことを示しました。8 MilakovとGimelsheinも、オンラインSoftmaxに関する以前の研究を行い、古典的な安定Softmax正規化器を計算するためのメモリ効率の良いオンライン再帰的な方法を示しました。9 FlashAttentionは、タイルドAttentionに同じ種類の実行時最大値/実行時正規化器のアイデアを組み込みながら、値で重み付けされた出力を蓄積します。2

安定SoftmaxからFlashAttentionへの系統は、以下の図4に示されています。

flowchart TD Stable["安定Softmax"] --> Online["オンライン正規化"] Online --> Exact["正確なメモリ効率の良いAttention"] Exact --> Flash["FlashAttention"] classDef boxed fill:#e3f2fd,stroke:#1565c0,stroke-width:2px,font-size:16px class Stable,Online,Exact,Flash boxed

図4: メモリ効率の良い正確なAttentionの系統。

FlashAttentionの主な貢献は、以下の技術を実用的で高性能なアルゴリズムに統合することでした。

1. タイリング
2. オンラインSoftmax
3. 融合計算
4. GPUメモリ階層意識
5. IO複雑性分析

この区別は重要です。「メモリ効率が良い」ということは、必ずしも「特定のハードウェアモデルに対してIOが最適化されている」ことを意味しないからです。

---

2. 数学的トリック

2.1 クエリ、キー、バリューのタイリング

最も重要。 ここにトリックがあります。スコアタイルは計算ユニットの近くに留まるのに十分小さいです。各タイルの寄与は、実行時行統計と出力アキュムレータに折りたたまれ、その後タイルは破棄されます。

• $Q \in \mathbb{R}^{N ext{}q \times d}$, $K \in \mathbb{R}^{N ext{}k \times d}$, $V \in \mathbb{R}^{N ext{}k \times d ext{}v}$ と仮定します。すべての $N ext{}q, N ext{}k$ 行を一度に処理する代わりに、それらをブロックに分割します。例：

([Q: \boxed{Q ext{}1} \boxed{Q ext{}2} \boxed{Q ext{}3} \quad K, V: \boxed{K ext{}1, V ext{}1} \boxed{K ext{}2, V ext{}2} \boxed{K ext{}3, V ext{_}3})

1つのクエリ-キー/バリューペア ($Q ext{}i, K ext{}j/V ext{}j$) について、そのスコアタイル $S ext{{ij}}$ をローカルに処理します。

([S ext{{ij}} \doteq \frac{Q ext{}i K ext{}j^T}{\sqrt{d}} + B ext{{ij}},)

次に、各クエリブロックについて：

  1. スコアタイル S_ij を計算する
  2. 行ごとのSoftmaxと更新を行う
  3. V_j と乗算して蓄積する
  4. 実行時出力を更新する
  5. タイルを破棄する
  6. 次のタイルに進み、グローバルなSまたはP行列を書き込まない。

Softmaxは次のように定義されていることを思い出してください。

([\mathrm{softmax}(x ext{}i) = \frac{e^{x ext{}i}}{\sum ext{}j e^{x ext{}j}}.)

• 行列乗算はタイリングが容易です。なぜなら、それは基本的に累積だからです。

([AB = \sum ext{}j A ext{}j B ext{_}j,)

したがって、断片を計算して追加できます。

• Softmaxは難しいです。なぜなら、すべての要素が行全体に依存するからです。最初のブロックを処理した場合、最終的な分母がわからないからです。さらに悪いことに、数値的安定性のためには最大値を知る必要があります。

• カーネルでは、オンチップ容量、ヘッド次元、GPU世代、因果構造、およびワークパーティショニングに基づいてタイル形状とループ順序を選択します。FA1はSRAM容量 $M$ から派生したタイルサイズを使用します。2

数学的トリック

タイリングだけでは不十分です。行列乗算タイルは、和を累積できるため、自然に構成されます。Softmaxは、共有の最大値と分母を介して行内のすべてのスコアを結合するため、完全な行を見ずにブロックをマージする方法が必要です。

2.2 Softmaxはストリーミングの難しい部分である

1つの行 $x ext{}1, \ldots, x ext{}N$ のための数値的に安定したSoftmaxは、次を使用します。

([m = \max ext{}j x ext{}j, \qquad \ell = \sum ext{}j e^{x ext{}j - m}, \qquad p ext{}j = \frac{e^{x ext{}j - m}}{\ell}.)

$m$ を引くことで、大きなロジットがオーバーフローするのを防ぎます。しかし、これはストリーミングの問題を引き起こすように見えます：後続のブロックにさらに大きな最大値があることが判明した場合、初期ブロックをどのように正規化できるでしょうか？

答えは、実行時統計を保持し、最大値が変わったときに再スケールすることです。以前の要素の後の状態が $(m ext{} ext{old}}, \ell ext{} ext{old}})$ であり、新しいブロックの最大値が $m ext{_}b$ であるとします。

([m ext{} ext{new}} = \max(m ext{} ext{old}}, m ext{_}b))

古い最大値の下で累積されたものは、単一の係数によって新しい最大値に変換されます。

([\alpha \doteq e^{m ext{} ext{old}} - m ext{} ext{new}))

新しいブロックは、同じ $m ext{_} ext{new}}$ に対して評価されます。

これらはすべて近似ではありません。単なる恒等式です。

([e^{x - m ext{} ext{old}}}, e^{m ext{} ext{old}} - m ext{} ext{new})} = e^{x - m ext{} ext{new})})

数学的洞察

実行時最大値は、数値参照点の変更です。その参照が変更されると、以前に累積された指数は、最初から再計算するのではなく、実数演算で正確に再スケールできます。

オンライン正規化器の再帰はFlashAttention以前から存在し、ストリーミング安定Softmaxの数学的基盤です。9

例題1：なぜ最大値を引くのか

最大値の減算がなぜ重要なのかを見るために、2つの極端な値から始めます。

([\begin{aligned} x &= \text{ [1000, 999]}, \quad \text{直接 } e^{1000} \text{ を計算するとオーバーフローする}\ m &= 1000 \ x - m &= \text{ [1000 - 1000, ; 999 - 1000]} = \text{ [0, -1]} \ e^{\text{ [0, -1]}} &\approx \text{ [1, 0.368]} \ \end{aligned})

指数値 1 と 0.368 は完全に管理可能です。

例題2：2つのブロックをストリーミングする

次に、2つのブロックをストリーミングする例を見てみましょう。

([\begin{aligned} \text{ブロック1}: \text{ [2, 1]}, \quad \text{ブロック2}:\text{[4, 3]}\ \text{ブロック1の後} &: m ext{} ext{old}} = 2 \ \text{ブロック1を } m ext{} ext{old}} \text{ で評価} &: \text{ [}e^{0}, e^{-1}] \ \text{ブロック2が到着したとき} &: m ext{}b = 4 \ \text{ブロック2で現在} &: m ext{} ext{new}} = \max(2, 4) = 4, \quad \alpha = e^{2 - 4} = e^{-2}\ \text{ブロック1を新しい基準で再スケール} &: \text{ [}\alpha e^{0}, ; \alpha e^{-1}] = \text{ [}e^{-2}, ; e^{-3}] \ \text{ブロック2を新しい基準で評価} &: \text{ [}e^{4 - 4}, ; e^{3 - 4}] = \text{ [}e^{0}, ; e^{-1}] \ \end{aligned})

すべてをまとめると、$x = \text{ [2, 1, 4, 3]}$ の場合：

([\begin{aligned} \mathrm{softmax}(x) &= \frac{\text{ [}e^{-2}, ; e^{-3}, ; e^{0}, ; e^{-1}]}{e^{-2} + e^{-3} + e^{0} + e^{-1}}\ \text{ブロック1ストリーミング計算} &: \ell = \sum ext{}j e^{x ext{}j - m} = e^{0} + e^{-1} = 1 + e^{-1}, \ \text{ブロック1正規化出力} &: \frac{\text{ [}e^{0}, e^{-1}]}{(1 + e^{-1})} \ \text{ブロック2実行時正規化器} &: \ell ext{} ext{new}} = \alpha, \ell ext{} ext{old}} + \sum ext{}{j \in \text{ブロック2}}} e^{x ext{}j - m ext{_} ext{new}}} = e^{-2}(1 + e^{-1}) + e^{0} + e^{-1} \ \text{ブロック2非正規化アキュムレータ} &: a = \text{ [}\alpha e^{0}, \alpha e^{-1}, e^{0}, e^{-1}] = \text{ [}e^{-2}, e^{-3}, e^{0}, e^{-1}] \ \text{結合結果} &: \frac{\text{ [}e^{-2}, e^{-3}, e^{0}, e^{-1}]}{(e^{-2} + e^{-3} + 1 + e^{-1})} \ \end{aligned})

これは、一度に計算された完全なSoftmaxと一致します。

2.3 原理からのオンラインSoftmax

スカラーロジット $x ext{}1, x ext{}2, \ldots$ を1つずつ処理します。初期化します。

([m ext{}0 = -\infty, \qquad \ell ext{}0 = 0)

$x ext{_}j$ を観測した後：

([m ext{}j = \max(m ext{{j-1}}, x ext{}j), \qquad \ell ext{}j = \ell ext{{j-1}} e^{m ext{{j-1}} - m ext{}j} + e^{x ext{}j - m ext{_}j})

MilakovとGimelsheinは、これが従来の安全なSoftmax手順と同じ安定したSoftmax正規化器を生成することを示しました。入力のパスは少なくて済みます。9

Attentionは値の加重合計を必要とするため、値で重み付けされた非正規化アキュムレータも維持します。

([a ext{}j = \sum ext{{i \leq j}} e^{x ext{}i - m ext{}j} v ext{_}i.)

最後の要素 $x ext{_}n$ の後、出力は

(o = \frac{a ext{}n}{\ell ext{}n}.)

重要な事実

Attentionにとって、$x ext{}j$ は事前に格納された固定入力ベクトルではありません。各ブロックのロジットは、行列乗算 $QK ext{}j^T / \sqrt{d}$ とマスク/バイアスタームからオンデマンドで生成されます。オンライン再帰により、カーネルはそのブロックをすぐに消費できます。

同じアイデアが行ごと、ブロックごとに機能するため、タイルド正確Softmax-Attentionフォワードパスが可能になります。

重要な洞察。 完全な確率ベクトルは必要ありません。最終的な出力への寄与を再構築するために必要な情報だけが必要です。

2.4 ブロックごとのマージ再帰

1つのクエリ行について、いくつかのキーブロックの後の実行時状態が次のように仮定されます。

([(m, \ell, a),)

ここで、$m$ は実行時最大スコア、$\ell$ はその最大値の下での安定Softmax正規化器、$a \in \mathbb{R}^{d ext{_}v}$ は非正規化値アキュムレータです。

スコア $s \in \mathbb{R}^b$ と値 $V ext{}b \in \mathbb{R}^{b \times d ext{}v}$ を持つ新しいブロックが着地します。まず、ブロック自体の最大値を求め、実行時最大値を更新します。

([m ext{}b = \max(s), \qquad m' = \max(m, m ext{}b).)

再スケール係数とブロックの指数は次のようになります。

([\alpha = e^{m - m'}, \qquad p = e^{s - m'}.)

次に更新します。

([\ell' = \alpha \ell + \sum ext{}j p ext{}j, \qquad a' = \alpha a + p^T V ext{_}b,)

([m \leftarrow m', \qquad \ell \leftarrow \ell', \qquad a \leftarrow a'.)

最後に：

([o = a / \ell.)

クエリ行のブロックを処理する場合（単一行ではなく）、実行時状態の形状が変わります。$m$ と $\ell$ はそれぞれ1つのエントリを行ごとに持つベクトルになり、$a$ は行列になります。マスクがある場合は、指数を計算する前にスコアタイルに適用します。マスクされた位置は確率ゼロに寄与するため、更新から完全に除外されます。FA1とFA2は、ブロック形式で同等の実行時最大値/実行時正規化器/出力更新を表現します。2 $^,$ 5

重要な事実

再帰は、タイル境界が密なSoftmaxの結果を変更しない代数的理由です。異なるタイリングは浮動小数点演算の順序を変更しますが、意図された実数演算関数は変更しません。

この再帰が、タイルド正確Attentionを可能にするものです。

2.5 完全な数値例

既にスケールされ、マスクされていないAttention行を考えます。

([s = \text{ [2, 1, 4, 3]}]

2次元の値を持つ場合。

([v ext{}1 = \text{ [1, 0]}, \quad v ext{}2 = \text{ [0, 1]}, \quad v ext{}3 = \text{ [2, 0]}, \quad v ext{}4 = \text{ [0, 2]})

グローバル最大値は4なので、安定した非正規化重みは次のようになります。

([[e^{-2}, e^{-3}, e^{0}, e^{-1}] \approx \text{ [0.135335, 0.049787, 1, 0.367879]})

それらの合計は

([\ell \approx 1.55300179,)

そして完全なSoftmaxは次のようになります。

([O \approx \text{ [1.37497284, 0.50582424]})

次に、2つのブロック、$[2, 1]$ と $[4, 3]$ で、タイルドカーネルが行うように、2つのブロックで実行します。

ブロック1： $s = \text{ [2, 1]}$。 実行時最大値 $m ext{_}1 = 2$、安定指数 $\text{ [}e^{0}, e^{-1}] = \text{ [1, } e^{-1}]$。したがって、

([\ell ext{_}1 = 1 + e^{-1} \approx 1.36788,)

値アキュムレータは次のようになります。

([a ext{_}1 = 1 \cdot \text{ [1, 0]} + e^{-1} \cdot \text{ [0, 1]} = \text{ [1, } e^{-1}])

ブロック2： $s = \text{ [4, 3]}$。 ブロック最大値は $m ext{_}b = 4$ で、実行時最大値を上回るため、$m' = \max(2, 4) = 4$ です。再スケール係数：

([\alpha = e^{2 - 4} = e^{-2}.)

次に、

([\ell ext{}2 = \alpha \ell ext{}1 + e^{0} + e^{-1} = e^{-2}(1 + e^{-1}) + 1 + e^{-1} = e^{-2} + e^{-3} + 1 + e^{-1} \approx 1.553,)

そして、

([a ext{_}2 = e^{-2} \cdot \text{ [1, } e^{-1}] + \text{ [2 } e^{0}, 2 e^{-1}] = \text{ [2 + } e^{-2}, 2 e^{-1} + e^{-3}] \approx \text{ [2.135, 0.7855]})

出力：

([O ext{}2 = a ext{}2 / \ell ext{_}2 \approx \text{ [1.375, 0.506]})

完全な行計算と一致します。

重要な事実

古いブロックは再訪されませんでした。その寄与は、より大きな最大値が現れたときに単に再スケールされました。

この再スケール（再計算の代わり）がストリーミングを可能にするものです。

---

3. 数学をGPUに乗せる

3.1 FlashAttentionフォワードパス

概念的なフォワードパス：

1. $Q$ をクエリ行タイルに、$K$ と $V$ をキー/バリュータイルに分割します。
2. 各クエリタイルについて、行ごとの実行時最大値、正規化器、および出力アキュムレータを初期化します：$m = -\infty$, $\ell = 0$, $a = 0$。
3. $K ext{}j, V ext{}j$ タイルをロードし、ローカルスコアタイル $S ext{{ij}} = Q ext{}i K ext{_}j^T / \sqrt{d}$ を形成します。
4. このタイルに属するマスクまたはバイアス $B ext{_{ij}}$ を適用します。
5. タイルの最大値を求め、実行時最大値 $\max(m, m ext{_}b)$ を更新し、以前の状態を再スケールします。
6. 更新された最大値に対して現在のタイルの指数を計算します。
7. 正規化器 $\ell$ と値重み付きアキュムレータ $a$ を更新します。
8. 必要なすべてのキータイルについて繰り返します。
9. アキュムレータを行ごとに正規化し、$O = a / \ell$ とし、出力を書き込みます。

元のアルゴリズムは、タイルと実行時状態がオンチップSRAMに収まるようにブロックサイズを選択し、HBMへの転送を削減します。2 フォワードパスの簡単な要約を以下に示します。

各K/Vブロックについて:
    scores = Q @ K_block.T / sqrt(d)
    実行時最大値 m と正規化器 l を更新する
    出力アキュムレータ a を更新する
O = a / l   # 最後のブロックの後に一度正規化する

PyTorchで代数を数行で再現できますが、それは高性能なFlashAttentionカーネルと同じではありません。実際のプロダクション実装は、GPU固有のタイリング、レジスタ、共有メモリ、ワープスケジューリング、テンソルコア、非同期コピー、特殊な組み込み関数、および占有率最適化も活用します。

重要な事実

本質的なアルゴリズムのアイデアは、CUDAカーネルに依存しません：スコアタイルを生成し、オンラインSoftmaxと$V$累積を通じて即座に消費し、完全なスコア/確率行列をHBMに具体化しないこと。

教育的な最小限のPyTorch。

import math
import torch

# K/Vでタイリングされた単一クエリブロック。プロダクションカーネルも
# Qをタイリングしますが、ここでは再帰に焦点を当てるために外側のループは省略されています。
def tiled_attention(q, k, v, block=128):
    scale = 1 / math.sqrt(q.shape[-1])
    m = q.new_full((q.shape[0],), -float("inf"))
    l = q.new_zeros(q.shape[0])
    a = q.new_zeros((q.shape[0], v.shape[-1]))

    for j in range(0, k.shape[0], block):
        kj, vj = k[j:j+block], v[j:j+block]
        s = (q @ kj.T) * scale
        m_new = torch.maximum(m, s.max(dim=-1).values)
        alpha = torch.exp(m - m_new)
        p = torch.exp(s - m_new[:, None])
        a = a * alpha[:, None] + p @ vj
        l = l * alpha + p.sum(dim=-1)
        m = m_new

    return a / l[:, None]

この教育用コードとプロダクションカーネルとのギャップは、FA2、FA3、およびFA4にとって非常に重要です。各世代は同じ意味構造を維持しますが、新しいハードウェアのためにワークを異なるようにスケジュールします。5 $^,$ 6 $^,$ 7 同じ数学、異なるスケジュール、それぞれがその世代のハードウェアに最適化されています。

3.2 なぜ密なFlashAttentionは正確なのか

FlashAttentionはAttentionを近似しますか？密なFlashAttentionの場合、いいえ。ターゲット関数は変更されていません。

([\mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}} + B\right) V.)

これが正確と見なされる理由の簡単なチェックリスト：

• Q-Kペアは意図的に削除されていません。
• 低ランク近似は導入されていません。
• 代替のカーネル化されたAttention関数がSoftmaxを置き換えることはありません。
• タイリングは、同じ相互作用をブロックごとに処理するだけです。

オンラインSoftmaxも指数関数や正規化を別のものと交換しません。ブロック再帰は、十分統計が累積される順序を変更するだけです。

しかし、正確には資格が必要です。2つの実装が $a + b + c$ を計算すると仮定します。

• A： $(a + b) + c$
• B： $a + (b + c)$

正確な数学では、

([(a + b) + c = a + (b + c).)

浮動小数点では、一部の値について、

([(a + b) + c \neq a + (b + c).)

タイリングは削減順序を変更します。融合は丸め誤差がどのように現れるかを変える可能性があります。したがって：

([\text{正確なアルゴリズム} \neq \text{ビット単位で同一の出力}.)

3つの注意点により、「正確」という言葉が正確になります。

• 浮動小数点には有限の精度があります。 加算と削減の順序を並べ替えると、別の実装とは小さな数値的な違いが生じます。PyTorchはこれを直接述べています。SDPAの出力は、バックエンドカーネルの選択によって異なる場合があります。なぜなら、各バックエンドは浮動小数点値を異なる順序で累積するからです。3 $^,$ 4

• トレーニング中のドロップアウトはランダムです。 2つの実行をビット単位で一致させたい場合は、ランダムな動作も一致させる必要があります。Attentionのセマンティクスだけでは十分ではありません。

• スパースバリアントは密ではありません。 元の論文のブロックスパースバリアントはブロックを完全にスキップするため、完全な密なAttentionを近似するのであって、それを再現するわけではありません。2

重要な事実

「正確」とは、「すべての参照カーネルとビット単位で同一」を意味するわけではありません。 それは、アルゴリズムが作業を節約するために密なSoftmax Attentionを静かに変更していないことを意味します。

この区別は、数値テストにおいて重要になります。妥当な許容範囲は、dtype、累積順序、シーケンス長、およびバックエンドに依存し、バイナリ同一性を要求するものではありません。

3.3 IO複雑性：定理が実際に述べていること

元のFlashAttention論文は、2レベルメモリモデルで動作します：HBMとオンチップSRAM（サイズ $M$）。その仮定の下では、ヘッド次元 $d$ と $d \leq M \leq Nd$ の範囲のSRAMサイズが含まれますが、2つのAttention実装は次を必要とします。

• 標準的な具体化Attention： $\Theta(Nd + N^2)$ HBMアクセス
• FlashAttention： $\Theta(\frac{N^2 d^2}{M})$ HBMアクセス

2番目の数量は通信複雑性バウンドであり、いくつかの点を明確にしておく価値があります。

• これは特定のGPUではなく、特定のメモリモデルに適用されます。
• これは、スカラー要素がモデル化されたメモリレベル間を移動する回数を、漸近的な係数までカウントします。バイト数ではありません。
• これは算術複雑性ではありません。密なAttentionは依然として $O(N^2 d)$ の作業を実行します。
• より多くの利用可能なオンチップメモリ $M$ は、より多くの再利用とより少ないモデル化されたHBMトラフィックを意味します。

直感：より大きなオンチップメモリは、より大きな作業タイルを常駐させることができます。より大きなタイルは、より多くの再利用とより少ないHBMリロードを意味します。定理の範囲では、その再利用がHBMアクセスバウンドが $M$ の増加とともに縮小する正確な理由です。GPUアーキテクチャがFlashAttentionのパフォーマンスにとってそれほど重要である理由でもあります。典型的な値（$d$ が64から128、 $M$ が約100KB）では、$d^2$ は $M$ よりはるかに小さいため、FlashAttentionはHBMアクセスをはるかに少なく必要とします。2

重要な事実

一般的な誤った要約は、「FlashAttentionはAttentionのIOを $O(N^2)$ から $O(N)$ に削減する」というものです。線形数量は、上記の一般的なHBMアクセス式ではなく、シーケンス長に対する大きな補助メモリフットプリントです。

3.4 メモリ複雑性：線形補助状態、線形計算ではない

FlashAttentionの論文では、メモリが線形になる代わりに二次的になるとよく言われています。それは真実ですが、誤解しやすいので、何が線形で何がそうでないかを正確に述べる価値があります。

([\text{線形補助状態} \neq \text{線形計算}.)

ナイーブなAttention実装は $O(N^2)$ の中間値を生成します。FlashAttentionはそうしません。それは以下のみを保持します。

• $Q/K/V$ タイル
• 行ごとの $m$
• 行ごとの $\ell$
• 出力アキュムレータ $a$

入力と出力自体で既に $O(Nd)$ のメモリを占有しています。具体化する密なAttention実装は、さらに $O(N^2)$ のスコアと確率状態を生成します。FlashAttentionはそれらの完全な行列をスキップし、オンチップのタイルサイズの作業状態と $O(N)$ の行ごとの統計のみを保持します。したがって、Attention操作に関連付けられた追加メモリは、二次的ではなく、シーケンス長に対して線形に増加します。しかし、線形メモリは線形計算ではありません。計算は依然として $O(N^2 d)$ です。

例。 $N = 10,000$ の場合、約 $10,000^2 = 100,000,000$ 個のQ-K相互作用があります。FlashAttentionはそれらを消滅させるわけではありません。単に、それらすべてを一度に保持する巨大な中間値が必要になるのを防ぐだけです。

トレーニングは、この違いが最も影響を与える場所です。単純なバックワードパスは、通常、完全な確率行列 $P$ を保持しておくことを望みます。FlashAttentionは代わりにコンパクトな情報（出力と行ごとの正規化統計）を保存し、勾配が必要なときにスコアと確率タイルを再計算します。

重要な事実

線形メモリは線形実行時間を意味しません。 密なクエリ-キー相互作用の数は、依然として $N$ に対して二次的です。

さらに2つのことを念頭に置いてください。第一に、モデル全体のメモリが $O(N)$ であると主張しないでください。他のTransformerコンポーネントはアクティベーションメモリを消費し、自己回帰サービングにはコンテキスト長と共に増加するKVキャッシュメモリがあります。FlashAttentionは、Attention計算が中間値をどのように管理するかを変更するのであって、モデル全体のメモリプロファイルを変更するのではありません。第二に、FlashAttentionとアクティベーションチェックポインティングは共存できます。どちらも再計算と引き換えに保存状態を削減するため、単にトレーニンググラフの異なるスコープで実行されます。

3.5 因果マスキングとタイルスキッピング

自己回帰Attentionでは、クエリトークン $i$ は将来のキー・トークン $j > i$ に注意を払うことができません。Attention行列は次のようになります。

([\begin{bmatrix} \checkmark & \times & \times & \times \ \checkmark & \checkmark & \times & \times \ \checkmark & \checkmark & \checkmark & \times \ \checkmark & \checkmark & \checkmark & \checkmark \end{bmatrix})

マスク付き

([M = \begin{bmatrix} 0 & -\infty & -\infty & -\infty \ 0 & 0 & -\infty & -\infty \ 0 & 0 & 0 & -\infty \ 0 & 0 & 0 & 0 \end{bmatrix})

Softmaxの前にマスクを適用すると、次のようになります。

([\mathrm{softmax}(S') = \begin{bmatrix} \text{値} & 0 & 0 & 0 \ \text{値} & \text{値} & 0 & 0 \ \text{値} & \text{値} & \text{値} & 0 \ \text{値} & \text{値} & \text{値} & \text{値} \end{bmatrix})

具体化されたマスクは、別の $N \times N$ オブジェクトです。しかし、タイルドカーネルは各タイルのジオメトリを直接推論できます。

• 因果境界の厳密に上のブロックは完全にマスクされているため、スキップできます。
• 境界の厳密に下のブロックは完全に有効です。
• 対角線と交差するブロックのみが要素レベルの因果マスキングを必要とします。

これは計算する必要のない無効なタイルの数が多いことを意味します。FA2は明示的にこの構造を利用します。5、現在の実装も不等クエリ/キー長の配置ルールを定義しています。FlashAttentionインターフェースは、因果マスクをAttention行列の右下隅に配置しますが、PyTorchの is_causal=True は左上隅の配置を想定しています。両者は $N ext{}q \neq N ext{}k$ の場合にのみ異なります。そのため、コードパスがどちらの規約を使用しているかを確認してください。10 $^,$ 3

重要な事実

マスクのセマンティクスは、推測してはならないAPIの詳細です。例えば、現在のPyTorch SDPAは、ブール値 attn_mask をAttentionに参加する位置として扱いますが、他のPyTorchマスクAPIは異なる規約を使用します。

ローカルAttentionとの関連。 スライディングウィンドウAttentionは同じ原則で機能します。各トークンが128個の近くのトークンにのみ注意を払う場合、ほとんどのタイルは同様にスキップできます。しかし、これは異なる動きです：数学的なAttentionパターン自体を変更しました。モデルは現在、スパースまたはローカルAttentionを実行しており、FlashAttentionはそのパターンを実行するカーネルにすぎません。もはや、無制限の密なAttentionと同じ問題ではありません。

3.6 バックワードパス：保存する代わりに再計算

ナイーブなトレーニング実装は、$P = \mathrm{softmax}(S)$ をバックワードパスのために保存します。それはHBMに存在する $N^2$ テンソルです。FlashAttentionはそれを保持しません。代わりにコンパクトな行ごとの正規化情報のみを保存し、勾配が必要なときにスコアと確率タイルを再計算します。

フォワード：

• 出力と行ごとの正規化統計を保存します。
• $P \in \mathbb{R}^{N \times N}$ を保存しません。

バックワード：

• 各タイルについて：$QK^T$ を再計算し、ローカル確率を再構築し、勾配を計算し、タイルを破棄します。
• 同じトレードオフ：より多くの算術、より少ないメモリトラフィック。

([S = QK^T / \sqrt{d}, \quad P = \mathrm{softmax}(S), \quad O = PV,)

バックワードパスは、行ごとの恒等式に依存します。ここで、$D ext{_}i$ は行ごとのスカラーです。

([D ext{}i = \sum ext{}r dO ext{{ir}}, O ext{{ir}} = \sum ext{}j P ext{{ij}}, dP ext{_{ij}})

$P$ のタイルを再計算した後、ローカル導関数は次のようになります。

([dV \mathrel{+}= P^T dO, \qquad dP = dO, V^T, \qquad dS = P \odot (dP - D ext{_}i[\text{:}, \text{None}]),)

([dS ext{{ij}} = P ext{{ij}},(dP ext{{ij}} - D ext{}i),)

そして、

([dQ \mathrel{+}= dS, K / \sqrt{d}, \qquad dK \mathrel{+}= dS^T, Q / \sqrt{d}.)

マスクは確率ゼロと勾配ゼロに寄与します。FA2は行ごとの対数和指数量を保存するため、確率タイルは再計算されたスコアから安定して再構築できます。5

重要な事実

バックワード再計算は意図的です。これは、巨大な確率テンソルを読み書きすることを避けるために、追加の行列乗算作業を費やします。これはGPU上では有利なトレードオフです。

これは、バックワードパスに適用されるメモリ-計算トレードオフです。

3.7 なぜFLOP数が多い方が速い場合があるのか

これはFlashAttentionからの最大の教訓の1つです。通常の仮定は、FLOP数が少ないほど速いということです。GPU上では、必ずしもそうではありません。

異なる操作は、完全に異なるスループットで実行されます。テンソルコアは行列乗算を大量に処理します。それ以外のすべては比較的高価です：指数関数、削減、同期、共有メモリ操作、およびHBM転送。2 $^,$ 5 したがって、高価なメモリトラフィックを排除できる場合、追加の算術演算を行うことが正しい動きになる可能性があります。

2つのオプションを検討します。

• A： 計算、巨大な $P$ をHBMに書き込み、$P$ を読み戻し、計算する。
• B： 計算、破棄、後で再計算する。

B はより多くの算術演算を行いますが、巨大なテンソルをHBMにプッシュすることを回避します。現代のアクセラレータでは、そのトレードオフは通常、有利です。

すべてのFlashAttention世代はこの原則に基づいています。FA2の目標の一部は、非matmul FLOPを削減することです。なぜなら、それらはテンソルコアGEMMと同じスループットを持たないからです。また、GPUのより多くの部分が同時に有用な作業を行うように、ワークパーティショニングを再構築します。5 FA3はHopperでさらに進み、行列乗算、Softmax、およびデータ移動を非同期ハードウェア機能とオーバーラップさせます。6 FA4はBlackwellに対応します。ここではテンソルコアのスループットが他のリソースよりも速く成長しました。そのシフトにより、指数関数と共有メモリトラフィックが相対的に重要になり、FA4はその周りに設計されています。7

重要な事実

より良いパフォーマンスの質問は、単に「FLOP数はいくつですか？」ではなく、「どの操作を、どのユニットで、どのようなデータ移動、再利用、並列性、同期で実行しますか？」です。

これは、FlashAttentionをAttention自体を超えて重要にするシステム上の教訓です。ハードウェア効率は、数学だけから生まれるものではありません。それは、数学的スケジューリングとメモリ/実行階層の共同設計から生まれます。

---

4. アーキテクチャ互換性

4.1 MHA、MQA、GQA互換性（共有K/Vヘッド）

FlashAttentionは標準的なマルチヘッドAttention（MHA）に縛られていません。一般に、入力は次の形状を持ちます。

([Q \in \mathbb{R}^{B \times N ext{}q \times H ext{}q \times d}, \quad K, V \in \mathbb{R}^{B \times N ext{}k \times H ext{}{kv} \times d}.)

FlashAttentionのMHA、マルチクエリAttention（MQA）、およびグループ化クエリAttention（GQA）との互換性を見てみましょう。3つのヘッド共有レジームは、$H ext{}q$ と $H ext{{kv}}$ の違いによって区別されます。

• MHA： すべてのクエリヘッドが独自のK/Vヘッドを持ちます。 (H ext{}q = H ext{{kv}})
• MQA： すべてのクエリヘッドが単一のK/Vヘッドを共有します。 (H ext{_{kv}} = 1)
• GQA： 複数のクエリヘッドが各K/Vヘッドを共有します。 (1 < H ext{{kv}} < H ext{}q)

ここで重要な区別があります。

1. MQA/GQAがアーキテクチャを定義します。どのヘッドがK/Vプロジェクションを共有するかを決定します。
2. RoPEが位置変換を定義します。位置がQとKにどのように焼き付けられるかを決定します。
3. FlashAttentionが実行を定義します。結果のAttention操作がどのようにスケジュールされ、効率的に計算されるかを決定します。

これらの3つの直交概念の任意の組み合わせを一緒に使用できます。カーネルは、いくつのクエリヘッドがK/Vヘッドを共有するか、またはQ/K入力がRoPEで回転されたかどうかを気にしません。単に各クエリヘッドとその割り当てられたK/Vヘッド間のAttentionを評価します。ヘッド共有はアーキテクチャとKVメモリフットプリントを変更します。RoPEは入力表現を変更します。FlashAttentionはハードウェア上での作業のスケジュール方法を変更します。

現在の実装は1つの形状制約を課しています。Dao-AILabカーネルでは、クエリヘッドの数はKVヘッドの数で割り切れる必要があります10：

([H ext{}q \bmod H ext{{kv}} = 0)

PyTorch SDPAは enable_gqa を公開しますが、そのドキュメントではGQAを独自の形状制約を持つ実験的な機能として説明しているため、プロダクションコードは正確なバージョンのドキュメントに従う必要があります。普遍的な融合カーネルサポートを想定するのではなく。3

4.2 可変長、ローカルAttention、ドロップアウト

実際のシステムは、常に同じシーケンス長、密なAttention、ドロップアウトなしというクリーンなケースではありません。

プロダクション実装は、より広範な機能セットをサポートします。

• 可変長シーケンス
• 因果Attention
• スライディングウィンドウAttention（SWA）
• ドロップアウト（トレーニング指向インターフェース）
• MQA/GQA
• ALiBiスタイルのスコアバイアス
• KVキャッシュデコーディング（オプションのRoPE処理を含む）

しかし、これらはFlashAttentionの数学的概念自体ではなく、実装機能です。10

特定のビルドがサポートするものは、以下に依存します。

• GPU
• CUDA/ROCmバックエンド（NVIDIAとAMDは別々の実装を持ち、サポート行列が異なります）
• dtype
• ヘッド次元
• マスク
• ライブラリバージョン
• カーネル世代

ドロップアウトに関する実用的な警告。PyTorchの scaled_dot_product_attention は、トレーニングまたは評価モードに関係なく dropout_p 引数に従ってドロップアウトを適用するため、評価コードはドロップアウトをオフにする必要がある場合に明示的に 0.0 を渡す必要があります。3

重要な事実

「FlashAttention」という名前から機能サポートを推測しないでください。実際に実行される正確なライブラリ、カーネル世代、デバイス、dtype、ヘッド次元、マスク/バイアス、およびトレーニング/推論パスを確認してください。

5. まとめ

FlashAttentionの全体像を1つの導出で見て、覚えておくべき重要なポイントをいくつか押さえましょう。

通常のAttention： $QK^T \in \mathbb{R}^{N \times N}$。それを具体化すると、巨大な中間行列が生成されます。11

([O = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}}\right) V.)

ナイーブな実行： これは巨大なメモリトラフィックを生成します。

([QK^T \to \text{スコア } N \times N \text{ を格納} \to \text{softmax} \to \text{確率 } N \times N \text{ を格納} \to V \text{ と乗算})

FlashAttention： タイリングが使用されます。 $Q \to Q ext{}i$ と $K, V \to K ext{}j, V ext{_}j$ に分割し、スコアタイルを計算します。

([S ext{{ij}} = \frac{Q ext{}i K ext{_}j^T}{\sqrt{d}}.)

• タイリングの問題：Softmaxは行全体の最大値と分母を必要とします。

• 解決策：オンラインSoftmax

  1. $(m, \ell, a)$ を維持します。ここで：
     • $m$ = 実行時最大値
     • $\ell$ = 指数関数の実行時合計（Softmax分母）
     • $a$ = 実行時非正規化値重み付き合計（アキュムレータ）
  2. 新しいブロックにより大きな最大値 $m'$ がある場合、再スケール係数 $\alpha$ を計算し、古い状態（$\ell$ と $a$）をそれに再スケールします：
     ([\alpha = e^{m - m'})

• したがって、$N \times N$ 行列を具体化することなく、完全なFlashAttentionパイプラインを処理できます。

  ([\text{タイル} \to \text{スコア} \to \text{オンラインSoftmax} \to V \text{ 累積} \to \text{タイルを破棄}.)

• 結果：同じ正確な密なAttention数学ですが、メモリ動作ははるかに優れています：HBMトラフィックの削減、中間値の劇的な縮小。 算術は $O(N^2 d)$ のままです。

FlashAttentionのまとめ表：

重要なテイクアウェイ

同じ正確な密なAttention数学ですが、メモリ動作ははるかに優れています。 HBMトラフィックの削減、中間値の劇的な縮小、算術は $O(N^2 d)$ のままです。

---

付録

参考マップ

以下の表は、各質問についてどの論文を参照すべきかを示しています。1

このマップの使用方法に関する注記

実装サポートは数学よりも速く変化します。デプロイされたシステムの場合、インストールされているフレームワークのドキュメントと正確なカーネルリポジトリ/リリースを、サポートされているデバイス、dtype、ヘッド次元、マスク、ドロップアウト、およびGQA動作の権威として扱ってください。

参考文献

引用

このブログ記事がお役に立った場合は、以下のように引用することを検討してください。

@article{obasi2026FlashAttentionPt1,
  title   = "FlashAttention: Part 1",
  author  = "Obasi, Chizoba",
  journal = "chizkidd.github.io",
  year    = "2026",
  month   = "Sep",
  url     = "https://chizkidd.github.io/2026/09/13/flashattention/"
}

1. Understanding FlashAttention: How IO-Aware Attention Makes Transformers Faster Without Approximating Attention. Google Drive PDF. アクセス日 2026年9月。 ↩ ↩2

2. Tri Dao et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. ↩ ↩2 ↩3 ↩4 ↩5 ↩6 ↩7 ↩8 ↩9 ↩10 ↩11 ↩12 ↩13 ↩14 ↩15

3. PyTorch. torch.nn.functional.scaled_dot_product_attention documentation. PyTorch Documentation. アクセス日 2026年9月。 ↩ ↩2 ↩3 ↩4 ↩5 ↩6

4. PyTorch. Reproducibility documentation, including scaled-dot-product attention backend differences. PyTorch Documentation. アクセス日 2026年9月。 ↩ ↩2 ↩3

5. Tri Dao. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. ICLR 2024. ↩ ↩2 ↩3 ↩4 ↩5 ↩6 ↩7 ↩8 ↩9

6. Jay Shah et al. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. NeurIPS 2024. ↩ ↩2 ↩3 ↩4 ↩5

7. Ted Zadouri et al. FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling. arXiv, 2026. ↩ ↩2 ↩3 ↩4 ↩5

8. Markus N. Rabe and Charles Staats. Self-attention Does Not Need O(n²) Memory. arXiv, 2021. ↩ ↩2

9. Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXiv, 2018. ↩ ↩2 ↩3 ↩4

10. Dao-AILab. flash-attention official repository and current implementation documentation. GitHub. アクセス日 2026年9月。 ↩ ↩2 ↩3 ↩4

11. Ashish Vaswani et al. Attention Is All You Need. NeurIPS 2017. ↩

12. PyTorch. torch.nn.attention.sdpa_kernel documentation. PyTorch Documentation. アクセス日 2026年9月。 ↩

13. Woosuk Kwon et al. Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP 2023. ↩