AMD GPUを“brrr”な速度で回す方法

 (hazyresearch.stanford.edu)
1 ポイント 投稿者 GN⁺ 2025-11-16 | 1件のコメント | WhatsAppで共有
  • HipKittensは、AMD GPUの潜在性能を引き出すために設計されたプログラミング・プリミティブ群で、メモリアクセス・スケジューリング・キャッシュ再利用を最適化する
  • AMD MI355X GPUは256基のコンピュートユニットと8つのチップレット(XCD) 構造を持ち、大規模なレジスタファイルと細粒度のマトリクスコア命令を提供する
  • NVIDIAと異なり、AMDにはレジスタ再割り当て・非同期行列命令・mbarrierがなく、wave specializationの代わりに8-wave ping-pongと4-wave interleaveスケジューリングが効果的である
  • HipKittensはチップレット認識(grid)スケジューリングによってL2およびLLCキャッシュの局所性を改善し、GEMMとAttention演算で最大帯域幅およびTFLOPSの向上を達成した
  • このアプローチはAMD GPUエコシステムにおけるソフトウェア成熟度の不足を補い、多様なハードウェアに基づくAIコンピューティングの拡張性を高める基盤を提供する

AMD CDNA GPUアーキテクチャと性能特性

  • AMD MI355X GPUは256基のコンピュートユニット(CU) を含み、各CUは4つのSIMDで構成される
    • 1つのSIMDは64スレッドからなるwaveを実行し、これはNVIDIAの32スレッドwarpと対比される
  • MI355XはB200比で70%水準のSRAM(165KB) を持ち、非同期行列積命令、レジスタ再割り当て、テンザーメモリアクセラレーション、mbarrier機能がない
  • 一方で、2倍大きいレジスタファイルと60%多いプロセッサ数(256 CU対160 SM) を提供する
    • 小さく細粒度なマトリクスコア命令をサポートし、直接グローバル→共有メモリロード(TMA類似) 機能が存在する
  • AMDは8つのチップレット(XCD) からなるチップレットアーキテクチャを採用しており、各XCDは独立したL2キャッシュを持ち、その上位にLLCキャッシュが存在する
  • 表によればMI355Xは、BF16 2.5 PFLOPs、MXFP8 5.0 PFLOPs、MXFP6 10.1 PFLOPsの演算性能と、288GBのメモリ容量、8TB/sの帯域幅を持つ

AMD向けカーネル設計の課題

  • メモリアクセス最適化: HIPCCコンパイラの制約と非公開のI/O動作により、データ配置とスウィズル(swizzle)パターンの設計が重要
  • プロセッサ内スケジューリング: AMDでは共有メモリの代わりにレジスタファイルと小型のマトリクス命令を活用する必要がある
  • プロセッサ間スケジューリング: チップレットベースの構造により、キャッシュレベルのNUMA効果を考慮した作業分配が必要

HipKittensのメモリアクセスパターン

  • HipKittens(HK)はタイル(tile) を基本データ単位として使用し、PyTorch類似の演算関数を提供する
    • タイルはデータ型、サイズ、レイアウトで定義され、C++テンプレートメタプログラミングによってさまざまな入力に対応する
  • レジスタスケジューリング: HIPCCは特定のレジスタをMFMA入力として使用できないため、HKは明示的レジスタ固定機能を提供する
    • 開発者が直接レジスタを指定して最大性能のカーネルを作成可能
  • レジスタレイアウト: AMDではデータ型・行列形状に応じてレイアウトが異なるため、単一のスウィズルパターンは不可能
    • 例として、16×16 bf16タイルと16×32 bf16タイルでは異なるスウィズルパターンが必要
  • 命令フェーズ構造: AMDの共有メモリ命令は非連続なフェーズグループを持ち、内部文書化が不足している
    • HKはこれに対してリバースエンジニアリングしたソルバーを提供する
  • アドレス生成: AMDは非同期HBM→共有メモリロードをサポートし、HBMアドレスのスウィズルで最適化を行う

