开源大模型上 TPU：SGLang/vLLM 部署与 GPU 对照

适用场景

手里有 Google Cloud TPU 配额，或正在评估把推理服务从 GPU 迁到 TPU 的成本与性能，需要在 TPU 上跑 Kimi、DeepSeek 系开源模型的对话/推理服务。这套流程覆盖容量规划、TPU VM 创建、vLLM 与 SGLang 的 TPU 后端启动、AI 词典：KV Cache">KV Cache 与批处理调参，以及和 GPU 做同模型对照压测的完整链路。适合已经会起 vLLM/SGLang 服务、但对 TPU 内存模型和并行方式不熟的人。

环境与前置条件

账号与工具

• Google Cloud 项目，已启用 Cloud TPU API，tpus 相关配额充足。
• 本地装好 gcloud CLI，并完成 gcloud auth login / gcloud config set project。
• 目标 zone 有对应加速器规格的库存与配额，具体型号与 zone 的对应关系以官方 TPU 区域可用性页面为准。

TPU VM 运行时

• 操作系统镜像用 TPU VM 官方镜像（例如 tpu-ubuntu2204-base 这一系列），运行环境版本号以官方文档当前版本为准。
• Python 3.10 及以上，建议在 TPU VM 内建虚拟环境，不要直接动系统 Python。
• 底层依赖 JAX 或 PyTorch/XLA；vLLM 与 SGLang 的 TPU 后端都构建在这两者之上。

资源建议

主机内存：加载权重时要先过主机 RAM，建议每台 TPU VM 主机内存不低于「该机所承载分片大小」的 3~4 倍。

磁盘：模型权重目录至少留出模型体积的 1.5 倍（下载缓存 + 可能的格式转换），优先本地 SSD，其次 GCS 挂载。

HBM 容量参照（具体参数以官方规格页为准）：

模型选择建议

• 单机 v5e-8 / v6e-8：DeepSeek-R1-Distill-Qwen-14B/32B、Qwen 系 7B~32B 这类模型比较合适。
• 单机 v5p-8：可以放下 70B 级 bf16，或 100B+ 的 int8 权重。
• DeepSeek-V3/R1 671B、Kimi K2 这类千亿级 MoE：需要多机 TPU（v5p-16/32/64 或多切片 v6e），MoE 的所有专家权重都要常驻 HBM，不能按激活参数量估算。

分步骤部署

第 0 步：先算清楚三笔账

不要建完机器再试模型能不能放下。先用模型 config.json 里的真实参数算：

1. 权重占用 = 参数量 × 每参数字节数（bf16 为 2，int8 为 1）。

2. 每 token KV 占用 = 2 × 层数 × KV 头数 × head_dim × 精度字节数。

3. KV 总量 = 每 token KV × 单请求长度 × 并发数。

举例：Qwen2.5-32B，64 层，8 个 KV 头，head_dim 128，bf16。

• 每 token KV = 2 × 64 × 8 × 128 × 2 = 262,144 字节 ≈ 0.25 MiB
• 单请求 16K 上下文 → 约 4 GiB
• 16 路并发 → 约 64 GiB
• 权重 32B × 2 = 64 GB

两者相加约 128 GB，正好卡在 v5e-8 的上限，必须把 max-model-len 或并发数压下来。这就是为什么先算再建。

第 1 步：创建 TPU VM

```bash

gcloud compute tpus tpu-vm create llm-tpu-demo \

--zone="<你的可用区>" \

--accelerator-type="<加速器规格，如 v5litepod-8 / v6e-8 / v5p-8>" \

--version="tpu-ubuntu2204-base"

```

这步在创建一台挂载 TPU 芯片的虚拟机。成功时会输出 STATE: READY。多机场景把 --accelerator-type 换成更大的规格（例如 v5litepod-16 对应两个 8 芯片主机），会一次创建多个 worker。

参数名在不同 gcloud 版本里可能是 --version 或 --runtime-version，以官方文档当前版本为准。

第 2 步：登录并确认 TPU 可见

