单块 TPU 自托管轻量级 agent 后端：Gemma 4 2B + vLLM（v5e-1）

AI技术与应用 2026-08-11 17:38:25 2

最近折腾了一下在单块 TPU 上跑轻量级 Agent 后端，踩了几个坑，这篇把问题说清楚。实测跑 google/gemma-4-E2B-it 在 vLLM 下能到 1,496 tokens/秒 的聚合吞吐，单流延迟 8.02 ms，还自带原生 tool-calling。算下来支撑 8-16 个并发轻量 Agent，每百万输出 tokens 成本大概 $0.107。

这是一篇带数据的复盘文，所有数字都是实测跑出来的。其中最有价值的部分是我预测错了的地方——四个我原本很有信心的判断被 benchmark 打脸了，每次打脸都比猜对有用。

测试环境： v5litepod-1（单块 v5e），us-west4-a，vllm/vllm-tpu:nightly，vLLM 0.26.1rc1.dev125+ga7a204cc6，tpu-inference JAX backend，google/gemma-4-E2B-it bf16 精度。

第一部分——环境搭建和启动

1.1 前置准备

gcloud auth login # 供 gcloud 子进程调用
gcloud auth application-default login # ADC，供 Secret Manager 客户端使用

Hugging Face token 要存到 Secret Manager，别写进脚本或者环境变量——TPU VM 的启动脚本存在实例元数据里，任何写进去的东西都能被读取：

printf '%s' "hf_xxxxxxxxxxxx" | gcloud secrets create hf-token --data-file=- --project=YOUR_PROJECT

区域限制，这个坑了我半天： flex-start v5litepod-1 只在 us-west4-a 可用。europe-west4-a 和 -b 会直接被 API 拒绝，返回 FLEX_START provisioning model is not supported for accelerator type "v5litepod-1"，跟配额没关系。区域里显示有非零配额没用，provisioning model 才是拦路虎。

1.2 申请 TPU 资源

三种 provisioning model，三套不同的命令。注意 v5e 在 gcloud 里要写成 v5litepod——"v5e-1" 在正文里写没问题，但 CLI 参数里永远不认这个拼法。

# Spot — 最便宜，会被抢占，大概 30 秒预警，不会自动停（会一直计费到你删掉）
gcloud alpha compute tpus tpu-vm create gemma4-v5e \ --zone=us-west4-a --type=v5litepod --topology=1x1 \ --provisioning-model=spot --version=v2-alpha-tpuv5-lite # On-demand — 全价，不会被抢占，同样不会自动停
gcloud alpha compute tpus tpu-vm create gemma4-v5e \ --zone=us-west4-a --type=v5litepod --topology=1x1 \ --version=v2-alpha-tpuv5-lite

Flex-start 要走 Queued Resource API，而且它是唯一支持 --max-run-duration 的模式——也就是唯一能自己停止计费的：

gcloud alpha compute tpus queued-resources create gemma4-qr \ --node-id=gemma4-qr-node --zone=us-west4-a \ --accelerator-type=v5litepod-1 --runtime-version=v2-alpha-tpuv5-lite \ --provisioning-model=flex-start --max-run-duration=4h

2026-08-09 实测：这个命令执行成功，QR 到达 ACTIVE 状态，provisioningModel: FLEX_START 和 maxRunDuration: 14400s 都对，然后正常删除了。注意没有 dry-run：create 之后要么排队要么分配，PROVISIONING 状态是删不掉的，所以如果立刻有资源的话至少要付几分钟费用。

Spot 和 on-demand 不会自动停止。 会一直计费到被抢占或者你手动删除。设个日历提醒，或者用 flex-start。看下面的成本部分——flex-start 只比 spot 贵 3.8%。

Spot 用的是独立配额（TPUV5sPreemptibleLitepodPerProjectPerZoneForTPUAPI），不是标准 TPU 配额。某个区域 on-demand 配额充足不代表 spot 也能用。

1.3 把 token 传到节点上并启动服务

从某些沙盒环境里 gcloud compute tpus tpu-vm ssh 会崩，报 ConnectionResetError——它在自己内部 API 调用时就挂了，但直接调 gcloud API 完全正常。直接用 SSH 总是靠谱的：

IP=$(gcloud compute tpus tpu-vm describe gemma4-v5e --zone=us-west4-a \ --format='value(networkEndpoints[0].accessConfig.externalIp)') # token 直接管道传进去，别经过 shell 变量或者日志行
gcloud secrets versions access latest --secret=hf-token \ | ssh -i ~/.ssh/google_compute_engine xbill@$IP 'umask 077; cat > ~/.hf_token' ssh -i ~/.ssh/google_compute_engine xbill@$IP 'sudo docker pull vllm/vllm-tpu:nightly'

