Google 与 Meta 发布 TorchTPU：PyTorch 原生 TPU 后端

TL;DR

TorchTPU 是 Google 与 Meta 合作开发的开源项目，为 Google TPU 引入了原生且符合设计原则的 PyTorch 后端。它支持开箱即用的无缝模型执行与 TPU 迁移，并通过 eager 执行优化、torch.compile 集成和原生硬件加速，提供具有行业竞争力的推理与训练性能。

推出 TorchTPU

Google TPU 在 Google 内外都已被非常成功地用于大规模训练和部署大型机器学习模型。它们专为机器学习而设计，因此在性能、能耗和成本方面具有显著优势。尽管 TPU 历史上一直与 TensorFlow 和 JAX 生态系统相关联，我们很高兴推出 TorchTPU——这是 Google 于 2025 年中启动、与 Meta 的 PyTorch 团队合作开展的新项目，旨在让 PyTorch 原生运行在 TPU 上。

为了在 Python 框架之下对 Google TPU 进行编程，开发者使用 XLA 编译器栈。2020 年，Google 推出了首个面向 TPU 的 PyTorch 编译版本 PyTorch/XLA，基于 Lazy Tensors。尽管具有奠基意义，这一早期集成难以跟上 PyTorch 由社区驱动的快速演进。关键在于，PyTorch/XLA 需要侵入式修改源代码以显式定义可编译区域，编译时延迟严重，不支持 PyTorch 原生的命令式（“eager”）执行模式，并且在支持 PyTorch 生态内新兴的分布式计算原语方面滞后。与 PyTorch/XLA 和 TorchAX 等以往将 PyTorch 强行塞入 JAX 运作方式的尝试不同，TorchTPU 是作为 PyTorch 的一个符合设计原则的 TPU 后端而构建的：eager 执行是首要优先级，现有 PyTorch 工作负载只需极少的代码改动即可开箱即用（例如，.to("cuda") → .to("tpu")）。

更具体地说，在启动 TorchTPU 时，我们刻意制定了一套遵循的原则：

• PyTorch Citizenship:
  • 运行时中不依赖 JAX；
  • 稳健的 eager 模式；
  • 基于 Torch Dynamo（而非 Lazy Tensor）的 Torch.compile()，性能与 JAX compile 相近；
  • 通过在 PyTorch 栈的最底层（例如 Process Group API）提供 TPU 集合通信，广泛支持 PyTorch 分布式惯用法（DDP、FSDP、DTensor 等）；
  • 无缝集成到 PyTorch 生态系统中；
  • 社区参与。
• Focus on the User:
  • 将工作负载移植到 TPU 应当像把 model.to("cuda") 改为 model.to("tpu") 一样简单。（仅靠这一简单改动无法达到 TPU 的峰值性能，因为模型的形状、分片和算子可能偏向 GPU 架构；但模型应当能够开箱即用。）
  • 为运行时/编译器错误提供用户友好且可操作的提示信息；
  • 与 PyTorch 社区已在使用的工具和框架集成（例如 torch.profiler、DTensor、Ray）。
• Improve the State of the Art:
  • 构建一个能够在大规模工作负载上与 GPU 上的 PyTorch 竞争的解决方案；
  • 为 eager 执行引入 JIT 编译方法，通过运行时优化弥合逐算子执行与编译执行之间的性能差距；
  • 通过引入有界动态性、避免因形状变化导致 XLA 重新编译，以及探索可提前预编译的动态 TPU 内核，改进 TPU 编程栈。

TorchTPU 正在积极开发中，它已在生产环境中用于部分最大 Google Cloud 客户的推理和训练工作负载，并且可以在该 GitHub 仓库 中找到其开源项目。撰写本文时，我们已：

• 新增对 eager 和编译模式的支持，包括新的 DeferAndFuse eager 模式；
• 新增对约 450 个 Aten 算子的支持（跳过具有结构化或复合委托的算子），包括主要深度学习架构所需的全部 Core ATen 算子的前向和反向传播；
• 新增通过 JAX/Pallas 和 Helion DSL 编写自定义内核的能力；
• Validated out-of-box execution on the 5,000 most downloaded Hugging Face models, including:
  • LLM：Qwen、DeepSeek、Gemma、LLama、Mistral、GPT、OPT、BLOOM、Falcon、StableML
  • 音频/语音：Whisper、Wav2Vec2、HuBERT、MMS
  • 视觉：DinoV2、SigLIP、CLIP、Swin Transformer、ConvNeXt、ViT、ResNet、DETR
  • 文本编码器与搜索：BERT、ALBERT 与 Electra、RoBERTa、DeBERTa、T5、MT5、Flan-T5
• 已验证基于 DTensor 的分片；
• 已验证多主机推理（vLLM 和 SGLang）以及微调/训练工作负载（TorchTitan），并采用基于 Ray 的编排；
• 已验证与 vLLM、SGLang、TorchTitan、Ray 和 TorchRec 的集成。

在本文档的剩余部分，我们将讨论 TorchTPU 架构的关键方面以及一些初步结果。

架构深入探讨

TPU 入门

Google TPU 是专门设计的 ASIC，用于执行作为深度学习基础的巨量张量运算。从 硬件架构角度 来看，一个 TPU 芯片 包含（取决于代际）：

• 一个或多个 TensorCore，每个包含一个或多个矩阵乘法单元（MXU）、一个向量单元（VPU）和一个标量单元（参见 此链接）。根据 TPU 代际，一个 MXU 可包含 256×256 个乘加器，它们以 脉动阵列 的形式排列。
• 一个或多个 SparseCore，它们是用于稀疏操作（例如 gather、scatter、embeddings）和集合通信的 专用单元。它们可以与 TensorCore 并行运行，从而实现计算-计算和计算-通信的重叠。
• 大容量 高带宽内存（HBM）——在 TPU 8i 中为 288 GB、8.6 TB/s，可用于存储 ML 数据并与 CPU 共享。
• 双向 芯片间互连（ICI）——在 TPU 8t 和 8i 中为 19.2 TBps，允许用户配置的 TPU Slice 内的 TPU 芯片之间快速通信，TPU Slice 是 TPU Pod 的逻辑子组，由预配置数量的 TPU 芯片组成，这些芯片以 3D torus（适用于 TPU 7、TPUv 8t）或 Boardfly（TPU 8i）配置连接。
• 双向 数据中心网络（DCN）——在 TPU 8t 中为 400 GBps，允许跨 Slice 的 TPU 芯片之间通信。

TPU 通过两个不同的入口点使用 XLA 编译器 栈进行编程：StableHLO（用于 ML 算子的较高级 IR）和 Mosaic（用于自定义内核的较低级 IR）。实际的 TPU ISA 是专有的、可能发生变化且不公开。实际上，可以将 StableHLO 以及（更是如此）Mosaic 视为等同于 PTX。参见下图的高层示意图。

C++ 程序可以通过实例化 PjRT Client 来访问本地 TPU，其 API 提供了将 MLIR（StableHLO 或 Mosaic）编译为 TPU 可执行文件、启动该可执行文件以及在 TPU 之间复制数据的方法。请注意，TPU 执行是异步的，CPU 可以通过等待关联的 PjRTBuffer 来等待 TPU 结果就绪。

TPU 编程模型的显著特点：

