inferact 发布 TPU v7 megakernels，Kimi K3 解码速度超 700 tokens/s

今天我们发布了 inferact/tpu-megakernels，这是一组面向 TPU v7 的 megakernel 集合。我们的 Kimi K3 实现在投机解码下达到每秒超过 700 个 token，而 GB200 上为每秒 452 个 token。在不使用投机解码的情况下，我们针对 K3 和 Qwen 3.8 27B 的 megakernel 在 batch size 1 到 8 时，解码吞吐量约为 GB200 基线的 1.4 到 2 倍。

解码速度受限于硬件的峰值内存带宽。我们的 megakernel 通过有效利用 TPU 的大容量片上 VMEM、顺序编程模型，以及显式的异步流水线来跨层预取权重，将端到端性能推向这一理论峰值。本文将解释我们的设计思路，以及为什么 megakernel 与 TPU 是很好的搭配。

仅看顶层规格，Google 的 TPU v7 Ironwood 芯片与 NVIDIA 的 GB200 GPU 相当。 来源：Google Cloud TPU7x 文档，NVIDIA GB200 NVL72。

这种对等的比较忽略了一个 TPU 最有趣的特征：其片上内存的形状和大小。解码速度主要取决于如何高效地将数据从 HBM 移动到快速的片上内存，这使得这些维度成为关键的设计考量。

与 GPU 由硬件管理的 L1/L2 缓存以及程序管理的 Shared Memory 和 Tensor Memory 空间构成的多级内存层次不同，TPU 的内存层次很简单。来自 HBM 的数据被传输到一个名为 Vector Memory（VMEM）的大型片上 SRAM 池；TPU 的向量寄存器（VREG）直接从 VMEM 中填充。重要的是，VMEM 中所有数据的生命周期完全由我们的 TPU 程序管理：当我们将某些数据移入 VMEM 后，它会一直留在那里，直到我们用其他数据显式覆盖它。

每个 TensorCore 拥有 64 MiB 的 VMEM，我们可以预取大多数模型权重张量，确保在任何操作需要使用它们时，它们已经暂存在快速的片上 SRAM 中。GPU 的 Tensor Memory（TMEM）的组织方式则不同。一整块 GPU 总共只有约 38 MiB 的 TMEM，但这一容量被分配给全部 152 个 SM，每个 SM 256 KB 的空间几乎没有余量来提前暂存权重。

大多数 kernel 都有大致相同的轮廓：从 HBM 加载权重和激活值，执行一些计算，然后将结果存回。只有加载阶段会从 HBM 拉取大量数据，导致其他阶段带宽利用率不足。设计良好的 kernel 通过细分和重叠各个阶段来缓解这些带宽气泡，但程度有限。单次解码步骤可能会依次调度数百个这样的 kernel，这些气泡会累积起来。

GPU 已经提供了应对这一问题的方法。CUDA Graphs 减少了主机启动开销，stream 允许独立分支并发执行，程序化依赖启动（PDL）允许依赖的 kernel 在其生产者完成之前就开始准备和加载权重。这些机制在保留独立 kernel 的同时恢复了一定的重叠。

megakernel 则消除了这一边界。当整个解码步骤都在一个程序内时，权重加载不再绑定到使用它的 kernel——只要片上存储允许，就可以尽早发出。约束条件变成了缓冲区生命周期和容量。在 TPU v7 上，大容量 VMEM 和显式异步传输让我们可以在当前层仍在计算或通信时，暂存下一层的权重。目标是在一个以搬运权重为主要开销的解码步骤中，让 HBM 接口始终保持忙碌。

现有的 TPU 实现依赖 XLA 来执行融合和调度，但找到最优的跨层流水线需要对存储复用、传输时机和同步进行联合决策。自动优化可能无法发现我们想要的调度方案。Pallas 让我们能够直接表达这些决策，同时将指令调度留给编译器。

Pallas 是 TPU 的 kernel 语言。据我们所知，这是第一个用 Pallas 编写的开源推理 megakernel，我们发现这门语言出奇地适合这一任务。以下是本工作中最重要的 Pallas 特性。

Pallas 语言可用于为 TPU 和 GPU 编写 kernel。因此，可以用任何 GPU kernel 工程师都感到熟悉的方式来编写 Pallas kernel：我们可以将 kernel 作为程序网格启动，每个程序接收并操作输入的特定 tile。

然而，TPU 遵循固有的单线程执行模型。即使我们将 kernel 作为网格启动，每个单独的程序仍然是顺序运行的，一次一个。虽然基于网格的编程模型有其优势，但我们发现，对于 megakernel 开发而言，拥抱 TPU 的单线程特性在可读性和细粒度控制方面都是更好的选择。我们的 megakernel 作为单个程序启动，没有显式的网格。

我们的 megakernel 必须在积极预取模型权重与每个 TensorCore 有限的 64 MiB VMEM 池之间取得平衡。有些数据，如激活值，在整个前向传播过程中常驻 VMEM，而其他数据，如 MLA 中使用的 KV cache，生命周期非常短。Pallas 的 run_scoped 原语可以对特定 VMEM 分配的生命周期进行细粒度控制。这简化了程序员的任务：我们无需显式管理 VMEM 中哪些字节范围在任何时间点可以安全读写，而是可以编写自包含的“迷你 kernel”，精确地定义它们所需的 VMEM 的形状和大小。

（原文较长，此处译出前 6000 字，完整内容请点「打开原文」。）