{"slug": "gemma-4-in-pure-jax-what-ports-from-tpu-to-gpu-and-what-doesn-t", "title": "Gemma 4 in Pure JAX: What Ports from TPU to GPU, and What Doesn't", "summary": "A developer's hand-written Gemma 4 port in pure JAX runs across TPU v5e/v6e and NVIDIA T4G GPUs, revealing that while most of the code is portable, two hardware-specific issues break the abstraction. The heterogeneous head dimensions that force Triton on other stacks are handled natively by XLA, but the fused W4A16 kernel, tiled for TPU VMEM, cannot run on GPUs due to shared memory limits, forcing the GPU rigs to serve a dense 16-bit checkpoint instead.", "body_md": "This article is about running a hand-written **Gemma 4** port in **pure JAX** on three different accelerators, and about the two places the abstraction leaks.\n\nThe code is here:\n\nThis project aims to serve one Gemma 4 checkpoint from one JAX port across every accelerator I can rent, and to find out — by measurement, not by reading docs — which parts of \"it's just JAX\" are true.\n\nThe port lives in `ports/gemma4/`\n\nand is driven by a generation loop behind an OpenAI-compatible server. No PyTorch, no vLLM, no `torch_xla`\n\n. The same code runs on Cloud TPU v5e and v6e, and on an NVIDIA T4G attached to an AWS Graviton2 host.\n\n\"Pure JAX\" is the whole experiment. If the port is really portable, the only thing that should change between those rigs is a config file.\n\nIt mostly is. Two things are not, and they are the interesting part.\n\nAny port has to carry four irregularities, and none of them are optional:\n\n`head_dim=256`\n\n, global layers use That first one is worth dwelling on, because it is what breaks other stacks. On the vLLM path, the heterogeneous head dims force the Triton attention backend:\n\n```\nGemma4 model has heterogeneous head dimensions\n{'sliding_attention': 256, 'full_attention': 512}. FA4 not available,\nforcing TRITON_ATTN backend.\n```\n\nAnd on a Turing GPU that backend then asks for shared memory the hardware does not have:\n\n```\ntriton.runtime.errors.OutOfResources: out of resource: shared memory,\nRequired: 98304, Hardware limit: 65536\n```\n\n**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. This is the clearest win of the whole exercise: the irregular geometry that is a special case everywhere else is just array shapes here.\n\nThis is the single most expensive lesson in the repo.\n\n**A wrong compute dtype does not raise. It emulates.** `bfloat16`\n\non a pre-Ampere GPU does not fail — XLA routes it through fp32 and you simply lose most of your decode to conversion. Nothing in the logs is red.\n\nSo the port does not take the dtype from a config file. It reads the live compute capability off the device and decides:\n\n```\nCOMPUTE_DTYPE = float16 if IS_PRE_AMPERE else bfloat16\n```\n\nOn TPU that resolves to `bfloat16`\n\n, which the MXUs run natively. On an SM 8.9 Ada card, `bfloat16`\n\n. On the SM 7.5 Turing card in this rig, ** float16** — Turing's only real 16-bit datapath, since it has neither bf16 nor fp8.\n\nThe first line the process emits states what it decided, so a misconfiguration is one grep away rather than a mystery in the throughput:\n\n```\nINFO ports.gemma4.jax_e_model: jax_e_model device policy: platform=gpu\ncompute_capability=7.5 compute_dtype=float16 pallas_interpret=False\n```\n\n`pallas_interpret=False`\n\nmatters just as much — it is the difference between serving and silently running a simulator.\n\nHere is the part that does not port, and it is not a bug — it is a real hardware difference wearing a portable API.\n\nThe fused **W4A16 kernel is written in Pallas**, and it is tiled for **TPU VMEM**, which gives you 16 MB per core. At this model's shapes the tiles want **550 KiB – 1.1 MiB per block**.\n\nOn 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.\n\nSo the same Pallas kernel that is the fast path on TPU **cannot run on either GPU**. The rig computes the requirement at startup and refuses with the arithmetic attached, rather than dying as a cryptic `OutOfResources`\n\nat the first token:\n\n```\ncheck_w4a16_fits_scoped_memory()\n```\n\nThe practical consequence: the GPU rigs serve the **dense reference checkpoint** at 16-bit, while the TPU rigs serve the `-qat-w4a16-ct`\n\nexport. Same port, same model family, different weights — because a kernel written against VMEM does not describe a GPU.\n\n**If you take one thing from this article, take that one.** Pallas is portable as an API and not portable as a memory model.\n\n`200 OK`\n\nA padding-eviction bug in the KV ring cache cost a week, and it is the kind only Gemma 4's geometry produces.\n\nThe invariant is: **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 512-slot ring violates it, and the failure mode is not a crash or a NaN. It is a **token loop** — a clean HTTP `200`\n\n, `status: \"success\"`\n\n, and output like `The The The The`\n\n.\n\nNothing 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.\n\nThe scariest bugs in this whole project all returned success.\n\nEnough that the exercise was worth it:\n\n`max_new_tokens`\n\nis a `static_argnames`\n\nentry, so `(bucket, max_tokens)`\n\nis the compiled shape on every backend. Warm up at the shape you measure — the same request took `pip`\n\nsupplies CUDA.`jax[cuda13]`\n\nmeans the GPU rig installs in Profiling decode with xprof on the Turing card gave this:\n\n```\nconversion   54.0%   <-- dtype conversion\nfp32 gemv    32.9%\nfusion       12.2%\nTensorCore    0.0%\n```\n\n**Zero.** 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.\n\nThe obvious hypothesis was bf16 weights being converted on a chip with no bf16 datapath. So I converted the checkpoint to float16 host-side and re-ran. Parameter dtypes read `{'float16': 541, 'uint8': 1, 'int8': 1}`\n\n— and **conversion stayed at 54.0%**.\n\nThe obvious explanation is wrong and I do not yet know the real one. I would rather publish the open question than a tidy story.\n\nWhat I can stand behind is that the measurement is real: 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.**\n\nThe next rig is the control for exactly this — the same port on an **Ada** card, where `_compute_dtype()`\n\nreturns `bfloat16`\n\nwith no code change and the conversion pressure is removed at the hardware level. If 54% survives onto a bf16-native chip, the cause was never dtype at all.\n\n| Model |\n`google/gemma-4-E2B-it` , dense reference build |\n| Weights resident | 6.155 GB |\n| Decode |\n13.10 tok/s on the T4G |\n| Decode vs context | flat: 12.9 / 13.0 / 12.9 tok/s at 41 / 521 / 2,057 input tokens |\n| Context |\n`MAX_MODEL_LEN=4096` , and that is the honest number — 4,105 prompt tokens serve, 5,120 fails on a prefill transient |\n\n**Quote the gauge, not end-to-end.** End-to-end throughput does fall with a longer prompt (12.43 → 8.22), but that is prefill being linear in the padded bucket, not decode degrading. They are two different claims and conflating them makes a benchmark a lie.\n\nOne JAX port, three accelerators. The model code, the compilation cache and the static-shape discipline all transferred untouched, and Gemma 4's awkward geometry — the thing that forces a special-case kernel on other stacks — turned out to be the easiest part, because in JAX it is just shapes.\n\nWhat did not transfer was the one piece written against a specific memory model. Pallas gives you a portable API on top of VMEM and shared memory, and those are not the same size. That boundary is worth knowing before you plan a port around a fused kernel.", "url": "https://wpnews.pro/news/gemma-4-in-pure-jax-what-ports-from-tpu-to-gpu-and-what-doesn-t", "canonical_source": "https://dev.to/gde/gemma-4-in-pure-jax-what-ports-from-tpu-to-gpu-and-what-doesnt-3m09", "published_at": "2026-08-29 00:39:38+00:00", "updated_at": "2026-08-29 00:48:15.941246+00:00", "lang": "en", "topics": ["machine-learning", "large-language-models", "ai-infrastructure", "developer-tools"], "entities": ["Gemma 4", "JAX", "TPU v5e", "TPU v6e", "NVIDIA T4G", "AWS Graviton2", "Pallas", "Triton"], "alternates": {"html": "https://wpnews.pro/news/gemma-4-in-pure-jax-what-ports-from-tpu-to-gpu-and-what-doesn-t", "markdown": "https://wpnews.pro/news/gemma-4-in-pure-jax-what-ports-from-tpu-to-gpu-and-what-doesn-t.md", "text": "https://wpnews.pro/news/gemma-4-in-pure-jax-what-ports-from-tpu-to-gpu-and-what-doesn-t.txt", "jsonld": "https://wpnews.pro/news/gemma-4-in-pure-jax-what-ports-from-tpu-to-gpu-and-what-doesn-t.jsonld"}}