# Gemma 4 E2B in JAX on One TPU v5e: bf16 or 4-Bit, Python or Rust, and How It Stands Against vLLM > Source: > Published: 2026-11-04 22:49:08+00:00 This article provides a step by step guide to serving Gemma 4 E2B on a single Cloud TPU v5e chip with JAX. A pure-JAX engine serves the model from Python, a second engine builds the same graph in Rust and hands it to the libtpu compiler that JAX itself uses, and vLLM serves the same checkpoint as the reference. Every number below was measured on one flex-start v5e-1 VM on 2026-11-03, and every run is committed with its logs. [https://github.com/xbill9/jax-dev](https://github.com/xbill9/jax-dev) A v5e-1 is the smallest Cloud TPU you can rent: one chip, 16.909 GB of HBM, in a `ct5lp-hightpu-1t` VM. Gemma 4 E2B in bf16 puts 9.257 GB of weights on it, which leaves room for the KV cache with no quantization at all. That makes it the place to ask three questions with nothing else in the way: `TPUV5sLitepodPerProjectPerZoneForTPUAPI`) in `us-west4-a` `gcloud` with the alpha components`tpu-jax-v5e1-2b` directory as the working directory No Hugging Face token is needed: `google/gemma-4-E2B-it` and both 4-bit builds are public. Each rig ships an MCP server, `server.py`, whose tools wrap the `gcloud` calls. The rig's `tpu.env` sets `ACCELERATOR_TYPE=v5litepod-1`, and the server derives the machine type from it. `workload="jax"` installs Python 3.13 and `jax[tpu]` on a bare VM with no Docker: ``` python $ python3 -c 'import asyncio, server as s; print(asyncio.run(s.create_tpu_vm_instance(instance_name="jax-v5e1-2b", zone="us-west4-a", accelerator="v5litepod-1", workload="jax", max_run_duration="4h", request_valid_for="30m")))' Executing gcloud command: gcloud compute instances create jax-v5e1-2b --project= --zone=us-west4-a --machine-type=ct5lp-hightpu-1t --provisioning-model=FLEX_START --request-valid-for-duration=30m --max-run-duration=4h --instance-termination-action=DELETE ... --boot-disk-size=200GB --scopes=cloud-platform 🚀 Flex-start TPU VM `jax-v5e1-2b` (ct5lp-hightpu-1t, 1 chip(s)) created in us-west4-a; installing Python 3.13 + `jax[tpu]`. NAME ZONE MACHINE_TYPE PREEMPTIBLE STATUS jax-v5e1-2b us-west4-a ct5lp-hightpu-1t RUNNING ``` `wait_for_jax_ready` watches the serial console until the startup script has seen a TPU device from JAX: ``` 🟢 JAX environment ready on `jax-v5e1-2b`. Verify with `verify_jax_tpu`; run workloads over SSH with `python3.13`. $ python3.13 -c "import jax; print(jax.__version__, jax.devices())" 0.12.0.1 [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0)] ``` The flex-start VM deletes itself after `--max-run-duration`, so a forgotten VM stops billing on its own. `jax_engine.py` is a Gemma 4 E2B decoder written directly in JAX: jitted prefill, a static KV cache, and one jitted decode step per token. It reads the Hugging Face safetensors itself, so no PyTorch is installed. `--quant-mode fp16` is its dense path, which keeps the bf16 weights as they are: ``` bash $ python3.13 jax_openai_server.py --model google/gemma-4-E2B-it --quant-mode fp16 --max-model-len 2048 JAX devices: [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0)] Loading tokenizer: google/gemma-4-E2B-it Loaded 9.26 GB of parameters on TPU_0(process=0,(0,0,0,0)) in 5.4s (KV cache: bf16) INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit) ``` The server speaks the OpenAI chat API, streamed or not, so any OpenAI client can talk to it. Two 4-bit builds of E2B load into the same engine with `--quant-mode w4a16`: Google's `google/gemma-4-E2B-it-qat-w4a16-ct`, and `xbill9/gemma-4-E2B-it-qat-q4_0-w4a16-ct-text`, a repack that stores the QAT grid exactly. `tools/ab_bench.py` runs each in its own process on 8 prompts, 128 greedy tokens each, one untimed run and three timed: ``` loaded google/gemma-4-E2B-it in 64.7s, weights 9.257 GB, HBM 9.260 / 16.909 GB loaded xbill9/gemma-4-E2B-it-qat-q4_0-w4a16-ct-text in 39.2s, weights 6.559 GB, HBM 6.561 / 16.909 GB loaded google/gemma-4-E2B-it-qat-w4a16-ct in 51.9s, weights 6.559 GB, HBM 6.562 / 16.909 GB arm tok/s x ref weights GB HBM GB google/gemma-4-E2B-it (fp16) 132.6 1.000 9.257 9.260 xbill9/gemma-4-E2B-it-qat-q4_0-w4a16-ct-text (w4a16) 103.0 0.777 6.559 6.561 google/gemma-4-E2B-it-qat-w4a16-ct (w4a16) 103.1 0.777 6.559 6.562 ``` bf16 decodes 1.29x faster than either 4-bit build and prefills in 7.4 ms against 23.2 ms. The 4-bit builds save 2.70 GB of HBM. This engine unpacks the int4 weights inside every matmul, and on a chip where bf16 already fits, that unpacking costs more time than the smaller weights save. The repack and Google's export tie on speed and size. Their outputs differ: none of the 8 completions match token for token, and they agree for 26.6 of 128 tokens on average. Which of the two is closer to the trained QAT weights is measured in the earlier article listed under References. JAX traces Python into StableHLO and hands it to the TPU's PJRT plugin, `libtpu.so`, to compile and run. The `tpu-jaxrust-v5e1-2b` rig writes that StableHLO directly. `rust/src/gemma.rs` builds the forward pass op by op, following `jax_e_model.py` down to where JAX promotes bf16 to f32: ``` php fn rms_norm(b: &mut Builder, x: &V, w: Option<&V>, eps: f64) -> V { let dt = x.dt; let last = x.rank() as i64 - 1; let d = x.dim(-1) as f64; let xf = b.convert(x, DType::F32); let sq = b.mul(&xf, &xf); let sum = b.reduce_sum(&sq, &[last]); let var = b.scale(&sum, 1.0 / d); let var = b.add_scalar(&var, eps); let r = b.rsqrt(&var); ... ``` `rust/src/hlo.rs` prints the module as MLIR text, and `rust/src/pjrt.c`, a small C shim over the vendored PJRT C API header, opens the same `libtpu.so` that JAX loaded and compiles it. A PyO3 module, `jaxrust_rs`, hands the result to Python for the tokenizer: ``` platform {'plugin': '/usr/local/lib/python3.13/dist-packages/libtpu/libtpu.so', 'header_api': (0, 117), 'plugin_api': (0, 116), 'devices': 1} build {'rustc': 'rustc 1.101.0-nightly (1c3a65069 2026-11-02)', 'pjrt_header': '0.117', 'crate': '0.1.0'} engine Engine(context=256, weight_bytes=9257336832, cache_bytes=4718592) in 8.4s, {'load_s': 2.318838774, 'compile_s': 5.73853671} ``` The weights upload in 2.3 s and the decode program compiles in 5.7 s. There is no JAX in this path, and the Rust crate has no ML dependencies: `pyo3`, `safetensors`, `memmap2` and `serde_json`. `tools/parity.py` runs each engine in its own process on the same 8 prompts, with the same chat template and the same leading ``, and compares the last-position logits after the first prompt and every greedy token: ``` logits after prompt 0: max |diff| 0.2306, mean |diff| 0.07346, argmax 7216 vs 7216, top-10 overlap 10/10, KL(ref||got) 1.31e-06 prompts tokenized the same: 8 of 8; completions identical: 7 of 8; matching prefix per prompt: [128, 128, 128, 128, 128, 128, 42, 128] (mean 117.2 of 128) decode tok/s (median over prompts): jax 131.6, rust 150.7, ratio 1.145 weights on device: jax 9.257 GB, rust 9.257 GB ``` Both engines pick the same next token with the same top 10, 1.31e-6 apart in KL divergence. Seven of eight 128-token completions match token for token; the eighth parts at token 42, where "Time until 4:00 pm" and "Next full hour: 4:00 pm" were nearly tied. Timed inside each engine, Rust decodes 1.145x faster on identical weights. The Rust engine gets an OpenAI-style server of its own, `rust_openai_server.py`, and vLLM 0.33 runs from the `vllm/vllm-tpu` image on the same VM. Each server takes port 8000 alone, and `tools/serve_bench.py` times all three the same way: streamed chat, one request at a time, temperature 0, exactly 128 tokens (`ignore_eos`), context 2048, 8 prompts with one untimed and three timed requests each: ``` bash $ python3 tools/serve_bench.py compare vllm.json jax.json rust.json engine decode tok/s range first token ms tokens same text as vllm vllm 148.7 148.7-148.7 13.0 128 8 of 8 jax 126.7 126.0-128.4 12.5 128 6 of 8 rust 148.4 147.8-148.6 8.9 128 6 of 8 prompt tokens per engine: {'vllm': [20, 23, 23, 21, 20, 24, 35, 17], 'jax': [20, 23, 23, 21, 20, 24, 35, 17], 'rust': [20, 23, 23, 21, 20, 24, 35, 17]} repeats identical: {'vllm': 8, 'jax': 8, 'rust': 8} ratio jax / vllm: 0.852 ratio rust / vllm: 0.998 ``` The Rust engine decodes at vLLM's rate, 148.4 against 148.7 tokens per second, and reaches the first token in 8.9 ms against 13.0. The Python JAX engine decodes at 0.852x vLLM with a first token level with it. All three tokenize every prompt to the same length, and each repeats its own text exactly. The JAX and Rust engines part from vLLM at the same two near-ties, in prompts 2 and 6. Engines report their own speed in their own way: inside the decode loop, per HTTP response, with or without the first token. One client gives every engine the same definition. In `serve_bench.py`, decode speed is completion tokens minus one, over the time from the first text chunk to the finish chunk, with the token count read from the server's usage chunk. The clock stops at the finish chunk because, past the end of the answer, `ignore_eos` tokens can decode to empty text. The same Rust engine reads 150.7 tokens per second timed inside the engine at context 256, and 148.4 through this client at context 2048. PJRT's compile call takes a serialized `xla.CompileOptionsProto`. With empty options, libtpu asks for zero cores and refuses to compile ("Requested total cores: 0"). Six bytes set one replica and one partition: ``` /// `xla.CompileOptionsProto { executable_build_options { num_replicas: 1, num_partitions: 1 } }` /// serialized by hand: field 3 (length 4) holding fields 4 and 5, varint 1 each. const COMPILE_OPTIONS: [u8; 6] = [0x1a, 0x04, 0x20, 0x01, 0x28, 0x01]; ``` Two rules of the MLIR text format surface at the same compile call: a float literal needs a decimal point (`2.0e0`), and an empty integer array is written `array`. | One v5e chip, E2B | JAX engine, bf16 | JAX engine, 4-bit | Rust engine, bf16 | vLLM 0.33, bf16 | |---|---|---|---|---| | Decode tok/s, timed in the engine | 🥈 131.6 | 🥉 103.0 | 🥇 150.7 | — | | Decode tok/s, one client | 🥉 126.7 | not run | 🥈 148.4 | 🥇 148.7 | | First token, one client | 12.5 ms | not run | 🥇 8.9 ms | 13.0 ms | | Weights on the chip | 9.257 GB | 🥇 6.559 GB | 9.257 GB | — | | Same text as vLLM | 6 of 8 | not run | 6 of 8 | — | | Concurrent requests | no | no | no | yes | For one user at a time, the Rust engine: it decodes at vLLM's rate, answers first, and puts the same 9.257 GB on the chip. For JAX work in Python, the bf16 JAX engine: 0.852x vLLM, with the whole model written in JAX you can read and change. For more than one user, vLLM, which batches requests and is the only one of the three that does. On this chip keep E2B in bf16: the 4-bit builds save 2.70 GB and decode at 0.777x its speed in this engine. The goal of this article was to measure Gemma 4 E2B in JAX on one TPU v5e chip: bf16 against 4-bit, Python against Rust, and both against vLLM. The key to the solution was holding everything else fixed: one VM, one checkpoint revision, the same 8 prompts through each server's own chat template, and one client timing every engine. The results were: Scope: one flex-start `ct5lp-hightpu-1t` VM in `us-west4-a` (one TPU v5 lite chip, 16.909 GB HBM) on 2026-11-03; JAX 0.12.0.1 and libtpu 0.0.51 for both JAX-side engines, vLLM 0.33.0rc1.dev430 from `vllm/vllm-tpu@sha256:f1f925f5…`; `google/gemma-4-E2B-it` at revision `4adf1574`. Every figure is the median over 8 prompts of the median of 3 timed runs, greedy, 128 tokens, one request at a time. The in-engine runs used context 4096 (bf16 against 4-bit: 132.6 tok/s for bf16) and 256 (JAX against Rust: 131.6 for bf16); the one-client run used 2048 for all three. Parts of the analysis and writing were done with AI assistance (Claude); every figure comes from the committed run files. The strategy for serving Gemma 4 E2B with JAX on a single TPU v5e chip was validated with an incremental step by step approach.