AMD GPU上のFP4混合精度推論最適化

はじめに

最先端の大規模言語モデル（LLM）は規模が拡大し続けており、GPUの計算能力とメモリ帯域幅への需要が日々高まっています。GPUベンダーとモデル開発者は低精度浮動小数点フォーマットに転向しており、その中でもFP4（4ビット浮動小数点）量子化が特に注目を集めています。例えば、FP4量子化されたLlama 3.3 70Bモデルは体積が3.5倍縮小され、MMLUなどのベンチマークテストでの品質損失はごくわずかです。

しかし、現在のハードウェアサポートには明らかな短所があります。NVIDIA GB200やAMD MI350などの新世代GPUはFP4行列乗算をネイティブサポートしていますが、広く展開されているAMD Instinct MI250およびMI300シリーズGPUにはまだこの能力がなく、ユーザーは既存のAMDハードウェア上でFP4モデルを効率的に実行できません。

このギャップを埋めるため、私たちはPetitを開発しました——AMD GPU向けに設計されたFP16/BF16 × FP4混合精度GPUカーネル集です。Petitはハードウェアのアップグレードを必要とせず、MI200およびMI300シリーズ上でFP4モデルをサービス提供でき、顕著な性能向上をもたらします：

• SGLang使用時、Llama 3.3 70Bのエンドツーエンド推論性能が1.74倍向上；
• 等価行列乗算の実行速度がAMDの最先端GEMMライブラリhipBLASLtより最大3.7倍高速。

PetitはBSDライセンスでオープンソース化されており、SGLang 0.4.10に統合されています。ユーザーは以下のコマンドでAMD MI250/MI300X上でLlama 3.3 70B FP4モデルサービスを起動できます：

python -m sglang.launch_server --model-path nvidia/Llama-3.3-70B-Instruct-FP4 --host 0.0.0.0 --port 30000

本稿では最適化の過程と技術的詳細を詳述します。PetitはAMDのオープンソースソフトウェアエコシステムを十分に活用し、オフライン再配置やハードウェア特有の低レベル最適化などのイノベーションを導入しています。

ハードウェアアーキテクチャとの協調設計による効率的なGPUカーネル

現代のGPUはコンパクトな計算ユニット（CU）を積み重ねることで大規模な計算スループットを実現していますが、ピーク性能を発揮するにはアプリケーションと基盤アーキテクチャの深い協調設計が必要です。図1に示すように、Petitの開発は複数の重要な協調設計原則に従っています。

前処理による効率的な逆量子化

PetitはAMD GPUの専用MatrixCoreハードウェアを効率的に活用して行列乗算を加速します。1つのwavefront（64スレッドグループ）は、2つのBF16/FP16 16×16行列を効率的に集合的に乗算できます。しかし、MI300X GPUにはネイティブなFP4重みMatrixCoreサポートがないため、メモリロードとMatrixCore準備の効率性を保ちながら、FP4重みをBF16/FP16に逆量子化する必要があります。

これが核心的な課題をもたらします：メモリロードとMatrixCore準備には異なるデータレイアウトが必要です。メモリ効率のためには1024バイトブロックの連続ロードが必要ですが、MatrixCoreは16×16タイルがwavefront全体に分散することを期待します。従来のGPU上での再配置はオーバーヘッドが大きすぎます。

NVIDIA GPUのMarlin実装は、ディスク上で要素を事前に配置することでこの問題を回避しています。私たちは8つの連続したFP4値を32ビット整数にパックし、逆量子化には31命令が必要でした。Petitはさらにビットパッキングフォーマットをカスタマイズしました：最初の4つのFP4をBF8レイアウトで再配置し、残りを残りのビットに格納します。AMDユニークなv_bfrev_b32とv_cvt_pk_f32_bf8命令（サブダブルワードアドレッシングSDWAサポート）を活用することで、わずか15命令で8つのFP4値を逆量子化し、乗算性能を30%向上させました。

メモリ階層構造の習得