然后启动服务。这就是本文后面要讲的所有配置：

sudo docker run -d --name vllm-gemma4 --privileged --net=host \ -v /dev/shm:/dev/shm --shm-size 10gb \ -v ~/.cache/vllm:/root/.cache/vllm \ -e HF_HOME=/dev/shm -e HF_TOKEN="$(cat ~/.hf_token)" \ vllm/vllm-tpu:nightly \ vllm serve google/gemma-4-E2B-it \ --dtype bfloat16 \ --kv-cache-dtype auto \ --max-model-len 32768 \ --max-num-batched-tokens 4096 \ --tensor-parallel-size 1 \ --gpu-memory-utilization 0.92 \ --enable-prefix-caching \ --disable-chunked-mm-input \ --limit-mm-per-prompt '{"image":4,"audio":1}' \ --enable-auto-tool-choice --tool-call-parser gemma4 --reasoning-parser gemma4

整篇文章里价值最高的就是这行 -v ~/.cache/vllm:/root/.cache/vllm。 JAX 编译缓存存在那里（实测 197 MB），否则是容器本地的，每次 docker rm 就没了。编译占冷启动 857 秒里的 685 秒。实测挂载之后重启只要 497 秒，省了 42%。不挂载的话，每次重启、每次改参数、每次 spot 被抢，都要再付一次完整编译的代价。

1.4 验证

冷启动 857 秒（14 分钟），其中 80% 是 XLA 编译，不是权重加载——9.54 GiB 的 checkpoint 下载大概 10 秒就完了。别盯着时钟看，看日志：

sudo docker logs -f vllm-gemma4 2>&1 | grep -E "Memory statistics|Init kv-cache|startup complete"

想看到这个——一行展示完整内存预算：

Memory statistics | total_hbm_limit_gb=15.75GiB | total_hbm_limit_cap_gb=14.49GiB
 | total_hbm_used_gb=8.97GiB | total_hbm_avail_gb=5.52GiB

然后做个冒烟测试。用 /v1/chat/completions，别用 /v1/completions——对 -it 模型用原始 completions 接口会返回空字符串，看起来像部署坏了但其实不是：

curl -s localhost:8000/v1/chat/completions -H 'Content-Type: application/json' -d '{ "model":"google/gemma-4-E2B-it", "messages":[{"role":"user","content":"Say hi in five words."}]}' | jq -r '.choices[0].message.content'

1.5 清理

gcloud compute tpus tpu-vm delete gemma4-v5e --zone=us-west4-a --quiet

第二部分——各参数的实际作用

芯片本身：单块 v5e 实际能给你什么

规格	v5e 单芯片	来源
HBM 容量	标称 16 GB · 运行时实际可见 15.75 GiB	官方 · 实测
在 `0.92` 下可用于权重 + KV	14.49 GiB	实测
HBM 带宽	800 GiBps	官方
峰值 bf16	197 TFLOPS	官方
峰值 int8	393 TOPS（正好是 bf16 的 2 倍）	官方
TensorCore	1 个，4 个 MXU（128x128）	官方
ICI	400 GBps 双向，4 端口	官方
机器类型	`ct5lp-hightpu-1t`
gcloud 参数写法	`v5litepod-1`，运行时 `v2-alpha-tpuv5-lite`	实测

两个单位陷阱，比较前要留意。Google 对 v5e 报的是 GiBps，对 v6e 报的是 GBps——要先换算再对比，真实的代际提升大概是 1.9 倍，不是整 2 倍。另外"16 GB"容量和运行时显示的 15.75 GiB 放一起很别扭（15.75 GiB = 16.9 GB），所以官方那个数字八成是把 16 GiB 写随意了。跟实测的 15.75 GiB 比，别跟营销数字比。

数据类型：矩阵单元实际支持什么精度

这张表决定了这块芯片上所有量化问题的答案。

格式	MXU 原生支持？	在 v5e 上的实际收益
bf16	✅	基准线——这里所有东西都用这个跑
int8	✅ 2 倍于 bf16 吞吐	唯一有低精度计算加速的格式
fp8	❌	只能省存储和带宽——实际矩阵运算前会扩回 bf16
int4 / fp4	❌	只能省存储和带宽，然后解包回 bf16

Google 发布了 v5e 的 bf16 和 Int8 峰值数据，但根本没给 fp8 的数字，这就是信号。实际后果：在这个芯片上测 fp8 没加速是正确结果，不是配置错误。 v7/Ironwood 是第一个在 MXU 里支持 fp8 的 TPU——别把这里的结论套到它身上。

