💻 computer science

Flashlight: PyTorch Compiler Extensions to Accelerate Attention Variants

本論文は、静的なテンプレートに依存せずに任意かつデータ依存の注意パターンに対して融合された FlashAttention 風のカーネルを自動的に生成する PyTorch ネイティブのコンパイラフレームワークである Flashlight を紹介し、これにより FlexAttention などの既存のソリューションと比較して優れた柔軟性と競争力のある性能を提供する。

原著者： Bozhi You, Irene Wang, Zelal Su Mustafaoglu, Abhinav Jangda, Angélica Moreira, Roshan Dathathri, Divya Mahajan, Keshav Pingali

公開日 2026-05-22

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Bozhi You, Irene Wang, Zelal Su Mustafaoglu, Abhinav Jangda, Angélica Moreira, Roshan Dathathri, Divya Mahajan, Keshav Pingali

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

以下は、FLASHLIGHT論文の解説を、日常的な言葉と創造的な比喩を用いて翻訳したものです。

全体像：AI における「交通渋滞」

あなたが今話しているような大規模言語モデル（LLM）を想像してください。それは複雑な製品（文章やタンパク質の構造）を組み立てようとする巨大な工場です。この工場の中で最も重要な部門は**「Attention（注意）部門」**です。ここは、機械がどの情報が重要で、どの情報を無視すべきかを決定する場所です。

問題は、この部門の現在の仕組みが**「交通渋滞」**のようになっていることです。

従来の方法： 工場労働者（コンピュータコード）は一度作業を止め、紙にメモを書き（データを低速なメモリに保存し）、別のステーションへ歩き、メモを取り出し、それから再び作業を開始しなければなりません。これが何度も繰り返されます。これは遅く、エネルギーの無駄です。
「Flash」による解決： 数年前、エンジニアたちはFlashAttentionを発明しました。彼らは、労働者がすべてのメモをポケット（高速メモリ）に入れて、作業全体を一度の連続したスプリントで完了させる方法を考案しました。これは劇的な速度向上をもたらしましたが、それは**「特殊なレーシングカー」**のようでした。特定の種類のトラック（標準的な Attention）でのみ機能したのです。もし異なる種類の車（新しい実験的な Attention 手法）を運転したい場合、手作業で最初から新しいレーシングカーを丸ごと作り上げなければなりませんでした。

問題：カスタムレーシングカーが多すぎる

最近、科学者たちは AI をより賢く、あるいは効率的にするために、「Attention」を行う多くの新しい洗練された方法を考案しました（「Differential Attention」や「Gated Self-Attention」など）。

ボトルネック： これらの新しい方法を高速化するには、通常、各手法ごとに人間の専門家がカスタムの「レーシングカー」（特殊なコンピュータカーネル）を手作業でコーディングする必要があります。これには永遠にかかります。
中間的な解決策（FlexAttention）： FlexAttentionというツールは、建設者にレゴのテンプレートのセットを提供することでこの問題を解決しようとしました。もしあなたの新しい車がテンプレートに合えば、それは高速に動作します。しかし、もしあなたの車が奇妙でユニークで（テンプレートに合わなければ）、FlexAttention はあなたを助けることができず、あなたは再び遅い手作業のコーディングに戻らなければなりませんでした。

解決策：FLASHLIGHT（汎用工場のアップグレード）

登場するのがFLASHLIGHTです。著者たちはこれをコンパイラネイティブなフレームワークと説明しています。簡単に言えば、それは最も多くの AI 研究者が使用する PyTorch ソフトウェアに直接組み込まれた、賢い「オートパイロット」です。

FLASHLIGHT がどのように機能するかを、比喩を使って説明します。

1. 「賢いシェフ」対「レシピ本」