```bash

gcloud compute tpus tpu-vm ssh llm-tpu-demo --zone="<你的可用区>"

```

进机器后先自检：

```bash

python3 -c "import torch_xla.core.xla_model as xm; print(xm.xla_device())"

```

预期输出 xla:0。走 JAX 路线的话：

```bash

python3 -c "import jax; print(jax.devices())"

```

多机时确认总芯片数：

```bash

python3 -c "import torch_xla.runtime as xr; print(xr.world_size())"

```

也可以安装监控小工具看每芯片 HBM 与利用率：

```bash

pip install tpu-info

tpu-info

```

这一步报错就别往下走，环境不对后面全是坑。

第 3 步：安装运行环境

```bash

python3 -m venv ~/venv-tpu

source ~/venv-tpu/bin/activate

pip install --upgrade pip wheel

pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html

```

libtpu 是与 TPU 运行时匹配的驱动层，版本必须和 TPU VM 的 runtime 对齐；装错会直接表现为「找不到设备」。安装源与参数以 PyTorch/XLA 官方文档当前版本为准。

随后按需安装推理框架的 TPU 后端：

```bash

pip install vllm-tpu # vLLM 的 TPU 后端，包名与安装方式以官方文档为准

pip install sglang # SGLang，TPU 后端依赖与安装方式以官方文档为准

```

第 4 步：准备模型权重

```bash

pip install -U "huggingface_hub[cli]"

export HF_HOME=/mnt/models/.hf_cache

hf download <模型组织>/<模型名> --local-dir /mnt/models/<模型名>

```

较老版本的 CLI 入口是 huggingface-cli download，两者功能一致。

要点：

• 权重放在本地 SSD 或大内存机器的 /dev/shm，避免每次重启服务都重新下载。
• 多机时每台都要有同一份权重，或统一从 GCS 桶挂载。
• 下完先打开 config.json 核对层数和 KV 头数，把第 0 步的账用真实数字再算一遍。

第 5 步：启动 vLLM TPU 后端

```bash

export PJRT_DEVICE=TPU

export VLLM_USE_V1=1 # 新版本通常默认开启，具体以官方文档为准

vllm serve /mnt/models/Qwen2.5-32B-Instruct \

--tensor-parallel-size 8 \

--max-model-len 8192 \

--max-num-seqs 16 \

--port 8000

```

说明：

• --tensor-parallel-size 通常等于本机芯片数；多机时等于总芯片数。
• 首次启动会做 XLA 图编译并加载权重，几分钟到十几分钟都属于正常，看到日志停在编译阶段不用急着重启。
• 显存相关的开关名称在 TPU 后端可能与 GPU 侧不同（例如 GPU 上的 --gpu-memory-utilization），以官方文档当前版本为准。

成功标志是日志出现 Application startup complete 和监听 0.0.0.0:8000 的记录。

第 6 步：启动 SGLang TPU 后端

```bash

export PJRT_DEVICE=TPU

python -m sglang.launch_server \

--model-path /mnt/models/Qwen2.5-32B-Instruct \

--tp-size 8 \

--device tpu \

--context-length 8192 \

--mem-fraction-static 0.8 \

--max-running-requests 16 \

--host 0.0.0.0 --port 30000

```

• --mem-fraction-static 控制静态分配给权重和 KV 池的 HBM 比例，要留出编译与临时缓冲的余量。
• --context-length 直接决定 KV 上限，作用等同于 vLLM 的 max-model-len。

两个后端都能起来之后，用同一个模型做一次单请求延迟对照，再决定主力用哪个。

第 7 步：批处理与 KV Cache 调参

三个旋钮互相牵制：

1. 并发上限（--max-num-seqs / --max-running-requests）：直接成倍放大 KV 占用。

2. 单请求长度上限（--max-model-len / --context-length）：决定单序列的 KV 峰值。

3. prefill 分块（--chunked-prefill-size / --max-prefill-tokens）：把长输入的 prefill 拆开，避免一个超长请求把 KV 池顶爆。

KV 池大小一般在启动时就静态确定，改并发上限往往要重启进程才生效。