量化：实际能用到什么

Gemma 4 在这个技术栈里只有 JAX 实现，所以 torch 路径里的任何量化都够不着，不管平台广告了什么。实测情况：

路径	本构建状态
KV 缓存 bf16（`auto`）	✅ 唯一值得跑的
KV 缓存 `fp8_e4m3` / `fp8_e5m2`	可以触达，容量倍率 1.000x——block 布局按 word 对齐，元素宽度缩小换来的是填充而不是更多空间。大概慢 2%
KV 缓存 `int8`	❌ CLI 枚举里没有——根本到不了引擎
KV 缓存 `int8_per_token_head`、`turboquant_*`、`nvfp4`、`fp8_inc`、`fp8_ds_mla`	CLI 接受，但启动时把服务炸掉
权重，compressed-tensors w4a16（Google QAT 格式）	❌ JAX 路径抛 `NotImplementedError`
权重，mxfp4	❌ 只支持 MoE；E2B 是 dense 的，没东西可以挂
权重，qwix PTQ int8/int4	❌ 启动失败——具体路径 OOM 在量化中间变量上，抽象路径报权重绑定错误
权重，AWQ / GGUF / q4_0	❌ torch 路径或不存在

所以 bf16 不是选择，是唯一能跑的东西。 这也是芯片最大的限制。权重占 14.49 GiB 预算里的 8.97 GiB（62%），而且一块都压不了。如果能跑 int8 权重，KV 池容量大概能翻倍，还能拿到真实 FLOPS 收益，因为 int8 是唯一有原生 MXU 路径的格式。但这条路被上游堵死了，不是配置问题，所以每次镜像更新都值得重测一下。

一个实测结论：每步 decode 要搬动 ~3.15 GiB，对应 3.94 ms 的带宽下限，实测 8.02 ms——大概 达到峰值带宽的 49%。芯片在本文所有场景下都不是瓶颈。

模型：关于 Gemma 4 E2B 的六件容易踩坑的事

E2B 是个奇怪的 checkpoint。几乎所有常规 decoder 的直觉在这儿都错，而且后面算内存的时候只有把这些都摆出来才能对得上。

字段	值
`num_hidden_layers`	35（28 个滑动窗口 / 7 个全注意力，`i % 5 == 4` 是全注意力）
`num_kv_shared_layers`	20——也就是说只有 15 层自己持有 cache
`num_attention_heads` / `num_key_value_heads`	8 / 1
`head_dim` / `global_head_dim`	256 / 512
`hidden_size` / `intermediate_size`	1536 / 6144
`vocab_size`	262,144（绑定 embedding）
`sliding_window`	512
bf16 精度下占用	8.97 GiB

1. "E2B" 不是 2B 的模型。 它有效参数 ~2B，但总参数 ~5B，占 8.97 GiB。"E" 前缀不是装饰——把 E4B 读成"一个 4B 模型"会把权重低估大概 2 倍，而 2 倍正好是能不能塞进 16 GB 芯片的区别。

2. 有两种注意力几何，不是只有一种。 滑动层用 head_dim 256；7 个全注意力层用 512，而且 K 和 V 也是这样，不只是 Q。用一个 head_dim 描述不了这个模型——任何读一个值然后套到全部 35 层上的逻辑都会把全注意力层低估 2 倍。算出来 KV 容量会差 17%，而且真有人这么算。

3. 35 层里有 20 层在用别人的 cache。 first_shared = 35 − 20 = 15，所以层 0–14 持有 KV，层 15–34 共享。共享规则是"同一注意力类型的前一个层"，所以在 0–14 里，全部 20 个共享层其实只依赖两个 cache 源——滑动层对应层 13，全注意力层对应层 14。20 层，2 个 tensor。

4. KV 每 token 18 KiB，但启动日志会在这方面骗你。

12 个滑动缓存层 x 1 个 KV head x 2 (K,V) x 256 x 2 B = 12,288 B 3 个全注意力缓存层 x 1 个 KV head x 2 (K,V) x 512 x 2 B = 6,144 B 合计 = 18,432 B = 18 KiB/token

乘以实测的 321,344 个驻留 token 得到 5.52 GiB——正好是引擎报告的数字。但描述 cache 的那行日志 regular_attn_shape=(num_blocks, (64, 1, 2, 256))，是从层 0 取的第一个命中样本——而层 0 是滑动注意力，所以是 256。它对层 4、9、14 什么也没说。在任何混合维度模型上它都会报小。从 config 几何算 KV 大小，拿这个数字跟 total_hbm_avail_gb 核对，别直接从那行抄。

