HazyResearch/HipKittens

HipKittens とは？

HipKittens は、AMD GPU 上で高性能な AI カーネル（CDNA3、CDNA4、CDNA5）を記述するための低レベルビルディングブロックを提供するオープンソースの C++ ライブラリです。ThunderKittens エコシステムに属し、Triton や Torch‑Inductor などのプロジェクトが使用するのと同じ種類のテンソルコアスタイルのプログラミングを対象としていますが、AMD の ROCm スタック（HIP、hipcc）向けに書かれています。

コアアイデア

アイデア	意味
タイルプリミティブ	AMD テンソルコアユニットに直接マッピングされる小さな固定サイズのブロック。ライブラリは、コアレス、バンクコンフリクト回避、ネイティブテンソルコアレイアウトを使用するメモリアクセスヘルパーを提供します。
Python 風のバルク計算	Python 関数（例：`gemm`、`attention`）のように見える薄い C++ ラッパーですが、インラインアセンブリと HIP 呼び出しにコンパイルされます。
非同期ロード/ストア	共有メモリへの非ブロッキングロードを発行し、ハードウェアがレイテンシを隠せるようにするヘルパー。
スケジューリングパターン	計算とメモリをオーバーラップさせて GPU パイプラインをビジー状態に保つ、8‑wave ping‑pong と 4‑wave interleave の 2 つの再利用可能なパターン。

何ができますか？

BF16 GEMM カーネル（例：8192×8192×8192 行列乗算）で、AITER や PyTorch とベンチマークできます。
Attention カーネル（MHA、GQA、因果/非因果）で、フォワードパスとバックワードパスの両方をサポートし、バッチサイズ、シーケンス長、ヘッド数、ヘッド次元を設定できます。
メモリバウンドカーネル（回転位置埋め込みや融合レイヤー正規化など）。
BERT または LLaMA 用のエンドツーエンドのトレーニングスクリプトで、HipKittens attention カーネルを PyTorch トレーニングループに組み込めます。

これらはすべて ROCm の hipcc でコンパイルされ、提供されている AMD Docker イメージ内で実行できます。

誰が使うのか？

研究者：AMD 専用またはマルチシリコン AI パフォーマンスを探求し、軽量で「メタルに近い」カーネルライブラリを必要とする人。
エンジニア：AMD GPU 用のカスタムトランスフォーマーバックエンドを構築する人。特に、AITER、Triton、HipBLASLT などの既存のバックエンドと比較したい場合。
学生：AMD ハードウェアでのテンソルコアプログラミングを学ぶ人。ライブラリの単体テストとサンプル make ファイルは、具体的な学習教材として役立ちます。

始め方

リポジトリをクローンし、一致する ROCm Docker をプルします（README に MI35x および MI30x GPU 用の 2 つのイメージがリストされています）。
コンテナに入り、env.src を source して環境変数を設定し、比較ベンチマークを実行したい場合は AITER ベースラインをインストールします。
カーネルをビルドします（例：cd kernels/cdna4/gemm/bf16fp32 && make && python bench.py）。bench.py スクリプトは、HipKittens カーネル、AITER リファレンス、PyTorch の実装を実行し、タイミング数値を出力します。
単体テストを実行して GPU 世代をテストします（cd HipKittens/tests/unit/cdna4 && make && ./unit_tests）。
分析スクリプトを HipKittens/analysis で調べて論文のパフォーマンスプロットを再現するか、training/ を参照して簡単な BERT/LLaMA トレーニング実行を行います。

いつ役立つか？

AMD GPU でのトレーニングまたは推論（MI300X、MI350X、MI455 など）で、一般的な ROCm ライブラリよりも高速なカーネルが必要な場合。
新しいスケジューリングやメモリアクセスパターンを実験したいが、生の HIP アセンブリを自分で書きたくない場合。
マルチシリコン AI スタックを構築しており、他の DSL（例：Triton）が使用するのと同じタイルベースの抽象化に従い、AMD ハードウェアで動作するリファレンス実装が必要な場合。

連絡先とコミュニティ：William Hu (willhu@stanford.edu) と Simran Arora (simran@cs.stanford.edu) が HipKittens ユーザー向けの Discord チャンネルを運営しています。リポジトリには、より深い背景情報として arXiv 論文、ブログ記事、YouTube 講義へのリンクもあります。