8月1日、北京で第9回Meet AI Compilerテクニカルサロンが開催されました。このイベントでは、AIコンパイル技術の最新の進歩に焦点が当てられ、業界や研究機関の複数の専門家がプログラミング言語、演算子開発、コンパイル最適化、推論実行に関する知見を共有し、高水準言語表現からハードウェア実行に至るまでのAIコンパイラの協調的な進化を紹介しました。

で、BAAI AI Compilerの研究者である郭慧氏と暁航氏は、「FlagTree：Triton-TLE言語拡張、タイルIRバックエンド、およびコンパイラ最適化手法」と題した発表を行い、FlagTreeチームがTritonに関して実施した一連の研究を紹介した。

ハードウェアアーキテクチャの複雑化とモデル演算子の不規則性の増大に直面し、チームは高レベルのDSLと低レベルのハードウェア制御の間に段階的なチャネルを確立するためにTriton言語拡張（TLE）を提案しました。同時に、FlagTreeコンパイラのコンパイル最適化を実施し、TileIRバックエンド、自動チューニング、データレイアウト、命令スケジューリングなどの側面に着目しました。

この取り組みの核心的な目標は、Tritonの使いやすさとコミュニティのエコシステムを維持しながら、開発者が必要に応じてハードウェアの詳細をより深く掘り下げられるようにすること、そして比較的統一された開発システムで、迅速なオペレータ最適化、アーキテクチャを考慮したチューニング、ネイティブコードレベルの最適化をカバーすることです。

HyperAIは、共有されたコンテンツを元の意味を変えることなく編集・要約しました。

WeChat公式アカウント「HyperAI」をフォローし、背景にキーワード「」を入れて返信してください。0801 AIコンパイラ認定講演者のプレゼンテーション用PPTは、「…」をクリックすると入手できます。

Tritonを基盤として、抽象化とパフォーマンスのバランスを再調整する。

近年、AIコンパイラはますます顕著な矛盾に直面している。ハードウェアアーキテクチャとモデル演算子はますます複雑化しているが、開発者は依然として、より高レベルのDSLを使用して、専門家が作成したカーネルに近いパフォーマンスを実現することを期待している。

抽象化レベルが高すぎると、コンパイラが十分な情報を取得できず、詳細な最適化ができない可能性があります。一方、抽象化レベルが低すぎると、開発者はCUDAやベンダー独自の言語といった複雑な開発モードに戻ってしまうでしょう。

Tritonの成功の鍵は、GPUオペレータの開発をスレッドレベルからタイルレベルへと引き上げた点にある。ユーザーはPython DSLを使用してデータブロック間の計算上の関係を記述し、スレッドマッピング、レジスタ割り当て、データレイアウト、パイプライン処理、同期などのタスクは主にコンパイラによって処理されます。

このアプローチは、高性能オペレーターの開発障壁を低くし、大規模なコミュニティエコシステムを育成してきました。しかし、次世代GPU、ドメイン固有アーキテクチャ、国産AIチップの継続的な開発に伴い、Tritonの当初の抽象化は限界に直面し始めています。

一方では、コンパイラのバックエンドが新しいハードウェアのストレージ構造、通信メカニズム、または演算ユニットをまだサポートしていない場合、フロントエンド開発者がこれらの機能を独自に利用することは困難になるでしょう。一方で、ますます多くの重要なオペレーターが、ストレージレベル、並列処理の粒度、CTAコラボレーション、通信と計算の重複について、より精密な制御を必要としており、オリジナルのTritonコードでは、これらの意図を完全に表現することが困難な場合がある。

Gluon、TLX、TileLangといった新しい言語やDSLの出現は、同じ傾向を反映している。AIオペレータの開発は、もはや単一のカーネルを記述するだけではなく、データレイアウト、並列階層、パイプライン、通信トポロジー、ハードウェア特性を表現することにも及ぶようになった。

TLEはTritonを置き換えることを目的としているのではなく、Tritonの構文とエコシステムを階層的に拡張することを目的としている。これは、TLE-Lite、TLE-Struct、TLE-Rawという3つのレイヤーで構成されており、それぞれ軽量なセマンティックヒント、アーキテクチャを考慮した制御、ネイティブコードレベルの最適化に対応しています。

