Repacked QAT Gemma 4 on One TPU v5e: 12B Serves at 675 Tokens per Second

This article provides a step by step guide to repacking Google's quantization-aware-trained (QAT) Gemma 4 weights for vLLM and serving them on one Google Cloud TPU v5e chip, with every build scored for classification, math, tool calling, throughput and long prompts. Every per-record output, log and script is committed.

On one v5e chip the repacked QAT builds serve every Gemma 4 size from E2B to 26B. Through 12B they read level with bf16 on a 3,880-record classification suite, they score up to 2.4 points above Google's own 4-bit exports at the same speed, and they make 12B the largest model on the chip: 11.31 GiB of weights, 675 output tokens per second at 16 requests, 0.964 on GSM8K and 0.955 on BFCL tool calling.

https://github.com/xbill9/gemma4-dev/tree/main/jev-tpu-v5e1

One TPU v5e chip (v5litepod-1) has 15.75 GiB of HBM. Gemma 4 E4B at bf16 is 14.9 GiB and 12B is 22.4 GiB, so everything above E2B needs 4-bit or 8-bit weights on this chip.

Google trained 4-bit versions of every Gemma 4 size and publishes the trained values as "unquantized" bf16 checkpoints (-qat-q4_0-unquantized). Every group of 32 weights in them already sits on a 16-level grid. A repack stores those values in a format vLLM serves, without re-rounding them:

v5e multiplies int8 by int8 natively, which makes the int8 builds the fast ones on this chip.

us-west4-a, and the gcloud CLI logged inhf-token git clone https://github.com/xbill9/gemma4-dev The 4-bit repack recovers each group's trained step and stores the group as int4, so every value keeps its trained place on the grid:

huggingface-cli download google/gemma-4-12B-it-qat-q4_0-unquantized \
  --local-dir ~/models/gemma-4-12B-it-qat-q4_0-unquantized
python3 ../jev-tpu-31b/repack_q4_0.py repack ~/models/gemma-4-12B-it-qat-q4_0-unquantized \
  ~/models/gemma-4-12B-it-qat-q4_0-w4a16-ct

The int8 builds take the same QAT values to int8 per channel with ../jev-tpu-31b/w8a8_from_qat.py. The published builds are on Hugging Face under xbill9/gemma-4-*-it-qat-*.

Serving them on the TPU backend uses three additions to vLLM's tpu_inference, in jev-tpu-v5e1/patches/: an int8 W8A8 method on the JAX path, int4 embedding tables that stay packed on the chip and unpack only the rows a step reads, and an int4 lm_head.

Each run is a flex-start queued resource that boots, applies the patches to the pinned vLLM image, serves a list of builds one after another and uploads every result:

gcloud alpha compute tpus queued-resources create jev-tpu-v5e1-$RUN --zone us-west4-a \
  --accelerator-type v5litepod-1 --runtime-version v2-alpha-tpuv5-lite \
  --node-id jev-tpu-v5e1-$RUN-node \
  --provisioning-model flex-start --max-run-duration 4h --valid-until-duration 2h \
  --metadata jev-code=<bundle>.tgz,jev-run=$RUN,jev-qr=jev-tpu-v5e1-$RUN \
  --metadata-from-file startup-script=tpu/startup_quant.sh,jev-arms=arms.txt,jev-patches=patches.txt

One line per build names the checkpoint, what to measure and its serving flags:

/work/models/gemma-4-12B-it-qat-w8a8-int8-emb4=12b-w8a8-emb4=read+load+gen=tools,gmu=0.92
READY /work/models/gemma-4-12B-it-qat-w8a8-int8-emb4 after 405s

The suite is 3,880 public records from Bespoke Labs' benchmark set, paired record for record against bf16 (points, 95% range). Throughput is output tokens per second at 1, 4 and 16 parallel requests:

Every build through 12B is level with bf16; the 26B repack reads 1.1 points below it. E4B, 12B and 26B at bf16 do not fit this chip, so their bf16 references ran on v6e.

12B w8a8emb4 holds 11.31 GiB of weights and gives the KV cache the rest:

TPU KV cache size: 9,728 tokens, Maximum concurrency for 4,096 tokens per request: 2.38x

The 26B mixture-of-experts repack serves at 13.58 GiB with a 2,176-token cache.

