谷歌 TPU 跑 Kimi 比英伟达 GPU 快 57%！用的还是 DeepSeek 推理框架

16 块谷歌 TPU v7 跑 Kimi K3，每秒跑出了 709 个 Token，比老黄的 GB200 快了 57%。

做出这个成绩的，既不是造 Kimi 的月之暗面，也不是造 TPU 的谷歌，是一家叫 Inferact 的推理创业公司。

Inferact 创始团队就是 vLLM 的原班人马，今年拿了 a16z 领投的 1.5 亿美元种子轮，估值 8 亿美元。

他们给 TPU 写了一种叫 megakernel 的推理内核，把模型推理时原本要调度的几百个小程序焊成一个大程序。

配合 TPU 独特的片上内存架构，megakernel 把内存带宽利用率逼到了硬件理论峰值附近。

配上 DeepSeek 提出的 DSpark 加速框架，K3 在 TPU 上跑出了前面提到的每秒 709 token 的成绩。

这套代码，现在已经开源。

同样 16 块芯片，TPU 快 57%

Inferact 把 TPU 和英伟达 GPU 放在完全相同的条件下跑了一轮 benchmark，TPU 赢了。

测试是用 16 块 TPU v7 Ironwood 对阵 16 块英伟达 GB200，两边都跑 Kimi K3，都用 vLLM 推理引擎，唯一的变量是芯片和底层的 kernel 实现。

最终，TPU 跑出的成绩是每秒 709 个 token，GB200 跑出每秒 452 个，TPU 的速度快了 57%。

这个速度里有 DSpark 推测解码的功劳。

DSpark 是由 DeepSeek 提出的推理加速框架，其原理是让一个小模型先快速猜出一串候选 token，然后大模型再对这串候选 token 做批量验证，猜对的直接跳过，猜错的回退重来。

Inferact 在 TPU 上跑通了这套流程，acceptance length 达到了 6，也就是每轮猜出的 token 里，平均有 6 个能直接通过大模型的验证，单步 decode 的耗时大约在 8.5 毫秒。

关掉推测解码看裸速，差距同样拉得开。

在 batch size 为 1 的情况下，TPU 每秒能跑 249 个 token，GB200 只有 127 个，差距接近一倍。

把 batch size 放大到 8，TPU 达到 865 个 token 每秒，GB200 只有 636 个。

另外在 Qwen 上，两种芯片的差距更大。

Inferact 用 4 块 TPU v7 跑 Qwen 3.8 27B，每秒跑出 1515 个 token，同样配置的 GB200 只跑出了 695 个。

而且这种提速并没有造成精度损失，Inferact 用 greedy decoding 做了验证，Kimi K3 在 TPU 上的 GPQA-Diamond 得分是 94.4%，GSM8K 是 97.2%，和 GPU 上的结果完全一致。

推理引擎一样，模型一样，芯片换一下，速度就差出了 57%。

这种程度的差距，按理说应该能从硬件规格表上找到解释，但实际情况恰好相反。

两款芯片的算力也在同一个量级，TPU v7 的 HBM 带宽是 7380 GB/s，而 GB200 的 HBM 带宽反而更高，达到 8000 GB/s，带宽更大的那块芯片，跑出来的速度反而更低。

这种速度的差距，实际上来源于芯片上面跑的那层软件。

把几百个小 kernel，焊成一个大程序

Inferact 用一个叫 megakernel 的方案解决了 TPU 上的推理效率问题，核心思路，是把模型推理时原本要调度的几百个独立小程序，合并成一个大程序，从而消除程序之间切换时的带宽浪费。

大模型推理的速度瓶颈不在计算，在数据搬运。

每生成一个 token，模型的全部权重都要从 HBM 主存搬进芯片内部的高速缓存里算一遍，算完这一轮，下一个 token 再重新搬一遍。

传统的做法是把一次推理拆成几百个独立的 kernel，每个 kernel 负责一小块计算。

这些 kernel 排着队一个接一个地执行，但问题出依然会在交接的间隙。

