谷歌 TPU 跑 Kimi K3,比 GB200 快 57%

vLLM 核心团队创办的 Inferact 公布了一组推理测试:在 16 片谷歌 TPU v7 Ironwood 上跑月之暗面的 Kimi K3,单路解码速度达到每秒 709 个 token,同样 16 片英伟达 GB200 的对照组是每秒 452 个,快了约 57%。

测试报告和代码 9 月 23 日一起放出,代码仓库 inferact/tpu-megakernels 以 Apache 2.0 协议开源。Inferact 是 vLLM 原班人马的创业公司,此前完成 1.5 亿美元种子轮,估值 8 亿美元,a16z 和 Lightspeed 领投。

硬件纸面上 GB200 占优

先看两边的纸面参数。按报告列出的数字,单片 TPU v7 的 BF16 峰值算力 2.31 PFLOPS、FP8 4.61 PFLOPS,配 206GB HBM、带宽 7380 GB/s;单片 GB200 对应是 2.5 PFLOPS、5 PFLOPS,186GB HBM、带宽 8000 GB/s。算力和带宽两项,GB200 都高出一截,TPU 只在显存容量上多 20GB。

Kimi K3 是 92 层的 MoE 模型,注意力部分混用 Kimi Delta Attention 和多头潜在注意力(MLA)。测试用张量并行 4 路、专家并行 8 路切分到 16 片芯片上。

不开推测解码的裸速对比:

批大小TPU v7(megakernel)GB200(vLLM)
1249127
2392227
4515373
8865636

单位是每秒 token 数。批大小为 1 时 TPU 接近两倍,批大小到 8,领先幅度收窄到三成多。打开 DeepSeek 提出的 DSpark 推测解码后,单路速度变成 709 对 452,单步解码约 8.5 毫秒。

快在把 92 层合成一个程序

Inferact 的办法叫 megakernel。常规推理框架里,每一层的计算拆成若干个独立内核依次启动,内核之间有空档,权重也只能等到对应内核开跑时才开始从显存往芯片上搬。

Inferact 用 Pallas 把整个解码步骤的 92 层写成一个程序。报告的说法是,一个 megakernel 抹掉了内核边界,权重加载不再和使用它的内核绑在一起,芯片片上存储放得下多少,就可以提前多少去取。单路解码时芯片大部分时间在等数据,提前预取正好补在这里;批大小一上去,计算占比变高,这一招的收益就变薄,这和表里的趋势对得上。

一个附带好处是编译时间:原先走 XLA 编译要 30 分钟以上,megakernel 不到 90 秒。为证明提速没有牺牲结果,报告附上 Kimi K3 在 TPU 上的 GPQA-Diamond 得分 94.4%、GSM8K 97.2%。

粗算一下单位算力的产出

用单路推测解码的结果粗算:TPU v7 的 BF16 峰值约为 GB200 的 92%,产出的 token 速度却是后者的 1.57 倍,折合到每单位峰值算力,TPU 这边约为 GB200 的 1.7 倍。按带宽折算,结果也在 1.7 倍左右。

这笔账要打一个折扣。GB200 一侧用的是 vLLM 公开发布的 Kimi K3 配方,属于常规多内核路线,没有做同等力度的 megakernel 改造。报告里没有 GB200 上的对等优化对照,差距里有多少来自芯片、多少来自软件投入,目前无法拆开。英伟达方面也没有回应这组数字。

对月之暗面来说,这组测试提供了 Kimi K3 在非英伟达硬件上的一条可行部署路径。Kimi K3 是 2.8 万亿参数的开放权重模型,这个体量让自行部署的门槛很高;有了开源的 TPU 内核,在云上租 TPU 跑 K3 多了一个选项。

参考来源:Inferact 技术博客、CocoLoop、inferact/tpu-megakernels 代码仓库、量子位;吞吐数字以 Inferact 报告的 16 片对 16 片测试口径为准,融资信息据公开报道。