dev.to | 📄 原文链接 | 2026-11-06 收录

单张 TPU v5e 上的 Gemma 4 E2B：bf16 还是 4bit，Python 还是 Rust

来源：dev.to — 2026-11-04

📋 概述

作者在最小的 Cloud TPU（v5e-1，单芯片 16.9GB HBM）上把 Gemma 4 E2B 用三种方式跑起来：纯 JAX 写的 Python 引擎、把同一张计算图用 Rust 重建后交给 JAX 也在用的 libtpu 编译、以及 vLLM 作为参照，三个服务由同一个客户端计时，数据全部来自 2026-11-03 的一台 flex-start v5e-1，日志随代码一并提交。因为 bf16 权重只占 9.26GB，剩下的空间足够放 KV 缓存，作者要回答三个问题：已经装得下 bf16 的模型量化到 4bit 划不划算、JAX 引擎的时间有多少花在 Python 上、手写 JAX 引擎在同一块芯片上和 vLLM 差多少。结论是 JAX 引擎里 bf16 优于 4bit，Rust 引擎产出的 token 与 JAX 引擎一致，走 HTTP 时解码速度追平 vLLM 且首 token 最快。

🔑 核心要点

bf16 权重只占 9.26GB，16.9GB HBM 剩余空间够放 KV 缓存，所以量化并非必需，正好用来验证 4bit 是否真划算。
三个引擎跑同一个 checkpoint：纯 JAX 的 Python 引擎、用 Rust 重建同一张图交给 libtpu、以及 vLLM 做参照。
纯 JAX 引擎自己读 Hugging Face safetensors，全程不安装 PyTorch，用 jitted prefill 加静态 KV 缓存，每 token 一次 jitted decode。
最终结果：JAX 引擎里 bf16 胜过 4bit，Rust 引擎与 JAX 引擎 token 一致，HTTP 下解码速度与 vLLM 持平、首 token 最快。

💡 金句

“bf16 beats 4-bit in the JAX engine, the Rust engine matches the JAX engine's tokens, and over HTTP it decodes at vLLM's rate with the fastest first token.”
译：在 JAX 引擎里 bf16 胜过 4bit，Rust 引擎产出的 token 与 JAX 引擎一致，走 HTTP 时解码速度追平 vLLM，且首 token 最快。

← 返回 dev.to 首页