Triton 3.7、プラグイン拡張で TLX を上流版から利用可能に

導入

PyTorch Blog が紹介した PyTorch-Triton 3.7 の更新で重要なのは、単なる性能チューニングではなく、Triton の拡張方法そのものが変わった点だ。Triton Plugin Extensions により、開発者は上流版 Triton を変更したり再コンパイルしたりせず、独自のコンパイラ Pass、MLIR 方言、DSL 拡張を実行時に読み込めるようになった。Triton のフォークを長期運用してきたチームにとって、これは保守コストを大きく下げる可能性がある。

主なポイント

• フォーク依存からの脱却：高性能 GPU カーネルでは、ハードウェア固有の命令、専用の lowering、細かなメモリ制御が必要になることが多い。従来はこれらを組み込むために Triton の分岐版を維持する必要があった。
• 実行時ロード：プラグインは .so 形式の共有ライブラリとして提供され、TRITON_PLUGIN_PATHS 環境変数で指定するだけで Triton が検出する。Triton 本体の再ビルドは不要だ。
• パイプライン全体への介入：TTIR、TTGIR、LLVM IR、さらに PTX や AMDGCN などのターゲット出力に至るまで、複数の段階にフックが用意される。Pass の挿入、無効化、置換、さらにはステージ全体の上書きも可能になる。
• 方言と DSL の拡張：変換 Pass だけでなく、独立してコンパイルされた MLIR 方言や、Python 側から使える新しい DSL 操作も追加できる。
• カーネル単位の制御：プラグインはカーネルごとに有効化・無効化でき、キャッシュ管理に必要なハッシュ戦略もプラグイン側で扱う。

TLX の位置づけ

Meta の Triton Language Extensions（TLX）は、この仕組みの最初の大きな利用例だ。TLX は共有メモリの明示的な確保、グローバルメモリから共有メモリへの非同期ロード、非同期行列積、ローカルロード／ストア、ソフトウェアパイプラインといった操作を提供する。これらは、現代の GPU を高い稼働率で使う永続化 GEMM カーネルを書くうえで重要な機能だ。

これまで TLX を使うには Meta の実験的な Triton フォークが必要だった。今回の仕組みにより、TLX は utlx という独立パッケージとして上流版 Triton と組み合わせて利用できる。NVIDIA H100 では Hopper の TMA 非同期ロードや WGMMA 命令に対応し、AMD MI350 でも同じプログラミングモデルを各バックエンドにマッピングする。

意義

この更新の意味は、Triton 本体を安定した上流プロジェクトとして保ちながら、研究者、ハードウェアベンダー、実運用チームが独自最適化を素早く試せるようにすることにある。新しい Pass、方言、DSL 抽象はコアリポジトリに入る前から実験・配布でき、古いフォークに閉じ込められる必要がない。GPU プログラミング基盤としての Triton は、よりモジュール化され、変化の速いハードウェア環境に適応しやすくなる。

出典：PyTorch Blog