• TensorCore 设备可以被视为一个单线程处理器，带有一个或多个大规模矩阵乘法单元（MXU）。因此，要达到峰值性能，关键在于模型的形状要足够宽，以充分发挥 MXU 的潜力。这与 GPU 不同，在 GPU 上并行启动多个较小的内核有助于最大化硬件利用率。
• XLA 编译器可以指示 TensorCore 将稀疏运算和集合运算卸载到 SparseCore，以实现更高效和/或并行的执行。这可以自动发生，也可以通过编译后的 StableHLO 模块中用户提供的 MLIR 注解来实现。
• 一个 TPU 设备只能被一个用户空间进程访问。遵循我们无缝集成到 PyTorch 生态系统的目标，对于分布式 TorchTPU 工作负载，我们采用每设备一进程的模型，这与 JAX 的每主机一进程模型形成对比。然而，由于一个 TPU 只能被一个用户空间进程访问，RL 共置无法开箱即用地支持，我们正在研究放宽或绕过这一 TPU 限制的方法。
• StableHLO 遵循 SSA 范式，因此不允许原地修改，即写操作总是产生一个新值（至少在概念上如此）。这对 TorchTPU 尤为重要，它确实通过视图分解和反转提供了熟悉的 PyTorch“视图”语义（包括修改），如本文档后面所述。
• XLA 编译器和 PjRT 运行时假定张量形状在编译时已知。认识到这是 PyTorch 工作负载的一个限制，我们一直在添加对有界动态性的支持，即当形状维度的最大边界在编译时已知，但实际维度仅在运行时已知时。
• 集合操作被降级为 StableHLO 操作传递给 XLA，而不是使用 NCCL 等单独的库。

可用性与自动加载机制

我们在设计 TorchTPU 时考虑了可用性。当用户在安装了 TorchTPU wheel 的 TPU 机器上导入 PyTorch 时，在 TPU 上运行 PyTorch 代码只需进行微不足道的更改，并且得益于 PyTorch 的自动加载机制，只需几行代码即可完成：

请参阅下表进行并排比较。

import torch
import torch.distributed as dist

def main():
    device = torch.device("cuda")
    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    word_size = dist.get_world_size()
    ...
    model = model.to(device)
    ...

import torch
import torch.distributed as dist

def main():
    device = torch.device("tpu")
    dist.init_process_group(backend="tpu_dist")
    rank = dist.get_rank()
    word_size = dist.get_world_size()
    ...
    model = model.to(device)
    ...

import torch
import torch.distributed as dist

def main():
    device = torch.accelerator.current_accelerator()
    dist.init_process_group()
    rank = dist.get_rank()
    word_size = dist.get_world_size()
    ...
    model = model.to(device)
    ...

数值精度

TorchTPU 为 TPU 硬件适配标准 PyTorch 精度 API。与所有高性能机器学习架构一样，TPU 针对其特定代际所偏好的原生数据类型进行了优化，以实现峰值性能和能效。对于 TPU，原生矩阵乘法格式为 bfloat16（BF16），它在高动态范围与高效吞吐之间取得平衡。由于 TPU 矩阵乘法单元（MXU）被设计为使用 BF16 输入执行算术运算，并以 32 位浮点进行中间累加，这一配置是 TPU 高性能计算的基石。对于非收缩运算（例如逐元素加法、激活、layernorm 和 softmax），TorchTPU 使用 TPU 向量执行单元（VPU）执行 FP32 张量运算，该单元原生执行完整的 IEEE-754 FP32 算术。

TorchTPU 将 PyTorch 精度模式映射到底层 XLA/StableHLO 配置，以暴露这一硬件能力：

• medium：乘法使用较低精度（将输入转换为 MXU 的原生 BF16），而累加以 FP32 执行以保持稳定性。TorchTPU 默认使用此模式以充分利用 TPU 硬件加速，因为单遍 BF16 乘法配合 FP32 累加对几乎所有深度学习工作负载都能提供足够精度。
• high：使用 3 遍算法，将输入拆分为高、低分量以计算三个交叉项，以约 3 倍计算成本恢复大部分 FP32 精度。
• highest：使用 6 遍算法计算所有交叉项，以近似完整的 IEEE FP32 精度，在 MXU 上产生约 6 倍计算开销。

对于绝大多数深度学习工作负载，我们建议使用 medium（默认）精度。在 TorchTPU 中显式选择 high 或 highest 会触发性能警告（TORCH_WARN_ONCE）。只有在明确识别到与精度相关的收敛问题时才切换到 high 或 highest，因为这些模式会带来显著的性能开销。除全局设置外，TorchTPU 还提供线程本地上下文管理器（torch.tpu.precision）用于细粒度作用域（例如针对数值敏感的层）。

并排对比见下表。

Eager 执行、StableHLO 降级和 JIT 编译

遵循这份 PyTorch 指南，我们通过一个新的 C++ 后端扩展了 PyTorch dispatcher，并将其注册为 PrivateUse1。此外，我们使用了标准的 at::Tensor 机制来定义张量存储的设备特定实现（通过在 c10::StorageImpl 内使用 c10::UniqueVoidPtr）。

TorchTPU 的 eager 执行的一个独特特性是，它可以以多种模式将 Aten 算子分发到 TPU：

• 立即，即有效地实现逐算子分发；
• 延迟，一直延迟到自动确定的物化点（例如，一条 .to("cpu") 语句），并与截至该点被延迟的所有其他算子一起聚合/编译；
• 介于两者之间的任何情况。

虽然这种灵活性确实给设计带来了额外的复杂性，但我们从一开始就刻意加入它，因为我们希望试验不同的策略，以在基本的逐算子分发模式之上提升 eager 性能。（此外，正如本文档后面所讨论的，这种灵活性在处理视图语义和 torch.compile 时非常方便。）

总之，TorchTPU 提供了 3 种 eager 模式，可以通过环境变量或上下文管理器来设置：

• Strict（默认）：我们一次分发一个 Aten 算子，它们在 TPU 上与 CPU 异步执行。这相当于在 GPU 上以 eager 模式运行 PyTorch 时的预期行为。
• Debug（export TORCH_TPU_LAUNCH_BLOCKING=1）：我们一次分发一个 Aten 算子，但 TPU 执行与 CPU 同步。这很慢，旨在用于调试，相当于在 GPU 上使用 CUDA_LAUNCH_BLOCKING，并额外提供了 TorchTPU 的边界检查。
• DeferAndFuse（export TORCH_TPU_DEFER_AND_FUSE=1）：我们将多个 Aten 算子延迟并融合成更大的块。这是 TorchTPU 引入的一种新的 eager/编译混合模式，比 Strict Eager 快 50% 或更多，同时降低峰值 HBM 内存（算子融合减少了 HBM 中的临时存储）。

现在我们已经描述了 TorchTPU eager 模式的基础知识，让我们勾勒一下当收到来自 PyTorch dispatcher 的回调调用时会发生什么。假设我们使用的是 DeferAndFuse 模式，并以 at::mul.out 为例：

• Create a Deferred Op instance that contains:
  • A StableHLO builder object, which when invoked returns a StableHLO function implementing the semantics of aten::mul, as shown below (note that any implicit broadcasting and dtype promotion are handled elsewhere)

    func.func @main(%arg0: tensor<4x16xbfloat16>, %arg1: tensor<4x16xbfloat16>) -> tensor<4x16xbfloat16> {
      %0 = stablehlo.multiply %arg0, %arg1 : tensor<4x16xbfloat16>
      return %0 : tensor<4x16xbfloat16>
    }

  • 一个可选的静态算子参数映射（例如，对于 aten::native_batch_norm 使用 training 和 momentum），用于计算缓存键（下文讨论）；
• 返回一个新的 at::Tensor，其 c10::StorageImpl（通过间接方式）指向该 Deferred Op 实例。

当收到来自 PyTorch dispatcher 的对应于 mytensor.to("cpu") 或 mytensor.item() 或类似内容的回调时，TorchTPU 将：

1. 分析到目前为止累积的、导致 mytensor 计算的延迟算子区域 R；
2. 应用活跃性分析来确定 R 中哪些算子需要物化，以及产生 mytensor 的那个算子；重要的是，具有副作用的算子（例如，集合通信）总是被保留（其调用顺序也是如此），而 R 之外不活跃的算子可以在不同的工作负载副本之间确定性地检测到，因为 Python 的内存管理使用引用计数。
3. 为要物化的算子图 G 创建一个编译缓存键，以便将来快速且唯一地识别该图；
4. Use the computed cache keys to check in a Compilation Cache if graph G has been already compiled.
   • 如果是，则在 TPU 上分发 G 的可执行文件；
   • 如果不是，则使用线程池（异步地）编译 G，然后在 TPU 上分发其执行。
