今回の論文
今回取り上げるのは、Juechu Dong、Boyuan Feng、Driss Guessous、Yanbo Liang、Horace Heによる「Flex Attention: A Programming Model for Generating Optimized Attention Kernels」です。2024年12月7日にarXivで初版が公開されました。研究分野は、TransformerのAttention実装、機械学習コンパイラ、GPU最適化です。
論文の書誌はarXivの公開ページで確認できます。本記事の技術説明と数値は、比較条件をそろえるためarXiv v1の本文を基準にします。
この論文を選んだ理由は、「試したいAttentionのルール」と「速く動く実装」の間にある壁を扱っているためです。短いデータをまとめた学習や、独自の参照範囲を持つモデルを開発するとき、専用GPUカーネルを一から書かずに性能を確かめられるという、実装に直結する考え方があります。
どんな技術か
FlexAttentionは、Attentionの振る舞いを小さな関数で定義し、その関数を組み込んだ高速なGPUカーネルを生成するプログラミングモデルです。新しいLLMや学習済み重みを提案する論文ではなく、モデル内部の計算を作るための技術です。
たとえば、「同じ文書のトークンだけを参照する」「直近の一定範囲だけを見る」「距離に応じてスコアを補正する」という変更を考えます。普通のPyTorch演算なら書きやすくても、巨大な中間行列やマスクを保存すると遅くなります。FlexAttentionでは、変更する部分を関数として渡し、行列積や正規化の高速な処理に融合します。
基礎にあるのは、FlashAttentionの中間行列を保存しないブロック計算です。FlexAttentionは、この高速な計算の枠組みにユーザー定義の変更を差し込めるようにしたものと捉えると理解しやすいです。
課題
Attentionは、各トークンがほかのトークンをどの程度参照するかを計算します。その参照範囲やスコアの調整方法には、多くのバリエーションがあります。ところが、既存の高速実装が対応している組み合わせは限られます。
たとえば、ある実装が「直近のトークンだけを見る」処理と「距離によるスコア補正」に個別対応していても、両者を組み合わせたときに高速な経路を使えるとは限りません。論文は、手書きの高速カーネルが用意されているかどうかで試せる技術が左右される状況を、software lotteryと呼んでいます。
一方、汎用的な演算の組み合わせへ戻すと、系列長Nに対してN×Nのスコア行列やマスクを扱うことがあります。長い系列では、このメモリ使用量と読み書きが問題になります。単純にコンパイルするだけで、Attention特有の行列積と正規化を十分に融合できるとも限りません。
実際のシステムでは、長さの違う学習サンプルを詰め合わせたいときや、複数のAttention変種を扱う推論基盤で問題になります。モデルの工夫を試すたびに、低レベル実装と性能調整をやり直す負担を減らす必要があります。
用語解説
- AttentionスコアとSoftmax
- QueryとKeyの内積から参照の強さを求め、Softmaxで重みに変換してValueを集約します。FlexAttentionでは、Softmaxに入る前のスコアを変更する部分が重要です。
- Attentionマスク
- どのトークン間の参照を許すかを指定するルールです。許可しない位置のスコアを負の無限大として扱い、Softmax後の重みをゼロにします。単に出力を隠す処理ではなく、情報が伝わる範囲を決めます。
- 融合カーネル
- 複数の演算を一つのGPU処理へまとめたものです。Attentionの行列積、スコア補正、正規化、集約をまとめることで、巨大な中間結果をGPUの大容量メモリへ何度も書き出す負担を減らします。
- ブロック疎性とBlockMask
- スコア行列を小さな区画に分け、参照不要な区画を丸ごと計算から外せる性質がブロック疎性です。BlockMaskは、処理すべき区画の個数や位置を記録するデータ構造で、マスクを実際の計算削減につなげます。
- KVキャッシュとデコード
- 生成済みトークンのKeyとValueを保持し、次のトークン生成に再利用する仕組みがKVキャッシュです。1トークンずつ生成するデコードでは、今回のQuery内の位置と系列全体での位置が異なるため、マスクの位置計算に注意が必要です。
技術の仕組み
FlexAttentionの中心は、「参照スコアをどう変えるか」「どこを参照できるか」を分離し、共通のAttentionカーネルへ組み込む設計です。モデルの学習方法を新しくするのではなく、学習時の順伝播・逆伝播と推論時の処理を支えます。
score_modとmask_modでルールを表す
score_modは、スコアと、バッチ・ヘッド・Query位置・Key位置を受け取り、補正後のスコアを返す関数です。距離に応じたバイアスや、極端なスコアを抑える非線形変換などを表現できます。
mask_modは、位置情報から参照の可否を返す関数です。因果マスクなら「Key位置がQuery位置以下か」、文書マスクなら「両位置の文書IDが同じか」と判定します。複数の条件はANDやORで合成できます。
たとえば、以下は「同じ文書内の、直近Wトークンだけを見る」条件を表す擬似コードです。実際のAPI呼び出しや実行環境を含むコードではなく、論文の抽象化を説明するものです。
参照を許す = 文書ID[q] == 文書ID[k]
AND 0 <= q - k
AND q - k < WここでWには自分自身を含みます。2つのルールを組み合わせても、それぞれに専用のGPU実装を書く必要を減らせます。ただし、既存モデルの参照範囲をこの条件へ変更する場合、出力品質は別途検証する必要があります。
小さな関数を高速なテンプレートへ埋め込む
論文の実装では、TorchDynamoがユーザー定義関数の計算グラフを取り出し、TorchInductorがTritonのコードへ変換します。そのコードを、手書きで最適化したAttentionのテンプレートへ組み込みます。
テンプレートは、順伝播、逆伝播、デコード向けに用意されます。ブロック単位の行列積、処理を進めながら正規化統計を更新するonline softmax、メモリアクセスの工夫を共通部分として持ちます。ユーザーは、その内部に入れるスコア補正などを記述します。
ポイントは、任意のプログラム全体を自動で最適化するのではなく、Attentionに適した既知の高速構造を使い、変更される部分だけをコンパイルすることです。学習では自動微分の仕組みを使って補正処理の逆伝播も生成します。
BlockMaskで「隠す」を「計算しない」へ変える
マスクをすべてscore_mod内の条件分岐として表しても、数学的には参照を禁止できます。しかし、そのままでは不要なスコアまで計算する可能性があります。
FlexAttentionは、mask_modからBlockMaskを事前生成し、各Query側ブロックについて処理するKey側ブロックの個数と位置を記録します。全要素が参照禁止のブロックは読み込みや計算を省き、すべて参照可能なブロックでは要素ごとのマスク判定を省きます。一部だけ参照可能な境界ブロックには、要素単位の判定を適用します。
つまり、マスクを別の関数として渡すことには、計算を省ける場所をコンパイラへ伝える意味があります。ただし、要素がまばらでも各ブロックに参照可能な位置が残れば、ブロック全体の省略にはつながりません。
ページ化と生成位置の扱い
論文はKVキャッシュのページ化にも応用しています。BlockMaskの論理ブロック番号を、ページテーブルに沿って実際のメモリ上の番号へ変換し、既存の間接参照の仕組みを再利用します。位置依存の補正には、物理位置から復元した論理位置を渡します。
また、デコードでQueryが1トークンだけになると、その配列内の位置は通常ゼロから始まります。系列全体で何番目なのかを示すオフセットを使い、学習時と同じ意味の因果条件や距離補正を適用します。位置の基準まで含めてルールを保つことが必要です。
実験と結果
論文は、主に処理速度と数値誤差を評価しています。新しいモデルの回答精度を競う実験とは区別して読む必要があります。以下の数値は、arXiv v1の6節に報告された値です。
何と比較したか
評価では、通常のAttention、因果マスク、ALiBIによる距離バイアス、スライディングウィンドウ、PrefixLM、スコアのsoft cap、文書マスクの7種類を扱います。MHAと、複数のQueryヘッドでKVを共有するGQAを対象に、FlashAttention-2、FlashAttention-3、FlashDecoding、PyTorch SDPAなどと比較しています。各実装が対応する変種は異なります。
単体カーネルの評価は、KVサイズ256 MiB、ヘッド次元64、bfloat16を基本条件に系列長などを変えています。実験環境にはH100、A100、A6000を使用し、H100には650W・メモリ帯域2.4 TB/s、A100には330Wの制限を設定しています。
単体カーネルでは得意な条件と劣る条件がある
FlashAttention-2が対応する変種での相対性能は0.68〜1.43倍でした。1倍未満の条件ではFlexAttentionのほうが遅く、すべての設定で最速という結果ではありません。因果マスクに限ると、順伝播は1.00〜1.22倍、逆伝播は0.86〜1.05倍でした。
一方、専用対応がなく要素単位のマスクを扱うSDPAとの比較では、対象条件で5.49〜8.00倍の高速化を報告しています。これは特殊なマスクを効率よく実行できる価値を示す結果であり、通常のAttentionを常に8倍速くするという意味ではありません。
Query長1のデコードでは、FlashDecodingに対して通常0.93〜1.45倍でした。GQAとALiBIの組み合わせでは5.37倍という例外的な改善がありますが、比較対象側にその組み合わせの手動最適化がない影響を含みます。
学習・推論全体への効果
torchtuneでは、AlpacaデータセットでLLaMA3-8Bをファインチューニングしています。長さの異なる例を固定長系列に詰め、文書間の参照をマスクする設定です。論文は学習全体で2.4倍超の高速化を報告しています。背景には、大きなブールマスクの読み込みを避けられることがあります。
gpt-fastでSDPAを置き換えた推論では、LLaMA3.1-8Bが1.22〜2.04倍、70Bが0.99〜1.66倍でした。長い文脈でAttentionの割合が増えるほど効果が大きくなっています。70Bの一部条件ではほぼ同等であり、単体カーネルの改善幅がそのままアプリ全体に出るわけではありません。
数値精度は、fp64の参照出力に対するfp16・bfloat16出力のRMSEで確認し、比較実装に比べて追加の数値誤差を導入しないと報告しています。これはビット単位の一致や、参照範囲を変えたモデルの回答品質を保証するものではありません。
何に使える？
以下は、論文の仕組みと実験から考えられる応用です。実験で直接示された範囲を超える部分は、開発上の考察として捉えてください。
短い学習データをまとめて処理する
問い合わせと回答、コードの短い関数、業務文書の断片などを長い系列に詰める学習で役立ちます。文書IDによるマスクを使えば、各例の間に情報を流さずに処理できます。巨大なマスクを保存する方式がボトルネックなら、改善候補になります。
ただし、マスクだけでは学習データの扱いは完成しません。例の区切りに応じたラベルや損失計算、位置情報なども、学習目的に合わせて設計する必要があります。
独自のAttentionを持つモデルの検証
音声の局所的な範囲と一部の大域的な参照を組み合わせるなど、参照関係を位置ベースのルールとして書けるモデルの実験に向いています。スコア補正と参照範囲を分けることで、比較したい条件だけを変更しやすくなります。
これはモデル内部を操作できる開発向けです。外部のLLM APIへ渡すプロンプトに関数を追加しても、API側のAttention実装は変えられません。
複数モデルを扱う推論基盤
モデルごとに違う距離補正やマスクを、共通のカーネル生成の枠組みで表せる可能性があります。ページ化したKVキャッシュと組み合わせるときも、論理的な参照ルールと物理的な保存場所を分けて考えられます。
ただし、RAGの文書間参照を安易に禁止すると、複数資料を比較する能力を損なう可能性があります。マスクは処理を速くするだけでなく、情報経路そのものを変える設計です。
開発や事業へのヒント
この論文から得られる大きなヒントは、変更しやすい仕様と高速な実行を、別々に設計できるということです。少人数の開発でも、モデルのルールを記述する部分と共通基盤を分ける発想は参考になります。
意味の分かるルールを最適化に使う
「参照禁止」という情報を単なる数値変換に埋め込むと、処理を飛ばせる事実が見えにくくなります。FlexAttentionは、それをmask_modとして明示し、BlockMaskによる省略につなげます。
自分でAI処理基盤を作る場合も、不要な候補や独立したデータのまとまりを明示する設計は有用です。ただし、その情報を使って実際に計算を省ける実装まで用意して初めて速度に反映されます。
小さな製品では、まず詰まっている処理を測る
自社運用モデルで可変長データの学習が遅いなら、マスクの保存・読み込みとAttentionの時間を測ることが出発点になります。すでに専用カーネルが効いている通常のAttentionなら、置き換えによる利益が小さい場合もあります。
導入評価では、同じマスクと入力を使った出力・勾配の比較に加え、コンパイル後の処理時間、初回コンパイル、マスク生成、ピークメモリ、処理全体のスループットを分けて測ると判断しやすいです。GPU単価だけでなく、同じ品質で何件処理できるかを見ることが重要です。
ルールの組み合わせを試しやすくする
独自モデルで参照範囲と距離補正を変えたい場合、それぞれを関数として持つと比較実験を組み立てやすくなります。専用実装の開発期間を短縮できれば、モデル品質と実行コストの両方を見ながら候補を絞れます。これは論文の抽象化から考えられる開発上の利点で、事業成果を直接検証したものではありません。
限界
FlexAttentionで表しやすいのは、基本的なAttentionの流れを保ち、スコア単位の補正や位置に基づくマスクで記述できる変更です。Softmaxや集約構造そのものを自由に置き換える万能なコンパイラではありません。
密なAttentionの演算量が系列長に対して二乗で増える性質も残ります。ブロックごとに計算を省けるマスクなら削減できますが、疎な要素が散在するだけでは大きな効果を得にくくなります。BlockMaskも補助メモリを使い、固定のブロックサイズではそのメタデータ量も系列長に応じて増えます。
コンパイルやBlockMask生成の負担にも注意が必要です。参照ルールが変われば対応するマスクを用意する必要があり、再利用できる場合と毎回更新する場合では採算が違います。論文の定常的な性能値を、そのまま短命な処理や動的な入力の初回応答へ当てはめることはできません。
実装の再現性はGPU、データ型、ヘッド次元、ライブラリの版に依存します。本記事は2024年の論文版の説明であり、比較表を現在のライブラリの対応状況として使うことはできません。実装時には導入する版で利用可能な機能を確認する必要があります。
ページ化の実験もGPUメモリ内を対象とし、ホストやディスクへの退避を含む運用を検証したものではありません。混雑時の要求処理やネットワークを含むSaaS全体の性能、日本語タスクの品質向上も、この論文の結果だけからは判断できません。
よくある質問
Q. FlexAttentionは新しいAttentionアルゴリズムですか？
A. 中心は、Attentionの変種を記述して高速に実行するためのプログラミングモデルです。参照範囲やスコア補正は利用側が指定し、その計算を融合カーネルへ組み込みます。
Q. score_modだけでマスクを書いてはいけないのですか？
A. 参照禁止位置を負の無限大にする表現は可能ですが、それだけでは不要なブロックの計算省略につながるとは限りません。mask_modとBlockMaskを使うと、丸ごと参照不要なブロックを処理から外せます。
Q. 既存モデルへスライディングウィンドウを入れれば、品質を保って速くなりますか？
A. 保証されません。同じ参照ルールの計算実装を変えることと、モデルが見られる範囲を狭くすることは別の変更です。後者では長距離の情報を失うため、長文QAや複数資料の比較など、用途に沿った品質検証が必要です。
Q. 論文の最大8倍という性能は、通常のLLM推論でも出ますか？
A. その数字は、専用対応がなく要素単位のマスクを使うSDPAと比較した単体カーネルの対象条件です。gpt-fastでのLLaMA3.1-8B全体の改善は1.22〜2.04倍であり、比較対象と測定範囲を分けて読む必要があります。
Q. 外部LLM APIを使う小さなRAGアプリにも導入できますか？
A. APIの内部カーネルを利用者が置き換えることは通常できません。直接の導入対象は、モデルの学習コードや推論実装を管理できる環境です。API中心の製品では、まず検索・通信・生成のどこが時間を使っているかを測るほうが実際の改善につながります。
今日の学び
この論文は、Attentionに新しい工夫を入れるたびに、高速な専用実装が必要になる課題を扱いました。スコア補正と参照ルールを小さな関数で表し、最適化済みテンプレートへのコンパイルとBlockMaskによる計算省略で解こうとしています。
開発へのヒントは、変えたい仕様を明確に記述し、共通の実行基盤へ渡せる形にすることです。モデルの工夫を試す速さと、実際に処理する速さを両立させる設計として参考になります。