従来の方法： あなたにはレシピ本（コード）があります。標準的なスープを作りたい場合、レシピ本は刻んで、煮込んで、提供するように指示します。奇妙なスープを作りたい場合、新しいレシピを考案し、新しいレシピ本のページを書くためにシェフを雇わなければなりません。
FlexAttention： レシピ本には「特別スープのセクション」があります。もしあなたのスープが「特別スープ」のテンプレートに合えば、それは高速です。そうでなければ、行き詰まります。
FLASHLIGHT： FLASHLIGHT は、あなたが料理する様子を見守る超賢いキッチンロボットのようです。標準的なスープを作っているのか、奇妙で実験的なスープを作っているのかは問いません。それはあなたの行動（コード）を見て、あなたが多くの刻みと煮込みを行っていることを認識し、「ねえ、冷蔵庫に行ったり戻ったりする必要がないように、これらのステップをすべてまとめて、超効率的な動きにできるよ」と言います。

それは、あなたが車の作り方を知らなかったり、テンプレートに合わせる必要がなかったりすることなく、あなたのコードを手作業で仕上げたレーシングカーと同じくらい高速になるように自動的に書き換えます。

2. 魔法の仕組み（3 つのトリック）

この論文は、FLASHLIGHT が速度向上のために使用する 3 つの主要な「トリック」を説明しています。

トリック A：「格下げ」（ステップの融合）
ケーキを焼いていると想像してください。ステップ 1 は小麦粉を混ぜること、ステップ 2 は卵を加えることです。通常、小麦粉を混ぜてボウルを置き、新しいボウルを取り出して卵を加えます。
FLASHLIGHT は、「なぜ止めるの？小麦粉をボウルに入れたまま、そこで卵を加えればいいじゃない」と言います。それは別々のステップを一つの連続した流れに統合し、データが行き来する必要がないようにします。
トリック B：「数学の魔法」（代数的変換）
安全のために、計算を二度行わなければならない場合があります（温度をチェックし、確信するために再度チェックするなど）。これは遅いです。
FLASHLIGHT は「ホモモルフィズム」と呼ばれる数学的なトリックを使って、別々に停止して行うのではなく、料理をしている間に同時に両方のチェックを行えることを認識します。オーブンの温度をチェックするために止まってオーブンを見、その後戻ってかき混ぜるのではなく、かき混ぜている間にオーブンの温度をチェックするようなものです。
トリック C：「タイル化」（トラックの積み込み）
引っ越しをしていると想像してください。一度に椅子を一つずつ運ぶと、永遠にかかります。家具をトラック（「タイル」）いっぱいに積んで一度運転すれば、それは高速です。
FLASHLIGHT はトラックの積み込み方を賢く考えます。運んでいるものに基づいて、トラックの最適なサイズを計算します。家具が小さければ、スペースを無駄にしないように隅に収めます。コンピュータの「トラック」（メモリ）が常に満杯で、効率的に移動していることを保証します。

彼らは何を証明したか

研究者たちは、強力なコンピュータチップ（NVIDIA H100 および A100）で FLASHLIGHT をテストしました。

標準的なタスクの場合： FlexAttention がすでに処理できる「標準的な」Attention 手法を使用した際、FLASHLIGHT は同じくらい速く、あるいはさらに速い性能を示しました。
奇妙な/新しいタスクの場合： FlexAttention が処理できなかった Attention 手法（タンパク質折りたたみ AI である AlphaFold で使用されているものなど）を試した際、FLASHLIGHT は標準的で最適化されていないコードよりも5 倍高速でした。
実世界でのテスト： 彼らは実際のタンパク質折りたたみモデル（AlphaFold）でテストを行いました。それはモデルの全体的な実行速度を**6% から 9%**向上させました。

結論

FLASHLIGHTは、AI 開発者がこれまで通り、普通で簡単な方法でコードを書けるようにするツールですが、そのコードを自動的に高速で超効率的なエンジンに変換します。

以前： 「書きやすいが遅い」か「速いが書きにくい」かのどちらかを選ばなければなりませんでした。
FLASHLIGHT により： 「書きやすく、かつ速い」が実現します。

それは、新しいアイデアごとに専門家が手作業で速度向上を実装する必要性を取り除き、科学者たちがコードが遅すぎることを心配することなく、新しいタイプの AI Attention モデルを実験することを可能にします。現在、これはオープンソースであり、PyTorch のフォークとして誰でもすぐに使用できます。

