# Gemma 4 in Pure JAX: What Changes Between Turing and Ada, and What Doesn't > Source: > Published: 2026-08-31 01:40:09+00:00 This article is a measurement report on running a hand-written **Gemma 4** port in **pure JAX** across two NVIDIA GPUs a generation apart, and on the two places the "it's just JAX" abstraction leaks. One of those leaks costs 87% of decode and nothing in the logs is red. The code is here: [https://github.com/xbill9/gemma4-dev](https://github.com/xbill9/gemma4-dev) One port, one build, one checkpoint, two cards. Everything below comes from two archived runs, named so you can check them. | G5g | G6 | | |---|---|---| | Chip | NVIDIA T4G — Turing, SM 7.5, 15,360 MiB | NVIDIA L4 — Ada, SM 8.9, 23,034 MiB | | Host | `g5g.2xlarge` spot — Graviton2, aarch64 | `g6.2xlarge` spot — x86_64, `us-east-1d` | | Checkpoint | `google/gemma-4-E2B-it` , dense reference | `google/gemma-4-E2B-it` , dense reference | | Compute dtype | `float16` (device-chosen) | `bfloat16` (device-chosen) | | Stack | jax 0.11.1, CUDA from pip | jax 0.11.1, Python 3.14 | | Run cited | `2026-08-28-full-run-cached-g5g` | `2026-08-28-first-serve-g6` | Build id `51bc52c9e2e9` on both, config `ple4 + int8_lm_head` , and `tpu_jax_weight_bytes` reads **6,155,450,950** on both cards — the same integer. Only the chip and its host differ. `g5g.2xlarge` and `g6.2xlarge` `google/gemma-4-E2B-it` The port lives in `ports/gemma4/` and is driven by a generation loop behind an OpenAI-compatible server. No PyTorch, no vLLM, no `torch_xla` . The premise under test is that the same source runs on both cards with nothing changed but a config file. It mostly holds. The interesting part is where it does not. Any port has to carry four irregularities, and none of them are optional. `head_dim=256` , global layers use That first irregularity is the expensive one. On the vLLM path the heterogeneous head dims force the Triton attention backend: ``` Gemma4 model has heterogeneous head dimensions (sliding=256, global=512); falling back to the Triton attention backend ``` On a Turing GPU that backend then asks for shared memory the hardware does not have: ``` triton.runtime.errors.OutOfResources: out of resource: shared memory, Required: 147456, Hardware limit: 65536 ``` **JAX never enters that conversation.** Attention is ordinary XLA rather than a hand-tiled kernel, so there is no per-block shared-memory ceiling in the attention path at all. The irregular geometry that is a special case everywhere else is just array shapes here. This is the single most expensive lesson in the repository. **A wrong compute dtype does not raise. It emulates.** `bfloat16` on a pre-Ampere GPU does not fail — XLA routes it through fp32 and most of decode disappears into conversion. Nothing in the logs is red. So the port does not take the dtype from a config file. It reads the live compute capability off the device: ``` COMPUTE_DTYPE = float16 if IS_PRE_AMPERE else bfloat16 ``` On the SM 8.9 Ada card that resolves to `bfloat16` . On the SM 7.5 Turing card it resolves to `float16` — Turing's only real 16-bit datapath, since it has neither bf16 nor fp8. The first line the server emits is the policy, so a misconfiguration is one `grep` away rather than a mystery in the throughput: ``` INFO ports.gemma4.jax_e_model: jax_e_model device policy: platform=gpu compute_capability=8.9 compute_dtype=bfloat16 pallas_interpret=False ``` `pallas_interpret=False` matters just as much. It is the difference between serving and silently running a simulator. Here is the part that does not port, and it is not a bug. It is a real hardware difference wearing a portable API. The fused **W4A16 kernel is written in Pallas**, and it was tiled for a device with 16 MB of scratchpad per core. At this model's shapes the tiles want **550 KiB to 1.1 MiB per block**. On a GPU, Pallas lowers through Triton, and those tiles become **shared memory**. Turing gives you 64 KiB per block. Ada raises the ceiling, but nowhere near a megabyte. So the fast path runs on **neither card**. The engine computes the requirement at startup and refuses with the arithmetic attached, rather than dying as a cryptic `OutOfResources` at the first token: ``` check_w4a16_fits_scoped_memory() ``` The practical consequence is that both GPU rigs serve the **dense reference checkpoint** at 16-bit. A padding-eviction bug in the KV ring cache cost a week, and it is the kind only Gemma 4's geometry produces. The invariant is that **a cache index is an absolute real position, and padding never occupies an index a real position uses.** A port that right-pads into the `200` , `status: "success"` , and output like`The The The The` .Nothing in the logs is red. Nothing in the metrics is red. The only thing that catches it is a degeneracy check on the output itself, which the server now runs on every response. The scariest bugs in this project all returned success. `jax[cuda13]` supplies CUDA as wheels, so the install needs no CUDA toolkit, no Rust, and no compiler on the box. ``` Install: 117 s, with the cache restore included ``` XLA's persistent compilation cache ports as-is. On the T4G rig it restores **805 files / 12 MB in 6 seconds** onto a fresh instance, from a box that had already `max_new_tokens` is a `static_argnames` entry, so `(bucket, max_tokens)` is the compiled shape on every backend. A harness that does not warm up misreports the rig badly. On the T4G the first request off a fresh engine took **18.06 s against 4.50 s warm** — a 4.0x whole-request ratio, measured in `2026-08-21-cuda13-py314-g5g` .That run also notes something worth repeating: the 56x figure from the first-serve baseline is **TTFT specifically**, not the same measurement as the whole-request ratio. They are not interchangeable. 64 output tokens, concurrency 1, 3 repeats per cell, median. "Decode, gauge" is the engine's steady-state counter. "End-to-end" is wall time over the whole request, prefill included. | Input tokens | T4G gauge | T4G end-to-end | L4 gauge | L4 end-to-end | |---|---|---|---|---| | 41 | 12.9 tok/s | 12.43 tok/s | 48.5 tok/s | 46.23 tok/s | | 521 | 13.0 tok/s | 11.28 tok/s | 48.4 tok/s | 42.87 tok/s | | 2,057 | 12.9 tok/s | 8.22 tok/s | 48.3 tok/s | 34.57 tok/s | | 3,593 | — | — | 48.3 tok/s | 27.55 tok/s | Decode moves 0.8% across a 50x context range on the T4G and 0.4% on the L4. End-to-end falls hard on both. That fall is prefill being linear in the padded bucket, not decode degrading. They are two different claims, and conflating them makes a benchmark a lie. **Quote the gauge.** A cost proportional to the **weights** rather than the context produces exactly this shape, which is why KV is not what sets decode speed on either card — despite Gemma 4's whole KV story. On context specifically: `MAX_MODEL_LEN=4096` is the honest number on the T4G. 4,105 prompt tokens serve; 5,120 fails on a prefill transient. Profiling decode with xprof on the Turing card, 20 decode steps with the service stopped: | 🥉 T4G (SM 7.5) | 🥇 L4 (SM 8.9) | | |---|---|---| | dtype conversion | 54.1% | 0.0% | fp32 `gemvx` | 32.8% | absent | | Tensor Core | 0.0% | 0.0% | | Total kernel time | 1,466.0 ms | 362.8 ms | | Decode, gauge | 12.9 tok/s | 48.4 tok/s | | Peak HBM bandwidth | 298.083 GiB/s | 279.441 GiB/s | | Share of bandwidth roofline | 26% | ~100% | 1,466 ms of kernels across 108 distinct kernels on a Tensor Core GPU, without one Tensor Core firing. More than half of decode went to converting numbers between formats before any math happened. The obvious hypothesis was bf16 weights being converted on a chip with no bf16 datapath. So the checkpoint was converted to float16 host-side and re-run. Parameter dtypes read `{'float16': 541, 'uint8': 1, 'int8': 1}` — and **conversion stayed at 54.0%**. The measurement itself is solid. The same profile on a different instance, a different AMI and a restored cache landed at 1466.0 ms against 1467.1 ms. **1.1 ms apart on 1467.** The Ada card resolves it. Converting the stored weights changed nothing because storage dtype was never the problem: Turing has no native bf16, and the fp32 `gemvx` line is the tell — XLA was round-tripping through fp32 regardless of what the file on disk said. Give it a card where storage and compute dtype actually match, and the 54% conversion and the 32.8% fp32 path vanish **together**. An 87% tax gone, for 3.7x the throughput, and a rig sitting at its bandwidth roofline instead of 26% of it. The `/health` endpoint on the L4 reports `weights=bfloat16 activations=bfloat16` — storage dtype and compute dtype matching for kv_cache=bfloat16 pre_ampere=false the first time on this engine. Tensor Core utilization is **0.0% on the Ada card too** — 100 distinct kernels, 362.8 ms of them, and not one Tensor Core firing. Removing the dtype pressure made the machine roughly four times faster without making it touch the hardware it was sold for. That is the open question now, and it is a better one than the question this started with. Both rigs run on spot capacity and are terminated after collection. The XLA cache is pushed to S3 before teardown, which is what makes the 6-second restore on a fresh instance possible. The goal of this article was to find out which parts of "it's just JAX" survive a move between GPU generations. The key to the solution was reading the compute dtype off the live device rather than a config file. The measured results were: Scope: two spot instances, one in `us-east-1d` , each measured once with 3 repeats per sweep cell and medians reported. The two boxes differ in host architecture (aarch64 against x86_64) and base image as well as in GPU, so this is not a single-variable experiment; the payload is byte-identical across them — same build `51bc52c9e2e9` , same config, same 6,155,450,950 bytes of weights — which is the basis for attributing the difference to the chip. The Turing profile was reproduced on a second instance at 1466.0 ms against 1467.1 ms; the Ada profile was measured once. The strategy for using MCP for Gemma 4 serving across GPU generations was validated with an incremental step by step approach.