プロセッサ内スケジューリング: Waveパターン

  • Wave specializationはNVIDIAでは効果的だが、AMDではレジスタ再割り当ての不在により性能が低下する
    • Producer waveが不要なレジスタを占有し、Consumer waveはレジスタ不足によってスピル(spill) が発生する
  • HKの実験結果では、wave specializationはAMDで演算強度の低下とメモリボトルネックを引き起こした
    • 例: GEMMでHK 0/8構成は1605 TFLOPs、CUTLASSは1570 TFLOPs
  • 代替スケジューリングパターン
    • 8-wave ping-pong: 2つのwaveが交互にメモリ/演算クラスタを実行する
    • 4-wave interleave: 1つのwaveが細かくメモリと演算を交互実行する
    • 8-waveはコードが簡潔で、4-waveは細かいがコードが長くなる
    • GEMMとAttention Forwardで8-waveがSoTA水準の性能を達成

プロセッサ間スケジューリング: チップレット認識アプローチ

  • AMD MI355Xは8つのXCDチップレットを持ち、各チップレットは独立したL2キャッシュを備える
    • スレッドブロックがラウンドロビン方式でチップレットに割り当てられるため、グリッド順序がキャッシュ再利用効率に直接影響する
  • 単純なrow-major配置ではL2キャッシュ再利用率が低く、帯域幅の損失が発生する
    • 例: L2 55%、LLC 95%、15.1 TB/s、1113 TFLOPs
  • HKはチップレット認識(grid)スケジューリングを導入し、L2・LLCキャッシュ局所性を同時に活用する
    • スレッドブロックを出力行列の隣接領域単位でグループ化して入力データの再利用を最大化する

実際のカーネル例

  • Attention ForwardおよびBF16 GEMMカーネルのホットループ(hot loop) は、HKの8-wave ping-pongスケジュールを使用する
    • 各ループはCompute–Memoryクラスタを交互に実行し、スケジュールバリアで同期する
    • コード例ではmma_AtB, load, exp2, col_sumなどのHK演算が繰り返し使用される

結論: Multi-silicon AI時代のAMD

  • HipKittensはAMD CDNA3・CDNA4で競争力のある性能を達成した
    • 3つの核心: 最適化されたメモリアクセス、AMD中心のwaveスケジューリング、チップレット認識グリッドスケジューリング
  • HKカーネルはAMD基準で最高性能を達成し、NVIDIA Blackwellカーネルとも競争可能な水準
  • AIコンピューティングの多様性のためにAMD GPUへのアクセシビリティ拡大が必要であり、HipKittensはそのための中核ソフトウェア基盤を提供する
  • AMDのHIPCCレジスタスケジューリング改善が今後の重要な発展領域として挙げられる

関連記事

