アテンションメモリの計算
入力
| シーケンス長 | 4,096 |
|---|---|
| アテンションヘッド数 | 32 |
| バッチサイズ | 1 |
| 精度 | FP16 / BF16（2バイト） |
アテンションメモリの計算
標準的なTransformerアテンションで実体化されるアテンションスコア行列のメモリを、シーケンス長・ヘッド数・バッチサイズ・1要素あたりのバイト数から推定します。これはFlashAttentionが取り除く二次の項です。
入力
ワークロード
結果
値を入力すると計算結果が表示されます。
アテンションメモリ
標準的なTransformerアテンションは、すべてのトークンを他のすべてのトークンと比較し、softmaxの前にスコア行列を作ります。この行列はシーケンス長について正方なので、そのメモリはコンテキストの二乗で増えます。これがアテンションの有名な二次コストです。この計算ツールは、シーケンス長・アテンションヘッド数・バッチサイズ・スコア1つあたりのバイト数から、実体化されたその行列のサイズを推定します。これはまさに、FlashAttentionが格納を回避するメモリです。
二次の項
トークンのシーケンスに対して、アテンションは のスコア行列を作ります。 行 列はトークン がトークン にどれだけアテンションを向けるかを保持します。これを格納するには に比例するメモリが必要です。これは、コンテキストに対して線形に増えるモデルの重みやKey-Valueキャッシュとは異なります。スコア行列はその二乗で増えるため、長コンテキストでは支配的な中間バッファとなり、素朴なアテンションがメモリ不足になる原因になります。
計算式
各ヘッドが独自のスコア行列を作り、バッチ内の各シーケンスが独自のコピーを持つため、バイト単位のメモリは次のとおりです。
Abytes=B⋅H⋅s2⋅eここで はバッチサイズ、 はアテンションヘッド数、 はシーケンス長、 は格納する1スコアあたりのバイト数です。 で割るとギガバイトになります。 の二乗こそが長コンテキストを高コストにする要因で、ヘッド数とバッチは線形に掛かるだけです。
このメモリは何か
この値は、query-key積とvalueに対するsoftmax加重和との間で標準アテンションがメモリに書き込む、実体化されたスコア行列を測ったものです。これは素朴な実装が保持しなければならないアクティベーションバッファです。これが残り続けるかどうかは状況によります。プレーンなフォワードパスでは一時的でレイヤー間で再利用できますが、トレーニングのフレームワークは、チェックポイントや融合カーネルが介在しない限り、バックワードパスのためにレイヤーごとのコピーを保持することがあります。
FlashAttention
FlashAttentionは、完全な 行列を一度も書き込まずに同一の出力を計算します。keyとvalueを小さなブロック単位でたどり、softmaxの統計量を逐次保持することで、一度にメモリ上に必要なのは1ブロックだけになります。これにより、格納をシーケンス長について二次から線形へと変えます。この計算ツールがサイズを測る行列こそ、FlashAttentionが格納しないと判断するものであり、ここでの結果は、ある特定のコンテキスト長で融合アテンションカーネルが節約するメモリの良い目安になります。
計算例
4,096トークンの単一シーケンスを、32個のアテンションヘッド・16ビット精度で扱う場合を考えます。
Abytes=1×32×40962×2=1,073,741,824これは1レイヤーぶんのスコア行列で約1.07 GBです。コンテキストを8,192トークンに倍増させると二次の項が効いてきて、同じ式はおよそ4.29 GBを与えます。長さが2倍で4倍のメモリです。同じアテンション予算のKey-Value側はKVキャッシュサイズの計算で、モデルの重みの部分はLLM推論のVRAM計算で扱っています。
よくある質問 (FAQ)
なぜアテンションメモリはシーケンス長の二次なのですか？
アテンションはすべてのトークンを他のすべてのトークンと比較し、縦横ともシーケンス長となるスコア行列を作ります。したがってこの行列全体を格納するには、シーケンス長の二乗に比例するメモリが必要です。コンテキストを2倍にするとスコア行列は4倍になり、これが長コンテキストで標準アテンションがメモリ律速になる理由です。二次の項が、重みやアクティベーションの線形コストを上回ってしまうのです。
FlashAttentionはどのようにこのメモリを削減しますか？
FlashAttentionは、完全なスコア行列を一度も実体化せずに同じ結果を計算します。keyとvalueの小さなブロック単位でアテンションをストリーム処理し、softmaxの統計量を逐次保持することで、一度に1ブロックだけを保持します。この計算ツールが測定する二次の行列こそ、FlashAttentionが格納を回避するメモリであり、これによりはるかに長いシーケンスにスケールできます。計算自体は変わらず、中間的な格納分だけが二次から線形へと縮小します。
このメモリはレイヤーごとですか、それともモデル全体ですか？
実装によります。ここで示す値は1レイヤーぶんのスコア行列のサイズです。素朴なフォワードパスでは、このバッファをレイヤー間で解放して再利用できるため、レイヤー数を掛けるのではなく一時的なものになります。ただしトレーニングでは、勾配チェックポイントやFlashAttentionを使わない限り、フレームワークがバックワードパスのためにレイヤーごとのアテンションテンソルを保持することがあり、その場合は合計がレイヤー数に応じてスケールします。
免責事項
この推定は、指定された精度で実体化されるスコア行列のみを対象とし、query・key・valueの各テンソルや出力射影、その他のアクティベーションは含みません。これは標準アテンションを表すものであり、FlashAttentionのようなカーネルはこの行列を一切格納しません。