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

A developer built a pure-JAX inference engine for Google's Gemma 4 E2B model after discovering that quantization-aware-trained (QAT) checkpoints fail to load in vLLM on TPU due to a loader bug. The engine, which fits on a single Cloud TPU v6e chip, handles KV-sharing correctly by encoding it in the layer definition rather than patching at load time. The developer argues that small agent workloads are latency-bound and best measured by goodput, contrasting with throughput-focused benchmarks.

Cloud TPU v6e-1 ct6e-standard-1t, one v6e chip, 32 GB HBM , GCE flex-start, europe-west4-a. vLLM baseline measured 2026-07-21. Serving benchmarks optimize for the wrong shape. They report throughput at concurrency 100 with 1,024-token prompts, because that's what a public inference endpoint looks like. A small agent workload looks nothing like that: Worth flagging that this premise is contested. Zimbres' measurement series argues the reverse for agent fleets: agents "exchange messages among themselves, multiplying token volume by an order of magnitude or more per user task, with no human waiting on any individual token. There, per-token latency loses most of its meaning and aggregate throughput becomes the binding constraint" DOI 10.5281/zenodo.21212010 https://doi.org/10.5281/zenodo.21212010 . Both are right, for different topologies — a serial agent loop with a human at the end is latency-bound; a parallel machine-to-machine fleet is throughput-bound. The metric that covers both is goodput : throughput subject to a per-token latency objective. I use that framing below, but I do not claim a production goodput result from a kernel harness. That workload fits on one chip — if you can serve the model at all. Which brings us to the problem. Quantization-aware-trained exports of Gemma 4 E2B do not load in vLLM on TPU : | Checkpoint | Backend | Failure | Verdict | |---|---|---|---| -qat-w4a16-ct | tpu-inference JAX | int4 compressed-tensors scheme unimplemented for per layer model projection | ✕ no load | -qat-q4 0-unquantized | tpu-inference JAX | demands self attn.k norm.weight for layers 15–34 | ✕ no load | -qat-q4 0-unquantized | torchax MODEL IMPL TYPE=vllm | identical missing-weights error | ✕ no load | google/gemma-4-E2B-it BF16 | tpu-inference JAX | — | ✓ serves | The checkpoint is right and the loader is wrong. Comparing safetensors headers: the BF16 export ships self attn.k norm for all 35 layers; the QAT export ships it only for the 15 non-KV-shared layers. Both configs declare num kv shared layers: 20 . Layers 15–34 reuse K/V computed by lower layers. They hold no K-side parameters, so a k-norm for them is meaningless — the QAT export is the architecturally honest one, and the plain checkpoint only loads because it happens to carry the unused tensors along. Filed as tpu-inference 3225 https://github.com/vllm-project/tpu-inference/issues/3225 ; the fix is to skip instantiating K/V-side parameters for shared layers rather than demanding them unconditionally. So: serve BF16 and forgo QAT, or write the inference path yourself. I wrote it. ports/gemma4/jax e model.py is a pure-JAX Gemma 4 E2B — no PyTorch, no torch xla , nowhere in the path. That's not an aspiration: the checkpoint goes safetensors → JAX PyTree directly via safetensors.flax , which handles bfloat16 natively , and tests/test jax engine.py asserts torch never even enters sys.modules during a load. If you're used to "JAX" inference that quietly loads weights through AutoModelForCausalLM and converts, this isn't that. Sharing is encoded in the layer definition instead of patched at load time: php @property def first kv shared layer idx self - int: return self.num hidden layers - self.num kv shared layers 35 - 20 = 15 def kv share map self - List int : """Maps each layer index to the source layer index for KV state sharing.""" first = self.first kv shared layer idx last of type = {} for i in range first : last of type self.layer types i = i return i if i < first else last of type self.layer types i for i in range self.num hidden layers Sharing is per attention type — E2B interleaves sliding and full attention, and a shared layer inherits from the last non-shared layer of its own kind. Only range first kv shared layer idx ever allocates K/V. W4A16 weights eight int4 values per int32, BF16 scale per 32-element group are unpacked on chip, read straight from safetensors into a JAX PyTree. Decoding runs against a static KV cache written with dynamic update slice , so each step attends to the full prefix. The way you know a decoder is correct is to compare it against the slow thing that obviously is — tests/test kv cache parity.py runs greedy generation both ways, cached decode versus re-running the full model over the growing sequence every step: step | max|dlogit| | ref tok | cached tok | ref top-2 gap 0 | 0.000000 | 62, 36 | 62, 36 | 0.009738, 0.118984 ... 7 | 0.000001 | 97, 95 | 97, 95 | 0.012065, 0.000438 Maximum divergence across every step: 1e-6 , float32 roundoff. The test asserts logit agreement within tolerance and token agreement only where the top-2 gap exceeds it — near a 0.0004 tie, an exact argmax assertion is flaky by construction. A companion test proves padding a prompt to a 128-aligned bucket is invisible to the output: RoPE follows logical position, not cache slot. Here is the result that actually settles the "can one chip hold my agent fleet" question, and it needs no benchmark at all — just the architecture. E2B's KV sharing means only 15 of 35 layers hold KV tensors. Twelve sliding-attention layers use a 256-wide KV head; three full-attention layers use a 512-wide one: 12 × 1 × 256 × K+V + 3 × 1 × 512 × K+V = 9,216 KV elements/token = 18.00 KiB/token at bf16 That 18.00 KiB figure is computed from the checkpoint shapes and reproduced by init kv cache . An earlier 15.0 KiB estimate treated every layer as 256-wide and under-counted the three full-attention layers. So the resident bf16 KV alone costs: | Fleet | Context each | Resident KV tokens | KV total bf16 | |---|---|---|---| | 4 agents | 4,096 | 16,384 | 302 MB | | 8 agents | 4,096 | 32,768 | 604 MB | | 8 agents | 8,192 | 65,536 | 1.21 GB | | 32 agents | 8,192 | 262,144 | 4.83 GB | The real QAT checkpoint loads as 6.56 GB of weights, so eight 8K contexts plus weights account for about 7.77 GB before activations, executable buffers and allocator overhead. Do not subtract that from 32 GB and call the remainder free: capacity measurements show non-KV temporaries can consume roughly 40% of HBM at large batch. The defensible conclusion is narrower: eight 8K contexts fit comfortably; the weight-plus-KV arithmetic is not the limiting admission check. At ctx 8,192, the donated decode path was bisected at 1,425,408 resident KV tokens for bf16 and 2,818,048 for int8: B=174 and B=344 respectively. Those are decode residency ceilings, not promises that the same number of full prompts can be prefilled together. Prefill has a separate measured rule: batch × chunk size <= 8,192 prompt tokens per pass. The binding constraint is compute and scheduling. Which is where it gets interesting. This is the honest core of the assessment. Three gaps, ranked by how much they'd hurt an agent workload specifically: 1. No prefix caching — and agent loops are the pathological case. Every agent turn resends the entire conversation plus every tool result so far. With prefix caching, turn 20 reprocesses only the new tokens. Without it, you re-prefill the whole growing context every single turn. Over a 20-turn loop with context growing 500 → 8,000 tokens, that's quadratic wasted prefill. vLLM ships this; raw JAX does not. For agent workloads this is the single biggest gap , and it isn't exotic to build — but you do have to build it. Measured from the other direction: prefix caching "is at its most effective in agent fleets, because agents share long system prompts and accumulated context, so the prompt-processing cost of inter-agent messages collapses" DOI 10.5281/zenodo.21212010 https://doi.org/10.5281/zenodo.21212010 . 2. No guided decoding. Tool calls have to be parseable JSON. vLLM has schema/grammar-constrained generation; this raw JAX server does not. On the measured Gemma 4 vLLM stack, schema enforcement was itself conditional: it worked with thinking enabled and failed or partially worked with thinking disabled. Client-side validation remains necessary on either path, but raw JAX still lacks the server-side constraint layer entirely. 3. Lockstep batching penalizes uneven turns. A fixed batch runs together and a finished sequence holds its slot until the whole batch drains. Agent turns are wildly variable — one agent emits a 12-token tool call while another writes 400 tokens of reasoning. vLLM's continuous batching swaps in new work per step; here, short turns wait on long ones. At the low concurrency agent fleets run at, this is a real efficiency tax. Plus the standing costs of the approach: static shapes mean any uncompiled batch, seq combination stalls on XLA compilation, so you bucket and spend FLOPs on padding; and every new architecture or quantization format is hand-written. Against that, what raw JAX genuinely buys you: it runs the checkpoint vLLM can't , and the whole stack is inspectable. A persistent JAX compilation cache For calibration, vLLM serving the BF16 checkpoint on the same chip, measured with vllm bench serve 1,024-token input, 128-token output : | Concurrency | Output tok/s | Per-stream tok/s | Median TTFT | Median TPOT | |---|---|---|---|---| | 1 | 209 | 213 | 16 ms | 4.7 ms | | 8 | 1,209 | 161 | 27 ms | 6.2 ms | | 32 | 1,636 | 57 | 155 ms | 17.5 ms | | 64 | 2,140 | 39 | 122 ms | 25.3 ms | | 100 | 2,215 | 27 | 833 ms | 36.8 ms | Note the shape at agent-scale concurrency: at 8 streams it holds 161 tok/s per stream at 27 ms TTFT . That is a good bar, and it's what you give up by leaving vLLM. The honest framing is not "raw JAX beats vLLM" — it's "vLLM can't load this checkpoint, and here's what the alternative costs you." This distinction matters: correctness was tested with the real checkpoint; performance was measured with deterministic, architecture-shaped synthetic parameters. Under static JAX shapes the values do not change the compiled work, but this is still a decode-kernel benchmark , not an end-to-end serving result and not a direct vLLM comparison. The final revalidation used the checkpoint's dimensions, one configuration per process, 15 timed samples after warmup, ctx 8,192 and B=32. The donation and cache-dtype effects exceeded twice the measured IQR; the PLE differences did not. The largest relative IQR was 1.32%. | PLE | KV | cache update | step | aggregate | |---|---|---|---|---| | bf16 | bf16 | copying | 21.262 ms | 1,505 tok/s | | bf16 | bf16 | donated | 13.143 ms | 2,435 tok/s | | bf16 | int8 | donated | 11.088 ms | 2,886 tok/s | | int4 | int8 | donated | 11.080 ms | 2,888 tok/s | The largest fix was not quantization. It was buffer donation. Without donate argnums , dynamic update slice created a second full cache to write one token, then attention read it again. Donation made bf16 1.62× faster and nearly doubled its resident-token ceiling. On the corrected donated path, int8 KV was 1.17–1.19× faster than bf16 and retained 1.82–1.98× the KV-token capacity. On the real checkpoint, int8 was also indistinguishable from bf16 in a 583-step teacher-forced quality check 28.41 versus 28.73 perplexity, 97.08% greedy match , and its error did not grow across a separate 968-step continuous decode. Quantizing the 4.70 GB per-layer embedding table to int4 reduced the measured weight tree from 6.618 to 3.113 GB, but did not move step time: 13.143 versus 13.099 ms with bf16 KV, a 0.3% difference. That table is gathered, not streamed in full every token. Quantizing it buys residency, not throughput. The roofline cross-check also changed the conclusion. A calibrated streaming reduction reached 1,417 GB/s. The donated decode step's marginal KV read reached 794 GB/s, about 56% of that; without donation it reached 276 GB/s, or 19%. This engine still has headroom, but eager attention was not the dominant mystery: attention in isolation measured in the same 39–59% range. The extra cache copies were. I then ran the real checkpoint through the OpenAI-compatible HTTP endpoint. With int8 KV and donation enabled, warm single-stream decode stayed nearly flat as context grew. The prompt repeated a fixed sentence to reach each token count, so this is a latency/scheduling test rather than a quality evaluation: | Prompt tokens | HTTP wall time | Prefill | Decode | |---|---|---|---| | 506 | 240.5 ms | 9.3 ms | 141.1 tok/s | | 2,045 | 267.7 ms | 32.0 ms | 140.5 tok/s | | 7,679 | 570.9 ms | 318.6 ms | 138.5 tok/s | The long-context penalty is almost entirely prefill. But simultaneous HTTP requests expose the more important limitation. At a 2K prompt and 32 output tokens, concurrency 2/4/8 produced 128.7/133.6/143.3 aggregate tok/s while median request latency rose from 497 to 952 to 1,775 ms. Every request succeeded; none formed a device batch. That is exactly what the code says should happen: generate stream is B=1, and FastAPI launches independent B=1 executions. The device queues them. The 2,888 tok/s kernel point proves batched decode work can be fast; the current server cannot reach it until it owns a real request batcher and batched KV state. jax openai server.py runs generation on the JAX engine: /v1/chat/completions , /v1/completions , /health , Prometheus /metrics , and SSE streaming that emits one chunk per decoded token. python3 jax openai server.py --port 8000 --max-model-len 4096 curl http://localhost:8000/v1/chat/completions \ -H "Content-Type: application/json" \ -d '{ "model": "google/gemma-4-E2B-it-qat-w4a16-ct", "messages": {"role": "user", "content": "Explain TPU MXU vectorization."} , "stream": true }' tests/test openai server.py drives the real endpoints, including an assertion that streaming and non-streaming produce identical greedy text — precisely the thing that silently rots when streaming is bolted on separately. tests/test jax engine.py additionally asserts the load path never imports torch . Is raw JAX viable for a small agent workload on one TPU chip? As an experimental serving path, yes. As a vLLM replacement, not yet. The strongest evidence is correctness and capacity: the QAT checkpoint that blocks vLLM loads and generates; cached decode matches a full re-forward; eight 8K bf16 contexts require 1.21 GB of KV; and donated decode has far more residency headroom than that. The kernel result at B=32 is promising, but it is not an OpenAI-serving benchmark. The real HTTP test establishes the current serving limit instead: roughly 139–141 tok/s aggregate , with latency growing almost linearly when 2–8 independent requests contend for the device. The deciding gaps are prefix caching, constrained decoding and continuous batching. If agents resend a growing history, the current server re-prefills it quadratically over the conversation. Chunked prefill makes large batches admissible, but it does not reuse a previous turn's prefix and currently runs about 5× slower than windowed one-shot prefill at B=8, ctx 2,048. Use this path when you specifically need the QAT checkpoint, can tolerate an experimental scheduler, and value a stack you can inspect end to end. For a production agent service today, the evidence still favours vLLM with the BF16 checkpoint while waiting for 3225. python3 benchmarks/queued/revalidate.py --all --out results.json This harness uses synthetic parameter values but checkpoint-matched dimensions, runs each configuration in a fresh process and measures both donated and copying paths. Keep performance and capacity separate: use it for step latency, and the dedicated capacity bisection in the run directory for HBM ceilings. A doubling ladder previously manufactured a false "constant" ceiling, and a non-donated step later measured exactly half the usable capacity. Worth recording, because the reasoning was plausible and the result wasn't. Decode reads the full weight set per token, so 4-bit weights ought to cut traffic ~4× — a fused Pallas kernel that unpacks int4 inside the VMEM tile instead of dequantizing to BF16 in HBM first. Predicted ~3.4× on the memory floor. Measured: 0.59× at B=1, 0.21× at B=2 , and CompileTimeScopedVmemOom on five of eight cells because the kernel loads all of x as one VMEM block rather than tiling the sequence axis. Quantizing the LM head to int8 — the single largest read in a step — bought 1.00×/1.04×/1.05× at B=1/2/8, for 0.8% logit error. Neither failure was novel. The same series reports speculative decoding costing a factor of six , a precision change that "looked certain on paper" costing 14% and being reverted, and the general rule that this stack is best modified above the compiler configuration or below it whole kernels "rather than through point edits to the compiled middle" — which is exactly what a fused dequant-matmul is DOI 10.5281/zenodo.21212010 https://doi.org/10.5281/zenodo.21212010 . It also brackets the ceiling I was chasing: the entire logits pipeline is 18% of a decode step, so an int8 LM head could never have paid much. Those failures do not prove weight traffic is irrelevant. The old ~3.6 ms floor omitted KV traffic, and the old 20.6 ms step was measured before the donation fix. The narrower conclusion survives: these two implementations did not pay. Profile attribution and a calibrated bandwidth bound were more useful than another speculative kernel. The uncomfortable part is that this is a known failure pattern, not a novel one. The same series ran the fashionable strategies and measured them losing: a team "following the mainstream sequence without profiling would have deployed speculative decoding, measured here at six times slower in this regime; pursued quantization, measured at zero functioning routes on this stack, one failing silently; or hand-optimized a genuine numeric inefficiency, measured at 14 percent slower when removed at the wrong layer." I did the second one, and mine failed silently too. There's a landmine in that kernel worth flagging if you write one: it originally unpacked nibbles plane-major while contracting activations in natural order — silently wrong — and nobody noticed because a bare except Exception fell back to the reference path on any host without a TPU. It only computes garbage where Pallas actually compiles. Verify kernels against a reference on the hardware you ship on, and never let a fallback swallow the exception that would have told you. The next engineering step is no longer another kernel benchmark. It is a request batcher with batched KV ownership, continuous admission and prefix reuse, followed by the same 2–8 stream HTTP test. Peak aggregate kernel throughput cannot answer that question by itself. ports/gemma4/jax e model.py ports/gemma4/jax e loader.py jax engine.py jax openai server.py tests/test kv cache parity.py tests/test jax engine.py tests/test openai server.py ports/gemma4/jax e benchmark sweep v2.py benchmarks/runs/2026-07-29-kv-quant-v6e1/REPORT.md benchmarks/runs/2026-07-29-real-http-v6e1/REPORT.md benchmarks/reports/2026-07-21-gemma4-e2b-v6e1.json Rubens de Almeida Zimbres' TPU inference measurement series CC-BY-4.0 is the closest published work to this, and several of its results are cited above. Notes on how each bears on the engine here are in docs/references/tpu-inference-measurement-series.md https://github.com/xbill9/tpu-jax/blob/main/docs/references/tpu-inference-measurement-series.md .