Gemma 4 E2B on a Single TPU v6e Chip: A Serving Deep Dive

Google's Gemma 4 E2B model serves efficiently on a single TPU v6e chip, achieving 213 tok/s for a single user and scaling to ~2,200 output tok/s across concurrent streams, while its QAT variants fail to load due to a loader bug. A developer found that the plain checkpoint works, but the int4 and bf16 QAT exports hit unimplemented quantization paths or missing weights, filed as issue #3225.

Measured 2026-07-21 on vllm/vllm-tpu:nightly vLLM 0.23.1rc1.dev1076 , a GCE flex-start ct6e-standard-1t one TPU v6e chip, 32 GB HBM in europe-west4-a. The plain google/gemma-4-E2B-it serves well on one v6e chip; none of its QAT siblings load at all. The 2-billion-parameter "efficient" Gemma 4 sustains 213 tok/s for a single user with a 16 ms first token, scales to ~2,200 output tok/s across concurrent streams, handles OpenAI-style function calling — including parallel calls and refusal to hallucinate calls — without a miss, and answers simple vision questions accurately in ~200 ms. The QAT variants are a different story: the int4 compressed-tensors export hits an unimplemented quantization path, and the bf16 QAT export trips a loader bug — the Gemma 4 implementation demands per-layer norms that E2B's KV-sharing architecture legitimately doesn't have. Filed upstream as tpu-inference 3225 https://github.com/vllm-project/tpu-inference/issues/3225 . One capability coupling to know about: with a reasoning parser configured, schema enforcement only engages when thinking is enabled — thinking-off requests sail through unconstrained with a 200 status. Config interaction, not TPU limitation; details below. The host is a GCE flex-start VM — capacity granted on request, billed until deleted, hard-stopped at a 4-hour max run, $1.35/chip-hour. A startup script installs Docker, pulls vllm/vllm-tpu:nightly , fetches the Hugging Face token from Secret Manager via the metadata server, and launches vLLM. Boot timeline: VM RUNNING at t+0 200 GB boot disk — the 10 GB default cannot hold the vLLM image → Docker installed ~t+1:00 → image pulled ~t+6:00 → weights downloaded, XLA compiled, health green ~t+8:30. Two environment quirks worth knowing: gcloud compute ssh --tunnel-through-iap rides over HTTPS and works; so does tunneling the API port with gcloud compute start-iap-tunnel <vm 8000 .Serving flags: --max-model-len 65536 --gpu-memory-utilization 0.9 --max num batched tokens 4096 --enable-auto-tool-choice --tool-call-parser gemma4 --reasoning-parser gemma4 , bf16 weights, tensor-parallel 1. | Checkpoint / path | Failure | Verdict | |---|---|---| -qat-w4a16-ct · JAX | int4 compressed-tensors scheme unimplemented for E2B's per layer model projection | ✕ no load | -qat-q4 0-unquantized · JAX | k norm.weight "missing" for layers 15–34 | ✕ no load | -qat-q4 0-unquantized · torchax | identical missing-weights error via MODEL IMPL TYPE=vllm | ✕ no load | gemma-4-E2B-it plain · JAX | loads and serves | ✓ serves | The forensics point at the loader, not the checkpoint. Reading the safetensors headers of both repos: the plain export ships self attn.k norm for all 35 layers; the QAT export ships it only for the 15 non-KV-shared layers. Both configs are identical — including num kv shared layers: 20 . Layers 15–34 reuse K/V from lower layers and have no k-norm of their own, so the QAT export is the architecturally honest one; the plain checkpoint only loads because it carries those unused tensors anyway. Proposed fix in 3225 https://github.com/vllm-project/tpu-inference/issues/3225 : skip instantiating K/V-side parameters for KV-shared layers. Until it lands: serve the plain checkpoint. At 2B parameters ~5 GB bf16 against 32 GB HBM , 4-bit weights buy little here anyway — memory pressure lives in the KV cache, which is already fp8. Same workload at every level — 1,024-token prompts, 128-token completions, vllm bench serve , random dataset. | Concurrency | Req/s | Output tok/s | Total tok/s | TTFT med | TTFT p99 | TPOT med | Per-stream tok/s | |---|---|---|---|---|---|---|---| | 1 | 1.64 | 209 | 1,884 | 16 ms | 17 ms | 4.7 ms | 213 | | 8 | 9.44 | 1,209 | 10,878 | 27 ms | 99 ms | 6.2 ms | 161 | | 32 | 12.78 | 1,636 | 14,721 | 155 ms | 189 ms | 17.5 ms | 57 | | 64 | 16.72 | 2,140 | 19,262 | 122 ms | 349 ms | 25.3 ms | 39 | | 100 burst | 17.31 | 2,215 | 19,938 | 833 ms | 1,573 ms | 36.8 ms | 27 | Reading the curve: Single run per configuration — more meaningful here than usual: a kernel study on this same stack, cited in §7, measured run-to-run cv ≤ 0.3% under greedy decoding with static shapes. Served with --tool-call-parser gemma4 --enable-auto-tool-choice , probed with two OpenAI-style tools at temperature 0. Five scenarios, five clean results: | Probe | Behavior | Latency | |---|---|---| | Simple call | correct tool, inferred the optional unit arg from phrasing | 166 ms | | Result synthesis | fed the tool result back → clean natural-language answer | 140 ms | | No-tool restraint | answered directly, no spurious call | 97 ms | | Parallel calls | both tool calls emitted in one turn, correct args each | 150 ms | | Underspecified | asked "What city are you interested in?" instead of hallucinating a call | 44 ms | A 2B model producing well-formed tool calls JSON, choosing correctly between calling and answering, batching parallel calls, and asking for missing arguments — at double-digit-millisecond latency. For high-volume, low-complexity agent steps, the quality floor is higher than the parameter count suggests. | Probe | Observed behavior | Verdict | |---|---|---| json schema , thinking off default | free prose with a 200 status; strict: true , guided json , structured outputs spellings all equally unenforced | ✕ silently skipped | json object , thinking off | fenced code block, array where an object was asked, invented enum value | ± prompt-level | json schema + enable thinking | exact schema conformance — bare JSON object, typed integer, "ASAP" correctly mapped into the high enum | ✓ enforced | | Reasoning, default | no reasoning traces on any prompt | off by default | Reasoning, enable thinking | parser cleanly splits thinking trace from a terse answer; ~2.4× completion tokens | ✓ works | The mechanism: with --reasoning-parser gemma4 configured, vLLM defers grammar enforcement until the reasoning section ends. Thinking off → no reasoning terminator → the grammar never engages, and unconstrained prose ships with a 200 status. Enable thinking "chat template kwargs": {"enable thinking": true} and the same request is enforced exactly. Operational guidance: pair structured output with enable thinking: true under a reasoning parser — or drop --reasoning-parser from servers that don't need it. And keep client-side validation regardless: the silent-skip failure mode means a trusting client can't tell an enforced response from a lucky one. One restart with --limit-mm-per-prompt '{"image":4,"audio":1}' makes it a vision server. COCO validation images, base64 data URIs, temperature 0: | Probe | Answer abridged | Latency | |---|---|---| | Describe two cats | "Two tabby cats… on a bright pink surface… a remote control visible" | 197 ms | | Count + attributes | 2 animals, both cats, remote + blanket identified | 421 ms | | Scene bear | "A bear lying down in a grassy outdoor environment" | 329 ms | | Room inventory | wall-mounted TV, shelving, furniture correctly enumerated | 874 ms | An image costs ~280 prompt tokens and adds almost nothing over a text request. Two notes: server-side fetching of external image URLs proved flaky intermittent 422s — base64 data URIs are the reliable path — and the first multimodal request after boot can 422 while the processor warms; retry once. fp8 vs bf16 KV: six greedy prompts explanation, code, listing, translation, arithmetic, summarization — 889 completion tokens run under the default fp8 e5m2 cache, then re-run after a restart with --kv-cache-dtype bfloat16 . Result: 6 of 6 outputs byte-identical. On this small, greedy sample the compression is genuinely free — take the fp8 cache. Where the 32 GB goes bf16-KV boot, 65,536 max context : usable HBM reports as 31.24 GiB; at 0.9 utilization vLLM works within 28.12 GiB — roughly 5.5–6 GiB weights, 16.3 GiB KV cache 8,713 blocks × 128 tokens × 15 layers × 128 KiB , ~6 GiB workspace. The interesting physics: E2B's KV sharing is directly visible in the allocator — only 15 of 35 layers hold KV tensors , each with a single 256-dim KV head, so a token costs ~15 KiB of KV in bf16 ~7.5 KiB under fp8 . That yields ~1.1 million tokens of resident KV — seventeen full 65K-context conversations on-chip — which is why the sweep saturated on compute , never memory. Also: of the 404-second engine init, 329 seconds is XLA compilation — the dominant term in the ~10-minute cold start. Related work: a kernel-substitution study by Zimbres DOI 10.5281/zenodo.21404069 https://doi.org/10.5281/zenodo.21404069 shows the RPA v3 kernel's decode block-size heuristic costs 27.7–68.7% of large-batch throughput on 27B/31B models on v6e. E2B sits at the low-exposure end of that effect: at our c=64 operating point, attention is ~9% of memory traffic vs ~41% in their regime precisely because of the KV-sharing design above. Testing the override on E2B — smallest model, one chip, KV sharing — is a natural follow-up. Rate verified against Google's published Dynamic Workload Scheduler pricing https://cloud.google.com/products/dws/pricing : $1.35 per chip-hour for v6e flex-start europe-west4, us-east1, us-east5, asia-northeast1 . | Operating point | Output tok/s | $ / M output tokens | |---|---|---| | Saturation burst | 2,215 | $0.17 | | Sweet spot c=64 | 2,140 | $0.18 | | Interactive c=8 | 1,209 | $0.31 | | Single stream | 209 | $1.79 | Serve on the TPU VM docker run --name vllm-gemma4 --privileged --net=host -d \ -v /dev/shm:/dev/shm --shm-size 10gb -e HF HOME=/dev/shm \ -e HF TOKEN=$ gcloud secrets versions access latest --secret=hf-token \ vllm/vllm-tpu:nightly vllm serve google/gemma-4-E2B-it \ --tensor-parallel-size 1 --max-model-len 65536 \ --gpu-memory-utilization 0.9 --max num batched tokens 4096 \ --disable chunked mm input --enable-auto-tool-choice \ --tool-call-parser gemma4 --reasoning-parser gemma4 \ --limit-mm-per-prompt '{"image":4,"audio":1}' {"image":0,"audio":0} for text-only KV comparison: add --kv-cache-dtype bfloat16 default is fp8 e5m2 on v6e Benchmark per concurrency level C vllm bench serve --backend vllm --model google/gemma-4-E2B-it \ --dataset-name random --num-prompts 100 \ --random-input-len 1024 --random-output-len 128 --max-concurrency C Reach the endpoint from a network that blocks direct traffic gcloud compute start-iap-tunnel vllm-gemma4-e2b 8000 \ --local-host-port=localhost:8000 --zone=europe-west4-a Environment: vLLM 0.23.1rc1.dev1076+g5c342876a vllm-tpu:nightly, tpu-inference backend · TPU v6e-1 ct6e-standard-1t, 32 GB HBM, GCE flex-start · bf16 weights, fp8 e5m2 KV cache. Single benchmark run per configuration; treat deltas under ~10% as noise.