单块TPU v5e实测：哪些小模型能与Gemma 4 E2B一战？

在单块 Google Cloud TPU v5e 芯片上，用 Apache-2.0 和 MIT 许可的小模型替换 Gemma 4 E2B 是否可行？这是一个对成本敏感的开发者和研究者都关心的问题——v5litepod-1 是 Google Cloud 最便宜的 TPU 实例，15.75 GiB HBM 中可用约 14 GiB 给模型和 KV 缓存。

近期一项系统性实测（代码已开源在 GitHub 的 model-sweep 仓库）在完全相同的 vLLM 镜像和配置下，对十余个小模型进行了速度与准确率的对比。所有候选模型的许可证均为 Apache-2.0 或 MIT，意味着它们可以被自由重新打包和分发。

测试设计：为什么选单块 v5e 芯片？

v5litepod-1 的 15.75 GiB HBM 中，vLLM 只取用 0.9，留出约 14 GiB 给模型权重和 KV 缓存。这直接将候选模型限制在 ~4B 参数以内。Gemma 4 E2B 在此芯片上占 8.97 GiB，已经是阵容中较大的模型。

模型筛选由 survey.py 自动完成：从 Hugging Face API 读取每个候选模型的许可证、架构和权重体积，许可证必须是 Apache-2.0 或 MIT，且至少留有 2 GiB 给 KV 缓存。筛选再按 vLLM TPU 的支持方式分为两层——原生 JAX 模型和 PyTorch 回退（torchax）。最终 Tier 1 有 8 个模型，Tier 2 有 7 个。

Qwen2.5-3B 因许可证不符（标记为 other）被排除，Qwen3-8B 因内存超出被排除。说个细节，DeepSeek^广告-R1-Distill-Qwen-1.5B 和 OLMo 2 1B 虽然符合许可证要求，却在此镜像上无法加载。

TPU v5e 与 v5litepod-1 的背景：TPU v5e 是 Google 第五代张量处理单元中面向推理与中等规模训练的版本，单芯片提供约 15.75 GiB 的 HBM（高带宽内存）。v5litepod-1 是 Google Cloud 上最小的 TPU 切片实例，仅含一块 v5e 芯片，因此是成本敏感场景的首选。

vLLM 的 gpu_memory_utilization 参数：vLLM 通过 --gpu-memory-utilization（在 TPU 上同样适用）控制单实例可占用的显存比例，默认 0.9。文章提到的“vLLM 只取用 0.9”即指该值。它预留了约 10% 内存用于激活、临时缓冲等运行时开销，剩余约 14 GiB 可供模型权重与 KV 缓存分配。

KV 缓存与 max-num-seqs：KV 缓存（Key/Value Cache）是自回归推理时缓存历史 token 的注意力键值，其大小决定了可同时处理的序列数和上下文长度。--max-num-seqs 16 限制了 vLLM 每批最多调度 16 个并发序列，配合 --max-model-len 4096 也框定了缓存的上限。

测试方法：统一镜像与代码评分

两个 Tier 都在同一镜像上运行：vLLM TPU sha256:19a1a052...，vLLM 版本 0.29.1rc1.dev468+g0b7f11a1e，确保 Tier 2 的模型能与 Tier 1 的 Gemma 结果配对比较。

每个模型以 --max-model-len 4096 --max-num-seqs 16 启动，开启前缀缓存，通过三项标准测量：

• 输出速度：每次请求 256 个输出 token，在 1、4、16 并发下各测三轮取中位数
• 长提示：使用 1,024 和 3,072 token 的不重复提示，使前缀缓存无法发挥作用
• 准确率：全部 1,319 道 GSM8K 题（贪心解码，2,048 token 答案上限），以及 400 条 BFCL 简单工具调用记录

所有答案由代码评分，与 Gemma 的差异在相同问题上配对比较，使用 10,000 次重采样得出 95% 置信区间。区间包含零意味着没有可测量的差异。

vLLM TPU 的两条执行路径：Tier 1 的“原生 JAX 模型”指模型已用 JAX/Flax 重新实现，可直接在 TPU 的 XLA 编译器下运行；Tier 2 的“PyTorch 回退（torchax）”指模型仍是 PyTorch 权重与代码，通过 torchax（PyTorch/XLA 的桥接层）在 TPU 上执行，通常有额外开销。同一镜像同时支持两条路径，保证了对比的公平性。