1件のコメント

 
GN⁺ 2025-11-16
Hacker Newsの意見
  • HipKittens 関連の議論も参照するとよい
  • 同じ研究を扱った HipKittens: Fast and furious AMD kernels という投稿もある。George Hotz と AMD 社員のコメントが付いている
  • 学界でこうした問題が扱われるのは歓迎だが、結局これは AMD 内部で解決すべき問題だと思う
    • 私はハードウェア企業はハードウェアだけを作るのがよいと思う。そうすれば インセンティブの純度 が保たれる。性能が 20% 落ちるとしても、そのほうがよいと思う
    • 完全に同意する。AMD は 10 年前にこの問題を先送りし、今になってようやく追いつこうとしている。ハードウェアは素晴らしいが、ファームウェア作成能力の不足のせいで潜在力を生かせていない
    • ただ、この研究チームは Nvidia GPU 向けにも似たソフトウェアを作ったことがある。優れた研究者たちが自分たちの専門性を発揮しているのだと思う
    • 私の知る限り、AMD はすでに複数のレイヤーでこの問題に取り組んでおり、tinycorp とも協力している
  • 記事を読むと、AMD GPU の アーキテクチャ上の複雑さ のせいで最適化が難しいという印象を受ける。ただ、長期的には AMD のアプローチのほうがうまくスケールする可能性もある。Nvidia が 2 チップレットなのに対し、AMD は 8 チップレット構成なので メモリ局所性の問題 がある。将来的にはチップレット数がさらに増えるだろうから、今の複雑さに対処する経験は長期的に役立つかもしれない
    • AMD は高性能化のために warp specialization を必要としないので、プログラミングはより単純だ
  • 多くの開発者が AMD GPU を一般開発者向けに “go brrr” にしようとして失敗してきた。AMD がなぜ自前で ソフトウェアの問題を解決しないのか 理解できない。今は資金も十分あるのに、開発者を雇わない言い訳にはならない。データセンター向け GPU も悪くはないが、個人が ML・AI の実験をするなら今でも Nvidia のほうがずっとよい。5 年前の自分の RTX 3090 のほうが、これまでに出た AMD の民生向け GPU より良いと感じる
    • AMD の開発者体験はひどい。ドライバークラッシュのバグ報告すら受け付けない
    • 私は最近、推論サーバーを NVidia 5090 から AMD R9700 32GB を 2 枚に切り替えたが、完全にポジティブな体験だった。Fedora カーネルで DKMS の設定なしにそのまま動き、ROCm でのコンテナ接続も簡単だった。Ollama と Storyteller の設定を変えるだけで済んだ。CUDA よりずっと快適な体験だった
    • Nvidia は Unreal Engine のフォーク まで自前で保守している。AMD は競争になっていないレベルだ
    • Nvidia はハードウェア企業の中で唯一、ソフトウェアエンジニアに競争力のある報酬を提供している。AMD には今でもソフトウェアを「本当の仕事」と見なさない文化が残っており、こうした惰性は変えにくい
  • Mojo が AMD GPU 上の開発者体験（devX）を改善しようというアイデアを持っていたが、その進捗が気になる
  • AMD がソフトウェア改善に 数十億ドルを投資しないのが理解できない。Nvidia は世界で最も価値の高い企業で、AMD は唯一の競争相手なのに
    • AMD も努力しているが、毎年ハードウェアを更新する組織文化を ソフトウェア中心の文化へ転換するのは難しいと思う。ソフトウェアはハードウェアのようにすぐ収益を生まないため、経営陣は優先度を低く置きがちだ。また、外部ベンダーがオープンソースでコードを提供することは短期的には良く見えても、長期的な品質には悪影響を及ぼす。ハードウェアのトレンドを一度でも見誤れば、競合に後れを取るリスクが大きい
    • 複数の GPU ベンダーで働いたことがあるが、ソフトウェアを 資産（asset） と見なして投資しているのは Nvidia だけだ。他社はコストとしか見ていない
  • “go brr” ミームは個人的には好きではないが、Stanford のような場所で使われているのを見ると面白い
    • 実際、すでに 1 年前の ThunderKittens の発表 で “go brr” を使っていた
    • こういうミームが大学の公式チャンネルに出てきたなら、すでに 流行が終わったサインなのかもしれない
  • プロジェクト自体は素晴らしいが、なぜ AMD がこういうことを自分でやらないのか疑問だ。成熟したソフトウェアスタックの重要性を AMD はいまだに理解していないように見える。CUDA のように、すべてのカードで使える統合スタックが必要だ。かつては AMD もいずれ追いつくと信じていたが、今ではほとんど諦めている
  • プロジェクトは良いが、記事そのものは どこか奇妙な書かれ方 をしているように感じる
    • 文章がかなりぎこちない。AI に過度に依存したか、AI の文体を真似たように見える。「part one を確認しろ」や「AMD GPU を go brr にする方法」のような文が繰り返される。技術的な部分でも、本来グラフで説明すべき内容を 100 行のコードで書き下している のが特に残念だ