TLE-Liteは、アルゴリズムエンジニアや迅速な最適化シナリオ向けに設計されています。開発者は基盤となるハードウェアについて直接的に考慮する必要はなく、代わりに、より明示的な構造情報をコンパイラに追加すればよい。例えば、テンソルにサブタイルからアクセスする必要がある場合や、分散メッシュ上で計算を実行する場合、あるいはCTAがプロデューサー・コンシューマーパイプラインを使用する場合などが考えられます。サブタイル操作を例にとると、開発者はより大きなテンソルから論理的なサブブロックを直接抽出し、活性化、正規化、統計計算を実行してから、オフセットを手動で計算したり、マスクを構築したり、境界を処理したりすることなく、それを書き戻すことができます。

コンパイラはこれが通常のタイルアクセスであることを認識できるため、データレイアウト、ベクトル化、バンク競合、レジスタ再利用などに関してさらに最適化を行うことができます。このアプローチは、スパースアテンション、ローカル正規化、ブロック統計、ルーティング演算子に特に適しています。分散環境では、TLEはデバイスメッシュを使用して、ノード、GPU、ブロッククラスタ、ブロックなどのさまざまなレベルを記述し、それらを統一された多次元トポロジに整理します。

開発者はメッシュベースのアプローチでプログラムを記述し、コンパイラとランタイムが論理トポロジーを実際のハードウェアと通信メカニズムにマッピングします。その結果、リング通信、バリア同期、シャーディングアクセスは、コード内で散在するランクや通信グループではなく、分析可能な構造化されたセマンティクスになります。通信関係が明示的に表現されると、コンパイラはトポロジーを考慮したスケジューリング、通信と計算のオーバーラップ、バリアマージ、デッドロックチェックを実行できるようになります。

TLE-Liteは、パイプラインプリミティブを通じて、CTAの内部連携をプロデューサー・コンシューマーモデルに抽象化します。開発者は主に、誰がデータを生成し、誰が消費するかを記述し、基盤となるバリア、バッファの再利用、および同期メカニズムはコンパイラによって処理されます。これは、基盤となる機能を隠蔽することを意味するのではなく、むしろそれらを分析、検証、最適化できるプログラム構造に変換することを意味する。

軽量なセマンティクスからネイティブなパススルーまで、さまざまなレベルの最適化深度を網羅しています。

TLE-Lite が主にクロスプラットフォームのセマンティック表現を扱っている場合、一方、TLE-Structは、アーキテクチャの理解と微調整に重点を置いている。

異なるGPU、DSA、およびAIアクセラレータは、ストレージ階層、実行ユニット、同期メカニズム、およびオンチップネットワークにおいて大きく異なります。 したがって、TLE-Structは階層的な並列処理およびストレージ構造を開発者に公開し、ベンダー独自のインターフェースに直接縛られることなく、データレイアウト、計算マッピング、およびメモリ階層を明示的に定義できるようにします。

例えば、同じローカルバッファをGPUバックエンドでは共有メモリに、DSAバックエンドではスクラッチパッドやオンチップSRAMにマッピングすることができます。ユーザーは構造化されたメモリの意図を表明し、コンパイラはそれをターゲットハードウェアに適したアドレス空間とメモリアクセス命令に変換する役割を担います。

MoEにおけるエキスパートカウントは典型的なシナリオです。この演算子は基本的に、異なるエキスパートにルーティングされたトークンの数をカウントしますが、共有メモリのレイアウト、同時更新、バンクの競合、ブロック間の集計などの要因の影響を受けやすいです。

TLE-Structを使用すると、開発者はカウンタのレイアウトを明示的に構成し、異なるエキスパートやトークンをローカルストレージの異なる領域にマッピングできます。コンパイラはこの構造情報を取得すると、適切なアクセスおよび同期メソッドを生成します。

TLE-Rawは、パフォーマンス最適化の専門家向けに設計されており、ベンダーネイティブコードのインターフェースを維持します。

パフォーマンス向上のためには、CUDA、アセンブリ言語、または専用の組み込み関数を直接使用する必要がある場合があります。これらをより高レベルのDSLに再パッケージ化すると、パフォーマンスが低下したり、移行コストが増加したりする可能性があります。TLE-Rawを使用すると、開発者はTriton/TLEアーキテクチャ内でネイティブコードをインライン化したり、ベンダーのコンパイルパイプラインに直接アクセスしたりできます。