5. 一个 KV head 意味着芯片越多越糟糕。 num_key_value_heads = 1 是全 MQA，单个 head 不能分片。运行时会把 num_kv_heads 填充到 tensor-parallel size 的倍数，所以 TP=4 时 KV 内存要乘 4 倍来存同一个 head 的复制。在这个模型上更大的拓扑是乘 KV 成本，不是除。动手前先看 num_key_value_heads。

6. Head 们不能整除 hidden size。 8 x 256 = 2048，但 hidden_size = 1536，所以 Q projection 是矩形的。任何用 head_dim = hidden_size / num_heads 算的代码会得到 192，然后默默出错。

再补充一个能解释性能的细节，不只是内存的：8.97 GiB 里有 4.38 GiB 是每层的 embedding 表（262,144 x 256 x 35），这些是按 token 聚集的，不是流式的。每步 decode 实际只搬 ~3.15 GiB——dense transformer 加上 lm_head 读的那 0.75 GiB 绑定 embedding。这就是为什么一个 8.97 GiB 的模型 decode 能这么快。

还有两个值得知道的坑：如果要去摸 checkpoint，下面的别漏了：文件里还有 audio_tower 和视觉层，它们有自己的独立层编号，所以正则 layers\.(\d+)\. 会默默跟它们冲突——总是锚定到 model.language_model.。另外 QAT 导出的模型（-qat-w4a16-ct、-qat-q4_0-unquantized）在这个技术栈里完全加载不了，部分原因是它们确实没有给 KV 共享层的 k_norm，但 loader 偏偏要。

心智模型：一块 HBM 预算，和一块没人管的区域

其他一切都是从这里推出来的：

 15.75 GiB 芯片上总 HBM
× 0.92 --gpu-memory-utilization
───────── 14.49 GiB 引擎在里面分配的硬上限
− 8.97 GiB 模型权重（bf16）
───────── 5.52 GiB KV 缓存 → 每 token 18 KiB，共 321,376 tokens 1.26 GiB 上限之外剩下的——编译好的 XLA 程序存在这里， gpu_memory_utilization 管不到它们

最后这行就是坑。我试了 --gpu-memory-utilization 0.95，本以为能免费多拿 8% 的 KV。KV 池大小算出来完全对——上限 14.96 GiB，KV 5.99 GiB，348,864 tokens，+8.6%——然后 XLA 在 第 691 秒 崩了，load jit_structured_decode_fn 的时候：

RuntimeProgramAllocationFailure: Attempting to reserve 384.11M at the bottom of memory.
That was not possible. There are 347.33M free, 0B reserved, and 347.33M reservable.

这个参数只管权重 + KV 缓存；编译好的程序镜像从它剩下的那部分里出。我测了两次，间隔几个月，page size 不同（32 和 64），都是同一个失败——要 384.11 M，剩 346.77 M 和 347.33 M。确定性的，不是竞态。

0.92 是上限，不是保守默认值。 而且注意那个失败的形状：发现它要付一次完整编译的代价——691 秒之后才告诉你不行。

`--gpu-memory-utilization 0.92`

引擎可以为权重 + KV 缓存分配的 HBM 比例。不管激活值，不管编译好的程序。0.95 在这个模型/芯片/构建版本上跑不起来。0.93 和 0.94 没测；估算下来大概差 291 MB 和 127 MB，而失败那次只差 37 MB，风险收益比很差。

还有一个参数 --kv-cache-memory（配置字段 kv_cache_memory_bytes），它是用字节数固定池大小，而不是比例，而且后续启动会跳过内存 profiling。本构建实测好的值：5923602432。

`--max-model-len 32768`

最大上下文长度。出乎意料的是它不占 KV 容量：

`max-model-len`	`block_size`	每个请求的 block 数	KV block 数	KV tokens
16,384	32	512	10,043	321,376
32,768	64	512	5,021	321,344

Pallas backend 会让 block_size 跟着 max_model_len 等比放大，保持每个请求的 block 数不变。

上下文翻倍则 page size 翻倍，block 数量减半，最终落在同样的 token 容量上。实测，两种配置都跑过。

而且每个请求的 block 数恰好能预测 decode 速度：

每个请求的 block 数	c=1 TPOT
512（`max-model-len` 16384）	8.05 ms
512（`max-model-len` 32768）	8.02 ms
2048（配置错了的一个测试分支）	8. Powered By Marico's space \| 晋ICP备2022011647号 \| 技术合作: look@marico.cc TEL: +86 18660050334