技術的サマリー：FLASHLIGHT

問題定義

大規模言語モデル（LLM）のトレーニングおよび推論を加速するためには、アテンション機構の最適化が不可欠である。FlashAttentionは、タイル分割とカーネル融合を活用してメモリ帯域幅と起動オーバーヘッドを削減することで、標準的なアテンションのパフォーマンス向上に成功したが、これはバニラアテンションに特化した手動最適化カーネルである。そのため、差分アテンション、AlphaFold の Evoformer における行/列方向のゲーテッド自己アテンション、不変点アテンション、または修正されたスパースアテンションなど、より複雑な新しいアテンション変種には、大幅な手動エンジニアリング努力なしには適用できない。

FlexAttentionやFlashInferのような既存のソリューションは、特定のアテンション変種のサブセットに対して静的な高次テンプレートを提供するか、またはジャストインタイムコンパイルを行うことで、このギャップを埋めようとしている。しかし、これらのアプローチは表現力の限界に悩まされている。すなわち、プログラマーが事前に定義されたテンプレートに適合するようにモデルを再記述する必要がある。その結果、これらのテンプレートに適合しない多くの新しいアテンション定式化は、遅く、融合されていない PyTorch 実装によって制約され、その採用が妨げられている。さらに、標準的な PyTorch コンパイラスタック（TorchInductor）は、メモリ境界を超えた削減融合や複雑な演算子融合といった重要な最適化を欠いており、これらの変種に対する生成コードのパフォーマンスを制限している。

手法

本論文は、PyTorch エコシステム内のコンパイラネイティブなフレームワークであるFLASHLIGHTを紹介する。これは、任意のアテンションベースのプログラムに対して、融合された FlashAttention スタイルのカーネルを自動的に生成するように設計されている。従来のアプローチとは異なり、FLASHLIGHT は静的テンプレートや事前定義されたカーネルの特殊化に依存しない。代わりに、torch.compile ワークフローに直接統合し、標準的な PyTorch コードを分析・変換して構成演算を発見し、融合させる。

FLASHLIGHT は、3 つの主要なメカニズムを通じてTorchInductorの中間表現（IR）と最適化パスを拡張する。

統合された削減 IR：FLASHLIGHT は、一般行列積（GEMM）などのテンソル縮約を、並列次元（ $p$ -次元）と削減次元（ $r$ -次元）を区別する統合された IR 内の一般化された削減としてモデル化する。これにより、行列乗算がエンドツーエンドの融合エンジンに完全に参加できるようになり、通常 GEMM 演算を周囲の計算から分離している人工的な境界が除去される。
次元降格を伴う構造的融合：このフレームワークは、生成カーネルからの並列ループ（ $p$ -次元）を、融合された消費カーネル内の削減ループ（ $r$ -次元）に「降格」させることができる融合規則を導入する。この変換は、並列性を犠牲にして、中間テンソルをグローバルメモリに高遅延で具体化するのを排除する。これにより、softmax におけるmax()のような演算を、先行する行列乗算（ $QK^\top$ ）と直接融合することが可能になる。
代数的変換を伴う意味的融合：softmax のような状態依存演算を融合するために、FLASHLIGHT は代数的性質（特に準同型）を活用し、マルチパスアルゴリズムをシングルパスのオンラインアルゴリズムに変換する。例えば、安定化 softmax（最大値を見つけ、次に指数関数の和を求めるという 2 パスが必要）を、オンライン softmax アルゴリズムに変換する。これは、最大値と和のRunning Accumulator（累積値）を維持し、Running 最大値が変化するたびに、 $e^{x-y} = e^x / e^y$ という性質を用いて和を動的に再スケーリングすることで達成される。
タイル認識型の次元排除と論理グリッド：このフレームワークは、連続したタイルを反復することでループ実行を洗練させる。また、問題の構造を物理的な GPU ハードウェアの制約から分離するために、論理グリッド次元を導入する。これにより、次元ごとの柔軟なタイルサイズが可能となり、単一のタイル内に完全に収まる小さな次元を排除できるようになる。その結果、以前はコンパイラにとって不可能だった複雑な多段階計算（例： $softmax(QK^\top)V$ の融合）のための融合機会が解放される。

