One TPU Chip, Eight Agents: Serving Small Agent Workloads with Raw JAX
Scored daily by a customisable AI persona to surface the most relevant engineering leadership news.
Deep technical dive into serving AI agents on TPU with JAX, highly relevant to agent orchestration.
Small agent workloads (2–8 agents, latency-sensitive serial loops, growing context, structured JSON output) are poorly served by LLM benchmarks optimized for high concurrency. The author encountered a vLLM bug (#3225) that fails to load Gemma 4 QAT checkpoints on TPU because the loader unconditionally expects K-norm weights for all 35 layers, ignoring KV-shared layers 15–34 where K-params are architecturally absent. To serve these workloads on a single TPU v6e chip, they built a pure-JAX inference engine (no PyTorch/torch_xla) that loads safetensors directly into a JAX PyTree, bypassing the broken loader and achieving the correct operating point for agent latency bounds.