All-Gather GEMMを例にとると、この演算子は、通信、行列計算、バッファリング、および同期管理を伴います。開発者は、構造化されたテンソルとタイルを使用して計算を表現しながら、基盤となる通信機能を再利用できます。そして最終的に、統一されたコンパイルパイプラインによって、さまざまな部分が呼び出し可能な演算子に整理されます。

この階層的な設計により、アルゴリズムエンジニア、オペレーター開発者、パフォーマンス専門家は、最も低いレベルのプログラミングから始めることなく、同じシステム内で異なる最適化レベルを選択できます。

性能テストでは、TLEの全体的な抽象化オーバーヘッドは制御可能である。

Radix Selectテストにおいて、TLEはTensorRT-LLMアルゴリズムを再現し、複数の形状にわたって約85%～97%のパフォーマンスを達成しました。クロスプラットフォームの保守と迅速なイテレーションを必要とするチームにとって、より統一されたコードでエキスパートレベルに近いパフォーマンスを実現できることは、エンジニアリング上の大きな価値があります。

128K のコンテキストを持つ SparseMLA シナリオでは、TLEはパイプラインプリミティブを使用して、異なる実行ロール間の連携を表現し、FlashMLAのベースラインと同等の約90%のパフォーマンスを実現しています。

チームはまた、8基のNVIDIA H100プロセッサを搭載した単一ノード上でAll-Gatherのテストを実施し、GEMMとAll-Gatherの統合についてさらに検討した。焦点は、単に通信ライブラリを置き換えることではなく、通信を演算子レベルの式やコンパイル最適化に組み込み、データ伝送、ローカル計算、同期、そしてその後の利用を可能にすることでパイプラインを構築することにある。

推論システムにおいては、このような通信と計算の融合は、単一のGEMMのピーク性能を個別に向上させるよりも、多くの場合、より意義深いものとなる。なぜなら、ユーザーが最終的に体感するのはエンドツーエンドの遅延だからである。

Flagtreeコンパイル最適化手法

FlagTreeチームは、TLE言語拡張機能に加えて、さらに2つのタスクにも取り組みました。まず、CUDA Tile IRバックエンドとの統合を行い、次に、実世界のモデルワークロード向けにTritonコンパイルの最適化をいくつか実装しました。

CUDA Tile IRの核心的な概念は、プログラムが基盤となるスレッドマッピングを事前に決定するのではなく、タイルを表現できるようにすることです。その入力は、タイルプログラムと3次元タイルグリッドで構成されます。ダイアレクト内部では、計算、データビュー、および必要な依存関係は、タイル計算、ビュー、およびトークン順序操作（TKO）を通じて表現されます。

これらのうち、TensorViewはグローバルポインタ、形状、ストライドを記述し、PartitionViewはこれにブロックマッピング機能を追加します。Tokenは関連するTKO間の依存関係を制約するために使用され、Memory Modelはメモリのセマンティクスとスコープを個別に指定します。

FlagTreeは、NVIDIAのネイティブバックエンドであるTritonを完全に置き換えるものではありません。既存のシステムにTileIRパスを追加し、TLEプリミティブを介してTileIRのビューおよびトークンインターフェースをオーバーライドします。要件を満たすカーネルはTileIRバックエンドにアクセスできますが、まだサポートされていないカーネルはネイティブCUDAにフォールバックします。コンパイラでは、両方のパスを共存させることができます。

自動最適化に関して、チームは、Triton Autotuneの実世界モデルにおけるカバレッジとコストの間の矛盾を解決するために、FlagOSTuneを提案した。

実際のモデル推論では、同じ演算子が多数の異なる形状に対応する可能性があります。チームの統計によると、6つのモデルと4種類の推論シナリオには1994種類の固有のMM形状が存在し、これは日々のベンチマークの網羅範囲をはるかに上回っています。候補構成を直接拡張すると、理論上の探索規模は急速に数百万セットにまで拡大します。

FlagOSTuneは、モデル予測と限定的な実世界テストを組み合わせることで、探索空間を拡大し、探索コストを削減することでパフォーマンスを向上させます。このシステムは、まずXGBoostを使用して候補構成をランク付けし、次に最も有望な構成のごく一部のみを実際のコンパイルとGPUテストに送り、その後、遺伝的アルゴリズムを使用して探索を続行します。

