vLLM TPU：让 Google TPU 成为大模型推理的新一代标准硬件

vLLM TPU（仓库 vllm-project/tpu-inference）是 vLLM 官方项目，专为 Google Cloud TPU 提供高性能大模型推理服务。它将 vLLM 的核心能力——PagedAttention、连续批处理、投机解码——完整移植到 TPU 硬件上，同时带来了一个重要创新：统一 JAX 与 PyTorch 双后端，让两种框架的用户都能在 TPU 上获得开箱即用的高性能推理体验。

图1：vLLM TPU 项目 logo（支持亮色/暗色模式自动切换）

为什么 TPU 推理值得关注？

在大模型推理领域，NVIDIA GPU 长期占据主导地位。然而 Google TPU 正在以几个独特优势改变这一格局：

第一，TPU 的内存带宽惊人。TPU v6e 提供 1.2 TB/s 的内存带宽，远超 H100 的 3.35 TB/s 带宽容量（注意这里的对比单位有差异，TPU 的高带宽对长上下文推理尤其友好）。对于需要处理超长上下文（128K+ token）的大模型推理场景，TPU 的内存子系统具有天然优势。

第二，TPU 的性价比在大规模推理场景下往往更优。Google Cloud 的 TPU 租赁价格相比同等算力的 A100/H100 实例通常有 20-40% 的成本优势，尤其是在长序列推理和批量处理场景。

第三，Google 正在全力推动 TPU 生态。2025 年 10 月 vLLM 官方博客专门发布了 vLLM TPU: A New Unified Backend，宣布 TPU 后端从实验性功能升级为官方支持的稳定后端，与 GPU 后端并列。

核心架构：JAX + PyTorch 统一降级

vLLM TPU 最大的技术创新在于它的 统一降级路径（Unified Lowering Path）。整个架构分为三层：

第一层：框架接口层（Model Definition）

vLLM TPU 支持两种模型定义方式：

• PyTorch 原生模型：用户可以直接使用 HuggingFace 格式的 PyTorch 模型，无需修改任何代码。torchax 层负责将 PyTorch 算子自动翻译为 TPU 上的等价实现。
• JAX 原生模型：JAX 用户可以使用 flax 定义的模型，享受 TPU 原生的 XLA 编译优化。

项目目录中可以看到两组并行实现：

• tpu_inference/models/jax/ — JAX 原生模型（Llama4、Gemma4、DeepSeek V3、Qwen2.5-VL 等）
• tpu_inference/models/vllm/ — PyTorch 模型封装（vllm_model_wrapper.py）

第二层：核心推理引擎

TPU Runner（tpu_inference/runner/tpu_runner.py，126KB，最大文件）是整个推理引擎的核心，负责模型加载与权重分发、请求调度与连续批处理（Continuous Batching）、KV Cache 管理以及分布式协调。

KV Cache Manager（tpu_inference/runner/kv_cache_manager.py）借鉴了 vLLM 在 GPU 上的 PagedAttention 思想，针对 TPU 内存特性进行了重新设计，支持页面级的 KV Cache 分配和动态回收。

DP Scheduler（tpu_inference/core/sched/dp_scheduler.py，66KB）是调度器，实现了详细的请求分配、token 生成和批处理逻辑。

第三层：TPU 专用内核（Kernels）

kernels/ 目录包含大量手工优化或自动调优的 TPU 内核，是性能的核心来源：

内核目录	说明	规模
ragged_paged_attention/v2/v3	专用于 TPU 的 PagedAttention 实现，含 HD64 变体	最大
mla/	Multi-head Latent Attention（DeepSeek V4 架构）	~107KB
fused_moe/	混合专家模型的融合矩阵乘法	~68KB
sparse_core/	Sparse Core 利用（Weight-Only 场景）	~73KB
quantized_matmul/	INT8/FP8 量化矩阵乘法	~41KB

值得注意的是，ragged_paged_attention/v3/tuned_block_sizes.py 是整个仓库中最大的单个文件（231KB），记录了大量经过自动搜索（Autotuner）得到的 TPU 最优分块大小参数，体现了项目对极致性能的追求。

部署体验：安装简单，硬件门槛高

安装方式

vLLM TPU 的安装非常现代化，提供了三种方式：

# 方式一：推荐，用 uv 安装（最快）
uv pip install vllm-tpu

# 方式二：标准 pip
pip install vllm-tpu

# 方式三：Docker 一键启动
sudo docker run -it --rm --name vllm \
    --privileged --net=host \
    -v /dev/shm:/dev/shm --shm-size 150gb \
    -p 8000:8000 vllm/vllm-tpu:latest

代码仓库中提供了多阶段构建的 Dockerfile（docker/Dockerfile），基于 python:3.12-slim-bookworm，会自动克隆 vLLM 主仓库并安装所有依赖。