MI300XなどのGPUは極めて高い算術密度（>500）を持ち、CUはピークFLOPSに到達するために1バイトあたり数百の操作を実行する必要があるため、有効メモリ帯域幅の最大化が極めて重要です。Petitはタイリングとダブルバッファリングなどの成熟した技術を採用し、AMD特性に対して最適化しています：

• LDSバンク競合の回避：AMD GPUのLDSは32バンクに分かれており、各サイクルで32個のユニークバンクへの同時アクセスが可能です。競合はシリアライゼーションボトルネックを引き起こし、特にwavefrontが64スレッドを持つ場合に顕著です。Petitはバンク設計に基づいて置換データレイアウトを実装し、競合のないLDS使用を実現しています。
• チップレットと相互接続：各MI300チップレット（XCD）には4MBのローカルL2キャッシュがあり、GPU全体で256MBのL3キャッシュを共有します。相互接続帯域幅は高いものの遅延が大きくなります。Petitはトポロジー認識ワークロード分割を実装し、相互接続トラフィックを削減し、プロファイリングで相互接続オーバーヘッドが大きいことが示された場合、グローバルストライプ分割よりも素朴なグリッド分割を優先します。

高品質機械語の生成

GPUは単純なアウトオブオーダー実行ユニットを使用してCU密度を最大化していますが、分岐とパイプライン停止のコストは高くなります。AMD GPUは条件移動と境界メモリ命令を提供し、分岐を完全に排除します。例えば、Petitは範囲指定のバッファロード/ストア命令を利用し、GPUが自動的に境界外アクセスを破棄します；LDSの64KB超アクセスも自動的に処理され、性能ペナルティはありません。さらに、Petitはコンパイラヒントを提供し、MFMA（行列融合乗加算）命令とメモリアクセスをオーバーラップさせ、メモリ遅延を効果的に隠蔽します。

標準コンパイラは高度なGPU ISA（意図的な境界外アクセスが未定義動作など）を十分に活用していない可能性があります。これらの最適化には手動での構築と検証が必要です。

性能結果

エンドツーエンド推論性能

FP4とBF16モデルのエンドツーエンド推論性能比較によりPetitの実際の効果を評価しました。テストではLlama 3.3 70Bの2つのバリアントとSGLang v0.4.10を使用し、バッチサイズ10と64の入力/出力トークンスループットを測定しました。環境はAMD開発者クラウドVM（1× MI300X GPU、240 GB RAM、5 TB SSD）で、Ubuntu 24.04.1上でROCm 6.4.2を実行しています。

図2はオフライン生成ベンチマークの結果を示しています（実際のShareGPTトレースを使用し、本番性能を反映）。全体として、PetitはLlama 3.3 70B FP4モデルをSGLangネイティブBF16モデルより1.74倍（バッチ10）および1.60倍（バッチ64）高速にサービス提供します。小バッチメモリ帯域制限の本番シナリオでは、Petitは3.5倍小さいFP4モデルを効率的に活用し、より高いスループットを実現します。以下のコマンドで再現可能です：

python -m sglang.bench_offline_throughput --model-path nvidia/Llama-3.3-70B-Instruct-FP4 --num-prompts 10
python -m sglang.bench_offline_throughput --model-path nvidia/Llama-3.3-70B-Instruct-FP4 --num-prompts 64

詳細な性能分析

PetitとhipBLASLt（AMDの最先端低レベルアセンブリGEMMライブラリ）の性能をさらに比較しました。2つのライブラリの目標は若干異なることに注意してください：

• Petit：BF16行列 × NVFP4行列（16要素が1つのFP8スケールを共有）。
• hipBLASLt：2つのBF16行列。

完全に同一ではありませんが、結果は量子化に関する洞察を提供します。Llama 3 70Bサービング時の実際の重み行列サイズを考察し、m=16（デコードワークロード）とm=256（プリフィルワークロード）の性能を測定し、100回実行の平均を取りました（50回のウォームアップ後）。両ライブラリとも最適な設定に調整されています。

図3aと3bはGEMM性能を示しています。m=16（デコード主導）の場合、PetitはhipBLASLtより最大3.7倍高速で、平均2.56倍の向上を達成しました。m=256（プリフィル）の場合、Petitは最大1.09倍高速でした。