AMD、MI355X 8基でGLM-5.2 FP8の64K推論を629K tok/sに高速化

AMDのROCmブログは、8基のAMD Instinct MI355X上でGLM-5.2 FP8推論を最適化し、EvalScopeで測った64K入力のピークスループットを163K tok/sから629K tok/sへ引き上げた結果を2026年11月11日に公開しました。1Mトークン入力・同時実行数1では7.16倍の改善です。

GLMはZ.aiが開発するAIモデルのファミリーで、SGLangはLMSYSのもとで公開されている推論フレームワークです。

何を最適化したか

原文は、GLM-5.2が大規模MoE、スパースアテンション、長文脈処理を組み合わせているため、長い入力ではGPU計算、メモリ帯域、GPU間通信がボトルネックになり、単一カーネルの調整では足りないと説明しています。著者らは推論スタック全体に次の最適化を適用しました。

• CP8プリフィル: TP8ではランクごとに64ヘッド中8ヘッドしか扱えず、Lightning Indexerが選んだ同じKVを各ランクが重複ロードし、AITERカーネルの16の倍数制約で8ヘッドを16ヘッドに複製する必要がありました。コンテキスト並列でクエリトークンを分割し、各ランクが64ヘッドを保持したまま1/8のトークンを処理することで、MI355Xネイティブのgqa:64 MLAカーネルを直接使えるようにしています。
• FP8 KVキャッシュ: fp8_e4m3を採用し、BF16クエリの量子化パイプラインを単一のTritonカーネルに融合しました。
• 64ヘッドFP8 Persistent MLA: AITERのpersistent MLA実装でH64形状を扱います。
• PyHIP MoE: router/top-kと活性化量子化を集団通信の前に移し、BF16ではなくFP8隠れ状態・スケール・ルーティング結果を集約します。1トークン当たりの通信ペイロードは12,288 bytesから6,408 bytesに減りました。
• GEMMチューニング: block-scaled FP8 A8W8とBF16経路を個別にオフライン調整し、AITERのチューニングテーブルに格納します。
• EAGLE MTP: 4ステップ、top-k = 1、ドラフト5トークンで、リクエスト単位のデータ並列とエキスパート並列と組み合わせています。

ほかにカスタムall-reduce、小カーネル融合、アテンションカーネル調整も含みます。

ベンチマーク結果

64K・128Kは90%共有プレフィックス、出力長1の構成で、リクエスト失敗0件のケースのみを集計しています。原文はキャッシュなし1K・8K入力より長文脈サービスのSLOを反映する構成としています。

条件	ベースライン	最適化後
64K ピークスループット	163K tok/s(同時実行4)	629K tok/s(同時実行50)
128K ピークスループット	137K tok/s(同時実行2)	553K tok/s(同時実行32)
1M・同時実行1 スループット	21,877 tok/s	156,598 tok/s
1M・同時実行1 平均TTFT	45,660 ms	6,340 ms

64Kではベースラインが同時実行4以降で頭打ちになる一方、最適化後は同時実行50までスケールしました。高速化率は64Kで1.71〜3.97倍、128Kで2.03〜4.25倍です。

精度はGSM8K(5-shot、全1,319サンプル)で確認し、ベースラインの94.16%(1242/1319)に対し最適化後は94.92%で、精度低下は見られませんでした。

再現手順と環境

項目	内容
ハードウェア	8× AMD Instinct MI355X(gfx950)
ROCm	7.2.4
モデル	zai-org/GLM-5.2-FP8
ベースラインイメージ	lmsysorg/sglang-rocm:v0.5.13.post1-rocm720-mi35x-20260618
最適化イメージ	sabreshao/sglang:glm_0831
API	OpenAI互換 /v1/completions

EvalScopeはpython3 -m pip install evalscopeで導入し、evalscope --versionで確認します。ブログと一緒に公開されたlaunch_baseline.sh、launch_optimized.sh、run_evalscope.shをscripts/に置いて実行権限を付けます。

ベースラインはMODEL_PATH=/path/to/GLM-5.2-FP8 ./scripts/launch_baseline.shで起動し、./scripts/run_evalscope.sh 64k 30000、128k 30000、1m 30000で測定します。最適化版はlaunch_optimized.shで起動し、ポート30001に対して同じ3プロファイルを実行します。スクリプトは各ケースの前にプレフィックスキャッシュをフラッシュします。