Gemma 4 E2B in JAX on One TPU v5e: bf16 or 4-Bit, Python or Rust, and How It Stands Against vLLM

7.7 relevance

Benchmarking Gemma 4 on TPU with JAX/Rust/vLLM, directly relevant to AI/ML infrastructure.

AI/ML dev.to

Answer the board on AI readiness

Arrive with evidence, not opinion. See how your foundation benchmarks against your peers, get a clear read, and leave with a ranked place to start.

Why It Matters

For platform engineers optimizing AI inference on GCP TPUs, this provides a direct comparison of JAX vs vLLM performance on the smallest v5e chip, with actionable data on quantization tradeoffs and Python vs Rust overhead.

Key Takeaways

Evaluate whether 4-bit quantization improves throughput on models that already fit in bf16, and consider Rust-based JAX engines to reduce Python overhead in latency-critical serving.

Gemma 4 E2B in JAX on One TPU v5e: bf16 or 4-Bit, Python or Rust, and How It Stands Against vLLM

Summary

A pure-JAX engine serves Gemma 4 E2B on a single TPU v5e chip (16.9 GB HBM) using bf16 or 4-bit quantization, with a Rust-based variant that bypasses Python overhead by handing the graph directly to libtpu. Benchmarks compare both against vLLM on the same flex-start VM, measuring whether 4-bit pays off when the 9.26 GB bf16 model already fits, how much latency comes from Python, and where a hand-written JAX engine stands relative to vLLM. The engine loads Hugging Face safetensors without PyTorch, speaks the OpenAI chat API, and uses jitted prefill with a static KV cache.

Author

xbill