NVIDIA、Moore Threads、Muxiなどの様々なコンピューティングパワーシステムにおいて、複数の演算子のパフォーマンスが向上し、1.21倍から7.35倍の高速化が実現した。NVIDIA H20 MMオペレーターを用いた複数の形状実験において、探索空間の構成数は62万以上から4,070に圧縮され、チューニング効率が120倍向上し、パフォーマンスの低下はほとんど見られなかった。

データレイアウトに関しては、チームはデータレイアウト変換のパフォーマンスオーバーヘッドの最適化に注力した。

Tritonでは、`convert_layout`は通常、スレッドをまたいだデータ再配置を伴い、共有メモリへの書き込み、同期、そして読み出しが必要となるため、コストのかからない操作ではありません。小規模な演算子やメモリを大量に消費する演算子の場合、レイアウト変換の回数が少なくても、パフォーマンスのボトルネックになる可能性があります。

FlagTreeは、コストモデル、バックプロパゲーション、ローカルジョイントソリューションなどの技術を用いて、不要なデータレイアウト変換を削減することで、レイアウト変換削除メカニズムを強化します。100人以上のオペレーターによるテストでは、正味変換排除率は約68%～79%であり、最大で71%の性能向上が見られました。

もう一つの最適化は、命令の並べ替えです。ループ展開では、ループ本体をコピーするだけで、複数のロード命令を事前に自動的に実行することはありません。コンパイラの並べ替え機能を有効にすると、独立したロード命令を事前に実行し、後続の計算をインターリーブすることで、ハードウェアパイプラインを利用してメモリアクセスの遅延を隠蔽できます。代表的な3つの演算子において、この最適化により平均で約1.19～1.61倍の高速化が実現し、最大で2倍の高速化が達成されます。

チームはまた、DeepSeek-V4-Flash の実際の形状に対する Fused Marlin MoE オペレーターのパフォーマンスを最適化し、フラグメントマージによってデータ読み込み時間を短縮しました。アルゴリズムレベルの最適化と組み合わせることで、vLLM CUDAと比較して、FlagGemsはテスト対象の53種類の形状すべてにおいて高速化を実現し、形状の出現頻度で重み付けした場合、平均で約1.208倍の改善が見られました。

NVIDIA H20のDeepSeek-V4-Flashエンドツーエンドテストでは、TP=4の場合、最初のトークンの遅延は20.221 TP3T減少し、総スループットは12.631 TP3T増加します。TP=8の場合、最初のトークンの遅延は12.871 TP3T減少し、総スループットは7.651 TP3T増加します。

将来の仕事の見通し

次の段階は、FlagTreeは、NUMAプログラミングモデルとMegaKernelコンパイラの発展に注力します。

マルチノード、マルチGPU構成、GPU内部クラスタ、ローカルSRAMなどでは、データアクセスにおける近接性に大きな違いが生じます。TLEは、テンソルがメッシュに沿ってどのように分割され、異なる分散方法間でどのように遷移するかを明示的に記述することを目的としており、コンパイラがどのデータを計算のために近接させるべきか、どの通信を計算と統合またはオーバーラップできるかを判断できるようにします。

MegaKernelコンパイラは、最適化の範囲を単一のカーネルからモデル実行プロセスへと拡張しようと試みます。コンパイラは、共同最適化可能な計算を特定し、それらを具体的なタスクに分解し、これらのタスクの割り当て、並列化、パイプライン化の方法を統一的に決定し、最終的に1つまたは複数のMegaKernelを生成します。

TLEは、フロントエンドのセマンティクスとハードウェアアーキテクチャの間の橋渡し役を果たします。タイル、メッシュ、パイプライン、メモリレイアウト、およびネイティブ機能インターフェースにより、コンパイラは単なる独立した演算子の集合ではなく、スケジューリング、配置、同期、およびマージが可能なタスクグラフを認識できます。

FlagTreeは、TLE-Liteの軽量なセマンティックヒントから、TLE-Structのアーキテクチャ認識制御、TLE-Rawのネイティブ機能パススルー、そして様々なコンパイル最適化技術を組み合わせることで、Tritonの使いやすさ、クロスプラットフォーム移行、そして究極のパフォーマンスの間で、より柔軟な階層型システムを確立し、新しいハードウェア機能をモデルや演算子の実際のパフォーマンスに迅速に変換できるようにすることを目指しています。