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

简而言之

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 芯片之间快速通信，这些 Slice 是 TPU Pod 的逻辑子组，由以 3D 环面（TPU 7、TPUv 8t）或 Boardfly（TPU 8i）配置连接的预配置数量 TPU 芯片组成。
• 双向 数据中心网络（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 工作负载的一个限制，我们一直在添加对有界动态性的支持，即当形状维度的最大边界在编译时已知，但实际维度仅在运行时已知时。
• 集合操作被降低为 XLA 的 StableHLO 操作，而不是使用诸如 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 张量操作，VPU 原生执行完整的 IEEE-754 FP32 算术。

TorchTPU 将 PyTorch 精度模式映射到 underlying 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执行，而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位整数，并使用该元组作为缓存键：

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

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

视图语义

PyTorch的操作语义与StableHLO语义形成对比的一个特定领域是可变性和共享内存视图的处理：

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

在存在写入的情况下，在 TPU 上实现正确的 view 语义并不简单；然而，TorchTPU 通过一些机制实现了这一点，涉及 view 分解、view 反转 和延迟操作。这一语义桥接使 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 提供了关键算子的高度优化实现，包括专家混合构建模块（如分组矩阵乘法和令牌收集/分散）、注意力变体（如分页注意力、门控 deltanet、Kimi delta 注意力和多头潜在注意力）以及 TopK 计算。算子组合将随时间演进，以支持在最新硬件加速器上运行的最新 OSS 模型（例如，发布时支持 TPU v8i）。这些内核由稳健的、集成的、基础基础设施支持，例如自动调优（例如，瓦片大小搜索）、低开销的基于 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 和 compile 模式下的集合操作，并无缝支持 PyTorch 分布式框架，如 DTensor、FSDPv2、SimpleFSDP、FairScale 等。与 GPU 上的 PyTorch 类似，我们遵循（并且实际上强制执行）每设备一进程的方法。从实现的角度来看，TPU 栈不提供用于 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 的数量（世界大小），以及 (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 需要两个标准 PyTorch torchrun 不知道如何设置的 TPU 专用变量：

• TORCH_TPU_TOPOLOGY：TPU 切片的物理网格几何结构（例如，对于 2x2x1 切片配置为“2,2,1”）。
• TORCH_TPU_SLICEBUILDER_ADDRESSES：用于构建 TPU 芯片间互连（ICI）网络网格的开放端口/地址列表（每台设备一个）（例如“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

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

最优设备网格映射

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

• ICI：直接的光/铜互连，以 Torus 网格连接芯片。提供数十到数百 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 Kernel

TorchTPU 支持自定义 Pallas kernel，可通过 ICI 上的硬件 RDMA 直接进行芯片间通信，无需主机 CPU 参与。这使你能够编写高效的点对点通信 kernel（例如 Ring Attention、自定义 All-to-All / Shuffles、流水线并行移位或 Halo Exchange），并将其作为原生 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 kernel。

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

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

• Rank 对齐映射：TorchTPU 现在将 Pallas kernel 直接映射到 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 的机制，使 kernel 能够继承 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 Transformers——以保证关键社区热门模型运行无误。第二层是从第一层中选取的重点子集，确保我们从 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⁻⁷)），要求 Relative 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，36% MFU，以及
  • 在序列长度为 10,240 时，每 TPU 约 4.2k TPS，29% MFU。

总之，所有这些示例都表明，纯 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（SemiAnalysis InferenceX）上使用 Qwen3.5 进行 LLM 推理

大型语言模型（LLM）推理——尤其是针对大规模混合专家（MoE）架构——是 AI 云基础设施中需求最高、成本最关键的工作负载之一。在 SemiAnalysis 发布的一项独立基准评测（TPU Inference Externalization Full Steam Ahead — InferenceX）中，Google 的 TPU v7 搭配 vLLM-TorchTPU 软件栈，交付了具有行业竞争力的每美元性能。该基准测试提供了针对大规模混合专家（MoE）工作负载上 TorchTPU 软件就绪度的第三方推理审计，表明面向 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 使工程团队能够使用标准 PyTorch 惯用法，在公开可用的 TPU 基础设施上部署现成的开放权重模型。

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

TorchTPU 与 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 差异对比工具和 @spmd_safe 注解来确保确定性执行并防止同步死锁，这些注解消除了 torch.compile 下不必要的图中断。

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

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

• SimpleFSDP：一种对编译器友好的 FSDP 版本。FSDP2 通过模块钩子注入 all-gather 和 reduce-scatter，这会把 torch.compile 拆成逐算子调度；而 SimpleFSDP 将它们表示为 traced graph 中的普通操作。这让集合通信对 XLA 可见，因此 TPU 编译器可以围绕它们进行融合、使其与 MXU 计算重叠，并可将其卸载到 SparseCore。
• Scan 高阶算子（torch._higher_order_ops.scan）：将重复的 transformer 块作为对堆叠权重的一次循环来运行。随着深度增长，这使编译时间和聚合参数内存保持恒定。每层的集合通信位于循环内部，因此 XLA 可以将其与计算重叠：延迟隐藏调度器将每层的 all-gather 隐藏在该层的计算之后，集合通信流水线器可以选择性预取下一层的 all-gather，并且在反向传播中，它可以让某一层的 reduce-scatter 在下一层计算时完成。
• Pallas SplashAttention 内核：用自定义 Pallas 内核替换缺乏 vector-tile 对齐的默认注意力实现，支持完整的 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.