前缀缓存（Prefix Caching）：vLLM 的 prefix caching 会复用相同前缀 prompt 的 KV 缓存，避免重复计算。文章特意使用“不重复的长提示”来测试，就是为了禁用这一优化，暴露模型在真实长上下文下的吞吐差距。

GSM8K 与 BFCL：GSM8K 是小学数学应用题数据集，用于衡量多步推理与数学能力；BFCL（Berkeley Function Calling Leaderboard）则评估模型调用外部函数/工具的结构化输出能力。两者结合覆盖了“推理型”与“工具型”两类典型负载。

准确率对决：Qwen3-4B 是唯一同级别选手

在 GSM8K 数学推理上，只有 Qwen3-4B（0.915）与 Gemma 4 E2B（0.911）处于同一水平，差异 +0.4 个百分点，置信区间跨零。其余模型从 1.7B 参数往下，除 Qwen3-1.7B（0.827）外，分数下降 18 到 45 个百分点：Granite 3.3 2B 为 0.817（-9.5），Qwen2.5-1.5B 为 0.729（-18.2），Qwen3-0.6B 仅 0.599（-31.2），SmolLM2-1.7B 更只有 0.461（-45.0）。

Phi-4-mini 在 GSM8K 上得 0.895，比 Gemma 低 1.6 个百分点，置信区间排除零，说明有微小的真实差距。

工具调用（BFCL）方面情况完全不同——五个模型与 Gemma 处于同一水平，最下探到 Qwen2.5-1.5B。SmolLM3-3B 甚至以 0.945 略高于 Gemma 的 0.927，但差异不显著。Qwen3-4B（0.940）、Qwen3-1.7B（0.935）、Qwen2.5-1.5B（0.910）都紧随其后。Granite 3.3 2B 为 0.920，比 Gemma 低 0.7 个百分点，同样在置信区间内。

这意味着，如果你的应用主要是工具调用而非数学推理，选择空间要大得多。

速度表现：Gemma 为何能后来居上？

在 16 并发、256 token 输出的标准测试中，Gemma 4 E2B 以 1887 tokens/秒的成绩位居中游偏上。比它快的只有几个小模型：Qwen2.5-0.5B（4176，2.21x）、Qwen3-0.6B（3210，1.70x）、Qwen2.5-1.5B（2780，1.47x）和 Qwen3-1.7B（2049，1.09x）。

但真正拉开差距的是长提示场景。当 16 个请求各携带 3,072 token 的提示时，Gemma 仍保持 1049 tokens/秒，而：

• Granite 3.3 2B 降至 732（0.70x）
• SmolLM3-3B 降至 725（0.69x）
• Phi-4-mini 暴跌至 404（0.39x）
• Qwen3-4B 更是只有 373（0.36x）

Gemma 拥有全表最大的权重（8.97 GiB），却仍能胜过所有超过 4 GiB 权重的模型。根本原因在于 KV 缓存设计。

架构密码：Gemma 的 KV 缓存为何如此高效

从各模型的启动日志中可以提取 KV 缓存容量数据，差距惊人：

模型	KV 头数	头尺寸	层数	KV 缓存 token 数
Gemma 4 E2B	1	256/512	35（20 层共享）	321,280
Granite 3.3 2B	8	64	40	125,184
SmolLM3-3B	4	128	36	120,320
Qwen3-1.7B	8	128	28	100,096
Qwen3-4B	8	128	36	50,688

Gemma 每层只存储一个 KV 头，其 35 层中有 28 层仅保留 512 token 的滑动窗口，20 层复用更早层的缓存。这让它的缓存容量达到下方四个模型的 2.6 到 6.3 倍。

架构上的"怪异"设计——双头尺寸、单 KV 头、层间 KV 共享——恰恰是它长提示速度领先的原因。这也解释了为什么 Qwen3-4B 虽然 GSM8K 和 BFCL 都与 Gemma 持平，长提示吞吐却只有 Gemma 的 36%。

多头注意力与 KV 头的区别：标准多头注意力（MHA）中每个注意力头都有独立的 Key 和 Value 投影，因此“KV 头数”等于“查询头数”。Gemma 4 E2B 每层只存 1 个 KV 头，意味着所有查询头共享同一份键值，这属于分组查询注意力（GQA）的极端形式——多查询注意力（MQA），可显著缩小 KV 缓存体积。

