# OpenAI Triton ## 概要 [[OpenAI]] が開発するオープンソースの GPU カーネル記述用ドメイン特化言語(DSL)兼 JIT コンパイラ。CUDA C++ を書かずに Python で高性能な GPU カーネルを記述できる。single-program, multiple-data(SPMD)モデルを採用し、CUDA の SIMT(スレッド単位)モデルより高水準な抽象化(スレッドブロック=「プログラム」単位)を提供する。Triton の JIT コンパイラは LLVM の NVPTX バックエンド経由で NVIDIA PTX を直接生成し、NVCC は経由しない。(Source: [[@2025__OReilly__AI Systems Performance Engineering - Chapter 14 PyTorch Compiler, OpenAI Triton, and XLA Backends]]) [[PyTorch]] の [[torch.compile]] スタックにおいて、TorchInductor は NVIDIA GPU 向けのデフォルトコード生成バックエンドとして Triton を利用する。`torch.library.triton_op`/`wrap_triton` API で Triton カーネルを PyTorch カスタム演算として登録すれば、autograd 対応・fake-tensor 対応の一級演算として `torch.compile` グラフに組み込まれ、融合・並べ替え・インライン化の対象になる。(Source: 同上) ## 主要機能 - **プログラミングモデル**: `@triton.jit` デコレータで関数を定義し、`triton.language`(`tl`)モジュールの `tl.program_id`・`tl.load`/`tl.store`(境界チェック用 `mask=` 引数付き)・`tl.arange` などのベクトル化プリミティブでカーネルを記述する。 - **オートチューニング**: `@triton.autotune` で `BLOCK_SIZE`・`num_warps`・`num_stages` などの候補構成(`triton.Config`)を列挙し、初回呼び出し時にベンチマークして最速構成を入力形状ごとにキャッシュする。 - **warp specialization**: `tl.range(..., warp_specialize=True)` でスレッドブロック内の warp をメモリ(producer)役と計算(consumer)役に分割し、メモリレイテンシと計算をオーバーラップさせる。`num_consumer_groups`・`num_buffers_warp_spec` で調整可能。 - **TMA テンソル記述子**: `tl.make_tensor_descriptor(...)` は Blackwell 世代 GPU では Tensor Memory Accelerator(TMA)ハードウェアへマップされ、`tl.dot` は `tcgen05`/UMMA へ下げられて TMEM にアキュムレータを保持する。persistent GEMM カーネル(1次元グリッドで複数出力タイルをストライド処理)と組み合わせて用いられる。 - **Proton プロファイラ**: Triton 専用のプロファイリングパッケージ。`proton.scope()` コンテキストマネージャで区間計測し、Nsight Systems の NVTX タイムラインと連携できる。書籍中の実測例では、persistent GEMM カーネル(159.352 ms)が cuBLAS(174.920 ms, M=N=8192, K=512)を上回った。 (Source: [[@2025__OReilly__AI Systems Performance Engineering - Chapter 14 PyTorch Compiler, OpenAI Triton, and XLA Backends]]) ## 位置づけ・競合 2025年、NVIDIA は Triton と競合する Python 中心の CUDA ライブラリ群(cuTile、CuTe Python DSL、CUTLASS Python DSL、cuPyNumeric)を発表したが、本章執筆時点では TorchInductor は依然として Triton を主要な GPU コード生成経路として使用している。(Source: 同上) ## 関連 - [[torch.compile]] / [[PyTorch]] / [[OpenAI]] - [[@2025__OReilly__AI Systems Performance Engineering - Chapter 14 PyTorch Compiler, OpenAI Triton, and XLA Backends]] - [[wiki/entities/AI Systems Performance Engineering|AI Systems Performance Engineering]] ## 出典 - Chris Fregly, *AI Systems Performance Engineering*, O'Reilly Media, 2025, Chapter 14「PyTorch Compiler, OpenAI Triton, and XLA Backends」.