OpenAI Triton 1.0 リリース

OpenAIは、研究者が広範なCUDAの経験を必要とせずに、非常に効率的なGPUコードを記述できるように設計された、オープンソースのPythonライクなプログラミング言語およびコンパイラであるTriton 1.0をリリースしました。Tritonを使用することで、開発者は、cuBLASのパフォーマンスに匹敵するFP16行列乗算のように、専門家によってチューニングされたCUDAカーネルと同等のハードウェアパフォーマンスを、はるかに少ないコード量で実現できます。

GPUプログラミングの複雑さの解決

ネイティブなフレームワークのオペレータにおいて、一時的なテンソルを作成・移動することに伴うパフォーマンスの低下を避けるために、特化型のGPUカーネルを記述することがしばしば必要になります。しかし、手動でのGPU最適化は、開発者が主に3つの主要なアーキテクチャ・コンポーネントを管理する必要があるため、困難です。

DRAM: メモリ転送は、現代的なインターフェースのバス幅を活用するために、大きなトランザクションに結合（coalesced）される必要があります。
SRAM: データは再利用前にSRAMに手動で格納（stashed）され、共有メモリのバンク衝突（bank conflicts）を避けるように管理される必要があります。
ALUs: 命令/スレッドレベルの並列性を、およびテンソルコアを活用するために、計算はStreaming Multiprocessors (SMs) の間および内部で慎重に分割およびスケジューリングされる必要があります。

Tritonは、メモリの結合（coalescing）、共有メモリ管理、およびSMs内でのスケジューリングを自動化する一方で、タイリング（tiling）やSM間同期（inter-SM synchronization）などの高レベルなアルゴリズムの検討事項は開発者に委ねます。

Tritonプログラミングモデル

Tritonは、カーネルがインスタンスのグリッド上で起動されるデコレータ付きのPython関数として定義されるプログラミングモデルを使用します。CUDAやNumbaで使用されるSingle Instruction, Multiple Thread (SIMT) モデルとは異なり、Tritonはブロック（2の累乗の次元を持つ小さな配列）に対する操作を通じて、インスタンス内の並列性を公開します。

主要な技術的相違点

ブロックベースの実行: 個々のスレッドではなくブロックに対して操作を行うことで、Tritonは、共有メモリの同期やテンソルコアのスケジューリングを含む、CUDAスレッドブロック内の並列処理の問題を抽象化します。
ポインタ演算: Triton JITは入力をテンソルではなくポインタとして扱い、block-sparse tensorsのような複雑なデータ構造に必要なメモリへの低レベルな制御を提供します。
Fused Kernels (結合カーネル): Tritonは、結合カーネル（fused kernels）の作成を簡素化します。例えば、Tritonにおける結合softmax実装は、正規化プロセス全体を通じて行をSRAMに保持し、データの再利用を最大化できます。このアプローチは、同等のPyTorch実装よりも最大2倍効率的になる場合があります。

行列乗算のパフォーマンス

Tritonは、ニューラルネットワークのコア操作である行列乗算において非常に効果的です。約25行のPythonコードで、V100テンソルコア上でピークパフォーマンスを達成できます。このアクセシビリティにより、開発者は、並外れたGPUプログラミングの専門知識を必要とせずに、結合変換（スライシングやLeaky ReLUなど）を含む行列乗算カーネルをカスタマイズできます。

システムアーキテクチャとコンパイラバックエンド

Tritonのパフォーマンスは、Triton-IRを中心としたモジュール式アーキテクチャから得られます。これは、多次元ブロックが第一級オブジェクト（first-class citizens）として扱われるLLVMベースの中間表現です。

コンパイルパイプライン

Python AST: @triton.jit デコレータは、Python関数の抽象構文木（Abstract Syntax Tree）を走査します。
Triton-IR: ASTは、静的単一代入（SSA）構築アルゴリズムを使用してTriton-IRに変換されます。
LLVM-IR & PTX: コンパイラバックエンドは、IRを簡素化および最適化し、自動的に並列化し、NVIDIA GPUでの実行のためにLLVM-IRおよび最終的にPTXへと変換します。

コンパイラ最適化

Automatic Memory Stashing (自動メモリ格納): コンパイラは、計算集約的なブロックレベルの操作（例: tl.dot）のオペランドを分析し、liveness analysisを使用してデータを共有メモリに自動的に格納します。
Automatic Parallelization (自動並列化): Tritonは、異なるカーネルインスタンスを同時に実行することでSMs間で実行を並列化し、ブロックレベルの操作のSIMDユニット間で反復空間を分割することでSMs内部で並列化を行います。