MLSys向けモダンGPUプログラミング

MLSys向けモダンGPUプログラミング# 機械学習システムは、現代のAIワークロードの中心に位置しています。これらのシステムにおいて、パフォーマンスは少数のGPUカーネルの品質に左右されることがよくあります。アテンションカーネル、LLMのプリフィルおよびデコードカーネル、低精度ブロックスケールGEMM、融合MoEレイヤー、その他の大規模融合カーネルはすべて、トレーニングとサービングの両方でエンドツーエンドの速度を直接左右します。しかし、これらのカーネルを高速化するには、最適化トリックのリスト以上のものが必要です。現代のGPUは、もはや古い設計の単純なバリエーションではありません。最近のアーキテクチャは、より豊富なメモリ空間、新しいアクセスパターン、そしてますます専門化された実行ユニットを導入しています。それらをうまくプログラミングするには、ハードウェアの明確な精神モデルと、高性能カーネルがどのように構築されるかという実践的な理解の両方が必要です。本書は、その両方を開発することを目的としています。本書はシンプルな進行に従います。まずGPUハードウェアを理解し、次に使用するプログラミングモデルを学び、最後に最先端のカーネルを段階的に構築します。主なターゲットはBlackwell世代であり、主な実行例は高速行列乗算（GEMM）とFlashAttentionです。その過程で、GPU最適化の背後にある核心的な要素である、データレイアウト、非同期データ移動、非同期協調についても学習します。この教材は、カーネギーメロン大学の機械学習システムコースシリーズから派生したものです。アイデアを研究しやすく、実行しやすくするために、本書ではTIRx Python DSLを使用して、実際のGPUカーネルの例を段階的に構築します。TIRxはハードウェアに密接であり、実行可能なコードを通じて学習しながら、低レベルの制御について推論することができます。 本書の構成# パートI、GPUの理解。このパートでは、GPUの全体的な構成、高速カーネルを記述するための一般的なレシピ、およびデータレイアウト、非同期メモリ操作、協調などの主要な概念を紹介します。本書の残りの部分が依拠するハードウェアの直感を構築します。パートII、TIRxの概要。このパートでは、本書全体のコード例の基盤となるTIRxの主要な要素を紹介します。パートIII、GEMM：タイル化からSOTAまで。TMAパイプライン、永続的スケジューリング、ワープ特殊化、2-CTAクラスターを通じて構築された、タイル化GEMMを最適化するための完全なガイドです。パートIV、Flash Attention 4。パートIIIの技術から構築された完全なアテンションカーネル：2つのMMAとそれらの間のソフトマックス、オンラインソフトマックスリスケーリング、因果マスキング、GQA。リファレンス。TIRx言語リファレンスとコンパイラの内部。 パートI、GPUの理解 GPU実行モデル カーネルを高速化する要因 データレイアウトとその表記法 GPU世代間でのTensor Coreオペランドレイアウト 非同期データ移動：TMA Tensor Cores：tcgen05 特殊メモリ：TMEM 非同期協調：mbarriers 高度：クラスタ起動制御 パートII、TIRxの概要 TIRxの紹介 TIRxレイアウトAPI パートIII、GEMM：タイル化からSOTAまで タイル化GEMMの構築 GEMM最適化パス ステップ1：逐次シングルタイルGEMM ステップ2：K-ループ蓄積 ステップ3：空間タイル化（マルチCTA） 演習 TMAによるGEMMのパイプライン化 ステップ4：TMA非同期ロード ステップ5：ソフトウェアパイプライン（PIPE_DEPTH=2） ステップ6：永続カーネル + タイルスケジューラ 演習 ワープ特殊化とクラスターによるGEMMのスケーリング ステップ7：ワープ特殊化 + パイプライン ステップ8：2-CTAクラスター ステップ9：マルチコンシューマワープ特殊化 エンドツーエンドの結果 演習 パートIV、Flash Attention 4 Flash Attention 4アルゴリズム 形状 タイルプリミティブグラフ ワープの役割とスコープ フラグメントの読み取り 2つのMMAフェーズ TMEMレイアウトと再利用 バリアが役割をどのように接続するか パイプライン構造 リスケーリングと書き戻し 因果マスキング GQAサポート タイルスケジューリング コンパイルと検証 GEMMとの違い 演習 リファレンス リファレンス ワープ特殊化カーネルのデバッグ コンパイラの内部 TIRx言語リファレンス