第 8 步：多机部署（千亿 MoE 走这一步）

```bash

gcloud compute tpus tpu-vm ssh llm-tpu-demo \

--zone="<你的可用区>" --worker=all \

--command="bash /mnt/scripts/setup_and_serve.sh"

```

• --worker=all 是多台并行执行，脚本必须幂等、可重复执行。
• 服务入口只在 worker 0 暴露，其余 worker 作为计算节点参与集合通信。
• 建议把安装、下载、启动写成一个脚本放 GCS，各机拉下来执行，避免逐台手敲。
• DeepSeek-V3/R1、Kimi K2 这类模型：先按「总参数量 × 字节数 ÷ 单芯片 HBM」算最小芯片数，再乘 1.5 留出 KV 与临时空间。

第 9 步：准备 GPU 对照环境

公平对照的关键是控制变量：

• 同一份权重、同一 revision。
• 同一精度（bf16 对 bf16，int8 对 int8）。
• 同一上下文长度与并发上限。
• 尽量同一推理框架版本。
• 算力规模按「能跑起该模型的最小可用集群」对齐，而不是 1 卡对 1 芯片。

```bash

vllm serve /mnt/models/Qwen2.5-32B-Instruct \

--tensor-parallel-size 4 \

--max-model-len 8192 \

--max-num-seqs 16 \

--port 8000

```

第 10 步：压测

用随机数据集固定输入输出长度，两侧跑同一套参数：

```bash

vLLM 侧

vllm bench serve \

--backend openai-chat \

--base-url http://<host>:8000 \

--model /mnt/models/Qwen2.5-32B-Instruct \

--dataset-name random \

--random-input-len 1024 \

--random-output-len 256 \

--num-prompts 200 \

--max-concurrency 16

```

压测入口在不同版本里可能是 vllm bench serve、python -m vllm.entrypoints... 或 benchmarks/benchmark_serving.py，以官方文档当前版本为准。

```bash

SGLang 侧

python -m sglang.bench_serving \

--backend sglang \

--host <host> --port 30000 \

--model /mnt/models/Qwen2.5-32B-Instruct \

--random-input-len 1024 \

--random-output-len 256 \

--num-prompts 200 \

--max-concurrency 16

```

要看的指标：

• TTFT：首 token 延迟，直接决定体感。
• TPOT / ITL：每输出 token 耗时，决定长回答的稳定性。
• Output token throughput：整体吞吐。
• 在 1 / 4 / 8 / 16 / 32 各并发档位分别跑，画出吞吐—延迟曲线，而不是只报一个峰值数字。

压测纪律：

• 先跑 2~3 轮预热，XLA 首次编译的结果不计入。
• 每档至少 200 条请求，取稳定区间。
• 同时记录 TPU 侧 HBM 占用与算力利用率，避免把「KV 池被撑爆」误判成「TPU 性能差」。
• 报告里写清模型 revision、精度、上下文长度、并发、芯片/GPU 型号与数量、框架版本。

验证部署是否成功

1. 健康检查

```bash

curl -s -o /dev/null -w "%{http_code}\n" http://localhost:8000/health

```

返回 200 即服务已就绪。

2. 模型列表

```bash

curl -s http://localhost:8000/v1/models | head

```

返回体里能看到模型 id。

3. 一次真实推理

```bash

curl -s http://localhost:8000/v1/chat/completions \

-H "Content-Type: application/json" \

-d '{

"model": "/mnt/models/Qwen2.5-32B-Instruct",

"messages": [{"role":"user","content":"用三句话解释什么是 KV Cache"}],

"max_tokens": 128,

"temperature": 0

}'

```

预期：choices[0].message.content 是一段通顺中文，usage 里有 prompt_tokens 与 completion_tokens。temperature 设为 0 便于多次请求比对一致性。

4. 确认 TPU 真在干活

另开一个 SSH 会话执行 tpu-info（或 watch -n1 tpu-info），压测时 HBM 使用率与算力利用率应当明显上涨。

5. 多机一致性

每台机器都跑一次：

