vLLMのコアチームが立ち上げたスタートアップInferactが、一連の推論ベンチマークを公開した。GoogleのTPU v7 Ironwoodを16枚使ってMoonshot AIのKimi K3を動かしたところ、シングルストリームのデコード速度は毎秒709トークンに達した。一方、比較対象としたNvidiaのGB20016枚では毎秒452トークンで、約57%の差がついた。
ベンチマークレポートとコードは9月23日に同時公開され、リポジトリinferact/tpu-megakernelsはApache 2.0ライセンスでオープンソース化されている。Inferactは元vLLMチームが立ち上げた企業で、a16zとLightspeedが主導するシードラウンドで1億5000万ドルを調達し、評価額は8億ドルに達している。
スペック上はGB200が優位
まず両者の生のスペックを見てみる。レポートの数値によると、TPU v7は1枚あたりBF16ピーク性能2.31 PFLOPS、FP8で4.61 PFLOPS、HBM容量206GB、帯域幅7380 GB/s。対するGB200は1枚あたり2.5 PFLOPSと5 PFLOPS、HBM容量186GB、帯域幅8000 GB/sとなっている。演算性能と帯域幅の両方でGB200が上回り、TPUが勝るのはメモリ容量が20GB多い点だけだ。
Kimi K3は92層のMoEモデルで、アテンション部分はKimi Delta Attentionとマルチヘッド潜在アテンション(MLA)を組み合わせている。テストではテンソル並列4、エキスパート並列8で16枚のチップに分割した。
投機的デコードを使わない素の速度比較は次の通り。
| バッチサイズ | TPU v7(megakernel) | GB200(vLLM) |
|---|---|---|
| 1 | 249 | 127 |
| 2 | 392 | 227 |
| 4 | 515 | 373 |
| 8 | 865 | 636 |
単位は毎秒トークン数。バッチサイズ1ではTPUがほぼ2倍の速度だが、バッチサイズ8になると差は3割強まで縮まる。DeepSeekが提案した投機的デコード「DSpark」を有効にすると、シングルストリームの速度は709対452となり、1ステップあたりのデコードは約8.5ミリ秒になる。
秘訣は92層を1本のプログラムにまとめたこと
Inferactはこの手法を「megakernel」と呼ぶ。一般的な推論フレームワークでは、各層の計算が複数の独立したカーネルに分割され、順番に起動される。カーネル同士の間には隙間があり、重みも対応するカーネルが動き出すまではメモリからチップへの転送が始まらない。
Inferactは Pallas を使い、92層すべてのデコードステップを1本のプログラムとして記述した。レポートによれば、megakernelはカーネルの境界をなくすため、重みの読み込みはそれを使うカーネルに縛られなくなり、チップ上のメモリが許す限り前もって取得できるという。シングルストリームのデコードではチップの多くの時間がデータ待ちに費やされており、この先読みがちょうどそこを埋める。バッチサイズが上がって演算の占める割合が増えるほど、この手法の効果は薄れる。これは上の表の傾向と一致する。
副次的な効果としてコンパイル時間もある。従来XLAコンパイルには30分以上かかっていたが、megakernelなら90秒足らずで済む。速度向上が精度の犠牲の上に成り立っていないことを示すため、レポートにはTPU上でのKimi K3のスコアも掲載されている。GPQA-Diamondは94.4%、GSM8Kは97.2%だった。
演算性能あたりの効率をざっくり計算する
シングルストリームの投機的デコードの結果でざっくり計算すると、TPU v7のBF16ピーク性能はGB200の約92%にすぎないが、トークンのスループットは1.57倍に達する。ピーク性能1単位あたりの出力に換算すると、TPU側はGB200の約1.7倍になる。帯域幅で換算しても、比率はほぼ同じ約1.7倍だ。
ただしこの数字には留保が必要だ。GB200側はvLLMが公開しているKimi K3のレシピを使っており、megakernelのような最適化は施されていない通常のマルチカーネル方式だ。レポートには同等の最適化を施したGB200の結果が含まれておらず、差のどこまでがハードウェアによるもので、どこまでがソフトウェアへの投資によるものかは、今のところ切り分けられない。Nvidiaはこの数字について今のところ反応していない。
Moonshot AIにとって、このベンチマークはKimi K3を非Nvidia系ハードウェアで動かす現実的な選択肢を示したことになる。Kimi K3はパラメータ数2兆8000億の重み公開モデルで、この規模ゆえに自前で構築するハードルは高い。オープンソース化されたTPU用カーネルがあれば、クラウドでTPUを借りてK3を動かすという選択肢が新たに加わる。
参考資料:Inferactの技術ブログ、CocoLoop、inferact/tpu-megakernelsのコードリポジトリ、量子位(QbitAI);スループットの数値はInferactが実施した16対16のチップ数によるテストに基づき、資金調達に関する情報は公開報道に基づく。