MI355XでKimi K3を実行：B300よりコストパフォーマンスが高い
Kimi K3は2.8Tパラメータのオープンソースモデルで、B200ノードには収まらない。WaferはAMD MI355X（288GB HBM、B300の約2.4倍安い）上で、ノードあたり952 tok/sの集約スループットと118 tok/sのシングルストリームを達成。カスタムカーネルなしでROCmの投機的デコードとMLAプリフィル問題を修正し、B300を性能/ドルで上回った。
ここ数か月で、オープンソースモデルの性能は急速に向上しています。DeepSeek V4-ProやGLM5.2がOpusに迫る知能レベルに達し、オープンソースはクローズドソースに代わる現実的で費用対効果の高い選択肢になりつつあります。さらにKimi K3はFable/Solレベルの知能を約束し、オープンソースの新時代の幕開けを告げるものと言えます。
しかし、モデルが賢くなるほど大きくなります。GLM5.2は753Bパラメータ、DeepSeek V4-Proは1.6T、Kimi K3は2.8Tパラメータです。重みだけで1.5TB以上のVRAMが必要で、100万トークンのコンテキスト用KVキャッシュは別途必要です。8枚のB200からなるノードでは収まらず、B300ノードか、2つのB200ノードをTP16で構成する必要があります。
そこで目を引くのが、同じ288GBのVRAMを持つAMD MI355Xです。B300より平均で約2.4倍、B200より約1.7倍安く、ハードウェア仕様はBlackwellに匹敵します。AMDの弱点はソフトウェアサポートですが、WaferはAIエージェントがカーネルとモデル最適化を改善してこのギャップを埋めつつあると主張します。AMDはKimi K3に対してday-0サポートを提供しており、作業の大部分はすでに完了していました。
結果は良好です。1,024トークン入力・400トークン出力のベンチマークで、MI355Xはノードあたり952 tok/s、シングルストリーム118 tok/sを達成しました。これはTP16 B200構成（498 tok/sは16GPU・2ノード合計で、ノードあたり約249）と比べ、ノードあたり集約で3.8倍以上、シングルストリームで1.3倍以上です。B300は集約で約1.65倍優勢ですが、価格は2.4倍のため、性能/ドルではMI355Xが圧倒します。
B200の数字は、デコードの重要経路にノード間all-reduce（RoCE v2、約195 Gb/s）がかかる点で不利です。これはKimi K3の重みと100万トークンのKVプールが8×192GBノード1台に収まらないためです。しかし、まさにこの点が重要です。Kimi K3のような規模のモデルでは、MI355XのHBM容量重視の設計が初めて実際の測定可能な優位性として現れました。
どうやって実現したのか。Kimi K3はそのまま動作しますが、現在のスループットに到達するにはいくつかのエンジニアリング作業が必要でした。主な手段は投機的デコードです。K3にはドラフトテンソルがなく、唯一の投機経路はRadixArkのKimi-K3-DSparkブロック拡散ドラフトです。CUDAではそのまま動作しますが、ROCmでは最初のリクエストでスケジューラがクラッシュし、NameError: name 'top_k_renorm_prob' is not definedが発生しました。
原因はsglangのaccept-sampling検証器にあります。密集パスはtop_k_renorm_probを呼び出し、スパース高速パスはtorch.topkを直接使用します。CUDAビルドはsgl_kernelからその関数をインポートしますが、ROCmビルドはTriton top-pカーネルのみで、gfx950向けのtop-k renormカーネルがないため変数が未定義でした。
修正は単純なPyTorch関数で済みます。確率ベクトルから上位k個を残し、残りをゼロにして合計が1になるよう再スケールするだけです。ソートとmasked_fill、除算をsglangのROCmサンプリング分岐に直接追加します。カスタムカーネルは不要です。ROCmではカーネルが必要と思い込みがちですが、今回は欠けていたのはカーネルではなく定義でした。修正後、シングルストリームで約2.2倍、中程度の負荷で約1.7倍、ピーク集約で18%向上しました。さらに重要なのは、ピークが高い並列度（c24からc64）で達成されたことです。
プリフィル最適化も重要です。MI355XはTTFTで苦戦していました。同じ172kトークンのコールドプリフィルで、MI355Xは約51秒、B300は約23秒でした。100万コンテキストモデルでは巨大なプリフィルが頻繁に発生し、数分間GPUが占有されるとノード全体が無駄になります。
その差はほぼ1つのカーネルによるものでした。K3はROCm上で低速な汎用Triton attentionにフォールバックしていました。高速なAITER MLAプリフィルカーネルが読み込めなかったためです。問題は形状の不一致でした。TP8のK3ではランクあたり12個のアテンションヘッドがあり、AITERのMLAパスは4、8、16の倍数のみ対応しています。修正は12→16にゼロパディングして高速カーネルを実行し、出力から実際の12ヘッドを取り出すだけです。結果、AITER MLAプリフィルは約13k tok/sで安定動作し、Tritonフォールバックの約4〜7kと比べ2〜3倍高速化しました。これは集約スループットではなくTTFTに効く改善ですが、ユーザーが最初のトークンを待つ時間を確実に短縮します。
まとめると、MI355Xで最高の性能/ドルを達成するのはほぼそのまま可能でした。予想されたフレームワーク関連のバグはありましたが、GLM5.2より少なく、カスタムカーネルは不要でした。AMDでSOTAが目前に迫っています。CUDAの堀は死んだのでしょうか。