AMD ROCm：4ウェーブインターリーブによるFP8 GEMMカーネル最適化でMI355Xのレジスタ予算を2倍に

AMD ROCmチームは、CDNA 4アーキテクチャを搭載したAMD Instinct MI355X GPUにおけるFP8 GEMM（汎用行列乗算）カーネルの新しい最適化戦略について詳細な技術解説を公開しました。この最適化は標準的な8ウェーブpingpongアプローチを超え、ウェーブあたりのレジスタ予算を劇的に増加させ、ハードウェアレベルでメモリ競合を排除する4ウェーブインターリーブ設計を採用しています。

FP8 GEMMで4ウェーブ方式が8ウェーブ方式に勝る理由

核心的な違いはVGPR（ベクトル汎用レジスタ）の割り当てにあります。「1 SIMD単位あたり2ウェーブ」を配置する8ウェーブ設計では、各ウェーブがレジスタ空間を共有しなければならず、処理できる出力ブロックのサイズが制限されます。4ウェーブ方式では1 SIMD単位あたり1ウェーブを割り当て、512 VGPRのフル予算へのアクセスを提供します。

これは処理可能なタイルサイズに直接影響します：4ウェーブ方式では8ウェーブ設計の64×128に対して128×128要素の出力タイルを使用できます。カーネル全体のグローバルタイルは256×256×128で、4つの各ウェーブは各ブロック内で16×16 MFMA（Matrix Fused Multiply-Add）演算を用いて64×64要素の4ブロックを処理します。

MI355Xにおいて4ウェーブインターリーブがメモリレイテンシを排除する仕組み

GPU行列乗算カーネルにおける中心的な課題はメモリレイテンシの隠蔽です。GPUはメモリからのデータロードを待ちながらMFMA演算で占有し続けなければなりません。4ウェーブ方式は、ハードウェアスケジューラに頼ることなくサブタイル粒度での命令の精密なインターリーブによってこれを解決します。

実装では__builtin_amdgcn_sched_barrier(0)を使用して命令順序を強制します。コンパイラはすべてのメモリ操作を演算の後にまとめることを禁止されます。各interleaved_block呼び出しは16個のMFMA命令と8個のLDSからレジスタへのロード（ds_read_b128）を重複させ、pingpong同期の複雑さなしに一貫した重複実行を実現します。

XORベースのLDSスウィズルがメモリバンク競合を排除する仕組み

MI355X LDS（ローカルデータ共有）メモリは4バイト単位の64バンクで構成され、256バイト周期を持ちます。問題は構造的なものです：16×16×128 MFMA命令は本質的に列優先（column-major）でマトリックスの列にアクセスするため、標準的な配置では8重バンク競合が発生します。8つのアクセスが同時に同じバンクを待機します。

AMDはオフセットの数学的変換でこれを解決しています：

row_bits = (offset % (16 * 128) >> 7) / 2
mask = row_bits << 4
swizzled_offset = offset ^ mask

XOR演算がアドレスを再マッピングし、ds_read_b128命令の全4フェーズがバンク競合ゼロで実行されます。結果として、8重競合時の実効12.5%に対してLDSメモリのフルスループットが得られます。

ハードウェアプラットフォームとチップレット対応最適化

すべての最適化はCDNA 4アーキテクチャを搭載したAMD Instinct MI355Xを対象としています。MI355Xは8チップレット（XCD）トポロジーを採用しており、各XCDが独自のL2キャッシュを持ち、最終レベルキャッシュはXCD間で共有されます。ROCmチームはモートン曲線に似たマッピングを使用してスレッドブロックの割り当てを再マッピングするチップレット対応グリッドスウィズルを追加し、1つのXCD内のL2利用率とXCD間のLLC利用率を最大化しています。