同时，.buildkite/ 目录下的 Kubernetes manifests（kubernetes/manifests/v7x/）说明项目也支持 K8s 部署，提供了 single_decode.yaml 和 single_prefill.yaml 等典型场景的配置模板。

硬件门槛

这是最大的限制——vLLM TPU 需要真实的 Google Cloud TPU 硬件实例：

• 推荐型号：TPU v7x（Ironwood）、TPU v6e（Trillium）、TPU v5e
• 实验性支持：TPU v3、v4、v5p
• 不支持：任何 GPU（AMD/NVIDIA）或普通 CPU

如果读者没有 TPU 访问权限，这个项目无法在本地进行任何实际推理测试。但 Google Cloud 提供 TPU 试用配额，研究人员可以申请免费额度。

启动推理服务

# 登录 HuggingFace
huggingface-cli login --token $HF_TOKEN

# 启动服务
python -m vllm.tpu.entrypoints.openai.api_server \
    --model google/gemma-3-27b-it \
    --host 0.0.0.0 --port 8000

API 兼容 OpenAI 接口格式，可以直接用 curl 调用 /v1/chat/completions 等标准端点。

功能特性：紧随 GPU 版本

作为 vLLM 的官方 TPU 后端，这个项目几乎同步跟进 GPU 版本的所有核心功能。根据仓库的 support_matrices/ 目录和 CI 配置，已验证支持的功能包括：

• ✅ 连续批处理（Continuous Batching）：最大化 TPU 利用率
• ✅ PagedAttention：页面级 KV Cache 管理（对应 TPU 内核 ragged_paged_attention）
• ✅ MoE 支持：混合专家模型（fused_moe 内核）
• ✅ MLA（Multi-head Latent Attention）：DeepSeek V4 架构核心（mla/ 内核 ~107KB）
• ✅ LoRA/Torch：PyTorch 风格的 LoRA 微调支持（CI 测试：.buildkite/features/LoRA_Torch.yml）
• ✅ 投机解码（Speculative Decoding）：Eagle3、Ngram 等多种策略
• ✅ 多模态输入：Qwen2.5-VL 等视觉语言模型（models/jax/qwen2_5_vl.py，52KB）
• ✅ KV Cache Offload：CPU 卸载以支持更大模型（offload/ 目录，105KB）
• ✅ 量化推理：INT8/FP8 量化矩阵乘法
• ✅ P-D 分离（Prefill-Decode Disaggregation）： disagg 示例（examples/disagg/）

代码质量：工程化程度极高

从仓库结构来看，vLLM TPU 体现了 Google 级别的工程水准：

• Apache-2.0 许可证，完全开源
• 完整的 CI/CD：.buildkite/ 下有大量自动化测试和 benchmark pipeline
• 详尽的测试覆盖：tests/ 目录包含 correctness test、performance test、offload test 等
• 自动化调优框架：.buildkite/xla_autotune/ — 对 TPU 分块参数进行自动搜索（genetic algorithm 风格）
• 支持矩阵透明化：support_matrices/ 下有自动生成的模型/功能支持状态文档
• 详细贡献指南：CONTRIBUTING.md 明确了 PR 流程和代码规范

开发文档也很完善：docs/developer_guides/ 下有 JAX 模型开发和 TorchAX 模型开发的专门指南，说明项目对社区贡献持开放态度。

局限与注意事项

1. TPU 硬件是硬性依赖：无法绕过。没有 TPU 访问权限的用户无法实际使用这个项目。
2. 仅支持 Google Cloud TPU：不兼容 Colab TPU 或其他 TPU 提供商。
3. 相对较新：虽然 vLLM 团队在快速推进，但某些 GPU 版本功能可能尚未移植。
4. 依赖 Google Cloud 环境：需要配置 GCP 项目、TPU 配额和 VPC 网络等，对不熟悉 GCP 的用户有一定学习成本。
5. 仅 Linux x86_64：README 中未提及 macOS 或 Windows 支持。

行业意义

vLLM TPU 的出现标志着大模型推理硬件多元化进程中的一个重要里程碑。在此之前，vLLM 仅支持 NVIDIA GPU，而 AMD ROCm 后端和 Intel 后端也处于不同成熟度阶段。Google TPU 的加入，使得：

• 用户有了更强的硬件选择权：不必绑死在 NVIDIA 生态
• 成本竞争加剧：TPU 的价格优势可能迫使 GPU 云厂商降价
• 技术标准统一：同一个 vLLM API 同时服务 GPU 和 TPU 两种硬件，降低了多硬件部署的工程复杂度

特别值得关注的是 2025 年 10 月 vLLM 官方博客宣布这一消息后，vLLM TPU 项目迅速获得了 369 个 GitHub stars 和 233 个 forks，说明社区对 TPU 推理的需求是真实存在的。