カスタム GPU カーネル用の Triton 言語

GPU プログラミングでは、多くの場合、メモリの移動を管理しながら、多くのスレッドにわたるパーティショニング作業が必要になります。 Triton は、カーネルが値のブロックに対する操作を記述する、Python ベースのドメイン固有言語を提供します。 Triton コンパイル用に装飾された関数は、プログラム ID などの構成要素を使用してブロックや範囲を識別し、要素オフセットを作成できます。コンパイラはこのブロックレベルのプログラムを GPU 実行にマッピングするため、開発者は従来の CUDA C++ のようにすべてのスレッド命令を指定しなくてもカスタム操作を作成できます。基本的なベクトル加算では、長いベクトルをブロックに分割します。各 Triton プログラムは 1 つのブロック ID を処理し、オフセットを計算し、両方の入力から値をロードし、それらを加算して出力を保存します。最後のブロックは有効なテンソル長を超える可能性があるため、マスクはロードとストアを保護します。マスキングが正しくないと、無効なメモリが読み取られたり、有効な値が省略されたりする可能性があります。ワープの数やブロック サイズなどの起動設定は、コンパイルと実行に影響します。行列演算の場合、タイリングによりデータを高速オンチップ メモリに保存し、計算全体で値を再利用できます。これによりデータ移動パターンが改善されますが、タイル サイズはハードウェア リソースに適合する必要があり、入力形状によって動作が異なる場合があります。 Triton のコンパイラは最適化と一部の自動チューニング ワークフローをサポートしていますが、カーネルは最適化されたライブラリ操作より自動的に高速になるわけではありません。小さな入力は起動オーバーヘッドによって支配される可能性があるため、コンパイル時間を定常状態の実行時間と混同しないでください。カーネルの開発には、シェイプ、ストライド、dtype、デバイスにわたる正確性チェックが必要です。浮動小数点演算に適切な数値許容誤差を使用して、信頼できる実装と比較します。ウォームアップ、同期、代表的なワークロードを使用したベンチマーク。パフォーマンスが期待と異なる場合は、生成されたコードまたはプロファイラ トレースを検査します。 Triton は、低レベルのボイラープレートの量を削減します。メモリの結合、占有、精度のトレードオフ、または競合状態を理解する必要性がなくなるわけではありません。カスタムの融合操作または特殊なパターンがメンテナンス コストを正当化する場合に使用し、ハードウェアまたはコンパイラのサポートが異なる場合でも信頼性の高いフォールバックを保持します。