Pallas TPU Kernel 开发教程

本教程已停止更新，只作存档。 本教程面向 TPU v4。工作组已将它改写为面向 TPU v7x 的版本，并转为内部开发；这里的内容保留归档时的状态，不再修订，其中的结论和数字也不会随新的实验更新。

前言

本教程面向熟悉 JAX、了解 GPU kernel 开发，希望进一步掌握 TPU kernel 开发的读者。你将从最基础的概念出发，逐步掌握在 TPU 上使用 Pallas 编写高性能 kernel 的能力。

为什么要学习 Pallas TPU Kernel 开发

XLA 是一个非常成功的编译器。它把 JAX 程序整体降低为 HLO 计算图，自动完成算子融合、布局分配、内存规划、指令调度，以及跨芯片的分片与通信插入。研究者只需写出数学意义上的计算，就能在 TPU 上获得相当可观的性能。长期以来，Google 内外大量模型的训练与推理都直接建立在 XLA 之上。

也正因为 XLA 足够成功，过去几乎没有人有动力绕开它做更底层的优化。人们提出的优化绝大多数在 XLA 内部实现，而不是由用户另写 custom kernel。XLA 自身就是一个可以持续优化的对象：某个模式的性能不理想时，可以在 XLA 中增加一条融合规则，或编写一个新的优化 pass，而一旦合入，所有基于 XLA 的模型都能一同受益。相比之下，为单个模型维护一个 custom kernel 显得成本高、收益小。在这样的分工下，模型开发者只负责编写 JAX 程序，编译器团队负责性能，TPU 的内存分块、片上布局与流水线调度全部封装在编译器内部，对使用者而言是一个“黑盒”。这个黑盒并不是缺陷，而是一种合理的抽象：它以极低的开发成本换来了“足够好”的性能，而在当时，“足够好”确实已经足够。

大模型时代改变了这一权衡。如今没有“能运行就够了”，只有充分优化。一次多余的 HBM 往返、一个无法融合的中间张量、一段没有重叠起来的通信与计算，乘上数万颗芯片和数月训练时间，都会变成真实而巨大的成本。“足够好”与“逼近硬件上限”之间的差距不再可以忽略，而这段差距恰恰是通用编译器最难填补的部分。XLA 必须服务广泛的模型，依靠通用分析和优化 pass 寻找解法，在保证所有模型正确的前提下做出保守的决策；而新的 attention、ragged tensor、paged KV cache、稀疏路由以及跨芯片融合，往往需要从算法、数据映射、执行调度到硬件特性逐层贯通设计，其中一些策略在原生 JAX/XLA 中根本无法完整描述。曾经让开发者免于操心硬件细节的黑盒，此时变成了阻挡开发者触及硬件能力的墙。

Pallas 正是为此而生。它打破了原生 XLA 的黑盒限制，将硬件的控制权真正交还给开发者：整体计算图仍然交给 XLA 编译，而在决定性能的关键算子上，开发者可以绕过编译器的自动融合与调度，直接控制内存分块、片上布局与算子调度。

与其等待某个模式被编译器识别、等待一个新的优化 pass 被设计和部署，不如由最了解模型与工作负载的人直接写出所需的 kernel。Pallas 给研究者的正是这一级控制权：决定如何切块，如何映射 program，数据以什么布局进入 TC VMEM，HBM 与 TC VMEM 之间的 DMA 如何同 MXU 计算重叠，何时同步，以及多个阶段如何组成软件流水线；同时，这些 kernel 仍然能够接受 JAX 程序变换，并进入统一的训练、推理与分布式执行系统。

掌握 Pallas，意味着能够把新的模型思想迅速变成高效实现。模型架构不会停在今天的 Transformer，硬件也不会停在某一代 TPU。库算子只能覆盖已经稳定下来的共性，而研究前沿恰恰发生在共性尚未形成的地方。理解 kernel 的分块、访存、累加、同步和布局，研究者就能在模型结构、数值格式或硬件约束变化时重新组织数据流，为训练和推理建立一条从算法原型到生产性能的短路径。

掌握 Pallas，是在参与 TPU 的软硬件协同设计。kernel 是模型需求与芯片结构相遇的地方：哪些数据流反复受制于带宽，哪些布局转换消耗过高，哪些同步原语缺失，怎样的片上存储、矩阵单元和互联才能支撑下一代工作负载，这些问题只有通过真实 kernel 才会暴露。今天对 Pallas kernel 的深入研究，不只是优化现有 TPU，也会为 TPU v11 乃至更远期 TPU 的指令、存储层级、互联与编程模型提供设计依据。

掌握 Pallas，是为自动调优的研发铺平道路。未来的 kernel 开发必然需要自动调优，而自动调优的研发与理解底层系统直接相关。搜索空间如何定义，tile、layout、pipeline stage 与并行策略如何参数化，哪些约束必须满足，开销模型观察什么信号，以及如何判断搜索结果是否真正逼近硬件上限，都建立在对 kernel 和 TPU 架构的深刻理解之上。Pallas 既是今天手工逼近峰值性能的工具，也是研发下一代 kernel 生成、开销模型和自动调优系统的实验场。

