Triton（GPU言語） とは？

通常のPython全体がGPUへ移るわけではないよ。Tritonのカーネルでは、Tritonが扱う演算でデータの読み込み・計算・書き込みを記述する。GPUで多くのデータをまとめて処理する部分を作る道具なんだ。

CUDAで個々のスレッドの処理を書くのに対し、Tritonではデータのブロックに対する計算を中心に書くよ。コンパイラーがスレッドへの割当てやメモリを使う処理を組み立てる。ただし、ブロックのサイズや読み書きの形などは設計が必要で、GPUの仕組みを知らなくても常に最適になるわけではないね。

torch.compileの既定バックエンドTorchInductorは、対応GPU向けのコード生成でTritonを使うよ。自分で書くカーネルと、コンパイラーが生成するカーネルは使い方が違う。選ぶバックエンドや機器によって処理が変わるので、すべてのモデルが全部Tritonになる、とは考えないでね。

公式チュートリアルには行列積やfused attentionの例があるよ。ただし、Attentionの実装が全部Tritonというわけではない。特定のモデルが何を使うかや何倍速いかは実装と計測条件による。既存のライブラリで十分なら、先にそれを使う選択肢もあるね。

使うTritonの版とGPU・ドライバー・OSなどの条件を確認しよう。公式リポジトリには対応プラットフォームが案内されている。速度は計算、データの形、メモリ転送、コンパイル時間などで変わるので、基準の実装と比べて測り、出力の正しさも確かめることが大切だよ。