TPU megakernel 让 Kimi K3 达到 709 tokens/s，超越 GB200

一个开源 megakernel 引发的性能话题

@inferact 团队最近开源了一个针对 Kimi K3 的 TPU megakernel，宣称在推理吞吐上达到 709 tokens/s，明显高于 GB200 平台上的 450 tokens/s。这个数字本身就足够引人注目——它意味着在特定工作负载下，经过深度优化的 TPU 内核有机会追平甚至反超英伟达最新的高端 GPU。

对于关注大模型推理成本与部署效率的团队来说，这类工程实践的价值往往不在于某个孤立的基准数字，而在于它揭示了如何把硬件特性压榨到极致。megakernel 的核心思路，是把原本零散的多个算子调用合并成一个大内核，从而减少启动开销、内存往返和调度延迟。

megakernel 的关键：92 层 MoE 塞进一个 Pallas 内核

根据该团队披露的信息，Kimi K3 的全部 92 个 MoE（混合专家）层 都运行在单一的 Pallas 内核中。Pallas 是 JAX 生态里用于编写自定义 TPU/GPU 内核的编程接口，允许开发者在较低层面控制内存布局与计算调度。

把 92 层都放进一个内核，意味着避免了逐层的内核启动开销，也让编译器和运行时有机会做更激进的全局优化。对于 MoE 架构而言，专家路由和稀疏计算本身就容易产生大量碎片化的小算子，逐个调度会带来可观的开销，将其融合成 megakernel 正是对症下药的做法。

跨层的权重预取

更值得关注的是权重预取（weight prefetching）机制。该实现让权重的预取「跨越层边界」——即在计算前一层的同时，提前把下一层所需的权重从慢速存储搬运到快速存储。这样一层的数据传输就能与前一层的计算重叠，隐藏掉本该串行等待的搬运延迟。

对于 MoE 这类参数量巨大、但每次前向只激活部分专家的模型，权重加载往往是瓶颈之一。通过计算与传输的重叠，内核能持续保持算力单元的高利用率，而不是频繁地在等待数据中空转。这也是它能把吞吐推到 709 tokens/s 的重要工程细节。

MoE（Mixture of Experts，混合专家）是一种稀疏激活的神经网络架构：模型整体参数量很大，但对每个输入 token，只有少数几个「专家」子网络被路由器选中并参与计算，其余专家的权重在该次前向传播中保持沉默。这一设计使得模型可以在不线性增加推理计算量的前提下大幅扩展参数规模。Kimi K3 采用的正是这一架构，92 层意味着每层都有一套专家集合，路由决策也在每层独立发生。稀疏性带来的挑战在于：被激活的专家分布不规则，容易产生大量细碎的矩阵乘法调用，在传统逐层调度模式下会频繁触发内核启动，造成 GPU/TPU 算力单元大量空转。

权重预取本质上是一种「双缓冲」或流水线技术，在体系结构领域有悠久的历史。其原理是将存储层次划分为慢速（容量大、带宽低，如 HBM 或片外 DRAM）和快速（容量小、带宽高，如 SRAM 或片上缓存）两级，在算力单元消费当前数据的窗口期内，异步地把下一批数据搬入快速存储，使计算与搬运两条流水线并行推进。对大型 MoE 模型而言，单个专家的权重矩阵可能高达数 GB，整个模型权重远超片上存储容量，每层都需要从 HBM 重新加载。如果不做预取，加载延迟会完全暴露在关键路径上，导致算力利用率骤降。跨层预取将这一思路延伸到层边界之外，只要编译器或运行时能静态或动态地预判下一层所需权重，就能将搬运时间几乎完全隐藏。

为什么这个对比有意义

709 vs 450 tokens/s 的对比之所以受关注，是因为它触及了当前推理硬件生态的一个核心议题：在英伟达 GPU 之外，TPU 等替代硬件能否在真实工作负载上提供有竞争力的性价比。

需要客观看待的是，单一吞吐数字受到诸多因素影响——batch 大小、精度设置、上下文长度、是否为端到端延迟等，都可能改变结论。因此这一结果更应被理解为「在特定条件下、经过极致优化的 TPU 内核可以达到的水平」，而非普适的硬件优劣判断。原始信息来自单一来源（@inferact 团队的推文与仓库），完整的 writeup 和代码仓库链接在其推文串中给出，感兴趣的开发者可自行复现验证。

GB200 指英伟达 Blackwell 架构的 Grace Blackwell 超级芯片，是目前英伟达面向数据中心推理场景的旗舰产品，集成了 NVLink-C2C 互联的 Grace CPU 与 B200 GPU。450 tokens/s 作为对比基线，代表的是在该平台上运行同一模型（Kimi K3）的吞吐水平，但需注意两者的硬件配置、内存带宽、互联拓扑乃至软件栈均存在根本差异，因此这组数字更适合理解为「TPU 经极致优化后的相对位置」，而非同等条件下的直接竞品对比。

对开发者的启示

这类开源工作对推理优化领域有几点直接价值：

• 算子融合是高价值方向：将大量小算子合并为 megakernel，能显著削减调度与启动开销，尤其适合 MoE 这类碎片化计算密集的架构。
• 内存搬运优化不可忽视：跨层权重预取展示了如何用计算掩盖 I/O 延迟，这是大参数模型推理的通用优化范式。
• 硬件多样性正在增强：TPU + Pallas 的组合证明，围绕非英伟达硬件的高性能推理栈正在逐渐成熟，为部署方提供了更多选择。

对于正在评估推理基础设施的团队，这样的开源实现既是可直接借鉴的工程样本，也是一个提醒：硬件能力往往被软件栈的成熟度所限制，真正拉开差距的常常是内核层面的精细优化。