主要な貢献

カーネルレベルの融合：FLASHLIGHT は、並列性が許容される場合に、matmul + softmax や matmul + matmul などの互換性のある計算ブロックを自動的に融合するメカニズムを実装し、手動でのカーネル記述を必要とせずにパフォーマンス向上を実現する。
コンパイラ駆動アプローチによる汎用性：静的テンプレートやカスタムカーネルを必要とせず、統一されたコンパイラ駆動アプローチを通じて、標準的かつ新興のアテンション機構の両方をサポートする。
柔軟性を備えたパフォーマンス：既存のテンプレートベースのソリューション（FlexAttention など）と同等かそれ以上のパフォーマンスを持つカーネルを生成しつつ、ネイティブ PyTorch コードの完全な柔軟性を維持する。

結果

著者らは、NVIDIA H100 および A100 GPU 上で、さまざまなアテンション変種について FLASHLIGHT を評価した。

FlexAttention がサポートする変種：FlexAttention がサポートする変種（バニラ、ALiBi、Softcap、Causal、スライディングウィンドウなど）において、FLASHLIGHT によって生成されたコードは、FlexAttention と同等かそれ以上の速度を記録した。特に score_mod 変種では、ブロックマスクの処理に伴うオーバーヘッドを回避する単純な融合カーネルにより、FLASHLIGHT は最大1.48 倍高速であった。block_mask 変種については、スパーススキップによるカーネル実行の速さから FlexAttention の方が速いが、FLASHLIGHT は実行時にスパースブロックマスクを構築する際の大きなオーバーヘッドを回避する。
複雑な変種（FlexAttention 未対応）：FlexAttention のテンプレートで表現できないアテンション変種（差分アテンション、AlphaFold の Evoformer における行/列方向のゲーテッド自己アテンション）において、FLASHLIGHT は標準的な torch.compile を大幅に上回る性能を発揮した。
- 差分アテンション：A100 に比べて H100 でより高い速度向上を示した。
- Evoformer：FLASHLIGHT は、両方の GPU アーキテクチャにおいて torch.compile に対して5 倍以上の速度向上を達成した。
エンドツーエンドの推論：
- AlphaFold2：FLASHLIGHT は、行/列方向のゲーテッド自己アテンション層において、推論レイテンシを**6% から 9%**改善した。
- LLaMa-3.2-1B：Mooncake の会話トレースを使用した vLLM システムにおいて、FLASHLIGHT は Softcap アテンションにおいて FlexAttention を上回った。Causal アテンション（ブロックマスク変種）については、最適化されたスパースブロックマスク処理により FlexAttention がわずかに高性能であったが、中間テンソルの具体化に起因する標準的な torch.compile で発生するメモリ不足エラーを FLASHLIGHT は回避した。

意義と主張

本論文は、FLASHLIGHT がアテンション変種の最適化方法において重要な転換点をもたらすと主張している。すなわち、手動エンジニアリングからコンパイラ自動化へと負担を移転させるものである。標準的な PyTorch コードから融合された FlashAttention スタイルのカーネルを自動的に生成可能にすることで、AI 科学者はパフォーマンスやスケーラビリティを犠牲にすることなく、新しいアテンションモデルを迅速に探索できるようになる。

著者らは、FLASHLIGHT がオープンソースであり、PyTorch のフォークとして利用可能であることを強調している。同フレームワークは、FlexAttention モデルで表現可能なすべての変種を支援するだけでなく、FlexAttention の能力を超えてより一般的でデータ依存の定式化を処理する汎用ソリューションとして位置づけられている。現在の評価は推論に焦点を当てているが、基礎となる拡張とパスは汎用的であり、将来的な作業においてトレーニングワークロードのサポートやブロックスパース性の最適化へ拡張可能であると著者らは指摘している。