{"slug": "one-tpu-chip-eight-agents-serving-small-agent-workloads-with-raw-jax", "title": "One TPU Chip, Eight Agents: Serving Small Agent Workloads with Raw JAX", "summary": "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.", "body_md": "*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.*\n\nServing 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.\n\nA small agent workload looks nothing like that:\n\nWorth 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.\n\nThat workload fits on one chip — if you can serve the model at all. Which brings us to the problem.\n\nQuantization-aware-trained exports of Gemma 4 E2B **do not load in vLLM on TPU**:\n\n| Checkpoint | Backend | Failure | Verdict |\n|---|---|---|---|\n`-qat-w4a16-ct` |\ntpu-inference (JAX) |\n`int4 compressed-tensors` scheme unimplemented for `per_layer_model_projection`\n|\n✕ no load |\n`-qat-q4_0-unquantized` |\ntpu-inference (JAX) | demands `self_attn.k_norm.weight` for layers 15–34 |\n✕ no load |\n`-qat-q4_0-unquantized` |\ntorchax (`MODEL_IMPL_TYPE=vllm` ) |\nidentical missing-weights error | ✕ no load |\n`google/gemma-4-E2B-it` (BF16) |\ntpu-inference (JAX) | — | ✓ serves |\n\n**The checkpoint is right and the loader is wrong.** Comparing safetensors headers: the BF16 export ships `self_attn.k_norm`\n\nfor all 35 layers; the QAT export ships it only for the 15 non-KV-shared layers. Both configs declare `num_kv_shared_layers: 20`\n\n.\n\nLayers 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.\n\nSo: serve BF16 and forgo QAT, or write the inference path yourself. I wrote it.\n\n`ports/gemma4/jax_e_model.py`\n\nis a pure-JAX Gemma 4 E2B — no PyTorch, no `torch_xla`\n\n, nowhere in the path. That's not an aspiration: the checkpoint goes safetensors → JAX PyTree directly (via `safetensors.flax`\n\n, which handles bfloat16 natively), and `tests/test_jax_engine.py`\n\nasserts `torch`\n\nnever even enters `sys.modules`\n\nduring a load. If you're used to \"JAX\" inference that quietly loads weights through `AutoModelForCausalLM`\n\nand converts, this isn't that.\n\nSharing is encoded in the layer definition instead of patched at load time:\n\n``` php\n@property\ndef first_kv_shared_layer_idx(self) -> int:\n    return self.num_hidden_layers - self.num_kv_shared_layers   # 35 - 20 = 15\n\ndef kv_share_map(self) -> List[int]:\n    \"\"\"Maps each layer index to the source layer index for KV state sharing.\"\"\"\n    first = self.first_kv_shared_layer_idx\n    last_of_type = {}\n    for i in range(first):\n        last_of_type[self.layer_types[i]] = i\n    return [\n        i if i < first else last_of_type[self.layer_types[i]]\n        for i in range(self.num_hidden_layers)\n    ]\n```\n\nSharing 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)`\n\never 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.\n\nDecoding runs against a static KV cache written with `dynamic_update_slice`\n\n, 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`\n\nruns greedy generation both ways, cached decode versus re-running the full model over the growing sequence every step:\n\n```\nstep | max|dlogit| | ref tok  | cached tok | ref top-2 gap\n   0 |    0.000000 | [62, 36] | [62, 36]   | [0.009738, 0.118984]\n   ...\n   7 |    0.000001 | [97, 95] | [97, 95]   | [0.012065, 0.000438]\n```\n\nMaximum 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.\n\nHere 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.\n\nE2B's KV sharing means only 15 of 35 layers hold KV tensors. Twelve sliding-attention\n\nlayers use a 256-wide KV head; three full-attention layers use a 512-wide one:\n\n```\n12 × 1 × 256 × (K+V) + 3 × 1 × 512 × (K+V)\n    = 9,216 KV elements/token\n    = 18.00 KiB/token at bf16\n```\n\nThat 18.00 KiB figure is computed from the checkpoint shapes and reproduced by\n\n`init_kv_cache`\n\n. An earlier 15.0 KiB estimate treated every layer as 256-wide and\n\nunder-counted the three full-attention layers. So the resident bf16 KV alone costs:\n\n| Fleet | Context each | Resident KV tokens | KV total (bf16) |\n|---|---|---|---|\n| 4 agents | 4,096 | 16,384 | 302 MB |\n| 8 agents | 4,096 | 32,768 | 604 MB |\n| 8 agents | 8,192 | 65,536 | 1.21 GB |\n| 32 agents | 8,192 | 262,144 | 4.83 GB |\n\nThe real QAT checkpoint loads as 6.56 GB of weights, so eight 8K contexts plus\n\nweights account for about **7.77 GB** before activations, executable buffers and\n\nallocator overhead. Do not subtract that from 32 GB and call the remainder free:\n\ncapacity measurements show non-KV temporaries can consume roughly 40% of HBM at\n\nlarge batch. The defensible conclusion is narrower: **eight 8K contexts fit\ncomfortably; the weight-plus-KV arithmetic is not the limiting admission check.**\n\nAt ctx 8,192, the donated decode path was bisected at 1,425,408 resident KV tokens\n\nfor bf16 and 2,818,048 for int8: B=174 and B=344 respectively. Those are decode\n\nresidency ceilings, not promises that the same number of full prompts can be\n\nprefilled together. Prefill has a separate measured rule:\n\n`batch × chunk_size <= 8,192`\n\nprompt tokens per pass.\n\nThe binding constraint is compute and scheduling. Which is where it gets interesting.\n\nThis is the honest core of the assessment. Three gaps, ranked by how much they'd hurt an agent workload specifically:\n\n**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)).\n\n**2. No guided decoding.** Tool calls have to be parseable JSON. vLLM has\n\nschema/grammar-constrained generation; this raw JAX server does not. On the\n\nmeasured Gemma 4 vLLM stack, schema enforcement was itself conditional: it worked\n\nwith thinking enabled and failed or partially worked with thinking disabled.\n\nClient-side validation remains necessary on either path, but raw JAX still lacks\n\nthe server-side constraint layer entirely.\n\n**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.\n\nPlus the standing costs of the approach: static shapes mean any uncompiled `(batch, seq)`\n\ncombination stalls on XLA compilation, so you bucket and spend FLOPs on padding; and every new architecture or quantization format is hand-written.\n\nAgainst that, what raw JAX genuinely buys you: **it runs the checkpoint vLLM\ncan't**, and the whole stack is inspectable. A persistent JAX compilation cache\n\nFor calibration, vLLM serving the BF16 checkpoint on the same chip, measured with `vllm bench serve`\n\n(1,024-token input, 128-token output):\n\n| Concurrency | Output tok/s | Per-stream tok/s | Median TTFT | Median TPOT |\n|---|---|---|---|---|\n| 1 | 209 | 213 | 16 ms | 4.7 ms |\n| 8 | 1,209 | 161 | 27 ms | 6.2 ms |\n| 32 | 1,636 | 57 | 155 ms | 17.5 ms |\n| 64 | 2,140 | 39 | 122 ms | 25.3 ms |\n| 100 | 2,215 | 27 | 833 ms | 36.8 ms |\n\nNote 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.\"\n\nThis distinction matters: correctness was tested with the real checkpoint;\n\nperformance was measured with deterministic, architecture-shaped synthetic\n\nparameters. Under static JAX shapes the values do not change the compiled work,\n\nbut this is still a **decode-kernel benchmark**, not an end-to-end serving result\n\nand not a direct vLLM comparison.\n\nThe final revalidation used the checkpoint's dimensions, one configuration per\n\nprocess, 15 timed samples after warmup, ctx 8,192 and B=32. The donation and\n\ncache-dtype effects exceeded twice the measured IQR; the PLE differences did\n\nnot. The largest relative IQR was 1.32%.\n\n| PLE | KV | cache update | step | aggregate |\n|---|---|---|---|---|\n| bf16 | bf16 | copying | 21.262 ms | 1,505 tok/s |\n| bf16 | bf16 | donated |\n13.143 ms | 2,435 tok/s |\n| bf16 | int8 | donated |\n11.088 ms | 2,886 tok/s |\n| int4 | int8 | donated |\n11.080 ms |\n2,888 tok/s |\n\nThe largest fix was not quantization. It was buffer donation. Without\n\n`donate_argnums`\n\n, `dynamic_update_slice`\n\ncreated a second full cache to write one\n\ntoken, then attention read it again. Donation made bf16 1.62× faster and nearly\n\ndoubled its resident-token ceiling. On the corrected donated path, int8 KV was\n\n1.17–1.19× faster than bf16 and retained 1.82–1.98× the KV-token capacity.\n\nOn the real checkpoint, int8 was also indistinguishable from bf16 in a 583-step\n\nteacher-forced quality check (28.41 versus 28.73 perplexity, 97.08% greedy\n\nmatch), and its error did not grow across a separate 968-step continuous decode.\n\nQuantizing the 4.70 GB per-layer embedding table to int4 reduced the measured\n\nweight tree from 6.618 to 3.113 GB, but did not move step time: 13.143 versus\n\n13.099 ms with bf16 KV, a 0.3% difference. That table is gathered, not streamed\n\nin full every token. Quantizing it buys residency, not throughput.\n\nThe roofline cross-check also changed the conclusion. A calibrated streaming\n\nreduction reached 1,417 GB/s. The donated decode step's marginal KV read reached\n\n794 GB/s, about 56% of that; without donation it reached 276 GB/s, or 19%.\n\nThis engine still has headroom, but eager attention was not the dominant mystery:\n\nattention in isolation measured in the same 39–59% range. The extra cache copies\n\nwere.\n\nI then ran the real checkpoint through the OpenAI-compatible HTTP endpoint. With\n\nint8 KV and donation enabled, warm single-stream decode stayed nearly flat as\n\ncontext grew. The prompt repeated a fixed sentence to reach each token count, so\n\nthis is a latency/scheduling test rather than a quality evaluation:\n\n| Prompt tokens | HTTP wall time | Prefill | Decode |\n|---|---|---|---|\n| 506 | 240.5 ms | 9.3 ms | 141.1 tok/s |\n| 2,045 | 267.7 ms | 32.0 ms | 140.5 tok/s |\n| 7,679 | 570.9 ms | 318.6 ms | 138.5 tok/s |\n\nThe long-context penalty is almost entirely prefill. But simultaneous HTTP\n\nrequests expose the more important limitation. At a 2K prompt and 32 output\n\ntokens, concurrency 2/4/8 produced 128.7/133.6/143.3 aggregate tok/s while\n\nmedian request latency rose from 497 to 952 to 1,775 ms. Every request succeeded;\n\nnone formed a device batch.\n\nThat is exactly what the code says should happen: `generate_stream`\n\nis B=1, and\n\nFastAPI launches independent B=1 executions. The device queues them. The\n\n2,888 tok/s kernel point proves batched decode work can be fast; the current\n\nserver cannot reach it until it owns a real request batcher and batched KV state.\n\n`jax_openai_server.py`\n\nruns generation on the JAX engine: `/v1/chat/completions`\n\n, `/v1/completions`\n\n, `/health`\n\n, Prometheus `/metrics`\n\n, and SSE streaming that emits one chunk per decoded token.\n\n```\npython3 jax_openai_server.py --port 8000 --max-model-len 4096\ncurl http://localhost:8000/v1/chat/completions \\\n  -H \"Content-Type: application/json\" \\\n  -d '{\n    \"model\": \"google/gemma-4-E2B-it-qat-w4a16-ct\",\n    \"messages\": [{\"role\": \"user\", \"content\": \"Explain TPU MXU vectorization.\"}],\n    \"stream\": true\n  }'\n```\n\n`tests/test_openai_server.py`\n\ndrives 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`\n\nadditionally asserts the load path never imports `torch`\n\n.\n\n**Is raw JAX viable for a small agent workload on one TPU chip? As an\nexperimental serving path, yes. As a vLLM replacement, not yet.**\n\nThe strongest evidence is correctness and capacity: the QAT checkpoint that\n\nblocks vLLM loads and generates; cached decode matches a full re-forward; eight\n\n8K bf16 contexts require 1.21 GB of KV; and donated decode has far more residency\n\nheadroom than that. The kernel result at B=32 is promising, but it is not an\n\nOpenAI-serving benchmark. The real HTTP test establishes the current serving\n\nlimit instead: roughly **139–141 tok/s aggregate**, with latency growing almost\n\nlinearly when 2–8 independent requests contend for the device.\n\nThe deciding gaps are prefix caching, constrained decoding and continuous\n\nbatching. If agents resend a growing history, the current server re-prefills it\n\nquadratically over the conversation. Chunked prefill makes large batches\n\nadmissible, but it does not reuse a previous turn's prefix and currently runs\n\nabout 5× slower than windowed one-shot prefill at B=8, ctx 2,048.\n\nUse this path when you specifically need the QAT checkpoint, can tolerate an\n\nexperimental scheduler, and value a stack you can inspect end to end. For a\n\nproduction agent service today, the evidence still favours vLLM with the BF16\n\ncheckpoint while waiting for #3225.\n\n```\npython3 benchmarks/queued/revalidate.py --all --out results.json\n```\n\nThis harness uses synthetic parameter values but checkpoint-matched dimensions,\n\nruns each configuration in a fresh process and measures both donated and copying\n\npaths. Keep performance and capacity separate: use it for step latency, and the\n\ndedicated capacity bisection in the run directory for HBM ceilings. A doubling\n\nladder previously manufactured a false \"constant\" ceiling, and a non-donated\n\nstep later measured exactly half the usable capacity.\n\nWorth 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`\n\non five of eight cells because the kernel loads all of `x`\n\nas 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.\n\nNeither 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.\n\nThose failures do **not** prove weight traffic is irrelevant. The old ~3.6 ms\n\nfloor omitted KV traffic, and the old 20.6 ms step was measured before the\n\ndonation fix. The narrower conclusion survives: these two implementations did\n\nnot pay. Profile attribution and a calibrated bandwidth bound were more useful\n\nthan another speculative kernel.\n\nThe 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.\n\nThere'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`\n\nfell 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.\n\nThe next engineering step is no longer another kernel benchmark. It is a request\n\nbatcher with batched KV ownership, continuous admission and prefix reuse, followed\n\nby the same 2–8 stream HTTP test. Peak aggregate kernel throughput cannot answer\n\nthat question by itself.\n\n`ports/gemma4/jax_e_model.py`\n\n`ports/gemma4/jax_e_loader.py`\n\n`jax_engine.py`\n\n`jax_openai_server.py`\n\n`tests/test_kv_cache_parity.py`\n\n`tests/test_jax_engine.py`\n\n`tests/test_openai_server.py`\n\n`ports/gemma4/jax_e_benchmark_sweep_v2.py`\n\n`benchmarks/runs/2026-07-29-kv-quant-v6e1/REPORT.md`\n\n`benchmarks/runs/2026-07-29-real-http-v6e1/REPORT.md`\n\n`benchmarks/reports/2026-07-21-gemma4-e2b-v6e1.json`\n\nRubens de Almeida Zimbres' TPU inference measurement series (CC-BY-4.0) is the closest published\n\nwork to this, and several of its results are cited above. Notes on how each bears on the engine\n\nhere are in [ docs/references/tpu-inference-measurement-series.md](https://github.com/xbill9/tpu-jax/blob/main/docs/references/tpu-inference-measurement-series.md).", "url": "https://wpnews.pro/news/one-tpu-chip-eight-agents-serving-small-agent-workloads-with-raw-jax", "canonical_source": "https://dev.to/xbill/one-tpu-chip-eight-agents-serving-small-agent-workloads-with-raw-jax-2cc4", "published_at": "2026-07-29 22:02:22+00:00", "updated_at": "2026-07-29 22:33:34.110572+00:00", "lang": "en", "topics": ["artificial-intelligence", "large-language-models", "ai-infrastructure", "developer-tools"], "entities": ["Google", "Gemma 4 E2B", "vLLM", "Cloud TPU v6e", "JAX", "Zimbres"], "alternates": {"html": "https://wpnews.pro/news/one-tpu-chip-eight-agents-serving-small-agent-workloads-with-raw-jax", "markdown": "https://wpnews.pro/news/one-tpu-chip-eight-agents-serving-small-agent-workloads-with-raw-jax.md", "text": "https://wpnews.pro/news/one-tpu-chip-eight-agents-serving-small-agent-workloads-with-raw-jax.txt", "jsonld": "https://wpnews.pro/news/one-tpu-chip-eight-agents-serving-small-agent-workloads-with-raw-jax.jsonld"}}