5. 等待 mytensor 的 TPU 值就绪，然后将其从 TPU HBM 复制到与目标 at::Tensor 关联的 CPU 存储。

表面上看，上述方法似乎效率不高，因为 CPU 只有在遇到 .to("cpu") 语句时才会启动 TPU 执行。对于真实工作负载，我们可以预期，在少数初始迭代/步骤填充编译缓存之后，CPU 执行应仅涉及 PyTorch 追踪和算子分发，因此会始终领先于 TPU 执行，而后者将忙于冗长的数值计算。尽管如此，为了加速对多次重复的算子序列（想想一堆 transformer 块）的物化，我们添加了重复算子启发式（Repeated Ops Heuristic），它会检测此类重复，并在到达 TPU 到 CPU 传输之前就触发物化点。

编译缓存

TorchTPU 采用编译缓存来避免不必要的 XLA 编译。与现代计算机存储系统类似，该缓存包含多个协同工作的层级：

• 第 1 层：进程本地内存缓存，用于在同一进程内共享编译结果。这是三层中最快的。
• 第 2 层：机器本地共享内存，支持同一主机上的跨进程共享。这对于多设备训练/推理很有用。
• 第 3 层：由远程文件系统（例如 Google Cloud Storage）支持的持久缓存，允许跨主机、程序和作业重启复用。这在 Colab 和数据中心等执行环境中运行 TPU 程序时很有用。

第 1 层缓存始终启用。默认情况下，第 2 层对分布式程序启用，对单设备程序禁用。用户可以根据需要覆盖此行为。第 3 层默认禁用，但可以通过指定共享存储路径显式启用。

缓存中的每个条目都是针对 StableHLO 计算图编译的 TPU 可执行文件。在第 1 层缓存中，每个可执行文件都是一个 C++ 对象，可以直接调度到 TPU 设备上执行。由于对象无法直接跨进程共享，第 2 层和第 3 层缓存将可执行文件保存为序列化的二进制字符串。

出于性能原因，我们不使用 StableHLO 图本身作为缓存键。相反，我们将计算、编译器选项和主要组件的版本指纹化为三个 64 位整数，并使用该元组作为缓存键：

• 一个无形状键，捕获计算图的结构，但不捕获张量输入/输出的大小，
• 一个维度键，捕获张量大小，以及
• 一个编译选项键，捕获传递给 XLA 编译器的选项以及 TorchTPU、libtpu 和 PjRt 运行时的版本。

无形状键和维度键合在一起也称为图键。虽然理论上不同图、编译器选项和组件版本之间可能发生缓存键冲突，但我们认为 192 位键空间足够大，实践中不会出现此类冲突。

视图语义

PyTorch 的操作语义与 StableHLO 语义之间一个特别不同的领域是可变性和共享内存视图的处理：

• PyTorch 假设每个 Tensor 都是指向线性字节数组的指针，并在 Tensor 层级带有 stride 元数据，用于描述逻辑索引元组到该数组中线性偏移量之间的转换关系。多个 Tensor 可以共享同一底层字节数组，因此可以从任意 Tensor 执行原地写入，并且该写入会从任何其他 Tensor 可见。得益于附加的 stride 元数据，视图操作可以在 dispatch 阶段在 CPU 上以较低开销完成，只需更新该元数据即可。虽然这些视图大多属于少数几种常见操作，但使用 as_strided 运算符可以定义任意复杂的模式。
• 相比之下，StableHLO 是一种纯函数式语言，其中所有张量值在构造后均不可变。物理布局对用户是抽象的，它反映编译器做出的分块决策，因此不支持直接操作 stride。此外，StableHLO 不支持原地写入；相反，可以通过 dynamic_update_slice 操作应用选择性更新，将一组新值合并到现有张量的子切片中。

在存在写入的情况下，在 TPU 上实现正确的视图语义并不简单；然而，TorchTPU 借助一些机制实现了这一点，包括视图分解、视图反转和延迟操作。这一语义桥接使 TorchTPU 能够原生支持原本不可能支持的 PyTorch 模式，例如用于分布式训练的拼接参数数组。

流与 SparseCore

在优化模型性能时，PyTorch 用户可能希望同时执行 3 类操作：

• 计算操作；
• 集合通信操作；
• 主机与设备之间的数据传输。

在 GPU 上，这通常通过使用单独的命令队列来实现，以便设备可以将工作分派到不同的硬件单元（SM、DMA 等）。在 PyTorch 中，设备命令队列以 Streams 的形式暴露，它们通常既用于用户代码，也用于框架代码（例如在 FSDPv2 中）。

TorchTPU 也实现了 PyTorch Stream API，但并非在所有情况下都能实现 GPU 上可用的相同功能。考虑到 TPU TensorCore 本质上是一个单线程执行处理器，我们在给定 TensorCore 上一次只能分派一个计算。相反，GPU 包含许多相同的计算核心（SM），可以并行运行。所幸，一个 TPU TensorCore 关联一个或多个 DMA 引擎以及一个或多个 SparseCore 引擎。不出所料，DMA 引擎可用于在主机与设备之间进行数据传输，并与 TensorCore 操作并行执行。更有趣的是，SparseCore 引擎是一个独立单元，针对稀疏操作（例如 gather/scatter）进行了优化，并且还可以支持集合通信操作。原则上它可以支持完整的 StableHLO 操作集，但需要注意的是，它在稀疏数学运算上会优于 TensorCore，而在稠密数学运算上会逊于 TensorCore，而后者正是 TensorCore 内 MXU 的强项。为了让用户完全控制工作运行的位置，可以将 TorchTPU 流手动固定到 SparseCore，并且从性能和操作支持完整性的角度来看，该流与运行在 TensorCore 上的默认流并不相同。流还可以用于协调异步数据传输（使用 DMA 引擎），方式与 GPU 上类似。

编译模式

我们的首要目标是保持与其他 torch.compile() 后端一致的用户体验，并因此共享相同的编译器前端，即 Dynamo。根据这份 PyTorch 指南，我们注册了一个新的 TorchDynamo 后端，其用法如下：

compiled_model = torch.compile(model, backend="tpu")

Dynamo 通过传入 torch.fx.GraphModule 以及一组示例输入张量来调用我们的后端。此时，我们像在 eager 模式下一样对所提供的图进行追踪，但所有算子都被延迟执行，且不会在 TPU 上实际运行。这使我们能够为整个 FX 图生成并编译单个 StableHLO 程序。最后，我们的 Dynamo 后端返回一个 Python 可调用对象，该对象封装了对所生成 TPU 可执行文件的调用。请注意，作为一项临时限制，在 TPU 机器上调用 torch.compile 时，我们会自动设置 dynamic=False，因为支持 SymInt 需要有界动态性（下文将讨论），但该特性目前尚未默认启用。

有界动态性

TPU 与 GPU 的一个主要区别在于，XLA 传统上基于静态形状运行，以生成优化的硬件布局和融合内核。在实践中，这并非严重限制；当张量维度（如批大小或序列长度）在工作负载执行期间发生变化时，开发者通常通过预分配更大的张量并在计算前后手动填充/切片来解决。

为避免手动修改代码和重新编译的开销，我们正在积极开发，以通过 TorchTPU 中的 XLA 有界动态性原生支持动态形状。在有界动态性下，动态缓冲区在主模型模块内按最大上限分配，图只编译一次。计算使用有界缓冲区在 TPU 上执行，而 XLA 会跟踪逻辑维度大小并自动处理填充和切片，不会触发重新编译。

内核

我们区分两类内核：

1. 由 PyTorch 框架提供的内核，例如 torch.nn.functional.scaled_dot_product_attention。
2. 用户提供的自定义内核，用于加速模型的特定部分；这些通常与具体模型相关；

对于第 1 类，我们最初实现了一种方法：用 Pallas 编写 TPU 内核，然后以符号形状导出到 Mosaic，再通过一个 TorchTPU 步骤在调用时特化内核的形状。该机制被证明不如预期灵活，因此我们正在试验一种新方法，以编写支持动态形状的 Mosaic 内核。

