One TPU Chip, Eight Agents: Serving Small Agent Workloads with Raw JAX

Short summary

An expert deep-dive into serving small AI agent workloads (2-8 concurrent agents) on a single Cloud TPU v6e chip using raw JAX instead of vLLM. The author identifies why standard benchmarks don't match agent workload patterns (low concurrency, serial latency, growing context), documents QAT checkpoint loading failures in vLLM, and implements a pure-JAX inference path for Gemma 4 E2B. Includes benchmark data and references to published measurement studies.

• •Standard inference benchmarks optimize for high-concurrency public endpoints, not low-concurrency agent loops
• •QAT exports of Gemma 4 E2B fail to load in vLLM on TPU due to missing K/V-shared layer parameters
• •Author built a pure-JAX inference path (no PyTorch/torch_xla) as a working alternative

Generated with AI, which can make mistakes.