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

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.

TL;DR

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.

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.

1. Getting to a serving endpoint

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:

• Direct SSH silently times out on some networks even when the VPC allows tcp:22 — the block is upstream of the VPC. IAP tunneling (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.
• vLLM auto-selects an fp8_e5m2 KV cache on v6e — the largest memory consumer is 8-bit before any weight quantization enters the picture.

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.

2. Three ways the QAT checkpoints fail to load

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: 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.

3. What one chip is worth: the concurrency sweep

Same workload at every level — 1,024-token prompts, 128-token completions, vllm bench serve, random dataset.

Reading the curve:

• Prefill is effectively free at low load. A 1,024-token prompt reaches first token in 16 ms — roughly 64K prefill tok/s for a single stream.
• c=8 is nearly free concurrency: six times the throughput of a single stream for +1.5 ms/token; each of 8 users still sees ~160 tok/s.
• The knee is between 32 and 64. Going 32→64 buys +31% throughput for +45% per-token latency; 64→burst buys +3.5% for another +45%.
• Capacity-planning number: run at ≤64 concurrent streams for smooth per-user experience; the ceiling is ~17 req/s at this workload shape.

(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.)

4. Function calling at 2B scale

Served with --tool-call-parser gemma4 --enable-auto-tool-choice, probed with two OpenAI-style tools at temperature 0. Five scenarios, five clean results:

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.

5. Structured output works — but only with thinking on

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.

6. Vision at 2B: accurate and nearly free

One restart with --limit-mm-per-prompt '{"image":4,"audio":1}' makes it a vision server. COCO validation images, base64 data URIs, temperature 0:

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.

7. fp8 KV cache, HBM anatomy, and a related kernel result

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) 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.

8. Cost breakdown

Rate verified against Google's published Dynamic Workload Scheduler pricing: $1.35 per chip-hour for v6e flex-start (europe-west4, us-east1, us-east5, asia-northeast1).

• Breakeven vs. the API: against Gemini 2.5 Flash-Lite at $0.40/M output tokens, self-hosting wins once you sustain ~940 output tok/s — roughly the c=8 operating point held continuously.
• Cold start (~8.5–10.5 min, mostly XLA compile) costs $0.19–0.24 per provisioning.
• A full 4-hour session costs $5.40 and, at saturation, delivers ~30M output tokens — about $12 worth at Flash-Lite prices. Flex-start fits batch bursts, not idle always-on endpoints.

9. Reproduction

# 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.