谷歌TPU跑Kimi比英伟达GPU快57%，代码已开源

16块谷歌TPU v7跑Kimi K3，每秒跑出709个Token，比英伟达GB200快57%。做出这一成绩的是推理创业公司Inferact，其创始团队为vLLM原班人马，今年获a16z领投的1.5亿美元种子轮，估值8亿美元。

Inferact为TPU编写了名为megakernel的推理内核，将模型推理时原本需调度的几百个小程序合并为一个大程序，配合TPU片上内存架构，把内存带宽利用率逼至硬件理论峰值附近。再配上DeepSeek提出的DSpark推测解码框架，K3在TPU上实现每秒709 token。该代码已开源。

测试在完全相同条件下进行：16块TPU v7 Ironwood对阵16块英伟达GB200，均跑Kimi K3、均用vLLM推理引擎，唯一变量是芯片和底层kernel实现。TPU每秒709个token，GB200每秒452个，TPU快57%。DSpark原理是让小模型先猜一串候选token，大模型批量验证，猜对跳过、猜错回退。Inferact在TPU上跑通该流程，acceptance length达6，单步decode耗时约8.5毫秒。

关掉推测解码看裸速，差距同样明显。batch size为1时，TPU每秒249个token，GB200仅127个，差距近一倍；batch size为8时，TPU达865个，GB200为636个。在Qwen上差距更大：4块TPU v7跑Qwen 3.8 27B，每秒1515个token，同配置GB200仅695个。提速未造成精度损失，greedy decoding验证下，Kimi K3在TPU上的GPQA-Diamond得分94.4%、GSM8K为97.2%，与GPU结果完全一致。

这一差距无法从硬件规格解释：TPU v7的HBM带宽为7380 GB/s，GB200反而更高，达8000 GB/s。差距实际来源于芯片上运行的软件。

大模型推理瓶颈不在计算，在数据搬运。每生成一个token，全部权重都要从HBM搬入芯片内部高速缓存算一遍。传统做法将一次推理拆成几百个独立kernel排队执行，交接间隙内存带宽空转。CUDA Graphs和英伟达PDL能缩短间隙，但kernel边界仍在。megakernel直接消除边界：Kimi K3共92层MoE计算，Inferact将其从头到尾塞进一个Pallas程序，一次调用走完整个前向传播。边界消失后，跨层权重预取顺理成章——第N层还在算MoE时，第N+1层attention权重已开始搬运，计算与搬运同时进行。

TPU硬件特性让这种编排更顺畅：每个TensorCore带64 MiB VMEM片上内存，生命周期完全交给软件管理，容量足够同时装下当前层计算数据和下一层预取权重。GPU方面，Blackwell架构片上内存分散在152个SM中，总量约38 MiB，由硬件自动调度，跨层编排限制更多。

Inferact用Pallas语言手写整个megakernel。TPU默认编译工具链XLA擅长优化单层计算，但跨92层协调数据搬运决策分支太多，找不到最优解。绕过XLA后，编译耗时从30分钟以上降至不到90秒。目前megakernel针对Kimi K3定制，换模型架构需重新适配，团队接下来将扩展支持更多模型架构。

Inferact去年11月成立，团队几乎就是vLLM原始开发者。CEO Simon Mo是vLLM原始维护者，伯克利EECS毕业，曾在Anyscale工作；联合创始人Woosuk Kwon是vLLM项目发起人，伯克利计算机科学博士，导师Ion Stoica，其提出的PagedAttention算法至今仍是vLLM核心技术；首席科学家游凯超曾获清华大学特等奖学金，主导vLLM分布式推理功能开发。a16z和Lightspeed领投，真格、红杉、Altimeter跟投。

此次TPU megakernel来自Inferact与Google Cloud的联合工程合作，目标是让TPU成为vLLM一流支持对象，所有优化成果回馈上游开源社区，tpu-megakernels代码仓库是首个公开成果。