```bash

python3 -c "import torch_xla.runtime as xr; print(xr.world_size(), xr.global_ordinal())"

```

world_size 等于总芯片数，global_ordinal 从 0 到 N-1 不重复。

6. 稳定性

用固定负载连续跑 30 分钟以上，观察是否出现 HBM 缓慢上涨、进程被 OOM kill、TPU 掉线重连。

常见报错与解决

报错：RuntimeError: No TPU devices found 或 XLA device not available

原因：PJRT_DEVICE 未设置，或 libtpu / torch_xla 与 TPU VM runtime 版本不匹配。

```bash

export PJRT_DEVICE=TPU

pip uninstall -y torch_xla libtpu

pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html

python3 -c "import torch_xla.core.xla_model as xm; print(xm.xla_device())"

```

仍不行就重建实例，把 runtime version 与 torch_xla 版本对齐。

报错：RESOURCE_EXHAUSTED: Out of memory / 加载到一半进程被 kill

原因：权重 + KV 池超过 HBM 上限。

```bash

vllm serve <model> --tensor-parallel-size 8 --max-model-len 4096 --max-num-seqs 8

```

调整顺序建议：先缩上下文，再降并发，再扩芯片数（提高 tensor parallel），最后才考虑 int8 量化。改完先回到第 0 步重算，不要靠反复试。

报错：CUDA error: no CUDA-capable device is detected，或找不到 nvidia-smi

原因：装成了 GPU 版推理框架，TPU 上没有 CUDA。

```bash

pip uninstall -y vllm

pip install vllm-tpu # 包名以官方文档为准

```

SGLang 则检查启动命令是否带了 --device tpu。

报错：Failed to connect to 10.x.x.x:port，多机启动卡在某台 worker

原因：--worker=all 是并行执行，某台机器上的脚本失败（权重没下完、依赖没装全），其余机器在等它。

```bash

gcloud compute tpus tpu-vm ssh llm-tpu-demo --zone="<zone>" --worker=all \

--command="ls /mnt/models/<模型名>/config.json && python3 -c 'import torch_xla.core.xla_model as xm; print(xm.xla_device())'"

```

逐台确认通过后再整体启动。

报错：创建实例时报 Quota exceeded 或资源不足

原因：所选 zone 没有该加速器规格的配额，或库存临时不足。

处理：换 zone、换规格，或在配额页面提交调整申请。可用型号与 zone 的对应关系以官方页面为准。

现象：启动长时间无新日志

原因：XLA 首次编译 + 大权重从磁盘或网络加载。

处理：第一次耐心等待；把权重放本地 SSD；启动加 -v 观察阶段日志；确认磁盘 IO 没有跑满。

后续维护

环境固化：把虚拟环境构建脚本、启动脚本、压测脚本放 GCS，新实例拉下来一条命令还原；也可以用自定义镜像或 --metadata-from-file startup-script=... 做开机初始化。

权重管理：统一放 GCS 桶，按「模型名 + revision」建目录，本地只做缓存。升级模型时改动目录名，方便回滚。

升级策略：TPU VM 的 runtime 版本通常要重建实例才能变更，不能原地升级。升级前锁定脚本；torch_xla 与推理框架一起升，不要只升其中一个；升完重跑一遍压测基准，和上次结果对比再上生产。

备份：真正需要备份的是启动脚本、压测脚本、模型 revision 记录、基准结果（JSON/CSV）。权重本身可以重新下载，不必单独备份。

日志：服务 stdout 重定向到文件并按 worker 编号分文件，再同步到 Cloud Logging。关键字段：启动耗时、编译耗时、OOM、连接断开、请求错误码。

监控：日常用 tpu-info 观察 HBM 利用率、算力利用率与温度；配合定时压测做回归，出现明显掉点时先排查是不是有人把并发上限或上下文长度改小了。

成本：不用的实例及时释放：

```bash

gcloud compute tpus tpu-vm delete llm-tpu-demo --zone="<你的可用区>"

```

对照实验做完，GPU 与 TPU 两侧机器一起关掉，避免闲置空转。