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

A developer repacked Google's quantization-aware-trained Gemma 4 weights for vLLM and served them on a single Google Cloud TPU v5e chip, publishing the scripts and per-record results on GitHub. The int8 builds (w8a8) run faster than bf16 on the chip's native int8×int8 path, with the 12B model fitting in 11.31 GiB of weights and reaching 675 output tokens per second at 16 concurrent requests, 0.964 on GSM8K and 0.955 on BFCL tool calling. The repacked 4-bit builds match bf16 on a 3,880-record classification suite through 12B and score up to 2.4 points above Google's own 4-bit exports at the same speed.

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 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: | Build | What is stored | |---|---| | q4w4a16 | int4 weights holding the QAT grid exactly, 16-bit activations | | q4w4a16emb4 | the same, with the vocabulary tables embed tokens , lm head , per-layer embeddings also int4 | | w8a8 | int8 weights per channel from the QAT values, int8 activations per token | | w8a8emb4 | the same, with int4 vocabulary tables | 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 in hf-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: | Build | Suite vs bf16 | Output tok/s | |---|---|---| | E2B bf16 | 0.683 | 144 / 560 / 2,008 | | E2B w8a8 | 0.686 +0.3, −0.7 to +1.3 | 220 / 841 / 2,872 | | E2B w8a8emb4 | 0.677 −0.6, −1.5 to +0.4 | 243 / 923 / 3,086 | | E4B q4w4a16 | 0.729 −0.2, −0.8 to +0.5 | 75 / 291 / 1,012 | | E4B w8a8emb4 | 0.728 −0.3, −1.0 to +0.5 | 133 / 509 / 1,747 | | 12B q4w4a16emb4 | 0.762 +0.2, −0.4 to +0.8 | 35 / 127 / 407 | | 12B w8a8emb4 | 0.761 +0.1, −0.6 to +0.8 | 57 / 219 / 675 | | 26B q4w4a16 | 0.754 −1.1, −1.8 to −0.3 | 31 / 102 / 201 | 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: | Size | Google export | Repack | Difference | |---|---|---|---| | E2B | 0.655 | 0.678 | +2.4 +1.4 to +3.4 | | E4B | 0.716 | 0.729 | +1.3 +0.7 to +2.0 | | 12B | 0.752 | 0.758 | +0.6 0.0 to +1.2 | 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. | Build | GSM8K | BFCL | |---|---|---| | 12B w8a8emb4 | 0.964 | 0.955 | | 12B q4w4a16emb4 | 0.958 | 0.948 | | E4B w8a8emb4 | 0.940 | 0.912 | | E4B q4w4a16 | 0.934 | 0.910 | | E2B q4w4a16 | 0.901 | 0.920 | | E2B w8a8emb4 | 0.895 | 0.920 | | E2B w8a8 | 0.889 | 0.915 | | E2B bf16 | 0.910 | 0.928 | 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 | Build | Output tok/s at 16 requests | Median first token | |---|---|---| | E2B bf16 | 906 | 1.10 s | | E2B w8a8emb4 | 1,249 | 0.82 s | | E4B w8a8emb4 | 649 | 1.43 s | | 12B w8a8emb4 | 95 | 22.0 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 | 12B w8a8emb4 , 16 requests | bf16 cache | fp8 cache | |---|---|---| | ~1,000-token prompts | 244 tok/s | 330 tok/s | | ~3,600-token prompts | 95 tok/s | 157 tok/s | | First token, ~3,600-token prompts | 22.0 s | 13.3 s | | Suite / GSM8K / BFCL | 0.761 / 0.957 / 0.953 | 0.754 / 0.956 / 0.945 | A single request runs at the same speed with either cache. The fp8 cache costs 0.7 suite points and holds GSM8K and BFCL. | Size | Build | Suite | GSM8K | BFCL | tok/s at 16 | |---|---|---|---|---|---| | 12B | 🥇 w8a8emb4 | 0.761 | 0.964 | 0.955 | 675 | | 12B | q4w4a16emb4 | 0.762 | 0.958 | 0.948 | 407 | | E4B | 🥇 w8a8emb4 | 0.728 | 0.940 | 0.912 | 1,747 | | E4B | q4w4a16 | 0.729 | 0.934 | 0.910 | 1,012 | | E2B | 🥇 w8a8 | 0.686 | 0.889 | 0.915 | 2,872 | | E2B | q4w4a16 | 0.678 | 0.901 | 0.920 | 1,906 | 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.