vLLM 模型代码不改一行就能跑上谷歌 TPU：新后端复用现有调度器与 KV 缓存
核心概要
Inferact 与谷歌发布 vLLM TorchTPU，让 vLLM 里已有的模型定义无需改动即可在谷歌云 TPU 上提供服务。此前在 TPU 上部署模型需要另写模型代码，而这一后端沿用 vLLM 原有的命令行、OpenAI 兼容服务器、调度器和 KV 缓存管理器，TPU 相关工作放在它们之下。它通过 torch.compile 的 TPU 后端编译到 XLA，并用 JAX/Pallas 自定义算子处理注意力、混合专家和量化矩阵乘等操作。
深度剖析
vLLM 中已有的模型定义现在可以直接在谷歌云 TPU 上提供服务，大多数当前能在 vLLM 里运行的模型不需要改动模型代码。 此前把模型搬到 TPU 上需要另写模型代码，而这一后端复用 vLLM 现有的模型代码。 该后端建立在谷歌 TorchTPU 栈之上，模型走 vLLM 标准模型运行器，并以 torch.compile 的 tpu 后端编译到 XLA。
使用者的操作方式保持不变：命令行、OpenAI 兼容服务器、调度器和 KV 缓存管理器都与在其他硬件上相同。 TPU 专属的工作被放在这些组件之下，而不是替换它们。 每个 TPU 设备拥有自己的 worker 进程，设备间通信走 torch.distributed，调度器、worker、KV 缓存管理器和 API 服务器均为未修改的上游 vLLM。
注意力、混合专家和量化矩阵乘等复杂操作由手写的 Pallas 算子承担。 这些算子以自定义算子形式注册，并在编译图中被调用。 算子覆盖 MHA/GQA 的 Ragged Paged Attention v3、DeepSeek、Kimi 与 GLM 的稠密与稀疏 MLA、混合线性注意力模型的 Gated DeltaNet 与 Kimi Delta Attention、带专家并行的融合混合专家，以及 FP8、INT8、MXFP4、NVFP4 量化矩阵乘。
启示与展望
已展示的范围是：在谷歌云 TPU 虚拟机上以 Python 3.12 安装该插件后，vLLM 通过 vllm.platform_plugins 入口自动加载它，大多数 vLLM 支持的模型可在 TPU 上服务。TPU v7x（Ironwood）是推荐硬件并在每次拉取请求上经 CI 测试，v6e、v5e、v5p 为实验性。首次启动需要为配置的形状编译 TPU 图，可能耗时数分钟，可通过 VLLM_XLA_CACHE_PATH 在重启间复用编译产物。面向的读者是已在用 vLLM 服务模型、希望把负载放到 TPU 上的工程师。
支持范围以已验证检查点清单和夜间支持矩阵为准，而非全部 vLLM 模型；v6e、v5e、v5p 仍标注为实验性。首次启动的图编译需要数分钟，缓存路径用于跨重启复用。路线图中的 24 小时内接入新开源权重模型的自动 CI、路由与扩散模型服务尚未列为已交付能力。