Google's -qat-w4a16-ct exports re-round every group of the QAT weights. The repack keeps them, and reads the suite higher:

Served one after the other on the same VM, the two run at the same speed:

e2b-google load at concurrency 16: 1910.8 tok/s, range 1884.7 to 1910.8
e2b-repack load at concurrency 16: 1910.2 tok/s, range 1878.3 to 1910.9

At 16 requests: E4B 1,020 against 1,009 and 12B 392 against 388.

gen_eval.py scores two tasks against the served model, in code:

JEV_GEN_MAX_TOKENS=2048 python3 gen_eval.py run http://localhost:8000 "$MODEL" gsm8k out/gsm8k.jsonl
python3 gen_eval.py run http://localhost:8000 "$MODEL" bfcl_simple out/bfcl_simple.jsonl
12b-w8a8-emb4 gsm8k: {"task": "gsm8k", "max_tokens": 2048, "n": 1319, "right": 1272, "accuracy": 0.9644, "errors": 0, "truncated": 4, "wall_s": 584.8}

--tool-call-parser gemma4, right when the call has the right name and accepted arguments.

Tool calling holds at every size: every E2B repack is within 1.3 points of bf16. On GSM8K the E2B repacks trail bf16 by 0.9 to 2.0 points, and the QAT weights themselves, stored at bf16, trail by 1.3 (−2.7 to 0.0), so the repack adds at most 0.8 to what QAT training left. At E2B the 4-bit repack holds math best. E4B and 12B have no bf16 reference for these tasks on this chip.

A third-party int8 build rounded directly from the bf16 release uses the same format and runs at the same speed. Built from the QAT weights instead, int8 reads the suite 1.5 points higher at E2B (+0.5 to +2.6). At 12B it scores 7.2 points higher on BFCL (0.955 against 0.882): the build rounded from bf16 writes string arguments wrapped in extra quote marks in 16 of the 400 records, and the QAT build writes none.

w4a16_client.py loadlong streams requests with prompts of about 3,600 tokens and a unique prefix on each, so none is served from cache:

12b-w8a8-emb4 long 3584 at concurrency 16: 3611 prompt tokens, 94.7 out tok/s, 1430.0 total tok/s, ttft 21.96 s

E2B w8a8emb4 stays 1.38x bf16 with long prompts and answers sooner. At 12B the cache holds about two such requests at once.

--kv-cache-dtype fp8 stores the cache at one byte per value:

TPU KV cache size: 18,944 tokens, Maximum concurrency for 4,096 tokens per request: 4.62x

A single request runs at the same speed with either cache. The fp8 cache costs 0.7 suite points and holds GSM8K and BFCL.

For the most capable model on one v5e chip, 12B w8a8emb4, with --kv-cache-dtype fp8 when many long prompts arrive at once. It has its own serving directory, tpu-vllm-v5e1-12b-w8a8emb4, with an MCP server, and serves there at 57 / 216 / 713 output tokens per second. For E4B, w8a8emb4. For E2B, w8a8 for speed, or the q4w4a16 repack for math closest to bf16.

Each run deletes its own queued resource when its last build finishes:

gcloud alpha compute tpus queued-resources list --zone us-west4-a

The goal of this article was to find what Google's QAT Gemma 4 weights can do on one TPU v5e chip. The key to the solution was repacking the QAT values into int4 and int8 formats without re-rounding them, and three additions to vLLM's TPU backend. The results were:

w8a8emb4 fits at 11.31 GiB and serves 675 output tokens per second, with 0.964 on GSM8K and 0.955 on BFCL Scope: every build ran on one TPU v5e chip (v5litepod-1, flex-start) in us-west4-a, vLLM 0.29.1rc1.dev468+g0b7f11a1e at vllm/vllm-tpu@sha256:19a1a052… with the three patches in jev-tpu-v5e1/patches/, one run per build. The E4B, 12B and 26B bf16 suite references ran on v6e, and E4B and 12B have no bf16 reference for GSM8K or BFCL. Throughput is the median of three passes of 256 output tokens. Pairings are record for record, and a difference counts only when its 95% range excludes zero. The repacked checkpoints are unofficial and derived from Google's release under Apache 2.0. Parts of the analysis and writing were done with AI assistance (Claude); every figure comes from the committed output files.

The strategy for serving Google's QAT Gemma 4 on one TPU v5e chip was validated with an incremental step by step approach.