滑动窗口注意力与层间共享：滑动窗口注意力（Sliding Window Attention）让每个 token 只关注最近固定长度（如 512 token）的上下文，超出部分无需保留 KV。Gemma 35 层中 28 层使用 512 token 窗口，并让 20 层复用更早层的缓存，因此实际缓存只需存储被窗口限制后的键值，总容量自然远高于同等权重规模的标准模型。

双头尺寸（256/512）：日志中的“头尺寸 256/512”说明 Gemma 的注意力头在不同维度上采用了非对称配置，这也是其非常规架构的一部分，论文中常被归类为“混合注意力”设计。

实践建议：根据场景选择

基于以上数据，选择逻辑可以归纳为：

追求单芯片最佳综合能力：Gemma 4 E2B。没有其他模型能同时匹配它的准确率和速度。

需要标准注意力架构且准确率不妥协：Qwen3-4B。GSM8K 和 BFCL 均与 Gemma 持平，但输出速度只有 0.57x，长提示下更是降至 0.36x。

以工具调用为主、需要最简洁架构：SmolLM3-3B。BFCL 与 Gemma 持平，数学低 5.3 分，速度约 0.76x，架构是标准的多头注意力。

工具调用优先、对速度敏感：Granite 3.3 2B 或 Qwen3-1.7B。前者 GSM8K 低 9.5 分但速度接近 Gemma（0.93x），后者数学低 8.4 分但速度略超 Gemma（1.09x）。

极致速度、可接受准确率损失：Qwen2.5-0.5B 或 Qwen3-0.6B，速度是 Gemma 的 1.7-2.2 倍，但 GSM8K 低 31-44 分。

值得注意的工程细节

两个 Tip 对实际部署很有价值。

Qwen3 和 SmolLM3 默认开启思考模式，会在 GSM8K 的答案预算中消耗推理文本。通过每次请求传递 {"chat_template_kwargs": {"enable_thinking": false}} 可以在不更换 checkpoint 的情况下关闭思考。测试中 Qwen3-4B 因此在 2,048 token 限制内没有留下 GSM8K 答案，SmolLM3-3B 则留下 16/1319 条未答。

DeepSeek-R1-Distill-Qwen-1.5B 在此镜像上加载失败，报错为 'Qwen2Config' object has no attribute 'text_config'。原因出在 tpu_inference/models/jax/qwen2.py 中一行代码：getattr(hf_config, 'hidden_size', hf_config.text_config.hidden_size)。Python 会先求值默认参数，所以任何 tie_word_embeddings: false 的 Qwen2 checkpoint 都会崩溃。Qwen2.5 Instruct 系列因绑定词嵌入而幸免。

Python getattr 的默认值求值问题：getattr(obj, 'attr', default) 中的 default 参数在调用前就会求值，因此即使 hf_config 已有 hidden_size，hf_config.text_config.hidden_size 仍会被执行，若 text_config 不存在则直接抛 AttributeError。这是 Python 中常见的“默认值陷阱”，修复方式是改用 getattr(hf_config, 'hidden_size', None) or hf_config.text_config.hidden_size 或先判断属性存在。

tie_word_embeddings 是什么：该配置项决定输入词嵌入矩阵与输出投影矩阵是否共享权重。设为 true（绑定）可节省参数并常见于小模型；设为 false（解绑）则两者独立，部分 Qwen2 架构在此配置下会走 text_config 分支，从而触发上述崩溃。Qwen2.5 Instruct 系列因默认绑定词嵌入，绕过了这一分支，因此能在该 TPU 镜像上正常加载。

总结

要在单块 TPU v5e 上替换 Gemma 4 E2B，Qwen3-4B 是唯一在准确率上完全对等的选择，但代价是速度减半。如果应用场景以工具调用为主，SmolLM3-3B、Granite 3.3 2B 和 Qwen3-1.7B 都是速度与准确率平衡不错的替代方案。Gemma 4 E2B 凭借其独特的 KV 缓存设计，在长提示场景下建立了难以逾越的优势——这提醒我们，模型架构的"非标准"之处有时正是其工程价值的来源。