# HipKittens: Fast and Furious AMD Kernels > [!info] Talk metadata > - **会議:** [[MLSys2026]] Day 4 (May 21 / Thu)、Grand Ballroom 2、17:45 - 18:00 PDT > - **登壇者:** William Hu ほか（Stanford University / HazyResearch, AMD） > - **URL:** https://mlsys.org/virtual/2026/oral/3735 > - **OpenReview:** https://openreview.net/forum?id=xxSSrndQrI > - **関連研究:** https://github.com/HazyResearch/HipKittens > [!abstract] 概要（MLSys サイト） > AMD GPU は最先端の演算性能とメモリ帯域幅を提供するが、ピーク性能を引き出すカーネルは生のアセンブリで書かれている。AI アルゴリズムをハードウェアへ効率的にマッピングする難しさに対処するため、近年 ThunderKittens (TK) のような C++ 組み込み型・PyTorch 風のドメイン固有言語が NVIDIA 向けに提案されている。本研究では、タイルベースのプログラミング・最適化されたメモリアクセス・ワーカー間の非同期実行といったプリミティブが NVIDIA 固有か汎用かを検証する。高性能な AMD AI カーネルに必要なプログラミングプリミティブの初の体系的研究を提示し、その知見を HipKittens (HK) フレームワークとして具現化した。タイルベースの抽象は AMD GPU に汎化可能だが、その具体化アルゴリズムは AMD 向けに再設計する必要がある。CDNA3 および CDNA4 プラットフォーム上で検証し、HK カーネルは GEMM およびアテンションにおいて AMD の手書きアセンブリカーネルと同等の性能を達成し、コンパイラベースラインを一貫して上回る。一部設定（d=64 アテンション、GQA 非因果逆伝搬、メモリ律速カーネル）では全ベースラインを 1.2-2.4 倍上回る。 ## 著者・所属 William Hu$^1$, Drew Wadsworth$^1$, Sean Siddens$^2$, Stanley Winata$^2$, Daniel Y. Fu$^3$, Ryan Swann$^2$, Muhammad Osama$^2$, Christopher Re$^1$, Simran Arora$^1$（$^1$Stanford University / HazyResearch、$^2$AMD、$^3$UC San Diego）。ACM Artifacts Available / Evaluated / Results Reproduced の三冠バッジを取得。 ## テーゼ・背景 AI は歴史的に NVIDIA GPU 単一ベンダーに依存してきたが、AMD CDNA4 GPU（MI355X）は BF16 で 2.5 PFLOPS、メモリ帯域幅 8.0 TB/s と NVIDIA B200（BF16 2.2 PFLOPS）に匹敵するスペックを持つ。しかし成熟したソフトウェアの欠如が「CUDA モート」を形成しており、ピーク性能カーネルは一握りの専門家が生アセンブリ（AITER ライブラリ）で書いている。たとえば AITER と PyTorch Llama の GQA 逆伝搬は MI355X のピーク性能のわずか 30% および 24% しか達成できていない。 ### NVIDIA B200 vs AMD MI355X スペック比較 | 項目 | NVIDIA B200 SXM5 | AMD MI355X OAM | |---|---|---| | BF16 マトリクス / テンソル | 2.2 PFLOPS | 2.5 PFLOPS | | MXFP8 マトリクス / テンソル | 4.5 PFLOPS | 5.0 PFLOPS | | MXFP6 マトリクス / テンソル | 4.5 PFLOPS | 10.1 PFLOPS | | MXFP4 マトリクス / テンソル | 9.0 PFLOPS | 10.1 PFLOPS | | メモリ容量 | 180 GB | 288 GB | | メモリ帯域幅 | 8.0 TB/s | 8.0 TB/s | 一方 NVIDIA 側では、ThunderKittens (TK) や CuTe DSL、Gluon といったタイルベースの C++ 組み込み DSL が高性能カーネル開発を簡素化してきた。しかしこれらは全て PTX/CUDA に依存し AMD では動作しない。Triton や Mojo などのコンパイラベースのアプローチも AMD 向けにコンパイルできるが、レジスタ寿命の追跡やメモリアクセスの最適な命令への低下に苦闘し、AMD でのピーク性能を達成できないことが多い。Mojo の MHA カーネルは MI355X でピーク性能の 50% にとどまり、共有メモリのバンクコンフリクトが原因である（`rocprofv3 --pmc SQ_LDS_BANK_CONFLICT` で測定）。 本研究の核心的な問いは、「AMD カーネル開発を簡素化するために全く新しいプログラミングプリミティブが必要か、それとも既存のプリミティブで十分か」である。 ## AMD GPU アーキテクチャ詳細 ### 演算階層 AMD GPU の演算階層はスレッド→ウェーブ（64 スレッド、NVIDIA のワープに相当）→ブロック（CU 上で共同スケジュールされるウェーブ群）→XCD（チップレット）→グリッド（GPU 全体）の階層構造を持つ。AMD MI355X は 256 CU を 8 つの XCD（アクセラレータコンプレックスダイ）に配置し、各 XCD は 32 CU（CDNA4）を含む。CDNA3（MI325X）では各 XCD が 38 CU を持つ。各 GPU は合計 1,024 マトリクスコアを搭載する。 ### メモリ階層 メモリは以下の階層で構成される。 1. **レジスタファイル:** 各 SIMD は 512 本の 32 ビットベクトルレジスタを持ち、CU あたりの合計は 512 KB である。シングルウェーブカーネルの場合、ハードウェアがこれを 256 本の VGPR（汎用ベクトルレジスタ）と 256 本の AGPR（アキュムレータレジスタ）に分割する。 2. **L1 キャッシュ + 共有メモリ（LDS）:** 各 CU は L1 キャッシュと共有メモリ（165 KB）を持ち、同一スレッドブロック内の複数ウェーブがアクセス可能である。 3. **L2 キャッシュ:** 各 XCD 専用の 4 MB。非プログラマブルであり、キャッシュミスのペナルティは最悪 300 ns である。 4. **LLC（ラストレベルキャッシュ）:** 全 XCD で共有される 256 MB。L2 と HBM の間に位置し、キャッシュミスのペナルティは 500 ns である。 5. **HBM（グローバルメモリ）:** 288 GB の HBM3E。帯域幅 8.0 TB/s。 XCD 間は Infinity Fabric（Advanced Package 版、5.5 TB/s 双方向）で接続される。 ### オキュパンシ スレッドは物理実行ユニット（ALU、FMA、マトリクスコア）で命令を実行する。各ユニットの命令は固定の発行レイテンシと限られた帯域を持つ。異なるウェーブは異なるユニットを同時に占有でき、単一ユニットの飽和を回避できる。 ### NVIDIA との用語マッピング | AMD | NVIDIA | |---|---| | ウェーブ（Wave） | ワープ（Warp） | | CU（Compute Unit） | SM（Streaming Multiprocessor） | | SIMD | サブパーティション | | XCD（Accelerator Complex Die） | GPC（Graphics Processing Cluster） | | LDS（Local Data Share） | 共有メモリ | | MFMA | テンソルコア | | AGPR | アキュムレータレジスタ | ## 提案手法：HipKittens フレームワーク HipKittens (HK) は AMD GPU 向けの最小限かつ主張を持った C++ 組み込みプログラミングプリミティブ集合である。TK と同じタイルベースのインターフェース（PyTorch 風の `mma`、`exp`、`add` 等の演算子）を維持しつつ、AMD 固有のハードウェア制約に対応する 3 つの柱を持つ。 ### 1. プログラマブル GPU メモリの最適化されたアクセスパターン **開発者制御のレジスタスケジューリング。** AMD CDNA では SIMD あたり 512 本のレジスタが全ウェーブに静的分割される。HIPCC コンパイラは AGPR（アキュムレータレジスタ）をマトリクスコア命令の入力オペランドとして使用することを妨げる。ハードウェア自体は AGPR をマトリクスコア入力として受け付けるにもかかわらず、HIPCC がこれを阻止するため、コンパイラ経由のカーネルでは AGPR から VGPR へのデータ移動命令（`v_accvgpr_read`）が余分に発生する。HK はコンパイラをバイパスし、各タイルに属するレジスタを開発者が明示的にピン留めする機能を導入する。ピン留めインターフェースは標準のコンパイラ管理タイルと同じ API を維持し、開発者が制御レベルを選択できる。 **Table 1: 明示的レジスタスケジューリングの効果**（バッチ 16、ヘッド 16、ヘッド次元 128 の MHA 非因果逆伝搬）: | 手法 | シーケンス長 | TFLOPS | |---|---|---| | HK（標準） | 4096 | 855 | | HK（ピン留めレジスタ） | 4096 | **1024** | | AMD Assembly (AITER) | 4096 | 1018 | | HK（標準） | 8192 | 909 | | HK（ピン留めレジスタ） | 8192 | 1091 | | AMD Assembly (AITER) | 8192 | **1169** | **異種マトリクスコア形状向けタイル。** AMD のマトリクス命令レイアウトは NVIDIA と異なり統一的な構造を持たない。NVIDIA は全形状が $16 \times 16$ のコアマトリクスを基盤とし、全体の命令形状はこのコアマトリクスを繰り返し配置して構成されるため、単一のスウィズル戦略（TK や Linear Layouts）で全形状に対応できる。一方 AMD の各マトリクス命令は独自のレイアウトを持ち、タイルレイアウトの組み合わせが爆発的に増加する。 さらに AMD の共有メモリアクセスではウェーブ内のスレッドが**フェーズ**単位で実行される。命令ごとにバンク数・フェーズ数・アクティブスレッドの組み合わせが異なる。たとえば `ds_read_b128` は 64 バンク・4 フェーズで実行され、`ds_read_b96` は 32 バンク・8 フェーズで実行される。これらのフェーズは CDNA ISA では文書化されておらず、著者らはソルバーを作成して特定した。 HK は頻出するレイアウトの組み合わせを特定し、行レイアウトと列レイアウトの両方でバンクコンフリクトフリーなスウィズルパターンを提供する。たとえば 16x32 の BF16 タイルでは、先頭 8 列と末尾 8 列（8 行目から開始）を入れ替えるスウィズルにより、`ds_read_b128` の行レイアウトと `ds_read_b64_tr_b16` の列レイアウトの両方で 2-way バンクコンフリクトを解消する。 **グローバルメモリの非同期ロードとスウィズリング。** AMD GPU は HBM から共有メモリへの直接非同期ロードをサポートする（NVIDIA の TMA に類似するが、レジスタファイルをバイパスする）。入力はスレッドごとの HBM アドレスである。TK のように共有メモリアドレスを直接スウィズルする代わりに、HK では HBM アドレス側でスウィズリングを行う。 HK のタイルメモリ管理は 3 階層に対応する。(1) **レジスタ:** 既定では最小の MFMA 命令形状に合わせて初期化し、スケジューリング制御を最大化する。エッジケースでは開発者が MFMA 命令形状をパラメータとして指定可能。(2) **共有メモリ:** 上述のレイアウト別スウィズルパターンを自動適用。(3) **グローバル:** HBM アドレスでのスウィズルにより共有メモリへの直接非同期ロードに対応。 ### 2. 演算とメモリの重複スケジューリング NVIDIA カーネルおよび DSL で主流のウェーブ特殊化パターン（専用のプロデューサーウェーブがメモリ操作、コンシューマーウェーブが演算を担当）は、AMD の CDNA3/CDNA4 では性能が劣化する。AMD のレジスタ静的割り当てにより、プロデューサーウェーブが演算に寄与しないままレジスタを消費し、スレッドブロックあたりの出力タイルサイズ（演算密度）が制限されるためである。MI355X ではウェーブ特殊化は BF16 GEMM ピーク性能の 80% にとどまる。 NVIDIA でウェーブ特殊化が有効な理由は、(1) TMA による専用メモリアクセスハードウェア、(2) 共有メモリまたはテンソルメモリから直接オペランドを受け取る非同期行列乗算（`wgmma`、`tcgnen05`）、(3) プロセッサあたりの大きな共有メモリ（B200 は MI355X より 40% 大きい SRAM）、(4) レジスタ再配置（TMA のレジスタ効率によりプロデューサーがコンシューマーにレジスタを譲渡）、(5) ハードウェア同期プリミティブ（`mbarriers`）の存在である。AMD にはこれらのアーキテクチャ機能がない。 **Table 2: プロデューサー・コンシューマー比較**（$M = N = K = 8192$ の BF16 GEMM）: | 構成 | MFMA 形状 | 出力タイルサイズ | TFLOPS | |---|---|---|---| | HK 4P / 8C | $16 \times 16 \times 32$ | $128 \times 256$ | 893 | | HK 4P / 12C | $16 \times 16 \times 32$ | $192 \times 256$ | 1278 | | HK 0P / 8C | $16 \times 16 \times 32$ | $192 \times 256$ | 1281 | | HK 0P / 8C | $16 \times 16 \times 32$ | $256 \times 256$ | **1610** | | TK (B200) | $256 \times 256 \times 16$ | $256 \times 256$ | 1538 | | CUTLASS (B200) | $256 \times 256 \times 16$ | $256 \times 256$ | 1570 | プロデューサー 0（ウェーブ特殊化なし）で $256 \times 256$ の出力タイルサイズを計算する場合にのみ同等性能を達成でき、プロデューサー数が増えるほど性能が低下する。共有メモリアトミクスを `mbarriers` の代替として使用してもオーバーヘッドは無視できる水準であり、出力タイルサイズが性能を支配する主要因である。 HK は代わりに 2 つの代替スケジューリングパターンを同定した。 - **8-wave ping-pong（均衡ワークロード向け）:** スレッドブロックあたり 8 ウェーブ（SIMD あたり 2 ウェーブ常駐）を使用し、2 グループ（Waves 0-3 と Waves 4-7）に分ける。各 SIMD 内で 2 つのウェーブが演算とメモリの役割を交互に切り替える。一方が演算命令（MFMA）を発行している間、他方がメモリプリフェッチ（HBM→共有メモリ→レジスタ）を行い、条件付きバリアで交替を制御する。条件付きバリアは `if (kittens::warpid() / 4 == 1) { __builtin_amdgcn_s_barrier(); }` の形で実装され、同一 SIMD 上の 2 ウェーブのうち片方のみをストールさせる。演算時間とメモリ時間がほぼ均衡するワークロードに適する。BF16 GEMM のホットループはわずか 48 行のコードでアセンブリカーネルと同等の性能を達成する。8-wave パターンは大きなタイルプリミティブを使用でき、TK のウェーブ特殊化と類似したプログラミングモデルである。 - **4-wave interleave（不均衡ワークロード向け）:** プロセッサの 4 つの SIMD にそれぞれ 1 ウェーブを配置し、各ウェーブが演算とメモリ命令を細粒度でインターリーブする。MFMA と LDS の両パイプラインを同時に飽和させ、演算偏重・メモリ偏重いずれのワークロードにも動的に適応する。SIMD あたり 1 ウェーブが命令ミックスを動的に調整できるため、小さなベースタイルプリミティブで細粒度のスケジューリングが可能だが、コードサイズは大きくなる。GQA 非因果逆伝搬では 8-wave の 1.8 倍に対し 4-wave は 2.3 倍のベースライン超えを達成する。 **Table 3: スケジューリングパターンの性能とコード量のトレードオフ:** | カーネル | パターン | コード行数 | TFLOPS | |---|---|---|---| | FP8 GEMM | 8-wave | 48 | 3222 | | FP8 GEMM | 4-wave | 183 | 3327 | | MHA 逆伝搬 | 8-wave | 331 | 894 | | MHA 逆伝搬 | 4-wave | 989 | 1091 | 注目すべきは、単純な 8-wave パターンだけで BF16 GEMM、FP8 GEMM、アテンション順伝搬のアセンブリカーネルと同等以上の性能を達成できる点である。4-wave は GQA 非因果逆伝搬のような不均衡ワークロードでさらに大きな高速化を実現する。 ### 3. 非プログラマブル GPU メモリのアクセスパターン最適化 AMD MI355X は 8 つの XCD（アクセラレータコンプレックスダイ）から成るチップレット構成を採用し、各 XCD が 32 CU と専用 4 MB L2 キャッシュを持ち、全 XCD が LLC（256 MB）を共有する。スレッドブロックはラウンドロビンで XCD に割り当てられるため、素朴な行優先スケジュールでは L2 ヒット率がわずか 36% に留まる。 実効帯域幅はキャッシュヒット率の関数としてモデル化される: $\text{Bandwidth} = \text{LLC Bandwidth} \times \text{LLC Hit\%} + \text{L2 Bandwidth} \times \text{L2 Hit\%}$ L2 帯域は LLC 帯域のおよそ 3 倍であるため、L2 ヒット率の最大化が性能上最も効果的である。 HK はチップレットスウィズリングアルゴリズム（Algorithm 1）を導入し、L2 再利用と LLC 再利用を同時に最適化する。アルゴリズムは 2 段階で構成される。 **ステップ 1: XCD グルーピング。** 2D グリッドを 1 次元シーケンスにフラット化し、ブロック ID をリマップする。ハードウェアのラウンドロビン XCD 割り当てを逆デインターリーブし、チャンクサイズ $C$ 個の連続ブロック ID が同一 XCD に常駐するよう再配置する。これによりクロスチップレットトラフィックが削減される。 **ステップ 2: 階層的ウィンドウ走査。** グリッドを行ごとに処理する代わりに、高さ $W$ の垂直ウィンドウ内を列方向に先に走査し、$W$ 行分を処理したら次の列に移動する。これにより入力ブロック ID 空間を矩形タイルに「折り畳む」効果が生じ、L2 キャッシュの再利用が最適化される。 - **L2 再利用:** 同一 XCD に割り当てられるスレッドブロックが出力行列の矩形領域（L2 タイル）をカバーするよう配置し、入力行列 $A$ の同一行・$B$ の同一列を再利用する。MI355X では各 XCD が 32 CU を持つため、$8 \times 4$ または $4 \times 8$ の L2 タイル形状が最適なハードウェア利用率を達成する。ただし L2 のみの最適化は各 XCD が $A$ と $B$ の互いに素な部分をフェッチするため、LLC レベルで冗長なロードが発生しうる。 - **LLC 再利用:** 複数 XCD が入力行列の近傍領域にアクセスするようチャンクサイズ $C$ を調整し、LLC に残留する共有データを活用する。理想的には全 XCD のアクセスフットプリント（LLC タイル）が $A$ と $B$ の両方で重複するよう配置する。 ウィンドウ高 $W$ とチャンクサイズ $C$ の 2 パラメータで L2/LLC 再利用のトレードオフを制御する。 **Table 4: キャッシュスケジュール比較（BF16 GEMM、マクロタイル $192 \times 256 \times 64$）:** | ブロック順序 | L2% | LLC% | 実効帯域幅 | TFLOPS | |---|---|---|---|---| | **$M=N=K=9216$** | | | | | | 行優先 | 55% | 95% | 15.1 TB/s | 1113 | | XCD ($W$7/$C$216) | 79% | 24% | 14.9 TB/s | 991 | | XCD ($W$5/$C$25) | 75% | 93% | 18.3 TB/s | **1145** | | **$M=N=K=14592$** | | | | | | 行優先 | 36% | 76% | 10.7 TB/s | 900 | | XCD ($W$8/$C$542) | 79% | 7% | 13.9 TB/s | 980 | | XCD ($W$8/$C$64) | 78% | 55% | 16.6 TB/s | **1068** | $M=N=K=14592$ の場合、出力行列の幅（タイル数）が XCD 数 8 と互いに素（57 タイル）であるため、デフォルトの行優先スケジュールでは最悪のキャッシュ再利用パターンとなる。この場合、L2/LLC 対応スケジュールはデフォルトに対し最大 15% 以上の性能向上を達成する。L2 帯域は LLC 帯域の約 3 倍であるため、$W$ は L2 ヒット率を最大化するよう選択すべきである。 ## 実験・結果 評価は AMD CDNA3 MI325X および CDNA4 MI355X 上で実施。比較対象は PyTorch (SDPA)、ROCm Library Triton、Composable Kernel、AITER（アセンブリ）、HipBLASLT。ベンチマークは AMD が最近公開したベータ版 ROCm 7.0 Docker（`rocm/7.0-preview:rocm7.0_preview_pytorch_training_mi35x_beta`）上で、$\mathcal{N}(0,1)$ からのランダム入力テンソルに対し 500 回ウォームアップ後 100 回平均の TFLOPS/s を測定。HK カーネルは Python バインディング経由でベンチマーク（FP8 は AMD PyTorch サポートが実験段階のため例外）。 ### BF16 / FP8 GEMM HK はアセンブリで書かれた AMD ベースラインカーネル（AITER、HipBLASLT）と同等の性能を達成する。Triton コンパイラに対しては 1.3-3.0 倍の高速化を示す。HK の GEMM カーネルのホットループは 100 行未満のコードであり、単一の 8-wave スケジュールが評価した全問題形状に汎化する。 具体的な BF16 GEMM 性能（MI355X、正方行列 $N \times N$）: - $N=4096$: HK 1413、AITER 1433、HipBLASLT 1114、Triton 329 TFLOPS - $N=8192$: HK 1591 / 1610、AITER 1579、Triton 504 TFLOPS - $N=16384$: HK 1564 / 1571、AITER 1579、Triton 478 TFLOPS FP8 GEMM では大規模行列（$N \geq 8192$）で HipBLASLT や Composable Kernel がメモリ不足（OOM）になる中、HK は安定して動作する: - $N=4096$: HK 2749、HipBLASLT 2741、CK 548 TFLOPS - $N=8192$: HK 3233、HipBLASLT 3327（HipBLASLT/CK は OOM なし） - $N=16384$: HK 3293 / 3378（HipBLASLT/CK は OOM） ### アテンション順伝搬 MHA および GQA カーネルを因果・非因果設定、ヘッド次元 64 および 128 で評価（バッチ 16、クエリヘッド 64、KV ヘッド 8）。HK は AITER のアセンブリカーネルを含む全ベースラインを平均で上回る。AITER 比 1.0-2.1 倍、PyTorch SDPA 比 1.3-4.5 倍、CK 比 1.0-1.4 倍、Triton 比 1.2-4.5 倍の高速化を達成する。 アテンション順伝搬カーネルは約 500 行のコードで 8-wave ping-pong パターンを採用し、演算クラスタ内では各ウェーブフロントがオンラインソフトマックスのベクトル演算（max / subtract / exp2 / accumulate）と MFMA 命令をインターリーブする。MI355X と NVIDIA B200 の間でスケジューリングとハードウェアに大きな差異があるにもかかわらず、FlashAttention-3 と同等の設定で競合する性能を達成する。 ### アテンション逆伝搬 GQA の因果・非因果逆伝搬で既存ベースラインを 1.8-2.5 倍上回る（バッチ 16、クエリヘッド 64、KV ヘッド 8、ヘッド次元 128）。 具体的な GQA 逆伝搬性能（MI355X）: - **GQA 非因果逆伝搬:** シーケンス長 16384 で HK 1176、AITER 471、CK 312、SDPA 389 TFLOPS（HK が全ベースラインの約 2.5 倍） - **GQA 因果逆伝搬:** シーケンス長 16384 で HK 995、AITER 334、CK 259、SDPA 541 TFLOPS 逆伝搬は AI ワークロードの中でも最もレジスタ消費が激しいワークロードであり、HK カーネルは複数の MFMA 命令形状（$16 \times 16 \times 32$ と $32 \times 32 \times 16$）、行・列レイアウトの異なる共有メモリアクセスパターン（同一共有タイルからレジスタへの行レイアウトと列レイアウトの両方のロード）、明示的レジスタピン留めを組み合わせる。MHA 逆伝搬では AITER のアセンブリカーネルと競合する（MHA 非因果逆伝搬: シーケンス長 16384 で HK 1119、AITER 1104 TFLOPS）。 AITER ライブラリやアセンブリカーネルの不一致な性能（たとえばヘッド次元 64 のアテンションや GQA 非因果逆伝搬でのカバレッジの欠如）は、アセンブリエンジニアリングのスケーリングの困難さを如実に示しており、簡潔なカーネルプログラミング抽象の価値を裏付ける。 ### メモリ律速カーネル 融合 dropout-residual-layernorm カーネルおよび回転位置符号化（RoPE）カーネルでは、HK は AITER および PyTorch compiled カーネルを 1.1-2.2 倍上回る。 ### 学習での検証 カーネルの安定性検証として Llama 1B と BERT 110M を The Pile で事前学習し、PyTorch および AITER で学習したモデルのパープレキシティと 100 億トークン後に一致することを確認している。 ## 関連研究 論文では関連研究を以下のカテゴリに整理している。 1. **手書きアセンブリライブラリ:** AITER および Composable Kernel (CK) は AMD の主要なアセンブリカーネルライブラリであり、HK の直接的な比較対象である。演算命令とメモリ命令を細粒度でインターリーブするアプローチはタイルベースプログラミングと直交する。 2. **タイルベース C++ DSL:** ThunderKittens (TK; Spector et al., ICLR 2024)、CuTe DSL (NVIDIA)、Gluon (Zhou et al., 2025)、TileLang (Wang et al., 2025) はいずれも NVIDIA GPU 専用で PTX/CUDA に依存する。HK は TK のタイルインターフェースを AMD に移植し、具体化アルゴリズムを再設計した最初の試みである。 3. **コンパイラベースアプローチ:** Triton (Tillet et al., 2019)、TVM (Chen et al., OSDI 2018)、LLVM/MLIR ベースのフレームワーク、Mojo は AMD GPU 向けにコンパイル可能だが、再利用可能な原則やプリミティブを提供しておらず、包括的な高性能 AMD カーネルスイートもリリースしていない。 4. **AI 設計カーネル:** Kevin (Baronio et al., 2025) は強化学習で CUDA カーネルを生成する手法、KernelBench (Ouyang et al., 2025) は LLM によるカーネル生成を評価するベンチマークである。現時点ではこれらのモデルも新しいハードウェア機能への対応に苦心しており、報酬ハッキングの問題もある。 ## 限界と今後の課題 論文が明示的に述べている限界・今後の課題は以下の通りである。 - **アセンブリのスケーラビリティ:** アセンブリカーネルは AI ワークロードの多様性に対応できず、多くの重要な設定（ヘッド次元 64 のアテンション、GQA 非因果逆伝搬、メモリ律速カーネル等）をカバーしていない。HK はこの問題に対する解だが、一部の設定（長いシーケンス長での MHA 逆伝搬等）ではなおアセンブリに及ばない。 - **CDNA ISA の文書化不足:** 共有メモリのフェーズ挙動（スレッドがどの順序でバンクにアクセスするか）は ISA ドキュメントに記載されておらず、著者らがソルバーを独自に作成して特定する必要があった。 - **チップレットスウィズリングのチューニング:** ウィンドウ高 $W$ とチャンクサイズ $C$ は問題サイズに依存し、最適値の自動選択は今後の課題である。 - **マルチシリコンの展望:** 抽象とフロントエンドインターフェース（タイルと PyTorch 風の一括操作）は NVIDIA と AMD で共通だが、具体化（スケジュール、メモリ移動、キャッシュ最適化）はハードウェアごとに異なる。ベンダーをまたぐ単一のタイルベースソフトウェアレイヤーの実現は今後の挑戦である。 ## 結論・オープン課題 HipKittens は、高性能 AMD AI カーネルの原則を体系化し、それを構成可能なオープンソースの C++ プリミティブとして具現化した初の取り組みである。主要な知見は 3 点にまとめられる。 1. **タイルは依然として重要なプリミティブである**が、ハードウェア差異により具体化方法は異なる必要がある。 2. **ウェーブ特殊化は NVIDIA 向けの命令スケジューリングに適するが、AMD では 4-wave インターリーブまたは 8-wave ping-pong が有効である。** 3. **チップレット構成と分離メモリ階層の NUMA 効果**は今後の GPU カーネル開発において一級の設計関心事となる。AMD だけでなく NVIDIA（Blackwell は 2 チップ構成）もチップレット化が進んでおり、キャッシュ階層を意識したスケジューリングの重要性は増す。 HK は AITER（https://github.com/ROCm/aiter）に統合されプロダクション利用されており、ROCm の TransformerEngine（https://github.com/ROCm/TransformerEngine）にも採用が進んでいる。OpenAI と AMD は 6 ギガワット規模の AMD GPU デプロイメントの戦略的パートナーシップを発表しており、マルチシリコンプラットフォームでのカーネルサポートのスケーリングが「AI の可能性を最大限に引き出すために必要な演算能力の解放」の鍵であると位置付けている。アセンブリはスケーラブルでなく、AI ワークロードの多様性に対応できない。HK はベンダーをまたぐ単一のタイルベースソフトウェアレイヤーへの道を拓く。