Triton（トリトン）
OpenAI が開発した GPU カーネル記述用の Python ベース DSL であり、CUDA の低レベル知識なしに高性能な並列演算カーネルを記述できるコンパイラ言語。
Triton とは
Triton は OpenAI が開発したオープンソースの GPU プログラミング言語およびコンパイラである。Python に似た高水準の構文で GPU カーネルを記述でき、コンパイラが自動的にメモリアクセスの最適化・ループタイリング・命令スケジューリングを行う。CUDA C++ では数百行必要な最適化カーネルを、Triton では数十行で同等以上の性能で実装できる。
PyTorch 2.0 以降、TorchInductor バックエンドが Triton をデフォルトのコード生成ターゲットとして採用しており、torch.compile 経由で自動的に Triton カーネルが生成・実行される。
Triton のアーキテクチャ
フロントエンド
Triton の @triton.jit デコレータを付与した Python 関数がカーネル定義となる。テンソル演算をブロック（タイル）単位で記述し、プログラマは論理的なブロックインデックスのみを意識する。
中間表現（Triton IR）
Python AST から Triton 独自の IR に変換される。この IR 上で以下の最適化パスが適用される。
- ブロックレベルのメモリ合体（Coalescing）: グローバルメモリアクセスを 128 バイト境界に整列
- 共有メモリ自動配置: レジスタ圧力とバンクコンフリクトを考慮した自動配置
- ソフトウェアパイプライニング: メモリロードと演算のオーバーラップ
バックエンド
Triton IR から LLVM IR を経由して PTX（NVIDIA GPU アセンブリ）を生成する。Hopper 世代では TTGIR（Triton GPU IR）から直接 NVPTX を生成するパスも追加されている。
Triton カーネルの記述パターン
基本的な行列乗算
Triton で行列乗算を記述する場合、プログラマは出力行列のブロック座標を tl.program_id で取得し、入力行列の対応するタイルを tl.load で読み込み、tl.dot でブロック行列積を計算する。
FlashAttention の実装
Triton 版 FlashAttention は以下の手順で動作する。
- Q 行列のブロックを共有メモリにロード
- K, V 行列を外側ループでストリーミング読み込み
- ブロック単位で Softmax を online アルゴリズムで計算
- 結果を累積してグローバルメモリに書き出す
この実装により、標準的な Attention 演算と比較してメモリ使用量が O(N^2) から O(N) に削減される。
| 特性 | Triton | CUDA C++ | CuPy |
|---|---|---|---|
| 記述量 | 少ない（数十行） | 多い（数百行） | 少ない |
| 性能上限 | cuBLAS の 90-100% | 理論値に近い | cuBLAS の 60-80% |
| 自動最適化 | タイリング・合体 | 手動 | なし |
| デバッグ | Python ベース | printf/NSight | Python ベース |
| 対応 HW | NVIDIA GPU | NVIDIA GPU | NVIDIA GPU |
| 動的シェイプ | JIT 再コンパイル | テンプレート | 対応 |
PyTorch との統合
torch.compile 経由の利用
torch.compile(model) を呼ぶと TorchInductor が FX グラフを解析し、融合可能な演算パターンを Triton カーネルに変換する。プログラマが明示的に Triton コードを書く必要はない。
カスタムカーネルの登録
torch.library.custom_op を使って Triton カーネルを PyTorch のオペレータとして登録できる。これにより torch.compile のグラフ内でカスタムカーネルが呼び出され、前後の演算との融合最適化も適用される。
Autograd との連携
Triton カーネルに torch.autograd.Function のフォワード/バックワード定義を与えることで、学習ループでも利用可能になる。FlashAttention の学習時バックワードパスも Triton で実装されている。
性能チューニングのポイント
BLOCK_SIZE の選択
Triton カーネルの性能はブロックサイズ（BLOCK_M, BLOCK_N, BLOCK_K）に強く依存する。一般的な指針は以下の通り。
- Compute-Bound: BLOCK_SIZE を大きく（64〜256）して演算密度を上げる
- Memory-Bound: BLOCK_SIZE を小さく（16〜64）してキャッシュヒット率を上げる
- 共有メモリ制約: SM あたりの共有メモリ容量（A100: 164KB, H100: 228KB）を超えない範囲で最大化
自動チューニング
triton.autotune デコレータでブロックサイズやパイプラインステージ数の候補を列挙し、実行時に最速の組み合わせを自動選択できる。
Triton の制約と注意点
- NVIDIA GPU 専用（AMD ROCm 対応は実験段階）
- Warp レベルの直接制御は不可（Warp Specialization は手動実装困難）
- コンパイル時間が長い場合がある（初回 JIT で数秒〜数十秒）
- デバッグツールが CUDA に比べて未成熟
よくある質問（FAQ）
Q1: Triton は CUDA の完全な代替になりますか？
いいえ、Triton は CUDA の上位レイヤーに位置する抽象化であり、Warp レベルの細粒度制御が必要な場合（例: Warp Shuffle 命令の活用、非同期コピーの手動制御）は CUDA C++ が依然として必要です。ただし、一般的な GEMM・Attention・ElementWise 演算では Triton で十分な性能が得られます。
Q2: Triton カーネルのデバッグはどうしますか？
TRITON_INTERPRET=1 環境変数を設定すると、Triton カーネルが Python インタプリタ上で逐次実行されるため、通常の Python デバッガ（pdb, breakpoint）が使えます。性能プロファイリングには NVIDIA Nsight Compute が利用可能です。
Q3: 自作 PC で Triton を使うにはどのような環境が必要ですか？
NVIDIA GPU（Compute Capability 7.0 以上、Volta 世代以降）と CUDA Toolkit 11.4 以上が必要です。PyTorch 2.0 以降をインストールすれば Triton は自動的に含まれます。RTX 3060 以上であれば、torch.compile 経由で Triton の恩恵を受けられます。