これはAMDが長く待ち望んでいた瞬間かもしれない。
最近、Wafer AI が AMD MI355X 上にデプロイされました Kimi K3。結果として、元々16枚のNVIDIA B200を必要とし、2台のサーバーにまたがって実行されていたモデルが、8枚のMI355Xを搭載したAMDサーバー1台でデプロイ可能になりました。
さらに重要なのは、単にモデルを組み込んだだけではないということです。
1024トークンの入力、400トークンの出力のテストにおいて、MI355Xは総スループット952トークン/秒、単一ユーザー生成速度118トークン/秒を達成しました。
単一ノードで計算した場合、そのスループットは約16枚のB200構成の3.8倍であり、コストパフォーマンスもB200およびB300を上回ります。
そして最も驚いたのは、ROCmが今回は特に手間をかけなかったことだ。
モデルが大きすぎると、VRAMが計算能力よりも重要になり始める
Kimi K3は2.8兆のパラメータを有し、モデルの重みのみで1.5TB以上のVRAMを必要とし、百万トークンのコンテキストに必要なKVキャッシュは含まれていません。
1台のB200サーバーには8枚のGPUが搭載され、各GPUには192GBのVRAMが搭載されており、合計容量は約1.5TBです。つまり、モデルの重みすら完全に収容するのが難しく、KV Cache用のスペースを確保することはさらに困難です。したがって、B200では2台のサーバー、合計16枚のGPUが必要です。
B300は1枚あたり288GBのVRAMを搭載しており、モデルを単一ノードに収めることができます。興味深いことに、AMD MI355Xも288GBのVRAMを搭載しており、8枚で合計約2.3TBとなり、1台のサーバーで十分です。
これは単に一台のマシンを減らすだけの話ではありません。モデルがノード間で実行されるようになると、1つのトークンを生成するたびにネットワークを通じてデータを同期する必要が生じます。約195 Gb/sのRoCE v2ネットワークを使用しても、ノード間通信はデコードを遅らせます。
MI355Xはより大きなメモリにより、モデル全体を1つのノードに保持します。
最終結果によると、8枚のMI355Xのピーク合計スループットは952 Token/sに達し、単一路の生成速度は118 Token/sです。
対照的に、16枚のB200を用いたデュアルノード構成の合計スループットは498 Token/sであり、単一ノードに換算すると約249 Token/sとなります。
つまり、MI355Xの単一ノードスループットは、B200のデュアルノードデプロイメントにおける平均単一ノードスループットの約3.8倍です。単一ユーザーの生成速度においても、MI355Xの118 Token/sはB200の90 Token/sを上回っています。
B300は依然として最も高いパフォーマンスを誇るソリューションです。8枚のB300を搭載したノードの合計スループットは1568 Token/sに達し、単一路の生成速度は172 Token/sです。全体のスループットはMI355Xの約1.65倍です。
しかし、価格が結論を変更しました。Waferは、MI355Xで1枚あたり毎時2.5ドル、B200で4.25ドル、B300で6ドルと計算されています。
この価格仮定下では、MI355Xは1ドルあたり約48 Token/sのピークスループットを提供し、B200は約7 Token/s、B300は約33 Token/sです。
B300はより高速ですが、MI355Xは単位コスト当たりの効率が優れています。大規模でオープンモデルを実行するデータセンターにとっては、単にパフォーマンスの頂点を争うよりも、これがより重要である可能性があります。
さらに驚きなのは、ROCmがほぼそのまま使用できることです。
長年にわたり、AMDのデータセンターグラフィックカードの最大の課題はハードウェアではなく、ソフトウェアであった。
同じモデルはCUDA上で直接実行できるが、ROCmではフレームワークを変更したり、演算子を補完したり、甚至は底层カーネルを再記述しなければならない可能性がある。
しかし Kimi K3の状況は異なります。
AMDは、ほぼリリース同期でのサポートを提供しました。Waferは、モデルをMI355X上で直接実行でき、後続の作業はわずかな互換性の問題とパフォーマンス最適化に集中していると述べました。
推測デコード段階で一つの問題が発生しました。 Kimi K3は、MTPまたはEAGLEに必要なドラフトモデルパラメータを提供していないため、Waferは外部のブロック拡散ドラフトモデルを使用しました。
このソリューションはCUDA上で直接実行できますが、ROCm環境では最初の実際のリクエストでスケジューラーがエラーを報告します。理由は、ROCmブランチにtop_k_renorm_probという関数が定義されていないためです。
この関数の処理はそれほど複雑ではありません：確率分布から上位k個の値を選択し、他の確率をゼロにし、残った確率を再規格化します。
Waferは、このロジックを通常のPyTorch関数で補完し、GPUカーネルを手書きしたり、推論デコードシステムを再設計したりする必要はありませんでした。
修正後、シングルスレッド性能は約2.2倍、中程度の並列処理下でのシングルストリーム性能は約1.7倍、ピーク総スループットは約18%向上しました。
さらに、システムはより高い並列処理能力でピークスループットを達成できます。
最初の文字が遅すぎて、最後にゼロを四つ追加しただけ
もちろん、スループットは推論サービスのすべてではありません。実際のユーザーにとって、もう一つの体験に直接影響する指標は、リクエストを送信して最初のトークンが表示されるまでの待ち時間であるTTFTです。
このタスクにおいて、MI355X の初期パフォーマンスはそれほど良好ではありませんでした。約17.2万トークンのコールドスタートプリフィルタータスクに対して、MI355Xは約51秒かかりましたが、B300は約23秒で完了しました。
百万トークンのコンテキストをサポートするモデルでは、プリフェッチタスクが非常に巨大になる可能性があります。長コンテキストを処理する際、ユーザーが毎回数十秒、あるいはそれ以上待たなければならない場合、解码速度がどれほど高くても、体験上の問題を補うのは困難です。
Waferは、性能の差がほぼすべてあるアテンションカーネルに由来することを最終的に発見した。 Kimi K3は8路テンソル並列構成下で、各GPUに12個のアテンションヘッドが割り当てられます。一方、AMD AITERで高速なMLAプリフィルリングカーネルは、4、8、または16の倍数などの形状のみをサポートしています。
12個のヘッドが一致しなかったため、システムはより遅い汎用Triton実装に戻りました。
解決策は非常にシンプルです：12個のアテンションヘッドをゼロパディングして16個にし、既存の高速カーネルを呼び出して計算を実行した後、必要な12個のヘッドだけを取り出します。モデル構造は変更せず、新しいアセンブリカーネルを記述することもなく、ただ4つのゼロを追加しただけです。
最適化後、AITER MLAカーネルの安定したプリフェッチ速度は約13,000トークン/秒に達し、従来のTritonフォールバックパスは約4,000～7,000トークン/秒だったため、コールドプリフェッチ時間は約2～3倍短縮されました。
この最適化は最終的なデコードスループットを変更しませんが、最初の文字が表示されるまでのユーザーの待ち時間を大幅に短縮します。
これは、AMDとNVIDIAの間で大きく見えるソフトウェアの差が、時として基盤能力の不足ではなく、既存の高速カーネルが新しいモデル形状をまだカバーしていないだけであることを示している。
CUDAの競争優位はまだ存在するが、穴が開き始めている
一回のテストでは、AMDがNVIDIAに完全に追いついたことを証明できない。
B200はVRAMが不足しているため、ノード間で実行を強制されている；B300の絶対パフォーマンスは依然としてリードしている；ROCmのツールチェーン、フレームワークサポート、開発者エコシステムも、依然としてCUDAには及ばない。
しかし、オープンモデルは急速に兆パラメータ時代に入っています。モデルが1台のサーバーに収まらなくなると、VRAM容量は単なるパラメータ表上の数字ではなく、通信コスト、デプロイの複雑さ、そして最終的なスループットに直接影響します。
AMDが単一GPUに更多のHBMを搭載する戦略は、実際のシステム上の利点となりつつある。
AMDがROCmの安定性をさらに向上させ、高速コアの形状サポートを拡大し、新モデルに対してより迅速な当日対応を提供できれば、データセンターはこれらのGPUを真剣に検討する必要がある。価格が低く、メモリ容量が大きく、性能は十分で、ソフトウェアの調整に数ヶ月を費やす必要もなくなった。
これについてどう思いますか？
参考リンク：
https://x.com/wafer_ai/status/2083628389903315406
https://x.com/ChiragAsarpota/status/2083864019870634151
本文は微信公众号「機械の心」（ID：almosthuman2014）より、著者：LLMに注目