上一个 kernel 干完了，下一个还没有启动起来，这段时间里内存带宽就这么空转着。

CUDA Graphs 和英伟达最新的 PDL 技术能缩短这个间隙，但 kernel 之间的边界还在。

Inferact 的 megakernel 直接消除了这个边界。

Kimi K3 一共有 92 层 MoE 计算，Inferact 把这 92 层的计算逻辑从头到尾塞进了一个 Pallas 程序里，一次调用就走完整个模型的前向传播。

边界消失之后，跨层权重预取就变得顺理成章了。

第 N 层还在算 MoE 的时候，第 N+1 层的 attention 权重已经开始从 HBM 往片上缓存搬了。

计算和数据搬运同时进行，内存带宽几乎没有空转的时刻。

TPU 的硬件特性让这种编排做起来格外顺畅。

TPU v7 的每个 TensorCore 带有 64 MiB 的 VMEM 片上内存，这块内存的生命周期完全交给软件来管理，工程师可以精确控制什么时候往里面装什么数据、什么时候把空间腾出来。

64 MiB 的容量足够同时装下当前层的计算数据和下一层预取的权重。

GPU 那边情况不同，英伟达 Blackwell 架构的片上内存分散在 152 个 SM 里面，总量大约 38 MiB，由硬件自动调度，要做类似的跨层编排限制更多。

Inferact 用 Pallas 语言手写了整个 megakernel 的代码，Pallas 是谷歌给 TPU 做的底层编程语言，角色类似英伟达 GPU 上的 CUDA。

TPU 默认的编译工具链 XLA 擅长优化单层计算，但跨 92 层协调数据搬运的决策分支太多，XLA 的自动调度找不到最优解。

Inferact 的做法是绕过 XLA，用 Pallas 语言手写整个 megakernel 的代码，由工程师自己来决定每一步搬什么数据、存在哪块缓存、什么时候释放。

这样做还带来了一个额外好处，那就是编译速度大幅提升，耗时从 XLA 的 30 分钟以上降到了不到 90 秒，工程师改完一版 kernel 一分半钟就能上机验证。

目前这套 megakernel 是针对 Kimi K3 的模型结构做的定制优化，如果换一个模型架构还需要重新适配。

Inferact 在博客里提到，团队接下来会把 megakernel 的支持扩展到更多模型架构上。

vLLM 原班人马的创业公司

Inferact 去年 11 月成立，团队几乎就是 vLLM 的原始开发者，一群靠给英伟达 GPU 写推理引擎成名的人。

他们今年一月出来创业，转头把 TPU 的推理速度做到了 GPU 前面。

vLLM 是目前开源推理引擎领域的重量级框架，支持超过 500 种模型架构，社区贡献者超过 2000 人，Meta、Google、Character.ai 都拿它来做线上的推理服务。

CEO Simon Mo 是 vLLM 项目的原始维护者，伯克利 EECS 毕业，之前在 Anyscale 工作。

联合创始人 Woosuk Kwon 是整个 vLLM 项目的发起人，伯克利计算机科学博士，导师是 Ion Stoica。

Kwon 提出的 PagedAttention 算法解决了 GPU 显存管理中的碎片化问题，这个算法至今仍然是 vLLM 的核心技术。

首席科学家游凯超曾获清华大学特等奖学金，他在 vLLM 中主导了分布式推理功能的开发。

Inferact 今年完成了 1.5 亿美元的种子轮融资，公司估值 8 亿美元。a16z 和 Lightspeed 领投，真格、红杉、Altimeter 跟投。

这次的 TPU megakernel 来自 Inferact 与 Google Cloud 的一项联合工程合作。

双方的目标是让 TPU 成为 vLLM 的一流支持对象，所有的优化成果都会回馈给上游的开源社区，tpu-megakernels 代码仓库是这次合作的第一个公开成果。

参考链接：

https://inferact.ai/blog/tpu-megakernels

一键三连「点赞」「转发」「小心心」

欢迎在评论区留下你的想法！

— 完 —

点亮星标