TorchTPU
也称为 TorchTPU、TorchTPU backend、PyTorch 原生 TPU、torch device tpu
先用大白话
TorchTPU 让 TPU 表现为一个普通的 PyTorch 设备，vLLM 和 SGLang 可以直接在 Google 芯片上运行现有 PyTorch 代码。
技术定义
TorchTPU 是 Google 基于 PrivateUse1 扩展点构建的 PyTorch 原生 TPU 后端，在 device tpu 上暴露真正的 torch.Tensor，并通过 StableHLO 和 XLA 下沉编译图。
工程细节
PyTorch dispatcher 把 ATen 操作直接路由到 TPU 后端，而不是翻译成 JAX。开发者可以 eager 执行用于 bring-up 和调试，也可以调用 torch.compile：TorchDynamo 和 AOTAutograd 生成 FX 图，TorchTPU 下沉为 StableHLO，XLA 生成 TPU 可执行文件，不使用 Inductor 和 Triton。原生范围覆盖张量、dispatch、eager 执行、编译入口和分布式 API，包括 DDP、FSDP2、DTensor 以及 SPMD 和 MPMD。kernel 层仍然是 TPU 专用的：TorchTPU 可以调用 Pallas 和 JAX 自定义 kernel，因此为 TorchAX 栈编写的 kernel 可以迁移而不必重写。
为什么重要
vLLM 和 SGLang 可以复用上游的模型代码、调度器、continuous batching 和功能逻辑，不必跨 PyTorch 到 JAX 的边界重建。这降低了新模型 bring-up 的成本，让 TPU 朝着与 NVIDIA 同步的 day-0 支持靠近。TorchTPU 将取代即将弃用的 TorchAX。
如何在 InferenceX 中解读
InferenceX Official Preview 中的 TPUv7 结果来自原生 TorchTPU vLLM 栈以 FP8 服务 Qwen3.5 397B。截至 2026 年 9 月初 TorchTPU 仍处于私有测试阶段，预计 10 月中旬 PyTorch Conference 前后开源，届时 SemiAnalysis 会把 TPU 基准测试从 fork 迁移到公开的 InferenceX 仓库。