Triton (OpenAI GPU コンパイラ)（トリトン）
OpenAI が開発した GPU カーネルコンパイラ兼プログラミング言語。Python ライクな構文で CUDA カーネルを記述でき、MLIR ベースの最適化パイプラインにより手書き CUDA に匹敵する性能を自動生成する。PyTorch 2.x の torch.compile バックエンドの中核技術。
Triton とは
Triton は OpenAI の Philippe Tillet が開発した、GPU 向けの高水準プログラミング言語およびコンパイラである。2021 年にオープンソース公開され、Python の関数デコレータ（@triton.jit） だけで高性能な GPU カーネルを記述できる。CUDA C/C++ の複雑なメモリ管理やスレッド同期を抽象化し、研究者やML エンジニアが GPU の計算能力を直接活用できるようにする。
CUDA との比較
| 比較軸 | CUDA C/C++ | Triton |
|---|---|---|
| 言語 | C/C++ 拡張 | Python デコレータ |
| メモリ管理 | 手動（shared memory 配置） | 自動（コンパイラが最適化） |
| スレッド管理 | 手動（threadIdx, blockIdx） | ブロックレベルの自動並列化 |
| 行コード量 | 100-300 行（典型的な GEMM） | 30-80 行 |
| 性能 | 最高（熟練者が書いた場合） | CUDA の 80-95%（自動最適化） |
| デバッグ | printf + CUDA-GDB | Python トレースバック |
| 習得コスト | 高（GPGPU アーキテクチャ理解必要） | 中（Python + テンソル概念） |
プログラミングモデル
Triton のプログラミングモデルは ブロック単位のデータ並列処理 に基づく。
- プログラム ID（program_id）: 各ブロックに一意に割り当てられるインデックス
- ブロックポインタ: メモリ上のテンソルブロック（タイル）へのポインタ
- マスク処理: 境界条件の自動ハンドリング（配列末尾の端数処理）
- アトミック操作: tl.atomic_add 等による競合のない集約
CUDA の warp/thread レベルの制御を隠蔽し、タイルサイズとメモリアクセスパターンのみをプログラマが指定する。残りの最適化（shared memory 配置、レジスタ割当、命令スケジューリング）はコンパイラが自動実行する。
コンパイルパイプライン
Triton のコンパイルは以下の段階を経る。
- Python AST 解析: @triton.jit デコレータが Python 関数を AST として取得
- Triton IR 生成: AST から Triton 独自の MLIR Dialect に変換
- 最適化パス: 演算融合・メモリ合体・パイプライニング等の最適化を適用
- PTX/AMDGCN 生成: NVIDIA GPU 向けには PTX、AMD GPU 向けには AMDGCN を出力
- JIT コンパイル: 初回実行時にカーネルをコンパイルしキャッシュ
torch.compile との統合
PyTorch 2.x では torch.compile のデフォルトバックエンドとして TorchInductor が使用され、その GPU カーネル生成に Triton が採用されている。
- ユーザーが
model = torch.compile(model)を呼ぶだけで自動的に Triton カーネルが生成される - FlashAttention 相当のアテンションカーネルも Triton で自動生成可能
- カスタムオペレータが必要な場合は @triton.jit で直接記述し torch.autograd.Function に統合
FAQ
Q1: Triton は NVIDIA GPU 以外でも動作するか？
AMD GPU（ROCm）向けのバックエンドが開発中で、2025 年時点で実験的にサポートされている。Intel GPU 向けの XPU バックエンドも Intel が開発に参加している。ただし NVIDIA GPU での安定性と性能が最も高い。
Q2: Triton で書いたカーネルは cuBLAS より速いか？
一般的な GEMM（行列積）では cuBLAS が最適化の蓄積で優位だが、カスタム演算（融合カーネル・特殊なアテンションパターン）では Triton が cuBLAS + 手動融合を上回るケースがある。FlashAttention の初期実装は Triton で書かれ、CUDA 実装に匹敵する性能を示した。
Q3: CUDA を知らなくても Triton は使えるか？
GPU プログラミングの基本概念（並列実行・メモリ階層・ブロック分割）の理解は必要だが、CUDA の詳細な API 知識は不要。Python と NumPy/PyTorch の経験があれば、Triton の公式チュートリアルで数時間で基本的なカーネルが書ける。