tpu-inference

Inference on Google Cloud: Precision, Memory and Capacity on TPU, from v5e to Ironwood --- title: "Inference on Google Cloud: Precision, Memory and Capacity on TPU, from v5e to Ironwood" published: true description: "The companion to the talk: what the number format, memory bandwidth, slice shape and capacity rules do to Gemma 4 on Cloud TPU v5e and v6e, against the L4, Blackwell, MI300X and Inferentia, with the cost of each." tags: googlecloud, tpu, machinelearning, gemma cover_image: raw.githubusercontent.com --- This article is the written companion to the talk "Inference on Google Cloud". It covers serving Gemma 4 on Cloud TPU v5e and v6e through vLLM, measured on the gemma4-dev Cloud TPU test setups between July and October 2026, and compares them with the NVIDIA L4 on Google Cloud, Blackwell list prices, an AMD MI300X and AWS Inferentia. Four ideas explain most of the results: the number format the chip computes natively, the memory bandwidth that feeds it, the slice shape, and the rules that decide whether a TPU is there to rent at all. xbill9/gemma4-dev The hands-on steps run from a second repository, which starts one TPU and serves Gemma 4 on it with a single command: xbill9/gcp-inference The slides: docs.google.com --- What Is the Problem? This article compares four TPU generations, the L4 and Blackwell GPUs beside them, and every way Google Cloud prices a TPU chip-hour. The same Gemma 4 E2B serves at 144 tokens a second on a v5e in bf16 and at 243 in int8 with 4-bit embedding tables, on the same chip. A v6e-1 serves 3.3 times the L4's tokens on the same 26B checkpoint. And a zone that reports quota for 1,536 chips can still return a stockout. --- Which Generations? | Chip | HBM | Peak bf16 | Native low precision | Max pod | On-demand | |---|---:|---:|---|---:|---:| | v5e | 16 GB | 197 TFLOPs | int8, 393 TOPs | 256 | $1.20/hr | | v6e Trillium | 32 GB | 918 TFLOPs | int8, 1,836 TOPs | 256 | $2.70/hr | | v5p | 95 GiB | 459 TFLOPs | int8; fp8 at bf16 rate | 8,960 | $4.20/hr | | TPU7x Ironwood | 192 GiB | 2,307 TFLOPs | fp8, 4,614 TFLOPs | 9,216 | $12.00/hr | Google's first TPU, in its datacenters since 2015, was an 8-bit inference chip with a 92 TOPS peak, and Google announced Ironwood on 2025-04-09 as "the first Google TPU for the age of inference". The gemma4-dev test setups measured v5e and v6e. v5p and TPU7x rows come from Google's documentation only. vLLM reaches the TPU through tpu-inference, a JAX backend. XLA compiles one program per input shape, which makes the compile the largest part of a cold start. --- At This Point You Should Have… - The gcloud CLI with its alpha and beta components, signed in with gcloud auth login and gcloud auth application-default login - A project with Cloud TPU quota for one v5e chip in us-west4-a - python3 with pip install 'google-adk[extensions]' google-genai mcp - A clone of xbill9/gcp-inference --- Step 1 — One TPU, One Command demo.sh asks for a TPU v5e-1 as a flex-start request in us-west4-a, waits for the node, opens an SSH tunnel to vLLM on the TPU's port 8000, waits for the model to load and compile, then runs a Google ADK agent against it. Flex-start stops the node by itself after 4 hours. On a TPU that is already up, ./demo.sh again skips the first two steps: ./demo.sh again == [0m01s] 3/5 Connecting to it over SSH (the demo reaches vLLM on the TPU's port 8000 through a tunnel) [0m05s] SSH works [0m11s] tunnel open: localhost:8005 -> TPU port 8000 == [0m11s] 4/5 Starting Gemma on vLLM (first start downloads and compiles: several minutes) [0m30s] vLLM is answering 20:21:25Z PASS tpu chat 0.198s; tool run filters=['id >= 10'] quoted=8 3.71s READY == [0m37s] 5/5 The demo: Gemma on the TPU counts ids with the count_ids tool from the MCP server Question: How many of the ids are 10 or more? Answer with the number. Ids: [0, 2, 3, 20, 21, 22, 23, 10, 11, 12, 13] True answer: 8 tools: mcp tpu run 1/3 ok quoted=8 filters=['id >= 10'] 3.87s tpu run 2/3 ok quoted=8 filters=['id >= 10'] 0.51s tpu run 3/3 ok quoted=8 filters=['id >= 10'] 0.49s backend runs correct right-filter errors median s filters sent tpu 3 3/3 3/3 0 0.51 id >= 10 The agent is an ADK LlmAgent whose model is LiteLlm pointed at vLLM's OpenAI-compatible API, and an ADK McpToolset starts mcp_server.py over stdio and offers it one tool. The model never counts. It calls count_ids on that MCP server, which returns the exact count, and each run is graded on the filter it sent as well as the number it quoted. Gemma 4 E2B on the v5e sent id >= 10 in all three runs. The node serves google/gemma-4-E2B-it in bf16 with vllm/vllm-tpu:v0.30.0, one chip, no sharding. --- Step 2 — What the 16 GB Chip Holds vLLM logs the chip's memory as it starts. In order: Init worker | ... | hbm=[(0.0, 15.75)]GiB Init model | hbm=[(8.97, 15.75)]GiB 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 TPU KV cache size: 321,280 tokens, Maximum concurrency for 8,192 tokens per request: 39.22x Init kv-cache | num_total_layers=15 | num_blocks=[1255, ...] | regular_attn_shape=(num_blocks, (256, 1, 2, 256)) | regular_attn_dtype=bfloat16 | hbm=[(14.51, 15.75)]Gb | HBM on one v5e chip | GiB | |---|---:| | Weights, bf16 | 8.97 | | vLLM's ceiling (92% of 15.75) | 14.49 | | Left for the KV cache | 5.52 | | In use after the KV cache is allocated | 14.51 | | Never allocated | 1.24 | Each of the 15 cached layers holds 1,255 blocks of 256 tokens, one KV head of dimension 256, in bf16. That is 1,024 bytes per token per layer, 15,360 bytes per token across all 15, and 321,280 token slots in all. The model's config lists 35 layers with num_kv_shared_layers: 20: 20 layers reuse another layer's keys and values, so the cache holds 15. --- Step 3 — The Full 131,072-Token Context Gemma 4 E2B's config allows 131,072 positions. vLLM sizes its KV cache from the memory left after the weights, so raising --max-model-len from 8,192 to 131,072 changes the concurrency and leaves the cache the same size. To change a flag on a running TPU, replace the container: gcloud compute tpus tpu-vm ssh gcp-inference-tpu-node --project <tpu-project> --zone us-west4-a --command \ 'sudo docker rm -f vllm && sudo docker run -d --name vllm --restart unless-stopped --privileged --net=host \ -v /dev/shm:/dev/shm --shm-size 10gb -e HF_HOME=/dev/shm vllm/vllm-tpu:v0.30.0 \ vllm serve google/gemma-4-E2B-it --served-model-name gemma --max-model-len 131072 --port 8000 \ --max-num-batched-tokens 4096 --enable-auto-tool-choice --tool-call-parser gemma4' 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 TPU KV cache size: 321,280 tokens, Maximum concurrency for 131,072 tokens per request: 2.45x init engine (profile, create kv cache, warmup model) took 503.54 s (compilation: 413.41 s) Application startup complete. | --max-model-len | 8,192 | 131,072 | |---|---:|---:| | KV cache, tokens | 321,280 | 321,280 | | Full-length requests at once | 39.22 | 2.45 | | Engine start | 502.84 s | 503.54 s | | of which XLA compile | 408.59 s | 413.41 s | One full-length request needs 1.875 GiB of KV cache. scripts/longctx.py in tpu-survey buries a passphrase at a chosen depth in up to 124,949 tokens of filler, starts each prompt with a unique line so vLLM's prefix cache cannot reuse an earlier run, and asks for the passphrase: $ python3 scripts/longctx.py $TPU_URL 8000 0.5 {"prompt_tokens": 7949, "completion_tokens": 10, "depth": 0.5, "seconds": 0.61, "prompt_tokens_per_second": 13110.4, "answer": "violet-harbor-7341", "found": true} $ python3 scripts/longctx.py $TPU_URL 120000 0.1 {"prompt_tokens": 124949, "completion_tokens": 10, "depth": 0.1, "seconds": 9.53, "prompt_tokens_per_second": 13116.9, "answer": "violet-harbor-7341", "found": true} $ python3 scripts/longctx.py $TPU_URL 120000 0.5 {"prompt_tokens": 124949, "completion_tokens": 10, "depth": 0.5, "seconds": 9.54, "prompt_tokens_per_second": 13102.3, "answer": "violet-harbor-7341", "found": true} $ python3 scripts/longctx.py $TPU_URL 120000 0.9 {"prompt_tokens": 124949, "completion_tokens": 10, "depth": 0.9, "seconds": 9.41, "prompt_tokens_per_second": 13274.6, "answer": "violet-harbor-7341", "found": true} $ python3 scripts/longctx.py $TPU_URL 140000 0.5 HTTP 400: {"error":{"message":"This model's maximum context length is 131072 tokens. However, you requested 16 output tokens and your prompt contains at least 131057 input tokens, for a total of at least 131073 tokens. ..."}} The passphrase comes back at 10%, 50% and 90% depth of a 124,949-token prompt, in 9.4 to 9.5 seconds each, about 13,100 prompt tokens a second as timed from the client. vLLM's own log reads 12,495 tokens a second, a 0.0% prefix cache hit rate, and 38 to 39% of the KV cache in use while the request runs. A prompt past the limit is refused before it reaches the chip. --- 🔎 Tip: Raise the Context on the Chip You Have Flex-start capacity is granted per request. Replacing the container keeps the node, and its 4-hour clock, and costs one engine start, 503.54 s on this v5e. Deleting the node to relaunch it means a new flex-start request, and a new wait for capacity. --- Step 4 — Precision: What the Chip Computes A smaller format saves memory and bandwidth on every chip. It saves compute only where the matrix unit multiplies that format natively. | Format | v5e | v6e | v5p | TPU7x | |---|---|---|---|---| | bf16 | native | native | native | native | | int8 | 2x bf16 | 2x bf16 | native | native | | fp8 | storage | storage | bf16 rate | 2x bf16 | | int4 | storage | storage | storage | storage | "Storage" means the weights widen to bf16 before every matrix multiply. On v5e that ordering shows up directly in E2B throughput at one stream: | E2B format on v5e-1 | tok/s | |---|---:| | bf16 | 144 | | 4-bit QAT + 4-bit embeddings | 150 | | fp8 | 182 | | fp8 + 4-bit embeddings | 203 | | int8 W8A8 | 220 | | int8 + 4-bit embeddings | 243 | Accuracy holds: 0.683 for bf16 against 0.677 for int8 with 4-bit embeddings, on a 3,880-record suite. At 12B, fp8 serves 518 tok/s at 16 streams against 624 for int8 at the same file size. On v6e the 4-bit path pays off through a fused kernel. Gemma 4 26B-A4B in 4-bit QAT W4A16 serves 1,283 tok/s at 16 streams against 668 for fp8, in 17.43 GiB of weights against 27.99 GiB, and holds 53,888 KV-cache tokens against 3,456. Accuracy is 75.3% against 76.0%. W4A16 needs tpu-inference PRs #3299, #3653 and #3660 on top of the stock image. --- Step 5 — Memory and Bandwidth Each generated token reads every weight once, so bandwidth sets the ceiling for one request. | Chip | Memory | Bandwidth | E2B bf16 tok/s, 1 stream | tok/s per TB/s | |---|---:|---:|---:|---:| | TPU v5e | 16 GB | 859 GB/s | 120 | 140 | | TPU v6e | 32 GB | 1,638 GB/s | 199 | 122 | | TPU v5p | 95 GiB | 2,765 GB/s | not measured | — | | TPU7x | 192 GiB | 7,380 GB/s | not measured | — | | NVIDIA L4 | 24 GB | 300 GB/s | 46 | 154 | The three measured chips land between 122 and 154 tokens a second per TB/s. Google quotes v5e bandwidth as 800 GiBps, which is 859 GB/s. The L4 figure was measured on an AWS g6.2xlarge, the same chip as Google Cloud's g2. --- Step 6 — Topologies and Slices Weight size picks the slice. Every Gemma 4 size serves on one host: | Model | v5e-1 | v6e-1 | v6e-4, TP=4 | |---|---|---|---| | E2B | bf16 or int8 | bf16 | — | | E4B | int8 or 4-bit; bf16 does not fit | bf16 | — | | 12B | int8 or 4-bit | bf16, 22.18 GiB | — | | 26B-A4B | 4-bit, 201 tok/s | 4-bit, 1,283 tok/s | bf16, 1,982 tok/s | | 31B | does not fit | 4-bit, 500 tok/s | bf16, 2,031 tok/s | Small models scale out. On a v5e-4, four one-chip E2B engines serve 1.39x the output of one TP=4 engine at 16 streams, 2.52x at 64 and 3.36x at 256. E2B has one KV head, so tensor parallelism keeps four copies of the KV cache. 12B gains from TP=4: 2.25x the v5e-1 output at one stream, and 122,624 KV tokens against 9,728. --- 🔎 Tip: Quota Lets You Ask, Availability Decides Quota is permission to request a TPU. Availability decides whether the request gets hardware, through three checks: hardware in the zone, the provisioning model offered there, and capacity at that moment. 37 zones report v6e quota; 18 offer v6e-1 hardware. Three of five zones holding 1,536 chips of quota returned stockout. us-central1-b has hardware, quota and a flex-start price, and rejects flex-start. Check capacity with a spot request: it fails fast with "reason: stockout", where a flex-start request with no quota queues as PENDING exactly like a stockout. Ask for fewer chips. In europe-west4-a a v6e-8 flex-start request expired after 26 minutes, and a v6e-4 was granted at once less than an hour later. --- Step 7 — Against Other Accelerators On the same 26B 4-bit checkpoint, with the same flags and client: | | NVIDIA L4 | TPU v6e-1 | |---|---:|---:| | Output tok/s, 16 x 256 tokens | 393.7 | 1,283.3 | | KV cache tokens | 17,990 | 53,888 | | 3,880-record suite, wall time | 139 s | 58 s | | 4-bit kernel | Marlin | gmm_v2 | | Software | stock vLLM 0.30.0 | three tpu-inference patches | | List price per hour | $0.71 (g2-standard-4) | $2.70 | Cloud Run serves the L4 as a managed service: one L4 per instance, weights mounted from Cloud Storage at /mnt/models, a startup check that waits 180 s before its first check, and a default that scales between zero and one instance, so an idle service means a cold GPU start. Blackwell on Google Cloud lists above v6e per accelerator-hour: an RTX PRO 6000 on g4-standard-48 at $4.50 on-demand and $2.25 flex-start, and a B200 at $8.06 per GPU flex-start ($64.44 for the 8-GPU a4-highgpu-8g, divided by 8). Both are list prices only, with no measurements in this article. Per dollar, on E2B bf16 with vllm bench serve at 1,024 prompt and 128 output tokens, v5e at its $0.58 spot price leads to 16 streams and the MI300X at $1.99 leads at 64. The v6e-1 is priced at the europe-west4 on-demand rate of $2.97: | tok/s per $/hr | MI300X ($1.99) | v5e-1 spot ($0.58) | v6e-1 europe-west4 ($2.97) | |---|---:|---:|---:| | 1 stream | 153 | 208 | 67 | | 4 streams | 487 | 711 | 233 | | 16 streams | 1,333 | 1,572 | 587 | | 64 streams | 2,713 | 1,972 | 719 | Each software stack fails in its own way: | | TPU, vLLM + JAX | CUDA, L4 | ROCm, MI300X | Neuron, inf2 / trn1 | |---|---|---|---|---| | Start to healthy, E2B | 510 s v6e, 986 s v5e | 334 s | about 160 s | 171 s on trn1 | | Gemma 4 support | 4-bit needed patches | 26B 4-bit loads stock | fp8 only as fnuz | vLLM-Neuron cannot load it | | Hazard | OOM after a long compile | sm_89 missing; runs sm_86 code | AITER flag 0.97 to 0.99x | large gather returns zeros, HTTP 200 | --- 🔎 Tip: Cache the Compile XLA compile is 81 to 91% of vLLM's engine start for E2B: 329 of 404 seconds on v6e-1 and 738 of 814 on v5e-1, inside 510 and 986 seconds from container start to healthy. On today's v5e-1 with vllm/vllm-tpu:v0.30.0 the engine start was 503.54 seconds, 413.41 of them compile. A mounted compile cache cut a v5e restart from 857 to 497 seconds, 42%. 31B weights load in 5.6 minutes from in-region Cloud Storage against 60 to 90 from Hugging Face. Serving flags that work on GPUs fail late on TPU, after the compile. --gpu-memory-utilization 0.95 failed on v5e after a 691-second compile, twice; 0.92 serves. --kv-cache-dtype fp8 logs fp8 on v6e and allocates bf16. The vllm/vllm-tpu:nightly tag broke two 4-bit builds on 2026-10-02, so pin an image digest. --- Step 8 — What Each Chip Costs | $ per chip-hour | On-demand | Flex-start | Calendar | 1-yr CUD | 3-yr CUD | Spot read | |---|---:|---:|---:|---:|---:|---:| | v5e | $1.20 | $0.60 | $0.84 | $0.84 | $0.54 | $0.58 | | v6e Trillium | $2.70 | $1.35 | $1.89 | $1.89 | $1.22 | $1.40 | | v5p | $4.20 | $2.10 | $2.94 | $2.94 | $1.89 | — | | TPU7x Ironwood | $12.00 | $6.00 | $8.40 | $8.40 | $5.40 | — | Flex-start halves on-demand on every generation and can undercut spot: v6e in us-east5 lists $1.35 flex-start against a $1.40 spot price. Flex-start runs 10 minutes to 7 days and bills until deleted. Utilization sets the price per token. A v6e-1 at the europe-west4 on-demand rate of $2.97 an hour serves E2B at $4.14 per million output tokens with one stream, $0.70 with eight and $0.55 with 32, at 1,024-token prompts. A v5e at the $0.60 flex-start rate goes from $1.31 per million at one stream to $0.09 at 64, at 128-token prompts. --- Compare and Contrast | Option | Precision served | Bandwidth | Measured result | List $/hr | Software | |---|---|---:|---|---:|---| | 🥇 TPU v6e-1 | 4-bit QAT W4A16 | 1,638 GB/s | 26B at 1,283 tok/s | $2.70 | tpu-inference, patched | | 🥈 TPU v5e-1 | int8 W8A8 + 4-bit embeddings | 859 GB/s | E2B at 243 tok/s, one stream | $1.20 | tpu-inference | | 🥉 NVIDIA L4 | 4-bit QAT | 300 GB/s | 26B at 394 tok/s | $0.71 | stock vLLM | | TPU7x Ironwood | fp8 native | 7,380 GB/s | not measured | $12.00 | allowlisted | | Blackwell B200 | fp4 native | — | not measured | $8.06 flex, per GPU | not measured | --- So, Which One? For E2B or E4B at the lowest cost, run v5e-1 in int8 with one engine per chip, on spot or flex-start. For 12B to 31B on one chip, run v6e-1 in 4-bit QAT. For bf16 26B or 31B with long context, run v6e-4 with TP=4. For stock software and the fastest boot, run the L4, on Compute Engine or on Cloud Run. TPU7x brings fp8 math and 192 GiB per chip, behind an allowlist. Whichever you pick: check capacity with a spot request, ask for the smallest slice that fits, pin the image digest and mount the compile cache. --- Summary The goal of this article was to serve Gemma 4 on a Cloud TPU in one command, then explain the results across TPU generations through native precision, memory bandwidth, slice shape and capacity. The key to the solution was measuring each claim on real hardware, with every ratio and price computed in code. The results were: - 🟢 One v5e chip holds Gemma 4 E2B's full 131,072-token context: a 321,280-token KV cache beside 8.97 GiB of bf16 weights, and a passphrase found at 10%, 50% and 90% depth of a 124,949-token prompt in 9.4 to 9.5 s. - 🟢 The agent sends the right filter: 3 of 3 runs on the v5e called count_ids with id >= 10 and quoted 8. - 🟢 int8 is the fast format on v5e: with 4-bit embedding tables it serves E2B at 1.7x bf16, with accuracy within a point. - 🟢 v6e-1 serves 3.3x the L4's tokens on the same 26B 4-bit checkpoint, with 3x the KV cache. - 🟢 Flex-start halves on-demand on every generation, and can undercut spot. - ⚠️ Cold start takes 8.5 to 16.4 minutes, most of it XLA compile; cache it. - 🟢 One-stream decode follows bandwidth: 122 to 154 tok/s per TB/s on v5e, v6e and the L4. - ❌ Quota is permission to request; availability decides the grant: 37 zones report v6e quota, 18 have v6e-1 hardware, and zones holding full quota still return stockout. Scope: the hands-on steps ran on one flex-start v5e-1 in us-west4-a on 2026-10-25, google/gemma-4-E2B-it in bf16 on vllm/vllm-tpu:v0.30.0, three demo runs and one long-context request per depth. The rest covers Gemma 4 E2B to 31B on Cloud TPU v5e and v6e, measured on the gemma4-dev test setups between July and October 2026 with vLLM and tpu-inference; L4 on Google Cloud g2 and AWS g6; MI300X on DigitalOcean; Inferentia2 and Trainium on AWS. The cross-vendor tokens-per-dollar table mixes vLLM builds, hosts and repeat counts, so it shows direction. v5p, TPU7x and Blackwell figures are Google's published specifications and list prices, read on 2026-10-12 and re-read on 2026-10-25 with the same values; spot prices come from the Cloud Billing Catalog on recorded dates. The strategy for serving Gemma 4 on Google Cloud from v5e to Ironwood was validated with an incremental step by step approach. --- References - gemma4-dev | GitHub - gcp-inference | GitHub - tpu-survey, long-context test and evidence | GitHub - Gemma 4 E2B-it | Hugging Face - Inference on Google Cloud slides | Google Slides - Inference on AWS slides | Google Slides - tpu-inference | GitHub - Cloud TPU documentation | Google Cloud - TPU7x (Ironwood) | Google Cloud - Cloud TPU pricing | Google Cloud - Accelerator-optimized machine pricing | Google Cloud - In-Datacenter Performance Analysis of a Tensor Processing Unit | arXiv

Small Apache-2.0 Models Against Gemma 4 E2B on One TPU v5e: Which Ones Keep Up? --- title: "Small Apache-2.0 Models Against Gemma 4 E2B on One TPU v5e: Which Ones Keep Up?" published: true description: "Fourteen Apache-2.0 and MIT open-weight models that fit one TPU v5e chip, served one after another through vLLM on the same pinned image and scored on GSM8K and BFCL against Gemma 4 E2B. Qwen3-4B matches Gemma's accuracy at 0.57x its output, SmolLM3-3B and Granite 3.3 2B match it on tool calls, and Gemma keeps the lead on long prompts." tags: tpu, vllm, gemma, llm cover_image: raw.githubusercontent.com --- This article provides a step by step guide to surveying small open-weight models on one Google Cloud TPU v5e chip through vLLM, with Gemma 4 E2B as the baseline. Every model is Apache-2.0 or MIT licensed, so any of them can be repacked and redistributed, and every one is served on the same pinned image with the same settings, then timed and scored by code. Gemma 4 E2B is small and workable on a v5e-1, but its attention is unusual: two head sizes, 256 for sliding-window layers and 512 for full-attention layers, one KV head, and 20 of its 35 layers reading another layer's KV cache. The question here is which plainer model can stand in for it. Qwen3-4B matches Gemma on both accuracy tests at 0.57x the output. SmolLM3-3B and Granite 3.3 2B match it on tool calls at 0.76x and 0.93x. With long prompts, Gemma stays ahead of every model of comparable accuracy. xbill9/model-sweep --- Why One v5e Chip? A v5litepod-1 is the cheapest TPU Google Cloud rents and is usually available as flex-start capacity in us-west4-a. It has 15.75 GiB of HBM, and vLLM gets 0.9 of that, so a bf16 model has to leave room for its KV cache in about 14 GiB. That rules the field down to roughly 4B parameters. Gemma 4 E2B loads in 8.97 GiB on this chip, which makes it one of the larger models in the lineup. --- At This Point You Should Have… - A Google Cloud project with the Cloud TPU API enabled and flex-start TPU v5e quota in us-west4-a - A GCS bucket for results, and a Hugging Face token saved in Secret Manager as hf-token - gcloud with the alpha components, and Python 3 with no extra packages - A clone: git clone https://github.com/xbill9/model-sweep --- Step 1 — Build the Lineup From the Hub survey.py reads each candidate's license, architecture and weight size from the Hugging Face API and decides whether it fits the chip. A model is in when its license is Apache-2.0 or MIT and at least 2 GiB is left for the KV cache. vLLM on TPU runs some architectures as native JAX models and the rest through a PyTorch fallback, so the survey splits the lineup into two tiers: python3 survey.py | `tpu-vllm-v5e1-qwen3_4b` | `Qwen/Qwen3-4B` | apache-2.0 | Qwen3ForCausalLM | 4.02 | 7.49 | 6.68 | JAX | ✅ | 1 | | `tpu-vllm-v5e1-gemma4_2b` | `google/gemma-4-E2B-it` | apache-2.0 | Gemma4ForConditionalGeneration | 5.12 | 9.54 | 4.63 | JAX | ✅ | 1 | | `tpu-vllm-v5e1-smollm3_3b` | `HuggingFaceTB/SmolLM3-3B` | apache-2.0 | SmolLM3ForCausalLM | 3.08 | 5.73 | 8.45 | torchax | | 2 | | `tpu-vllm-v5e1-qwen25_3b` | `Qwen/Qwen2.5-3B-Instruct` | other | Qwen2ForCausalLM | 3.09 | 5.75 | 8.43 | JAX | | excluded: license | | `tpu-vllm-v5e1-qwen3_8b` | `Qwen/Qwen3-8B` | apache-2.0 | Qwen3ForCausalLM | 8.19 | 15.26 | -1.08 | JAX | | excluded: memory | Tier 1: 8 models. Tier 2: 7 models. Excluded: 2. Qwen2.5-3B is out on its license and Qwen3-8B on memory. Everything else goes into two arms files, one per tier. --- Step 2 — Describe Each Model in One Line An arms file lists one model per line: the checkpoint, its rig name, what to measure, and any serving flags it needs. The tool-call parser and the thinking switch differ by family: google/gemma-4-E2B-it=gemma4_2b=load+long+gen=parser=gemma4,mm=image+audio+video Qwen/Qwen3-4B=qwen3_4b=load+long+gen=parser=hermes,nothink HuggingFaceTB/SmolLM3-3B=smollm3_3b=load+long+gen=parser=hermes,nothink ibm-granite/granite-3.3-2b-instruct=granite33_2b=load+long+gen=parser=granite nothink sends chat_template_kwargs {"enable_thinking": false} with every request, so Qwen3 and SmolLM3 answer directly like the others. mm= turns off the image, audio and video inputs of multimodal checkpoints. A unit test checks every line against the lineup and against the tool-call parsers registered in the pinned vLLM build. --- Step 3 — Launch One Chip per Tier tpu/launch.sh bundles the committed tree, uploads it, and creates a flex-start queued resource whose startup script serves every arm in turn, uploads results after each one, and deletes the TPU at the end: tpu/launch.sh arms-v5e1-tier2.txt --go arms: tpu/arms-v5e1-tier2.txt (7 arms) bundle: gs://aisprint-491218-bucket/model-sweep/inputs/sweep-86d97ba02-0837b648.tgz results: gs://aisprint-491218-bucket/model-sweep/2026-10-12-v5e1-tier2/ Created queued_resource [ms-2026-10-12-v5e1-tier2-1349]. created ms-2026-10-12-v5e1-tier2-1349; it deletes itself when the last arm finishes, or at 4h Both tiers ran on vllm/vllm-tpu@sha256:19a1a0526476f902eb83e1057f3d8938f35b457dd7716a30e9ab4f7bee90d507, vLLM 0.29.1rc1.dev468+g0b7f11a1e, so the tier 2 models pair with tier 1's Gemma run: [sweep 2026-10-11T17:14:28Z] start 2026-10-11-v5e1-tier1 on us-west4-a, arms arms-v5e1-tier1.txt [sweep 2026-10-11T18:24:30Z] DONE: all arms attempted [sweep 2026-10-12T13:51:26Z] start 2026-10-12-v5e1-tier2 on us-west4-a, arms arms-v5e1-tier2.txt [sweep 2026-10-12T16:02:12Z] DONE: all arms attempted --- Step 4 — Measure Speed and Accuracy the Same Way Every model is served with --max-model-len 4096 --max-num-seqs 16 and prefix caching on. Each one gets three measurements: - Output speed: 256 output tokens per request at 1, 4 and 16 requests at once, median of three passes - Long prompts: the same with 1,024- and 3,072-token prompts that never repeat, so the prefix cache cannot help - Accuracy: all 1,319 GSM8K problems, greedy, with a 2,048-token answer limit, and the 400 BFCL simple tool-call records for models with a tool-call parser Answers are scored in code. Differences against Gemma are paired on the same problems, with a 95% range from 10,000 resamples; a range that includes zero means no measured difference. --- Step 5 — File the Results and Build the Report python3 sweep_report.py file 2026-10-12-v5e1-tier2 python3 sweep_report.py report 2026-10-12-v5e1-tier2 --baseline-run 2026-10-11-v5e1-tier1 filed smollm3_3b -> tpu-vllm-v5e1-smollm3_3b/benchmarks/runs/2026-10-12-v5e1-tier2 wrote results/2026-10-12-v5e1-tier2/REPORT.md (6 served, 1 failed) file moves each model's logs and scored answers into its own directory. report refuses to pair two runs that served different image digests. --- How Accurate Is Each Model? Against Gemma 4 E2B, in percentage points, with the 95% range: | Model | GSM8K | vs Gemma | BFCL | vs Gemma | |---|---|---|---|---| | Gemma 4 E2B | 0.911 | — | 0.927 | — | | Qwen3-4B | 0.915 | +0.4 (-1.2 to +2.0) | 0.940 | +1.2 (-1.8 to +4.2) | | Phi-4-mini | 0.895 | -1.6 (-3.2 to -0.1) | not scored | — | | SmolLM3-3B | 0.858 | -5.3 (-7.2 to -3.5) | 0.945 | +1.7 (-1.0 to +4.8) | | Qwen3-1.7B | 0.827 | -8.4 (-10.5 to -6.4) | 0.935 | +0.8 (-2.2 to +3.8) | | Granite 3.3 2B | 0.817 | -9.5 (-11.6 to -7.4) | 0.920 | -0.7 (-3.8 to +2.2) | | Qwen2.5-1.5B | 0.729 | -18.2 (-20.6 to -15.8) | 0.910 | -1.7 (-5.2 to +1.8) | | Qwen3-0.6B | 0.599 | -31.2 (-34.0 to -28.5) | 0.825 | -10.3 (-14.2 to -6.5) | | Qwen2.5-0.5B | 0.469 | -44.3 (-47.0 to -41.4) | 0.818 | -11.0 (-14.8 to -7.2) | | SmolLM2-1.7B | 0.461 | -45.0 (-47.8 to -42.2) | no parser | — | Five models are level with Gemma on BFCL tool calls, down to Qwen2.5-1.5B. Only Qwen3-4B is level on GSM8K, and from 1.7B down, Qwen3-1.7B aside, the score falls by 18 to 45 points. --- How Fast Is Each Model? Output tokens per second with 16 requests at once, and against Gemma: | Model | Weights GiB | 16 requests | vs Gemma | 16 × 3,072-token prompts | vs Gemma | |---|---|---|---|---|---| | Qwen2.5-0.5B | 0.92 | 4176 | 2.21x | 2596 | 2.48x | | Qwen3-0.6B | 1.40 | 3210 | 1.70x | 1301 | 1.24x | | Qwen2.5-1.5B | 2.88 | 2780 | 1.47x | 1570 | 1.50x | | Qwen3-1.7B | 3.78 | 2049 | 1.09x | 926 | 0.88x | | Gemma 4 E2B | 8.97 | 1887 | 1.00x | 1049 | 1.00x | | Granite 3.3 2B | 4.92 | 1761 | 0.93x | 732 | 0.70x | | SmolLM3-3B | 6.22 | 1440 | 0.76x | 725 | 0.69x | | Phi-4-mini | 8.31 | 1136 | 0.60x | 404 | 0.39x | | Qwen3-4B | 7.49 | 1080 | 0.57x | 373 | 0.36x | Gemma carries the largest weights in the table and still out-serves every model with more than 4 GiB of weights. With 3,072-token prompts the gap widens: Granite and SmolLM3 drop to about 0.7x, and Qwen3-4B to 0.36x. --- Why Does Gemma Hold Up on Long Prompts? The KV cache. Each model's boot log reports how many tokens of cache fit in the HBM left after its weights: | Model | KV heads | Head size | Layers | KV cache tokens | |---|---|---|---|---| | Gemma 4 E2B | 1 | 256 / 512 | 35, 20 sharing | 321,280 | | Granite 3.3 2B | 8 | 64 | 40 | 125,184 | | SmolLM3-3B | 4 | 128 | 36 | 120,320 | | Qwen3-1.7B | 8 | 128 | 28 | 100,096 | | Qwen3-4B | 8 | 128 | 36 | 50,688 | Gemma stores one KV head per layer, 28 of its 35 layers keep only a 512-token sliding window, and 20 layers reuse an earlier layer's cache. That leaves it 2.6 to 6.3 times the cache of the four models below it, even though its weights are the largest. The same features that make Gemma awkward to port are the ones paying for its long-prompt speed on this chip. --- 🔎 Tip: Turn Thinking Off Per Request Qwen3 and SmolLM3 think by default, which spends the GSM8K answer budget on reasoning text. Passing the chat template switch with each request turns it off without a different checkpoint: {"chat_template_kwargs": {"enable_thinking": false}} With it, Qwen3-4B leaves no GSM8K answer at the 2,048-token limit, and SmolLM3-3B leaves 16 of 1,319. --- 🔎 Tip: DeepSeek-R1-Distill-Qwen-1.5B Needs a One-Line Fix on TPU The MIT-licensed R1 distill stops at load on this image: FAILED deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B (exited) after 75s: AttributeError: 'Qwen2Config' object has no attribute 'text_config'. Did you mean: 'get_text_config'? tpu_inference/models/jax/qwen2.py builds an untied lm_head with getattr(hf_config, 'hidden_size', hf_config.text_config.hidden_size). Python evaluates the default argument first, so any Qwen2 checkpoint with tie_word_embeddings: false fails. The Qwen2.5 Instruct checkpoints tie their embeddings and load. --- Compare and Contrast | Model | Size | GSM8K | BFCL | 16 requests | Long prompts | |---|---|---|---|---|---| | Gemma 4 E2B | 5.12B | 0.911 | 0.927 | 1.00x | 1.00x | | Qwen3-4B | 4.02B | 0.915 | 0.940 | 0.57x | 0.36x | | SmolLM3-3B | 3.08B | 0.858 | 0.945 | 0.76x | 0.69x | | Qwen3-1.7B | 2.03B | 0.827 | 0.935 | 1.09x | 0.88x | | Granite 3.3 2B | 2.53B | 0.817 | 0.920 | 0.93x | 0.70x | Long prompts are 16 requests of 3,072 tokens; every ratio is against Gemma on the same image. --- So, Which One? For the best accuracy per chip, Gemma 4 E2B: no other model matches both its scores and its speed on a v5e-1. For accuracy with standard attention, Qwen3-4B, level with Gemma on both tests at a bit over half the output. For tool calling at Gemma's size with the plainest architecture, SmolLM3-3B: level on BFCL, 5.3 points lower on math, three quarters of the speed. For tool calling at close to Gemma's speed, Granite 3.3 2B or Qwen3-1.7B, if math matters less. The models from 1.7B down run faster than Gemma but, apart from Qwen3-1.7B, give up 18 to 45 points of GSM8K. --- Teardown Each queued resource deletes itself after its last model. Confirm nothing is left: gcloud alpha compute tpus queued-resources list --zone us-west4-a --- Summary The goal of this article was to find which small Apache-2.0 or MIT model can stand in for Gemma 4 E2B on one TPU v5e chip. The key to the solution was serving every model on one pinned vLLM image with identical settings, scoring every answer in code, and pairing each model against Gemma on the same problems. The results were: - 🟢 Qwen3-4B is level with Gemma on GSM8K and BFCL, at 0.57x the output - 🟢 SmolLM3-3B, Granite 3.3 2B, Qwen3-1.7B and Qwen2.5-1.5B are level with Gemma on BFCL tool calls - 🟢 Gemma has the largest weights in the lineup and still out-serves every model with more than 4 GiB of weights - ⚠️ Gemma's shared, single-head, sliding-window KV cache holds 2.6 to 6.3 times the tokens of Granite, SmolLM3, Qwen3-1.7B and Qwen3-4B, and keeps its lead with 3,072-token prompts - ⚠️ Every model under 2B parameters except Qwen3-1.7B scores 18 to 45 points below Gemma on GSM8K - ❌ DeepSeek-R1-Distill-Qwen-1.5B and OLMo 2 1B do not load on this image Scope: one TPU v5e chip (v5litepod-1) in us-west4-a, flex-start, two runs on 2026-10-11 and 2026-10-12 on the same image digest, vLLM 0.29.1rc1.dev468+g0b7f11a1e. Speed is the median of three passes; accuracy is one greedy pass. Gemma's figures come from the first run and pair with the second through the shared image. OLMo 2 1B fails to compile (it asks for 133.56 MB of the chip's 128 MB vector memory). Ministral 3 3B and Qwen3.5 0.8B and 2B served but are left out of the tables: on this image Ministral returned server errors on part of each test and Qwen3.5 returned unusable text, so their scores describe the serving path. Phi-4-mini produced no parsed tool calls with the phi4_mini_json parser and is reported on GSM8K and speed only. Parts of the analysis and writing were done with AI assistance (Claude); every figure comes from the committed report files. The strategy for surveying small open models against Gemma 4 on a TPU v5e was validated with an incremental step by step approach. --- References - Code, logs, reports and scored answers: xbill9/model-sweep - vLLM TPU: vllm-project/tpu-inference - Gemma 4 E2B: huggingface.co - Qwen3-4B: huggingface.co - SmolLM3-3B: huggingface.co - Granite 3.3 2B: huggingface.co - GSM8K: huggingface.co - Berkeley Function Calling Leaderboard: gorilla.cs.berkeley.edu - Cloud TPU v5e: cloud.google.com