HazyResearch/ThunderKittens

ThunderKittens – 高性能AIカーネル用CUDA-DSL

何であるか – ThunderKittensは、ヘッダーオンリーのC++/CUDAライブラリであり、深層学習カーネル（例：GEMM、FlashAttention）をコンパクトでタイル指向のコードとして書くことを可能にします。低レベルのGPUの詳細（テンソルコア呼び出し、非同期コピー、共有メモリバンキング）を抽象化しつつ、ハードウェアに近い状態を保ち、最新のNVIDIA GPU（H100、Blackwell、Vera Rubin）上で理論上の速度に近い実行を実現します。

なぜ重要か – 大規模言語モデルの学習と推論は、行列乗算やアテンションカーネルの実行速度に制限されています。こうしたカーネルを手作業で実装するのはエラーを引き起こしやすく、NVIDIAのプログラミングモデルに関する深い知識が必要です。ThunderKittensは、次のような特徴を持つ小型で拡張可能なDSLを提供します：

• 多くのカーネルでソースコードが100行未満に収まる。
• コンパイル時に正しいレイアウト処理を保証する。
• アセンブリを書かずに最新のテンソルコア命令（WGMMA、TCGEN05、MXFP8、NVFP4）を直接利用可能。
• C++から直接使用可能であり、PyBind11を介してPyTorch用にラップすることも可能。

キーポイント

コンセプト	機能
タイルプリミティブ	16×16以上のブロック（レジスタタイル、共有タイル、ベクトル）を操作し、テンソルコアのレーンに自然にマッピングされる。
ワープ / ワープグループ	関数はデフォルトで1つのワープ（32スレッド）向けに記述される。4ワープの協調グループ（ワープグループ）は非同期行列乗算加算命令を暴露する。
TMA / 非同期コピー	NVIDIAのTMA（テンソルメモリアクセス）機構を通じて、読み込み/書き込みを隠蔽し、レイテンシを回避するための組み込みヘルパー。
静的レイアウトチェック	テンプレートがデータ型、形状、メモリレイアウトをエンコードするため、不一致な操作はコンパイル時に検出される。
Load-Store-Compute-Finishテンプレート	メモリ移動と計算を重ねる推奨パターンで、オキュパシーを最大化する。

一般的なワークフロー

1. リポジトリをクローンし、CUDAソースに kittens.cuh を含める。
2. 提供された matmul_layout / matmul_template（または他のプリミティブ）を使ってカーネルを定義する – producer、consumer、およびオプションの common_setup コールバックのみを埋める必要がある。
3. 提供されたMakefileでカーネルをコンパイルする（CUDA 12.8+、C++20）。各カーネルは kernels/ 下の個別のフォルダにあり、独立してビルド可能。
4. Pythonから呼び出し（オプション） – make 後、小さなPyBind11ラッパーにより、PyTorch 2.8+からカーネルを呼び出せる。
5. ベンチマーク / テスト – 正確性テストとパフォーマンススクリプトは、各カーネルと同一フォルダに配置されている。

例：H100で855 TFLOPsの行列乗算

#include "kittens.cuh"
#include "prototype.cuh"
using namespace kittens;
using namespace kittens::prototype;

// レイアウト定義（タイル、グローバルポインタなど）
template<int M_BLOCK, int N_BLOCK>
struct matmul_layout { … };

// producer、consumer、common_setupを接続するカーネルテンプレート
template<int _M_BLOCK=2, int _N_BLOCK=4, int _SUPER_M=12>
struct matmul_template { … };

完全なソース（約100行）はREADMEに記載されており、H100の理論ピークの約86％に達するカーネルにコンパイルされる。

対応ハードウェア – 主にNVIDIA Hopper（H100）およびBlackwell（B200）GPU。2026年9月以降、新しいVera Rubin GPUも対応。Ampereは動作するが、今後は更新されない。AMDユーザーは姉妹プロジェクト HipKittens を参照。

インストールチェックリスト

• CUDA 12.8+（CUDA_HOME を設定し、PATH/LD_LIBRARY_PATH を更新）。
• C++20コンパイラ（gcc-11 または clang-11）。
• （オプション）Pythonバインディングが必要な場合、PyTorch 2.8+ と PyBind11。
• リポジトリをクローンし、ヘッダーを含め、提供されたMakefileでカーネルをコンパイルする。

使用事例 – Together AI、Jump Trading、Cursorなどの企業でのプロダクションスケールの学習・推論。また、スタンフォード大学のHazy Research Labでも内部的に使用されている。

学習リソース

• ThunderKittensマニュアル – タイル、スコープ、API規約をカバーする短いガイド。
• 教育用カーネルシリーズ – kernels/gemm/educational_h100 はステップバイステップのGEMM実装を説明。
• ディープダイブブログ – Hamza Elshafieによる解説（2026年5月）でDSLの内部構造を詳細に分析。
• オンボーディングドキュメント – 新規ユーザー向けのGoogleドキュメント（README内のリンク）。

デモ – demos/ フォルダには、LLM学習・推論用の実行可能な例（例：Qwen、Llama、LoLCATS）が含まれている。ThunderKittensカーネルをPyTorchワークフローに統合し、Hugging Faceの8Bモデルを実行する方法を示している。

---

結論 – NVIDIAのテンソルコアの性能を最大限に引き出したい、カスタムアテンション、GEMM、またはその他の行列中心の演算に特化したカーネルを書く必要がある場合、ThunderKittensは、CUDAの上に直接乗る小型で型安全なDSLを提供し、従来のアセンブリレベルの手間なくプロダクション品質のカーネルを書けるようにします。