Torch-TPU：让SGLang与vLLM原生跑在TPU上

PyTorch 原生 TPU 推理：一条不走 JAX 的新路径

在 PyTorch 大会（San Jose）上，来自 Google 的 Qi 以及来自 Meta 的 Colin 和 Angela 共同介绍了一项针对大模型推理的新工作：基于全新的 PyTorch 原生 TPU 后端 Torch-TPU，让主流的 SGLang 与 vLLM 推理引擎能够原生运行在 Google TPU 上。

这项工作的核心思路与过往的 TPU 适配方案截然不同。以往要在 TPU 上跑大模型，通常意味着两种选择：要么构建一套全新的服务引擎，要么把模型移植到 JAX 生态。这两条路都成本高昂，也让 TPU 与主流 GPU 工具链之间产生了割裂。而 Torch-TPU 的目标，是让现有的 SGLang 和 vLLM 基础设施原生跑在 TPU 之上，而不是另起炉灶。

为什么选择复用现有引擎，而非重写

SGLang 和 vLLM 是当前 LLM 推理领域事实上的两大主流开源引擎，围绕它们已经形成了成熟的调度器、批处理系统、服务 API 和编译工作流。大量工程团队已经熟悉这套工具链，并在 GPU 上形成了稳定的生产实践。

团队强调，他们保留了这些让 GPU 用户倍感熟悉的组件——调度器（schedulers）、批处理系统（batching systems）、服务端 API（server APIs）以及编译工作流（compilation workflows）。换句话说，从使用者的视角看，迁移到 TPU 的心智负担被降到最低：你依然在用熟悉的 SGLang / vLLM，只是底层换成了 TPU。

这种“复用而非重写”的策略，背后是对工程现实的务实考量。重写一套服务引擎不仅耗时，还意味着要重新追赶上游社区在调度、连续批处理（continuous batching）、前缀缓存等方面的持续演进。与其追赶，不如直接站在上游的肩膀上。

连续批处理（continuous batching）是现代推理引擎的核心调度技术之一，由 Orca 论文（2022）提出并被 vLLM 等引擎广泛采用。其核心思想是：不等一批请求全部完成后再处理下一批，而是在每个迭代步骤中动态地将已完成的请求替换为新请求，从而大幅提升 GPU/TPU 利用率。前缀缓存（prefix caching）则是另一项关键优化——对于共享相同系统提示或上下文前缀的多个请求，引擎可以复用已计算的 KV Cache，避免重复计算，显著降低首 token 延迟（TTFT）。这两项技术都需要深度嵌入调度器逻辑，重新实现成本极高。Torch-TPU 选择直接复用 SGLang/vLLM 的调度层，意味着这些精细调优的机制可以直接在 TPU 上生效。

上游化带来的关键红利：模型与特性零成本流转

Torch-TPU 方案最值得关注的一点，是服务引擎保持在上游（upstream）。这意味着 SGLang 与 vLLM 社区中新增的模型支持和功能特性，可以直接在 TPU 上运行，而无需为 TPU 单独编写一套实现。

这一点在实践中价值巨大。大模型生态迭代极快，新模型架构、新量化方案、新采样策略几乎每周都在涌现。如果 TPU 需要为每个新特性做专门适配，就永远处于“落后半拍”的状态。而当引擎保持上游一致时，TPU 用户理论上可以与 GPU 用户几乎同步地享受到社区最新成果。

从硬件生态竞争的角度看，这也是降低 TPU 使用门槛、吸引 GPU 用户迁移的关键一步。长期以来，CUDA 生态的护城河很大程度上来自软件工具链的丰富度；而让主流推理引擎原生支持 TPU，正是在软件层面缩小这种差距的尝试。

Google 与 Meta 的跨厂商协作

值得玩味的是这项工作的参与方构成：Google 提供 TPU 硬件与后端，Meta 则是 PyTorch 以及相关推理引擎生态的重要推动者。两家公司在这一项目上的协作，反映出 PyTorch 作为跨厂商中立框架的生态价值——它既能承载 NVIDIA GPU，也能通过原生后端扩展到 TPU 等异构硬件。

不过需要说明的是，本次公开的信息仅为大会演讲的预告，具体的性能数据、支持的模型范围、以及 Torch-TPU 后端的技术实现细节，团队并未在预告中披露，而是留待 PyTorch 大会现场详细展开。

PyTorch 的跨硬件扩展能力来源于其后端抽象机制。PyTorch 通过 torch.distributed 与设备抽象层，允许第三方为新硬件注册自定义算子和编译后端。早期的 XLA（Accelerated Linear Algebra）是 Google 将 TPU 接入 PyTorch 的主要路径，但 XLA 的编译模型与 PyTorch 的动态图风格存在较大阻抗，调试困难、编译开销大。Torch-TPU 所指的"原生 PyTorch TPU 后端"，意味着更深层的集成——直接通过 PyTorch 的 Dispatcher 和 TorchInductor/torch.compile 体系驱动 TPU，而非绕道 JAX/XLA 的编程模型。这使得 PyTorch 用户无需学习新的编程范式，就能将模型部署到 TPU 上。

小结

Torch-TPU 代表了一种清晰的工程哲学：不与成熟的上游生态对抗，而是把硬件适配做成底层后端，让上层应用无感迁移。对于希望在 TPU 上部署大模型推理、又不想放弃 SGLang / vLLM 工具链的团队来说，这是一个值得密切关注的方向。完整的技术细节，有待大会现场揭晓。