TPUがKimiを57％高速化、差を生んだのは推論ソフトウェア

導入

大規模モデルの推論でGoogle TPUがNvidia GPUに対抗できるかどうかは、ピーク演算性能だけでは決まりません。Inferactが公開したテストでは、16基のTPU v7 IronwoodでKimi K3を実行した場合、毎秒709トークンに達しました。一方、16基のGB200は452トークンで、TPUが57％上回ったとされています。重要なのは、チップ単体よりも、その上で動くソフトウェア設計です。

テストの要点

• 両システムは同じKimi K3とvLLM推論エンジンを使用し、主な違いはアクセラレーターと低レベルKernelでした。
• TPU側ではDSpark推測デコードを利用しました。小さなモデルが候補トークンを生成し、大きなモデルがまとめて検証する方式で、平均受理長は6と報告されています。
• 推測デコードを無効にしたBatch Size 1では、TPUが毎秒249トークン、GB200が127トークンでした。Batch Size 8では865と636でした。
• Qwen 3.8 27Bでは、4基のTPUが毎秒1515トークン、4基のGB200が695トークンでした。
• Inferactによれば、TPU上のKimi K3はGPQA-Diamondで94.4％、GSM8Kで97.2％となり、GPU側と同じ結果でした。

これらはInferactが提示したベンチマークです。実装の可能性を示す一方、すべてのモデルや運用条件での優位性を証明するものではありません。

Megakernelの狙い

自己回帰推論では、演算そのものよりデータ移動がボトルネックになることがあります。従来方式では処理を多数の小さなKernelに分割するため、Kernel間の起動や受け渡しでメモリ帯域が遊びます。InferactのMegakernelは、Kimi K3の複数層の処理を一つのPallasプログラムにまとめました。

これにより、ある層を計算している間に次の層の重みを先読みできます。TPUのVMEMはソフトウェアで明示的に管理できるため、何を読み込み、保持し、解放するかを細かく制御できます。InferactはXLAに任せるのではなく、モデル全体のデータ移動を手作業で構成したと説明しています。コンパイル時間は30分超から90秒未満に短縮されたとされます。

意義と限界

この結果は、推論競争の焦点がピーク演算性能からデータフローとメモリ利用効率へ移っていることを示します。TPUにとっては、オープンソース推論基盤への対応拡大が新しいハードウェアと同じくらい重要になる可能性があります。

一方、現在のMegakernelはKimi K3向けに調整されています。別のモデル構造へ展開するには再適応が必要です。今後、より多くのモデルを支援し、vLLMに統合され、実際の遅延・コスト・同時実行条件でも優位性を維持できるかが評価の焦点になります。

出典：量子位