vLLM 发布 vLLM-TorchTPU，打造 TPU 原生 PyTorch 推理服务

vLLM 全新推出 vLLM-TorchTPU 组件，直接实现了 TPU 硬件上原生 PyTorch 推理服务，开发者不用改动 PyTorch 模型代码，就能在谷歌 TPU 集群里直接部署 vLLM 推理服务，完整保留 vLLM 的分页注意力、连续批处理、KV 缓存等核心调度能力，同时依托 TorchTPU+XLA 编译链路充分释放 TPU 算力，彻底打通 PyTorch 生态与 TPU 硬件之间的部署壁垒，让 PyTorch 模型在 TPU 上的部署流程和 GPU 环境保持一致，大幅降低 TPU 大模型线上服务的开发成本。

一、vLLM-TorchTPU 诞生的核心原因

1. 旧版 TPU 部署方案存在框架割裂问题。此前 vLLM 在 TPU 上分为两条技术路线：vllm-jax 面向 JAX 模型，而 PyTorch 模型只能依靠 PyTorch/XLA 间接编译，两条代码路径互相独立，无法共用同一套调度与内核，维护成本很高vLLM。
2. PyTorch 开发者缺少原生 TPU 部署入口。大量大模型都是用 PyTorch 编写，想要跑在 TPU 上，必须修改模型代码适配 XLA 静态图，还需要手动处理张量形状约束，部署门槛很高。
3. 原有编译链路无法完整兼容 vLLM 调度能力。早期 PyTorch/XLA 方案很难稳定运行 vLLM 的分页 KV 缓存、分片张量并行等高级特性，很难把 GPU 端成熟的推理能力完整迁移到 TPU 集群。
4. 硬件算力没有充分利用。旧方案编译链路冗余，无法调用 TorchTPU 生态优化后的专用内核，TPU 硬件的吞吐与延迟表现达不到最优水平。

二、vLLM-TorchTPU 整体技术实现步骤

步骤 1：拆分独立项目，实现框架解耦

vLLM 把 PyTorch+TPU 的推理能力独立拆分出 vLLM-TorchTPU 项目，和面向 JAX 的 vllm-jax 做业务拆分，后续所有 PyTorch 原生 TPU 推理迭代都集中在 vLLM-TorchTPU 仓库，两套后端互不干扰，各自迭代优化vLLM。

步骤 2：搭建 Torch.compile+TPU 编译链路

1. 模型加载保持原生 PyTorch 写法，直接调用 vLLM 标准模型运行器，不需要做模型改写。
2. 启用 torch.compile 并指定 backend 为 tpu，自动把 PyTorch 计算图 lowering 到 XLA 编译器，自动适配 TPU 硬件指令集vLLM。
3. 每个 TPU 设备对应独立的 worker 进程，沿用 vLLM 原有的多进程调度架构，张量并行、流水线并行可以直接复用原有分布式逻辑。

步骤 3：适配 vLLM 核心推理内核

1. 移植分页注意力（RPA v2）到 TorchTPU 编译链路，实现 TPU 硬件上的动态 KV 缓存管理。
2. 对接 TorchTPU 生态里的优化算子，包含 RoPE、RMSNorm、混合精度矩阵运算等专用内核，提升算子执行效率。
3. 增加输入形状分桶机制，对不同批次、序列长度预先编译 XLA 静态图，把首次启动的编译开销固化到本地缓存中，后续启动直接加载缓存文件，大幅缩短服务启动耗时vLLM。

步骤 4：对齐 OpenAI 兼容服务接口

推理服务对外完全沿用 vLLM 原有 REST 接口，开发者不用修改调用代码，就能把原本部署在 GPU 上的 PyTorch 模型服务无缝迁移到 TPU 集群，日志采集、监控指标、负载均衡能力和 GPU 版本保持统一。

三、vLLM-TorchTPU 支持的硬件与运行环境

1. 支持硬件机型：TPU v4、v5e、v5p、v6e 全系列云 TPU 机型，支持单机多卡与多主机分布式集群部署vLLM。
2. 软件环境要求：Python 3.10 及以上版本，谷歌云 TPU 专用虚拟机，依赖 TorchTPU 与 PyTorch XLA 配套组件。
3. 部署方式：支持源码编译构建，也可以使用官方 Docker 镜像一键拉起推理服务，容器配置包含共享内存、网络权限等 TPU 运行必要参数vLLM。

四、vLLM-TorchTPU 最终落地效果

1. 代码零改动：原生 PyTorch 模型无需修改，直接在 TPU 上完成编译与推理，PyTorch 开发者可以直接上手 TPU 部署。
2. 能力完整继承：vLLM 分页注意力、连续批处理、张量并行、LoRA 微调推理、批量解码等全部高级功能都能在 TPU 上稳定运行。
3. 性能充分释放：依托 XLA 静态编译 + TorchTPU 优化内核，TPU 硬件的 token 吞吐能力接近硬件理论上限，长上下文场景下的延迟表现优于旧版 PyTorch/XLA 方案。
4. 运维保持统一：服务接口、监控指标、缓存机制和 GPU 版本完全一致，企业可以统一维护一套 vLLM 服务代码，同时跑 GPU 与 TPU 两套硬件集群。

【相关 FAQ】

问题 1：vLLM-TorchTPU 和 vllm-jax 有什么区别，该怎么选择？

回答 1：vLLM-TorchTPU 主打 TPU 原生 PyTorch 推理服务，专门面向 PyTorch 编写的大模型，不需要切换到 JAX 框架；vllm-jax 面向 JAX 模型，使用 JAX→XLA 编译链路。如果你手上是 PyTorch 模型，直接选用 vLLM-TorchTPU；JAX 模型则使用 vllm-jax，两个项目互相独立，各自维护迭代vLLM。

问题 2：使用 vLLM-TorchTPU 部署 PyTorch 模型，是否需要手动改写模型代码适配 XLA 静态图？

回答 2：使用 vLLM-TorchTPU 搭建 TPU 原生 PyTorch 推理服务，不需要手动改写模型代码，vLLM 会自动通过 torch.compile 把 PyTorch 计算图转为 XLA 静态图，同时内置输入形状分桶策略，自动处理静态张量约束，开发者只需要加载原生 PyTorch 权重文件即可完成部署vLLM。

问题 3：vLLM-TorchTPU 首次启动编译耗时很长，有什么办法优化吗？

回答 3：vLLM-TorchTPU 依托 XLA 编译会产生静态图缓存，首次启动需要预编译多种输入形状的计算图，耗时在 20 至 30 分钟；可以提前完成模型预热，把编译缓存保存到共享存储，并且配置 VLLM_XLA_CACHE_PATH 环境变量指定缓存目录，后续服务重启直接读取缓存，编译耗时可以缩短到 5 分钟以内，大幅提升 TPU 推理服务上线效率vLLM。

【数据来源】

[2026-10-20]（Introducing vLLM TorchTPU）vLLM 官方博客vLLM

[2026-10-20]（vLLM TPU 后端文档）vLLM 项目官方文档vLLM

[2025-10-16]（vLLM TPU 统一后端技术介绍）vLLM 官方技术博客vLLM

[2025-08-10]（TPU 编译缓存优化指南）vLLM 官方运维文档vLLM

免责声明：本文所有技术内容均来自 vLLM 官方公开文档与技术博客，仅作技术学习交流使用；硬件运行效果会受 TPU 机型、集群拓扑、模型大小、输入序列长度等条件影响，实际业务性能请以线下实测结果为准；项目版本持续迭代，部署步骤请以对应版本最新官方文档为准。