CUDAに狂わずにBlackwellを飼い慣らす方法

GPUカーネルの記述と聞くと、大抵の場合はCUDA C++の延々と続くコードページ、ポインタの操り合い、メモリリークやレースコンディション的地獄のデバッグを想像するでしょう。NVIDIA Blackwellアーキテクチャ（sm_100a）のような化け物が市場に出てくると、状況はさらに複雑になります。内部には非同期的処理や新しいデータエンジンが山ほど詰まっており、従来の「コアにスレッドを投げるだけ」というアプローチでは、パフォーマンスの半分も引き出せなくなっているのです。

最近、MLC-AIチームのmodern-gpu-programming-for-mlsysプロジェクトを見つけました。本質的には、リポジトリであると同時に、モダンGPUプログラミングに関するライブインタラクティブチュートリアルです。彼らは全力投入狀態：PythonとTIRxという新しいDSLを使ってSOTA（state-of-the-art）カーネルの書き方を教えてくれるのです。

なぜこれは単なるチュートリアルではないのか

通常、GPU курсは2010年代の古典的な行列積に足を止めがちです。ここでは、明日のハードウェアに焦点が移されています。中心はBlackwellアーキテクチャ、そのメモリ階層、テンソルコア、非同期データ移動メカニズムです。

著者は、ハードウェアの理解から実際のモダンMLシステムで動作するカーネル記述までの道を提案しています。低レベルC++と格闘する代わりに、TIRx（Tensor IR next）で記述します。これはApache TVMの拡張機能で、レジスタ、キャッシュ、テンソルエンジンへの制御を維持しながら、Pythonでカーネルロジックを記述できます。

このチュートリアルの内容

すべてのコンテンツは、難易度を引き上げていく論理的なブロックに分割されています。

まず、実行モデルとメモリモデルをカバーします。退屈な定義はなく、TMA（Tensor Memory Accelerator）の実践的な操作と、データが実際に計算ユニットに届く仕組みの理解が中心です。興味深いことに、著者はRooflineモデルに詳細に踏み込んでいます—これにより、コードがメモリ帯域幅かコアの純粋な計算能力のどちらにボトルネックがあるか理解できます。

次にTIRxの深掘りに入ります。タイル（TileLayout）、軸、スワズルメカニズム（bank conflictを避けるためのメモリ内の巧妙なデータシャッフル）の操作を学びます。そして最後には動作するGEMM（行列積）カーネルができあがります。

その後、本題の「コア」部分に移ります。著者は通常のGEMMをハイパフォーマンスな化け物に変える方法を示します。TMAパイプライン、ウェープ特化化、CTAのクラスタリングが活躍します。これらの用語が呪文のように聞こえるとしても、チュートリアルが丁寧に整理してくれます。

最終章では、Flash Attention 4を作成します。これはモダン大規模言語モデルが実際に動作しているものです。カーネルには2つのMMA（Matrix-Matrix Accumulation）ステージ、オンラインソフトマックス、GQA（Grouped Query Attention）サポートが含まれています。

技術スタックと特徴

このプロジェクトはApache TVMエコシステムと密接に連携しています。サンプルを実行するには以下が必要です：

1. ライブラリapache-tvm（TIRxは中に含まれています）。
2. 結果の検証とデータ準備用のPyTorch。
3. 実際のハードウェアでコードを実行したい場合はBlackwell GPU（B200など）。

手元にBlackwellがない場合（今の私たちの大半にとって当然のことです）、それでもリポジトリは理論的な基盤として非常に有用です。メモリオプティマイゼーションと計算スケジューリングのロジックは古いアーキテクチャにも適用できますが、TMAのような一部機能は新しいNVIDIA世代に固有です。

ローカルの Sphinxでビルドし、読みやすい形式で読むには、標準的なSphinxで十分です：

pip install -r requirements-docs.txt
sphinx-build -b html . _build/html

時間をかける価値はあるか

カーネル記述に飛び込むのを怖がっているMLインフラ開発者に经常出会います。それを「C++の神々」の特権と考えています。このプロジェクトは、適切な抽象化（TIRxのような）とアーキテクチャの理解があれば、Pythonエコシステムの中でワールドクラスのコードを記述できることを証明しています。

このプロジェクトは、ニューラルネットワークの「内部」最適化に取り組む人やすぐにFlash Attentionが通常のアプローチより高速な理由を理解したい人にとって、特に有用です。明日にでも自分のBlackwellライブラリをリリースする予定がなくても、テンソルメモリと同期キューの仕組みを追いかけ見ると、本当に頭を明晰にしてくれて、ハイパフォーマンスコンピューティングへの見方が変わります。

唯一の欠点は、入門障壁がまだ高いことです。これは「GPU初心者向け」ではなく、「GPUプロになりたい人向け」です。しかし、グラフィックカードの内部でバイトが実際にどのように移動するかを理解する準備ができているなら、今のところこのテーマに関するより優れた資料を見つけるのは難しいでしょう。