DEV Community

Self-hosting a lite agent backend on one TPU: Gemma 4 E2B + vLLM on a v5e-1

Self-hosting a lite agent backend on one TPU chip A single Google Cloud TPU v5e chip - 16 GB of HBM, about $0.58/hour on spot - will serve google/gemma-4-E2B-it under vLLM at 1,496 output tokens/sec aggregate, with 8.02 ms per-token latency at single stream and native tool-calling. That is enough to back a fleet of 8-16 concurrent "lite" agents for roughly $0.107 per million output tokens. This is a build log with numbers. Everything here was measured on the hardware, and the sections that say "I was wrong about this" are the ones worth your time - four of my confident predictions were falsified by the benchmark, and each falsification was more useful than the guess. Setup under test: v5litepod-1 (one v5e chip), us-west4-a , vllm/vllm-tpu:nightly , vLLM 0.26.1rc1.dev125+ga7a204cc6 , tpu-inference JAX backend, google/gemma-4-E2B-it at bf16. Part 1 - Scaffold and run 1.1 Prerequisites gcloud auth login # for gcloud subprocess calls gcloud auth application-default login # ADC, for the Secret Manager client Put your Hugging Face token in Secret Manager rather than in a script or an env file - the TPU VM's startup script is stored as instance metadata, and anything you bake in is readable from the instance: printf '%s' "hf_xxxxxxxxxxxx" | gcloud secrets create hf-token --data-file=- --project=YOUR_PROJECT Zone constraint that will waste your afternoon if you miss it: flex-start v5litepod-1 is only accepted in us-west4-a . europe-west4-a and -b reject it at the API with FLEX_START provisioning model is not supported for accelerator type "v5litepod-1" , regardless of quota. Non-zero quota in a zone tells you nothing - the provisioning model is the blocker. 1.2 Provision the chip Three provisioning models, three different commands. Note v5e is spelled v5litepod to gcloud - "v5e-1" is fine in prose and is never valid in a CLI argument. # Spot - cheapest, preempted with ~30s notice, NO run limit (bills until you delete it) gcloud alpha compute tpus tpu-vm create gemma4-v5e \ --zone=us-west4-a --type=v5litepod --topology=1x1 \ --provisioning-model=spot --version=v2-alpha-tpuv5-lite # On-demand - full price, no preemption, also unbounded gcloud alpha compute tpus tpu-vm create gemma4-v5e \ --zone=us-west4-a --type=v5litepod --topology=1x1 \ --version=v2-alpha-tpuv5-lite Flex-start goes through the Queued Resource API instead, and is the only model that accepts --max-run-duration , i.e. the only one that stops billing on its own: gcloud alpha compute tpus queued-resources create gemma4-qr \ --node-id=gemma4-qr-node --zone=us-west4-a \ --accelerator-type=v5litepod-1 --runtime-version=v2-alpha-tpuv5-lite \ --provisioning-model=flex-start --max-run-duration=4h Verified 2026-08-09: this exact command was run - the QR reached ACTIVE withprovisioningModel: FLEX_START andmaxRunDuration: 14400s , then deleted cleanly. Note there is no dry-run: a create either queues or provisions, andPROVISIONING state cannot be deleted, so you will pay for at least a few minutes if capacity is immediately available.Spot and on-demand have no automatic stop. They bill until preempted or deleted. Set a calendar reminder, or use flex-start. See the cost section - flex-start is only 3.8% more than spot. Spot draws on a separate quota (TPUV5sPreemptibleLitepodPerProjectPerZoneForTPUAPI ), not the standard TPU quota. A zone with plenty of on-demand quota can still refuse spot. 1.3 Get the token onto the node and start the server gcloud compute tpus tpu-vm ssh crashes with ConnectionResetError from some sandboxed environments (it fails inside its own internal API call, while plain gcloud API calls work fine). Direct SSH always works: IP=$(gcloud compute tpus tpu-vm describe gemma4-v5e --zone=us-west4-a \ --format='value(networkEndpoints[0].accessConfig.externalIp)') # Pipe the secret straight in - never through a shell variable or a log line gcloud secrets versions access latest --secret=hf-token \ | ssh -i ~/.ssh/google_compute_engine xbill@$IP 'umask 077; cat > ~/.hf_token' ssh -i ~/.ssh/google_compute_engine xbill@$IP 'sudo docker pull vllm/vllm-tpu:nightly' Then start it. This is the configuration the rest of the article defends: sudo docker run -d --name vllm-gemma4 --privileged --net=host \ -v /dev/shm:/dev/shm --shm-size 10gb \ -v ~/.cache/vllm:/root/.cache/vllm \ -e HF_HOME=/dev/shm -e HF_TOKEN="$(cat ~/.hf_token)" \ vllm/vllm-tpu:nightly \ vllm serve google/gemma-4-E2B-it \ --dtype bfloat16 \ --kv-cache-dtype auto \ --max-model-len 32768 \ --max-num-batched-tokens 4096 \ --tensor-parallel-size 1 \ --gpu-memory-utilization 0.92 \ --enable-prefix-caching \ --disable-chunked-mm-input \ --limit-mm-per-prompt '{"image":4,"audio":1}' \ --enable-auto-tool-choice --tool-call-parser gemma4 --reasoning-parser gemma4 That -v ~/.cache/vllm:/root/.cache/vllm line is the highest-value thing in this article. The JAX compile cache lives there (197 MB measured) and is otherwise container-local, destroyed on every docker rm . Compilation is 685 s of the 857 s cold start. Measured: mounting it cuts a restart to 497 s, a 42% saving. Without the mount, every restart, flag change and spot preemption repays the full compile. 1.4 Verify Cold start is 857 s (14 min) and 80% of it is XLA compilation, not weight loading - the 9.54 GiB checkpoint downloads in about 10 seconds. Be patient, and watch the log rather than the clock: sudo docker logs -f vllm-gemma4 2>&1 | grep -E "Memory statistics|Init kv-cache|startup complete" You want to see this, which is the whole memory budget in one line: Memory statistics | total_hbm_limit_gb=15.75GiB | total_hbm_limit_cap_gb=14.49GiB | total_hbm_used_gb=8.97GiB | total_hbm_avail_gb=5.52GiB Then smoke-test it. Use /v1/chat/completions , not /v1/completions - raw completions return an empty string on -it models, which looks exactly like a broken deploy and isn't: curl -s localhost:8000/v1/chat/completions -H 'Content-Type: application/json' -d '{ "model":"google/gemma-4-E2B-it", "messages":[{"role":"user","content":"Say hi in five words."}]}' | jq -r '.choices[0].message.content' 1.5 Tear down gcloud compute tpus tpu-vm delete gemma4-v5e --zone=us-west4-a --quiet Part 2 - The flags, and what they actually do The chip: what one v5e actually gives you | spec | v5e, one chip | source | |---|---|---| | HBM capacity | 16 GB nominal ยท 15.75 GiB visible to the runtime | vendor ยท measured | Usable for weights + KV at 0.92 | 14.49 GiB | measured | | HBM bandwidth | 800 GiBps | vendor | | Peak bf16 | 197 TFLOPS | vendor | | Peak int8 | 393 TOPS (exactly 2x bf16) | vendor | | TensorCore | 1, with 4 MXUs (128x128) | vendor | | ICI | 400 GBps bidirectional, 4 ports | vendor | | Machine type | ct5lp-hightpu-1t | | | gcloud spelling | v5litepod-1 , runtime v2-alpha-tpuv5-lite | measured | Two unit traps worth knowing before you compare anything. Google quotes v5e bandwidth in GiBps and v6e in GBps - normalise before dividing, the real generational ratio is ~1.9x, not a clean 2x. And the "16 GB" capacity figure sits awkwardly next to the 15.75 GiB the runtime reports (15.75 GiB is 16.9 GB), so the vendor number is almost certainly 16 GiB loosely written. Size against the measured 15.75 GiB, never the marketing figure. Data types: what the matrix units can actually compute in This single table decides every quantization question on this chip. | format | native in the MXU? | what it buys on v5e | |---|---|---| | bf16 | โœ… | the baseline - everything here runs in it | | int8 | โœ… 2x bf16 throughput | the only low-precision compute win | | fp8 | โŒ | storage and bandwidth only - values widen back to bf16 before the matmul | | int4 / fp4 | โŒ | footprint and bandwidth only, then unpack to bf16 | Google publishes bf16 and Int8 peaks for v5e and no fp8 figure at all, which is the tell. The practical consequence: a benchmark showing no speedup from fp8 on this chip is the correct result, not a misconfiguration. v7/Ironwood is the first TPU with fp8 in the MXU - do not carry any conclusion here forward to it. Quantization: what is actually reachable Gemma 4 exists only as a JAX implementation in this stack, so anything in the torch path is unreachable no matter what the platform advertises. Measured state: | route | status on this build | |---|---| KV cache, bf16 (auto ) | โœ… the only one worth running | KV cache fp8_e4m3 / fp8_e5m2 | reachable, 1.000x capacity - the block layout is word-aligned, so narrowing the element buys padding, not room. ~2% slower | KV cache int8 | โŒ rejected by the CLI enum - never reaches the engine | KV cache int8_per_token_head , turboquant_* , nvfp4 , fp8_inc , fp8_ds_mla | accepted by the CLI, then kill the server at boot | | Weights, compressed-tensors w4a16 (Google's QAT format) | โŒ NotImplementedError on the JAX path | | Weights, mxfp4 | โŒ MoE-only; E2B is dense, so there is nothing to attach to | | Weights, qwix PTQ int8/int4 | โŒ does not boot - the concrete path OOMs on quantization temporaries, the abstract path raises binding weights | | Weights, AWQ / GGUF / q4_0 | โŒ torch path or absent | So bf16 is not a choice here, it is the only thing that runs - and that is the single biggest constraint on the chip. Weights are 8.97 of the 14.49 GiB budget (62%), and none of it can be compressed today. Working int8 weights would roughly double the KV pool and buy real FLOPS, since int8 is the one format with a native MXU path. It is blocked upstream, not by configuration, so it is worth re-testing on every image bump. One measured consequence of all this: decode moves ~3.15 GiB per step against a 3.94 ms bandwidth floor, and measures 8.02 ms - about 49% of peak bandwidth. (The 3.15 GiB is derived from the model's layer geometry; the 8.02 ms is measured.) The chip is not the bottleneck at any point in this article. The model: six things about Gemma 4 E2B that will catch you out E2B is a strange checkpoint. Almost every intuition from a conventional decoder is wrong here, and the memory arithmetic later in this article only makes sense on

Read on DEV Community ↗ ← Back to News

Comments

No comments yet. Start the discussion.