Gemma 4 in Pure JAX: What Ports from TPU to GPU, and What Doesn't

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. The code is here: github.com/xbill9/gemma4-dev What is this project trying to Do? This 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. 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 same code runs on Cloud TPU v5e and v6e, and on an NVIDIA T4G attached to an AWS Graviton2 host. "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. It mostly is. Two things are not, and they are the interesting part. Gemma 4 E2B is not a stock transformer Any port has to carry four irregularities, and none of them are optional: Two attention geometries. Sliding layers use head_dim=256 , global layers use 512 . Most inference stacks assume one head dimension per model. 8:1 MQA , so the KV budget is nothing like the parameter count would suggest. A KV-share map that collapses 35 layers onto 15 caches . A 512-slot sliding ring , plus per-layer embeddings (PLE) held in a 4.70 GB table that gets quantized to 4 bits on load. 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: Gemma4 model has heterogeneous head dimensions {'sliding_attention': 256, 'full_attention': 512}. FA4 not available, forcing TRITON_ATTN backend. And 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: 98304, Hardware limit: 65536 JAX never enters that conversation. Attention is ordinary XLA rather than a hand-tiled kernel, so ther