对于第 2 类，我们开发了三种解决方案：

• (Production ready) The user writes their TPU kernel in Pallas and integrates it with PyTorch via a TorchTPU API, as follows

  # Define a simple vector addition Pallas kernel
  def add_vectors_kernel(x_ref, y_ref, o_ref):
    o_ref[...] = x_ref[...] + y_ref[...]
  
  # Wrap the Pallas kernel into a JAX function
  @torch.tpu.pallas.jax_op("pallas::add_vectors")
  def add_vectors(x: jax.Array, y: jax.Array) -> jax.Array:
    out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype)
    return pl.pallas_call(add_vectors_pallas_kernel, out_shape)(x, y)
  
  # Use it like a standard PyTorch function
  x = torch.tensor([1.0, 2.0, 3.0], device="tpu")
  y = torch.tensor([4.0, 5.0, 6.0], device="tpu")
  z = add_vectors(x, y)

• （生产就绪）用户可利用来自 Tokamax 的内核，这是一个高性能 TPU 内核库，现已可用于与 PyTorch 无缝互操作。Tokamax 提供关键算子的高度优化实现，包括混合专家构建模块（如分组矩阵乘法和 token gather/scatter）、注意力变体（如分页注意力、门控 deltanet、Kimi delta 注意力和多头潜在注意力）以及 TopK 计算。算子组合将随时间演进，以支持在最新硬件加速器上运行的最新 OSS 模型（例如，发布时支持 TPU v8i）。这些内核由稳健的集成基础基础设施提供支持，例如自动调优（如 tile 大小搜索）、低开销的基于 XProf 的基准测试和回归测试，以降低运维复杂性。虽然内核本身用 JAX/Pallas 编写，但该库的基础设施采用 torch.tpu.pallas.jax_op API 来透明地管理互操作，并暴露原生 PyTorch API 以方便开发者。
• (Experimental) The user rewrites their GPU kernel in Helion, which enables targeting multiple accelerators with the same PyTorch-based DSL, as shown below

  # Define the kernel directly in Helion
  @helion.kernel(backend="pallas")
  def add_vectors(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(x)
    # Device execution loop over automatically partitioned tiles
    for tile in hl.tile(out.size()):
      out[tile] = x[tile] + y[tile]
    return out
  
  # Use it like a standard PyTorch function
  x = torch.tensor([1.0, 2.0, 3.0], device="tpu")
  y = torch.tensor([4.0, 5.0, 6.0], device="tpu")
  z = add_vectors(x, y)

分布式支持

为确保与现有 PyTorch 分布式 API 和框架兼容，我们决定在 Process Group API 层面实现 TPU 集合通信，通过继承 c10d::Backend；这使我们能够在 eager 和编译模式下拦截集合操作，并无缝支持 PyTorch 分布式框架，如 DTensor、FSDPv2、SimpleFSDP、FairScale 等。与 GPU 上的 PyTorch 类似，我们遵循（并实际强制执行）每设备一进程的方法。从实现角度看，TPU 栈不提供单独的 TPU 间通信库，而是将 TPU 集合操作表示为 StableHLO 操作（stablehlo.all_gather、stablehlo.all_scatter、stablehlo.all_to_all 等），这与上述机制自然契合。

在 Process Group 层面拦截集合操作几乎让我们到达终点，除了与 GPU 相比 TPU 的一些显著特性：

• 不同的初始化序列。考虑一个连接了多个设备的主机。在 GPU 的情况下，由 torchrun 启动的进程可以看到所有连接的设备，即使它只会使用其中一个（例如，使用 device=f"cuda:{local_rank}"），并且可以在调用 torch.distributed.init_process_group() 之前将模型移动到该设备。然而，在 TPU 的情况下，一个进程只能看到一个设备，并且不知道主机连接了多少设备。要正确初始化 TPU 运行时，进程需要在调用 torch.distributed.init_process_group 之前知道 (a) 运行分布式的意图，(b) TPU 数量（world size），以及 (c) TPU 类型（因为它影响拓扑）。
• Default Process Group Knowledge. Consider two subgroups (4, 0, 1, 3) and (5, 2, 6, 7) running an all_reduce among themselves. Note that torch.distributed.all_reduce takes as input only the subgroup in which the collective op must execute. In contrast, stablehlo.all_reduce requires information about all subgroups (i.e., PyTorch’s default process group), even those the current process is not part of, as shown in the following MLIR snippet:

  %0 = "stablehlo.all_reduce"(%arg0) <{replica_groups=dense<[[4,0,1,3], [5,2,6,7]]>: tensor<2x4xi64>}> ({...

• MPMD Support. Historically, the TPU stack was built with SPMD in mind. Thanks to tools such as Shardy, distributed JAX programs are written for a single host, with no communication code included, and then are sharded at massive scale with user-provided directives, while the XLA compiler is fully responsible for inserting collective ops at the right points while maximizing the overlap between computation and communication. Unfortunately, the PyTorch ecosystem doesn’t adopt a clean SPMD approach (a notable exception is SimpleFSDP) and different processes may execute different ML ops leading to a collective op for reasons such as:
  • 用户插入的以进程 rank 为条件的语句；
  • 当张量维度不能被设备数量整除时，DTensor 对共享张量的不均匀拆分。

  当包含集合操作的 TPU StableHLO 模块不一致时，可能会发生不可预测的死锁或内存损坏（这是由于生成的 TPU 代码中跨 rank 的同步标志 ID 或缓冲区偏移可能不匹配）。TorchTPU 对此的处理如下：

  • (默认) Skinny Collectives。默认情况下，TorchTPU 会将每个集合通信操作编译到其独立的“skinny” StableHLO 模块中，这保证了各 rank 之间缓冲区和同步标志 ID 的一致性。这样做虽然可行，但相比下一条中讨论的方案会带来明显的性能损失（例如端到端 30-40%），因为它不允许 XLA 编译器将计算与通信重叠。目前正在进行的一项工作是将 Skinny Collectives 运行在绑定到 SparseCore 的独立流上，从而使其能够与运行在 TensorCore 上的数学计算重叠。我们预计该优化将挽回 Skinny Collectives 的大部分性能损失。
  • (可选) @spmd_safe 装饰器。如果用户知道某段代码是 SPMD 的（例如，不存在以 torch.distributed.get_rank() 为条件的 if 语句），他们可以用 TorchTPU 的 @spmd_safe 装饰器标记该段代码。在 torch.compile 下，这会消除 TorchTPU 原本会在集合通信周围插入的图断点（参见上文 Skinny Collectives），并因此生成一个内嵌集合通信的大型 StableHLO 模块，XLA 编译器可以通过移动集合通信位置来高效优化该模块，从而最大化计算与通信的重叠；这已被证明在 SimpleFSDP 或手动放置集合通信操作时能带来显著的性能收益。这种方法的一个挑战是 MPMD 代码可能不易察觉，用户可能在代码并非严格 SPMD 的情况下使用 @spmd_safe 装饰器，从而导致脆弱的情况。
  • (Optional) Compilation Handshake (RFC). This mechanism runs when a given torch.compile statement triggers compilation of a module containing communication operations. It initiates a consensus round across all replicas to verify that the compiled FX graphs are identical, which in turn allows replicas to disable graph breaks around collectives for the given FX graphs, similarly to the @spmd_safe mechanism described above.The handshake stage is controlled by an environment variable TORCH_TPU_SPMD_HANDSHAKE_STAGE which can be set to:
    • OFF – 从不执行握手（当前默认）；
    • COMPILE_STAGE – 在开始 FX 图编译之前执行握手；
    • DISPATCH_STAGE – 在开始已编译的 FX 图执行之前执行握手。仅当用户预期会出现特定于 rank 的重新编译时才应使用此模式，因为它比上一个选项涉及更大的开销。

分布式初始化

在 TorchTPU 中初始化分布式工作负载遵循标准 PyTorch 分布式语义，但需注意 TorchTPU 需要两个 TPU 特有的变量，而标准 PyTorch torchrun 不知道如何设置它们：

• TORCH_TPU_TOPOLOGY：TPU 切片的物理 mesh 几何结构（例如，对于 2x2x1 切片配置为“2,2,1”）。
• TORCH_TPU_SLICEBUILDER_ADDRESSES：用于构建 TPU 芯片间互连（ICI）网络 mesh 的开放端口/地址列表（每个设备一个）（例如，“localhost:8471,localhost:8472,…”）。

另一种方法是在 Python 中使用 torch.tpu.distributed.set_tpu_launch_env()，它可以在 Google Cloud 环境中设置这些变量。下面我们展示启动分布式 TorchTPU 工作负载的示例代码。

A. Worker 实现（train.py）

import os
import torch
import torch.distributed as dist

def worker_main():
  # 1. Initialize distributed process group on TPU using the tpu_dist backend
  dist.init_process_group(backend="tpu_dist")

  # 2. Query rank, world size, and construct target TPU device
  rank = dist.get_rank()
  world_size = dist.get_world_size()
  device = torch.device("tpu")

  print(f"[Rank {rank}/{world_size}] Initialized on {device}")

  # 3. Perform a distributed collective operation across TPU workers
  x = torch.ones(4, 4, device=device) * (rank + 1)
  dist.all_reduce(x, op=dist.ReduceOp.SUM)

  print(f"[Rank {rank}] Reduced sum: {x[0, 0].item()}")

  # 4. Clean up process group resources
  dist.destroy_process_group()

if __name__ == "__main__":
  worker_main()

B. 使用 torchrun 启动工作负载（CLI）

# Single host launch (spawns 1 process per TPU chip on an 8-chip node)
eval $(python -m torch.tpu.distributed | sed 's/^/export /')
torchrun --nproc_per_node=8 train.py

最后，请注意给定节点的默认进程组是该节点的切片。

最优设备 Mesh 映射

由于 TPU 独特的连接性，将模型的并行策略（TC、CP、FSDP、DP）与物理 TPU 硬件互连拓扑对齐非常重要，以最大化带宽并最小化网络延迟。回顾一下，TPU 使用分层网络架构：

• ICI：直接的光纤/铜缆互连，以 Torus mesh 连接芯片。提供数十到数百 GB/秒的带宽。
• DCN：用于连接不同 TPU 切片的标准主机级以太网。其带宽显著低于 ICI，延迟显著高于 ICI。

一个好的规则是使用 TorchTPU 的 topology_aware_mesh() API，并将通信最密集的并行维度映射到最快、最连续的物理 ICI 维度，也就是将 mesh 维度从网络强度最低到网络强度最高进行排序。示例见下表：

以下是一个代码示例：

from torch.distributed.device_mesh import DeviceMesh

# Single-Slice 2D Mesh Example (DP + TP)

# Logical mesh shape: (dp=4, tp=2) on an 8-chip slice ordered from least
# communication-intensive to most communication-intensive:
mesh_shape = (4, 2)
dim_names = ("dp", "tp")

# Computes optimal rank layout aligned with physical ICI rings
optimal_rank_tensor = torch.tpu.topology_aware_mesh(
    mesh_shape=mesh_shape,
)

# Instantiate the DeviceMesh
mesh_2d = DeviceMesh("tpu", mesh=optimal_rank_tensor,
                     mesh_dim_names=dim_names, preserve_rank_order=True)

dp_mesh = mesh_2d["dp"]
tp_mesh = mesh_2d["tp"]

作为未来工作，我们将探索一种 API，允许用户更直接地自定义其 mesh，从而在特殊场景下实现最优的集合通信性能。

分布式 Pallas 内核

TorchTPU 支持自定义 Pallas 内核，可通过 ICI 上的硬件 RDMA 直接进行芯片间通信，无需主机 CPU 参与。这使你可以编写高效的点对点通信内核（例如 Ring Attention、自定义 All-to-All / Shuffles、流水线并行移位或 Halo 交换），并将它们作为原生 PyTorch 函数调用。示例如下：

# 1. Pallas Kernel Body: Asynchronously sends local buffer to neighbor rank (rank + 1)
def pallas_send_to_next_kernel(x_ref, y_ref, send_sem, recv_sem):
  rank = jax.lax.axis_index("x")
  world_size = jax.lax.axis_size("x")
  dest_id = (rank + 1) % world_size
  # Trigger hardware Remote DMA over TPU interconnect
  async_copy = pltpu.make_async_remote_copy(
      src_ref=x_ref,
      dst_ref=y_ref,
      send_sem=send_sem,
      recv_sem=recv_sem,
      device_id=(dest_id,),
  )
  async_copy.start()
  async_copy.wait()

# 2. JAX Wrapper with ShardMap and DMA Semaphores
def build_jax_fn(mesh):
  def local_pallas_call(x):
    return pl.pallas_call(
        pallas_send_to_next_kernel,
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
        grid_spec=pltpu.PrefetchScalarGridSpec(
            num_scalar_prefetch=0,
            in_specs=[pl.BlockSpec(x.shape)],
            out_specs=pl.BlockSpec(x.shape),
            scratch_shapes=(
                pltpu.SemaphoreType.DMA,
                pltpu.SemaphoreType.DMA,
            ),  # send & recv semaphores
        ),
    )(x)
  return jax.shard_map(
      local_pallas_call,
      mesh=mesh,
      in_specs=P("x"),
      out_specs=P("x"),
      check_vma=False,
  )

# 3. Create the TorchTPU callable wrapper
def get_send_to_next_op(world_size: int):
  mesh = torch.tpu.pallas.get_pallas_mesh(axis_names=("x",), mesh_shape=(world_size,))
  jax_fn = build_jax_fn(mesh)
  return torch.tpu.pallas.jax_op(
      "pallas::send_to_next",
      build_jax_fn(mesh),
      mesh=mesh,
      input_partition_specs=[P("x")],
      output_partition_specs=[P("x")],
  )

下表总结了如何编写通信型 Pallas 内核。

解决设备映射不一致问题。此前在 TorchTPU 中，使用底层芯片间操作（如 lax.axis_index 和 async_remote_dma / make_async_remote_copy）的 Pallas 内核基于原始硬件设备 ID 而非 PyTorch rank 来路由通信，与 JAX 保持一致。当硬件设备 ID 与 rank 顺序不一致时，环通信和跨设备 DMA 会被路由到非预期的物理设备，而在隔离子组上的执行会因分区数量不匹配而失败。

TorchTPU 通过强制逻辑寻址解决了这些不一致问题：

• Rank 对齐映射：TorchTPU 现在将 Pallas 内核直接映射到 PyTorch rank，而非原始物理设备 ID，通过 get_pallas_mesh API 确保路由一致，该 API 将 pallas 中的分区和轴直接映射到 PyTorch Process Group 中的 rank 顺序。
• 逻辑设备寻址：TorchTPU 现在强制使用 DeviceIdType.LOGICAL 和规范的设备分配（da(0, p) = p），保证每个逻辑 rank 始终映射到其正确的物理目标。

统一 ProcessGroup 集成。TorchTPU 现在提供获取 rank 对齐 JAX Mesh 的机制，使内核能够继承 PyTorch ProcessGroup 拓扑。这可以在全量世界和隔离子组配置下高效执行，无需手动修补状态或中断计算图。

开发模型

TorchTPU 以开放方式开发，是 Google 与社区协作的成果，优先考虑透明度和参与度。该开发模式鼓励更广泛的 PyTorch 生态和行业合作伙伴贡献力量。

我们预计 TorchTPU 将服务于多样化的用户群体，包括研究机构和行业领导者。我们还计划将其与 Hugging Face 等平台集成，为广泛的开放源代码模型提供开箱即用的 TPU 支持。

为支持开发者，TorchTPU 提供全面的文档，包括 （参见 http://torchtpu.readthedocs.io/）：

• 教程：引导式学习路径，包括虚拟机配置和模型移植。
• 配方：面向问题的解决方案，适用于性能优化和内核实现等任务。
• 参考：技术规范和 API 映射。
• 说明：深入探讨架构和硬件基础。

文档使用可直接运行的 notebook 构建。

该项目利用 Bazel 进行稳健的构建和测试管理，测试框架仍在演进，并计划采用 Pytest。CI/CD 基础设施正在集成到官方 PyTorch CI HUD 中，以保持高可靠性并促进社区贡献。

TorchTPU 设计为可扩展，能够与关键 PyTorch 生态组件无缝集成。它支持通过 Pallas 和 Helion 进行自定义内核开发，并通过 vLLM、SGLang 和 TorchTitan 等项目为分布式训练和推理工作负载提供一等支持。这种模块化方法允许用户利用现有的 PyTorch 惯用法，同时针对 TPU 特定的硬件架构进行优化。

案例研究

案例研究 1：Hugging Face 模型覆盖

TorchTPU 提供与 PyTorch 生态的开箱即用兼容性，使开发者无需修改代码或重新编写即可在 Google Cloud TPU 上运行开放源代码模型。为验证这一承诺，我们构建了一个自动化测试框架，使用 Hugging Face 模型覆盖 NLP、视觉、音频、多模态和生成式扩散领域。该持续集成流水线以 CPU 和 GPU 基线为基准，对 TorchTPU 的覆盖范围、稳定性和数值等价性进行基准测试。我们在从 v5e 开始的 TPU 硬件代际上，以 eager 和编译模式持续评估约 5,000 个 Hugging Face 模型，涵盖数百种架构。

我们采用基于语料分层和均衡架构选择的稳健测试方法。我们的方法依赖于两级选择流程，旨在平衡实际采用与广泛的架构多样性，避免对细微变体进行冗余测试，同时确保真实生产权重得到充分评估。在第一级中，我们根据社区采用度、下载速度和趋势评分优先选择排名靠前的模型——例如 Qwen、Gemma、Mistral、DeepSeek、GLM、Kimi、Llama、Whisper、Stable Diffusion、FLUX 和 Vision Transformer——以确保关键社区热门模型完美运行。第二级是从第一级中选出的聚焦子集，确保我们从超过 500 个不同的模型家族中至少选择 1 个代表性模型。我们使用第二级在全面的 PyTorch 算子和编译器模式集合上验证相对于 GPU 基线的数值准确性。

验证过程从第 1 阶段开始，该阶段侧重于冒烟测试和降低验证，以保证运行时稳定性。在此初始阶段，流水线使用合成权重实例化模型，验证 eager 模式的前向执行，并确认图追踪、FX 到 StableHLO 的图生成以及 XLA 编译均能成功完成且无错误（例如，无崩溃或异常）。一旦确认运行时稳定性，第 2 阶段便使用领域特定的评估指标，对照既定基线评估数值正确性。对于因果 LLM，数学保真度通过困惑度（Perplexity = e^L = exp((1/N) Σ_i=1^N L_i)）来衡量，若 TPU 与基线困惑度之间的相对差异保持在 0.10 以内，则测试通过。对于视觉、音频和编码器架构中的高维特征空间，我们采用双阈值相对 MSE 指标（Relative MSE = MSE(Y_TPU, Y_Baseline) / (Var(Y_Baseline) + 10⁻⁷)），要求相对 MSE ≤ 0.10 或绝对 MSE < 0.05，以防止基线输出接近常数。

下表总结了截至本文撰写时的模型覆盖结果：

对于有界动态性支持（针对 SymInts），我们在 839 个模型（来自 5K 语料库的因果 LLM 列表）上评估了正确性，并确认了 834 个（99%）模型的数值正确。

案例研究 2：开箱即用的性能

TorchTPU 致力于成为 PyTorch 的真正原生 TPU 后端。在我们看来，这一期望包括在运行纯 PyTorch、设备无关代码时具有良好的性能。我们通过两个示例来展示这一点。

第一个示例包含在 TorchTPU 仓库中，是一个自包含的、纯 PyTorch 实现的 Muse Glimmer 30B 模型，我们仅使用标准 PyTorch 模块和算子编写，例如 nn.Linear、nn.RMSNorm、F.scaled_dot_product_attention 和 torch.distributed 集合通信。我们有意没有使用任何模型特定的自定义内核（但仍使用了 PyTorch 内部内核）。在此示例中，我们展示了：

• Inference (prefill) performance on a single TPUv7 chip at
  • 序列长度为 5,120 时，~25.7k TPS，59% MFU，以及
  • 序列长度为 10,240 时，~22.4k TPS，51% MFU。

  请注意，此示例使用 FSDP（通过原生 PyTorch 集合通信实现），world size 为 2，因为单个 TPUv7 芯片向 PyTorch 暴露两个设备（对应两个包含的 TensorCore）。

• Training on 4 TPUv7 chips (i.e. a full single host), using FSDP, at
  • 每个 TPU 在序列长度为 5,120 时达到 ~5.6k TPS，38% MFU，以及
  • 每个 TPU 在序列长度为 10,240 时达到 ~4.6k TPS，32% MFU。

  在第二个示例中，我们启用了激活检查点，它仅在 HBM 中存储前向传播激活的一个子集，并在反向传播计算需要时重新计算/重物化未存储的中间激活。

所考虑示例的一个决定性特征是，同一份代码在除指定加速器类型（例如编写 device="tpu" 而非 device="cuda"）之外没有任何实质性改动的情况下，即可在 TPU 和 GPU 上运行，并且在标准 PyTorch 优化与分片技术（如 torch.compile 和 FSDP）下在 TPU 上表现非常出色。

为进一步验证 TorchTPU 在纯 PyTorch 代码上的性能，我们还从 TorchTitan 项目中选取了一个外部产出的 Muse Glimmer 30B 实现，该实现最初是为 GPU/CUDA 编写的。在对其做最小限度适配到 TorchTPU 后（我们仅将 TorchTPU 尚不支持的 FlexAttention 替换为原生 F.scaled_dot_product_attention，并应用了 torch.compile 和 SimpleFSDP），TorchTPU 移植版实现了：

• Training on 4 TPUv7 chips:
  • 在序列长度为 5,120 时，每 TPU 约 5.3k TPS，MFU 为 36%，以及
  • 在序列长度为 10,240 时，每 TPU 约 4.2k TPS，MFU 为 29%。

总之，所有这些示例表明，纯 PyTorch 模型实现，甚至最初为 GPU 编写的实现，只需简单的 PyTorch 改动并配合标准 PyTorch 优化，即可在 TPU 上表现足够出色，而无需借助自定义内核等高级的、设备特定的爬山技术。我们认为这是 TorchTPU 采用的一项重要特性。当然，对于追求峰值性能的生产工作负载，我们绝对建议采用基于 Pallas 内核的最优方案，以及 vLLM-TorchTPU、SGLang-TorchTPU 或 TorchTitan-TorchTPU 等生产级技术栈作为示例。

案例研究 3：将训练扩展到 Pod 规模

我们使用 TorchTitan 验证了 TorchTPU 可扩展至 1 个 TPU v7 POD。对于此测试，我们选择了 Qwen3 MoE 235B。作为初步冒烟测试，我们确认在 64 个 TPU v7 配置上，TorchTitan 配合 TorchTPU 的模型性能与 MaxText 上的 JAX 相当（PyTorch/TorchTitan 为 1828 TPS，JAX/MaxText 为 1937 TPS）。

下文讨论的数据基于以下设置计算得出：

• 模型：qwen3_moe_235b_scan
• 本地批大小（即每个 TensorCore）：16，序列长度：4096
• 并行策略：单切片运行使用 FSDP（通过 SimpleFSDP），多切片运行使用 HSDP
• 精度：BF16 配合 FP32 混合精度归约（梯度归约：f32，RoPE：f32，优化器状态：f32）
• 内核（Pallas）：Tokamax Splash Attention（带来大幅性能提升）、融合交叉熵损失内核，以及 MoE 层的分组矩阵乘法
• 扫描优化：94 个解码器层作为单个扫描循环在堆叠权重上运行
• 解码器层输入的激活卸载（作为扫描优化的一部分）
• 集合通信流水线化，以将集合通信与相邻层的计算重叠（扫描的一部分）
• 完整激活检查点
• 使用 TorchTitan 的 Graph Trainer 进行联合前向和反向传播
• 专家不均衡：使用 TorchTitan 的辅助损失防止 MoE 不均衡
• 训练数据集：Hugging Face C4 en 数据集，配合自定义数据集分片

在所有拓扑下，我们的数据证实了弱扩展：整体计算量呈线性扩展，如下所示，总每秒 token 数线性增长，而 MFU 基本保持平稳。

我们的数据还显示，在所考虑的各种并行级别下，损失曲线几乎完全相同。

在多切片 HSDP 配置中，前向传播不涉及任何跨切片通信，每个切片处理自己那份批次数据。但在反向传播中，一个切片先计算第 t 层的梯度，然后通过 all-reduce（在 DCN 上，使用 Megascale-XLA）对各切片的所有这些梯度求和。计算随后继续处理第 t-1 层的梯度。得益于 XLA，TorchTPU 能够将第 t+1 层的 all-reduce 与第 t 层的计算完全重叠，如下图所示。实现这种重叠是每个 TensorCore 的 TPS 从一片到八片几乎保持平稳的主要原因，如上表所示。

案例研究 4：在 vLLM 上使用 Qwen3.5 进行 LLM 推理（SemiAnalysis InferenceX）

大语言模型（LLM）推理——尤其是针对大规模混合专家（MoE）架构——是 AI 云基础设施中要求最高、成本最关键的工作负载之一。在 SemiAnalysis 发布的一项独立基准评估中（TPU Inference Externalization Full Steam Ahead — InferenceX），Google 的 TPU v7 搭配 vLLM-TorchTPU 软件栈，交付了具有行业竞争力的每美元性能。该基准提供了对 TorchTPU 软件在大规模混合专家（MoE）工作负载上就绪程度的第三方推理审计，表明面向 TPU 的新 vLLM 软件栈在社区公开可用的 TPU 系统上提供了强大的成本效率和高吞吐量。

在此基础上，Google 最近提交了 Qwen 3.5 在 Ironwood TPU 上的首批 AgentX 性能结果。AgentX 捕捉了智能体工作负载的四个要素——多轮、长上下文、高前缀复用和子智能体突发。TPU 已能够证明其即将推出的 vLLM-TorchTPU 软件栈能够成功处理智能体工作负载。仅用 4 块芯片，TPU v7 就实现了帕累托前沿的全面领先。

• 在 InferenceX 8K/1K 上，每百万 token 总体成本最低，为 2.01 美分（低于我们 9 月首次亮相时的 3.6 美分）；
• 在 AgentX 上，每百万 token 总体成本最低，为 0.83 美分，并具备 SOTA 交互性（用户响应能力）。帕累托前沿全面领先。

这些基准结果反映了在 TPU 系统上实现的强大、真实的在线服务效率，突显了在 Google 定制设计的 ML 芯片上部署大规模模型的巨大经济可行性。

要实现这些有竞争力的性能和 TCO 结果，需要通过自定义 Pallas 内核和有针对性的运行时调优来解决关键的软硬件对齐瓶颈，下表列出了几个示例：

这些结果突显了现代服务基础设施中的 TPU 优化，表明 TCO 和软硬件对齐决定了真实世界的服务效率。最终，带有 TorchTPU 后端的新 vLLM-TPU 使工程团队能够在公开可用的 TPU 基础设施上使用标准 PyTorch 惯用法部署现成的开放权重模型。

案例研究 5：使用 TorchTitan 进行 LLM 微调

TorchTPU 与 TorchTitan 原生集成——TorchTitan 是 PyTorch 面向大规模分布式训练的参考平台——在稠密大语言模型（MuseGlimmer 30B、Llama 3 8B–405B）、稀疏混合专家模型（Qwen 3 235B MoE）以及生成式扩散 Transformer（FLUX.1-dev 12B）上提供统一的高性能训练环境。在 TPU v6e（Trillium）和 TPU v7x（Ironwood）集群上以 BF16 精度进行评估，多切片执行通过 torchrun 在 Google Kubernetes Engine（GKE）上无缝编排。TorchTitan 作为开源项目可在此 GitHub 仓库中找到。

分布式策略与分片架构。高效的多加速器训练要求模型并行策略与 TPU 的物理互连拓扑之间实现无缝对齐。为此，我们采用 2D+ DeviceMesh 映射，将通信密集的维度——例如张量并行和专家并行——与连续的高带宽芯片间互连（ICI）对齐，同时将复制的数据并行流量路由到较慢的跨切片数据中心网络（DCN）。这一策略通过我们与 SimpleFSDP 的集成来实现，后者提供符合 SPMD 规范的操作。为进一步降低延迟，我们将 XLA 运行时配置为把集合通信操作（如 All-Gather 和 Reduce-Scatter）直接卸载到片上 SparseCore 引擎，从而有效地将通信与计算密集的 TensorCore MXU 数学运算重叠。最后，我们通过使用 HLO diffing 工具和 @spmd_safe 注解来确保确定性执行并防止同步死锁，这些工具和注解消除了 torch.compile 下不必要的图中断。

性能与多节点扩展。通过系统性的逐层优化，TorchTPU 上的 TorchTitan 为 FLUX 模型实现了卓越的性能和多节点扩展效率。从开箱即用的 eager 基线步进时间 11.49s（0.26 samples/s）开始，应用 DeferAndFuse eager 模式将步进延迟改善至 9.95s（1.15× 加速）。引入 SimpleFSDP 分片进一步将执行加速至 6.52s（1.76× 加速），而集成自定义 Pallas SplashAttention 内核则将步进时间降至 5.55s（2.07× 加速）。启用完整的 torch.compile(backend=”tpu”) 带来显著的性能飞跃，达到 3.02s（3.80× 加速），而应用目标模型窥孔重写则产生最终优化后的步进时间 2.78s（1.08 samples/s）。在多主机集群中，TorchTitan 展现出接近理想的弱扩展和强扩展：在 FLUX.1-dev 上，使用 SimpleFSDP 从单个 v6e-8 节点（8 芯片）扩展到 v6e-256 切片（64 个节点共 256 芯片），每样本步进时间仅从 0.92s 略微增加到 1.00s（总步进时间 3.01s），在光互连 ICI 网络上展现出接近线性的效率。在更大规模下，TorchTitan 在 TPU v7x（Ironwood）上以 BF16 精度、8K 上下文窗口在 1,024 个芯片上运行 Llama 3 405B，实现 500 tokens/sec/chip 和 50+% MFU。同样，在 8,192 个 TPU v7x 芯片上评估 Qwen3 235B MoE 也展现出高效扩展，在 BF16 精度、4K 序列长度下达到 28% MFU。

关键架构优化。为在多样化的模型架构中维持峰值计算密度并消除硬件瓶颈，我们采用了多项系统和内核优化：

• SimpleFSDP：一种对编译器友好的 FSDP 版本。FSDP2 通过模块钩子注入其 all-gather 和 reduce-scatter，这会把 torch.compile 打断成逐算子分发；SimpleFSDP 则将它们表达为被追踪图中的普通操作。这样集合通信对 XLA 保持可见，TPU 编译器便能围绕它们进行融合、与 MXU 计算重叠，并可将其卸载到 SparseCore。
• Scan 高阶算子（torch._higher_order_ops.scan）：将重复的 transformer 块作为对堆叠权重的一次循环来运行。随着深度增长，这使编译时间和聚合参数内存保持恒定。逐层的集合通信位于循环内部，因此 XLA 可以将它们与计算重叠：延迟隐藏调度器将每一层的 all-gather 隐藏在该层的计算之后，集合通信流水线器可选择性地预取下一层的 all-gather，而在反向传播中，它可以让某一层的 reduce-scatter 在下一层计算时完成。
• Pallas SplashAttention 内核：将缺乏向量瓦片对齐的默认注意力实现替换为自定义 Pallas 内核，支持完整的 Grouped-Query Attention（GQA）、偏置和 dropout，在扩散注意力块上带来 +17.5% 的加速。
• Fused Pallas Cross-Entropy Loss: Avoids out-of-memory (OOM) faults when materializing [B×S×V] logits for large vocabularies by porting Tokamax’s online reduction kernel to compute loss in-place.
• Graph trainer: Captures the full training step, forward, loss, backward and optimizer update, as a single graph that compiles into one XLA program. This gives the TPU compiler a global view for fusion, collective scheduling and memory planning, and lets TPU-specific passes decide per activation whether to save, recompute or offload to host, and donate optimizer buffers so state is updated in place rather than double-buffered.
• Activation Offloading: The activations that the backward pass needs can be moved to host memory after the forward and copied back just before backward uses them, which frees HBM for larger models and longer sequences. The graph trainer inserts the offload and reload ops into the traced graph. Under scan, they sit inside the loop body, so saved activations move one layer at a time. XLA turns these copies into async transfers, and in backward copies the next layer’s activations back while the current layer computes. This is essential for the largest models, such as Llama3-405B and Qwen3-235B, to fit at their target sequence lengths.

Ongoing development for TorchTitan on TorchTPU targets four strategic priorities: expanding list of models optimized for TPU; integrating FP8 and FP4 low-bit quantization via torchao; and supporting end-to-end Reinforcement Learning (RL) by coupling training loops with vLLM/SGLang sampling backends in VeRL.

Case Study 6: Hardware-Accelerated Recommendation Models with TorchRec and SparseCore

Recommendation and ranking architectures—such as Meta’s Deep Learning Recommendation Model (DLRM)—face critical memory and compute bottlenecks due to high-cardinality embedding tables. On standard GPU clusters, managing these tables may require complex software-based sharding via DistributedModelParallel (DMP) and separate optimizer updates, which often introduce latency. To bridge this gap, Google collaborated with the Meta PyTorch team to upstream a native TPU backend directly into the TorchRec repository, allowing recommendation models to bypass host-side compute bottlenecks and tap directly into TPU SparseCore (SC) hardware.

• Architectural Integration. The resulting implementation creates a seamless data pipeline where high-cardinality lookups are accelerated by dedicated hardware. A multi-threaded C++ PyBind backend runs on the host CPU to parallelize partition mapping and index padding, outputting structured batches that map directly to the TPU SparseCore units. Within the TPU, we implemented In-Place Optimizer Fusion: rather than materializing gradient tensors (.grad) in High Bandwidth Memory (HBM)—a process that traditionally complicates memory management—we map backpropagation directly to SparseDenseMatmulGrad. This ensures parameter updates (SGD, AdaGrad, or Adam) occur in-place, eliminating gradient memory overhead entirely. Finally, low-level ATen operators, such as SparseDenseMatmulOp (for forward lookups) and RaggedAllToAll (for optical ICI cross-chip collectives), provide hardware-accelerated kernels that execute directly on the silicon.
• Performance & Results. When evaluated on an industry-standard DLRM v2 model across a distributed 32-chip TPU v7x (Ghostfish) cluster, this integration achieved 100% performance parity with Google’s highly optimized JAX baseline, reaching an identical 18.3ms step latency. Notably, this performance matching was established using baseline SparseCore kernels before applying advanced optimizations. Convergence and mathematical fidelity were equally robust, yielding identical loss trajectories compared to the JAX implementation.

Active development focuses on further saturating bandwidth through TensorCore/SparseCore pipelining—interleaving MLP execution with embedding lookups—and advanced table-stacking techniques to unify memory buffers and reduce kernel overhead. Additionally, ongoing work aims to expose lower-level SparseCore embedding lookup operations (e.g. SparseMapRow, SparseGather etc) via ATen kernels, enabling users to implement custom lookup algorithms and advanced sharding strategies. This integration solidifies the TPU as a high-efficiency engine for the next generation of recommendation architectures.

Future Work

With TorchTPU now publicly available on GitHub, this initial open-source release represents only the beginning of our development journey. As we engage directly with the broader PyTorch community, our roadmap for the next 6 to 12 months focuses on production hardening, ecosystem expansion, and continually integrating community feedback to bridge the gap between TPU hardware capabilities and PyTorch-native developer workflows. We are prioritizing four strategic pillars to drive this evolution:

• Foundational Hardening & User Velocity: We are refining the “True Eager” experience by building a robust library of precompiled, dynamic kernels for ATen operators. We are also making XLA Bounded Dynamism fully production-ready, allowing the torch.compile stack to handle SymInts without extra recompilation, and streamlining our GKE packaging to lower the barrier for deployment across PyTorch versions. Furthermore, we are continuously expanding operator coverage and advancing our quantization support—specifically targeting FP8 and FP4 formats for next-generation TPUv8 architectures, and implementing nested compiler regions to reduce compilation time.
• Pod-Scale Distributed Training & Advanced Training: Building on our successful POD-scale training validation of Qwen3, we are expanding to additional model architecture and sharding techniques at and beyond that scale. We are enabling end-to-end Reinforcement Learning (RL) training loops by tightly coupling scalable training with high-throughput vLLM/SGLang sampling backends orchestrated via Ray Core. To enhance distributed stability, we are also validating a consensus-based compilation handshake to safely optimize collective operations with compute without fragile @spmd_safe annotations.
• Emerging Architectures, Inference & Serving: We are committing to proactively support cutting-edge model architectures as they emerge in the open-source ecosystem—including reasoning models, multimodal and generative diffusion models (e.g., Wan 2.2, Flux Klein), and novel Mixture-of-Experts (MoE) variants. Our efforts focus on optimizing the serving stack in vLLM and SGLang for low-latency production, incorporating native FP8/FP4 quantization across future TPU generations, and ensuring day-one TPU support for new architectural paradigms.
• High-Efficiency Communication & Diagnostics: We are optimizing collective operations within and across TPU slices and adding zero-copy DMA to enable hardware-assisted memory offloading. To support debugging at scale, we are transitioning to a WebGPU-based Trace Viewer 2.0 and deploying “flight recorder” continuous profiling (XProfz) to identify intermittent bottlenecks in production environments.

TorchTPU is built on the belief that a successful ecosystem thrives through transparent, collaborative engineering. As we open-source our four foundational repositories (TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU, and TorchTitan) at the October 2026 PyTorch Conference, we actively seek feedback on model architecture prioritization and workload sequencing. We invite developers, researchers, and partners to collaborate with us via GitHub or our active TPU Builders program.

†The TorchTPU Team

Claudio Basile, Michael Voznesensky, Gunhyun Park, Ben Wilson, Parker Holloway, Zhanyong Wan, Vlad Belous, Manish Khanna, Paul Ganssle, Kevin Gleason, Andrew Dame, Brendan Folie, Bhavya Bahl, Haifeng Jin, Pedro Goncalves Mokarzel, Arian Arfain, Luke Boyer, Chi Zhang, Yiling Yuan, Adam Mainz, Aditya Patel, Huangze Tang, Jay Gan, Lizzie Liang, Will Froom, Tushar Agarwal, Karlo Basioli, Theotime Combes, Emilio Cota, Minos Papadopoulos, Tudor Pipernea, Yarong Mu, Qi Zhou, Yaoshiang Ho, Aleksey Vlasenko, Alexander Brauckmann, Ziyin Huang, Jack Oh, Jerry Hsu, Tanvi Bajpai, Jialei Chen, Wenjun Gu, Chintan Shah, Chris Chan, Sayce Falk, Chris Jones, Rob Torres, Kat Ko, Lee Howes, Freya Azad, Andi Gavrilescu, Joe Pamer, Bill Jia

Plus many, many others…

The post Introducing TorchTPU appeared first on PyTorch.