遗憾的是，有些用户给 Pallas 提出一些莫名其妙的 issue，根源往往只是想当然地要求 Pallas 像 JAX 一样工作。Pallas 不是另一套 JAX，它所表达的是 TPU 硬件已经具备的操作与执行能力，而不是凭空创造硬件并不支持的语义。系统学习 Pallas，正是为了正本清源，建立对其能力范围的基本认识，并尽可能减少这类由误解产生的无意义 issue。

学习的目的是服务 TPU 硬件本身

我们学习的根本目的只有一个：最大化 TPU 的实际性能。我们选择 Pallas，是因为当前 Pallas 是最贴近直接控制硬件的入口，而不是因为我们认同或依赖 Pallas 这个项目本身的能力范围。

必须说明一个现实：Google 的 JAX 团队首先服务的是内部的工作负载，对外部使用者的诉求并不积极响应，在资源投入上也相当保守。当前，外部个人开发者能够稳定获取并使用的最高级 TPU 世代是 TPU v4，而讽刺的是，在 JAX 团队的视角里，这已经被称为“早已落后数代”的硬件，不再是功能开发的重点。这意味着，如果我们的优化需求依赖官方在 Pallas、Mosaic 或底层 runtime 上的主动投入，我们大概率会长期等待，或者根本等不到。

既然等不来官方，我们就只能自己向下走。一个 Pallas kernel 要先经过 Mosaic 降低，再由闭源的 libtpu 编译为机器指令。我们不满足于停留在 Pallas 的抽象之上，而是持续追问编译器最终生成了哪些指令、各条数据路径实际需要多少周期、硬件还有哪些能力没有被现有软件栈暴露出来。只要 TPU 硬件具备某种能力，并且利用这种能力对我们的优化确实有益，我们就应当有方法把它实现出来，无论这条路径是否被 Pallas 的公开接口覆盖，或者是否被 JAX 的编程模型认可。

XLA 是通用编译器，它采用的算法未必最好；即便算法合适，其切块、布局与流水线参数也不会是为当前任务专门挑选的。所以我们总是先读 XLA 生成的反汇编，弄清它的具体做法，至少写出一个与之相当的版本，再针对具体任务继续优化。原生 XLA 是必须越过的底线，一个比 XLA 还慢的 kernel，没有存在的理由。

遗憾的是，许多人仍在用 GPU 的习惯开发 TPU kernel：先写出一个版本，再打开 XProf 时间轴判断瓶颈。这套方法在 GPU 上行之有效，因为 GPU 本质上是动态的处理器：warp 由硬件调度器在运行时切换，缓存是否命中、分支是否分化、访存能否合并，都要到执行时才能确定，无法在编译期算出。TPU 则完全不同，TensorCore 按编译器排好的 VLIW 指令包顺序执行，除了 DMA 完成、跨核同步等需要等待外部事件的少数操作，一段代码需要多少周期，可以直接从反汇编中算出来。而 XProf 在 TPU 上的逐指令时间轴只在稀疏的位置插入 trace 点，中间的时间按编译器估计的周期插值得到，看似精细，实则并非实测；插入的 trace 指令本身还可能打断原本紧密的指令发射，引入观察者效应。在一台本可以静态分析的机器上，用不准确、又会干扰执行的时间轴去猜测瓶颈，是舍近求远。

学习本教程中的 Pallas 用法时，请始终记住：你在学习的不是“Pallas 允许你做什么”，而是“TPU 允许你做什么”。

目录

• 第一章：基本操作的 TC 实现：一颗芯片、一个 TensorCore 上每种基本操作由哪些指令完成、要多少周期。
• 第二章：多 TC 与 DMA：两个 TensorCore、Megacore Shared CMEM、多芯片与主机内存之间的数据通路，以及由它们推导出的流水线与集合通信。
• 第三章：静态分析与计时：主机时间与设备周期数的关系、公平比较的边界、两级发射的规则、从清单推算周期数的发射模型，以及跨 TensorCore、跨芯片的全局时钟。
• 第四章：随机数：硬件随机数生成器、计数器式生成器与各种分布的开销。
• 附录：术语中英文对照：正文使用的术语、对应的英文和首次详细介绍的位置。

开发环境

本教程使用 TPU v4、JAX 0.12.0.dev20261002+7fc69a22c2、jaxlib 0.12.0.dev20261002、libtpu 0.0.49 和 tpuasm 0.2.1。

虽然 TPU 已经有了更新的世代，但 TPU v4 足以承载本教程要介绍的核心思想。后续世代的具体指令、存储容量、带宽与计算吞吐会变化，但从硬件约束出发、检查编译结果、分析瓶颈并验证优化的方法仍然适用。在 TPU v4 上建立起这套方法，迁移到新的 TPU 时，就能根据新的硬件特性重新推导实现与调优策略。

许可证

本仓库的源代码与代码示例采用 MPL-2.0，教程正文、图表等非代码内容采用 CC BY-SA 4.0，详细说明见 LICENSE。

致谢

本项目的 Cloud TPU 计算资源由 Google TPU Research Cloud (TRC) 提供支持。