{"slug": "gemma-4-in-pure-jax-what-changes-between-turing-and-ada-and-what-doesn-t", "title": "Gemma 4 in Pure JAX: What Changes Between Turing and Ada, and What Doesn't", "summary": "A developer's hand-written Gemma 4 port in pure JAX runs on both Turing and Ada NVIDIA GPUs with identical weights, but performance differs drastically due to hardware-specific compute dtype handling and memory constraints. The port avoids Triton's shared-memory ceiling by using XLA for attention, yet a wrong compute dtype on pre-Ampere GPUs silently emulates through fp32, costing 87% of decode speed. The developer's solution reads the device's compute capability to select the correct dtype, ensuring the fast path runs only on Ada.", "body_md": "This article is a measurement report on running a hand-written **Gemma 4** port in\n\n**pure JAX** across two NVIDIA GPUs a generation apart, and on the two places the\n\n\"it's just JAX\" abstraction leaks. One of those leaks costs 87% of decode and\n\nnothing in the logs is red.\n\nThe code is here:\n\n[https://github.com/xbill9/gemma4-dev](https://github.com/xbill9/gemma4-dev)\n\nOne port, one build, one checkpoint, two cards. Everything below comes from two\n\narchived runs, named so you can check them.\n\n| G5g | G6 | |\n|---|---|---|\n| Chip | NVIDIA T4G — Turing, SM 7.5, 15,360 MiB |\nNVIDIA L4 — Ada, SM 8.9, 23,034 MiB |\n| Host |\n`g5g.2xlarge` spot — Graviton2, aarch64\n|\n`g6.2xlarge` spot — x86_64, `us-east-1d`\n|\n| Checkpoint |\n`google/gemma-4-E2B-it` , dense reference |\n`google/gemma-4-E2B-it` , dense reference |\n| Compute dtype |\n`float16` (device-chosen) |\n`bfloat16` (device-chosen) |\n| Stack | jax 0.11.1, CUDA from pip | jax 0.11.1, Python 3.14 |\n| Run cited | `2026-08-28-full-run-cached-g5g` |\n`2026-08-28-first-serve-g6` |\n\nBuild id `51bc52c9e2e9`\n\non both, config `ple4 + int8_lm_head`\n\n, and\n\n`tpu_jax_weight_bytes`\n\nreads **6,155,450,950** on both cards — the same integer.\n\nOnly the chip and its host differ.\n\n`g5g.2xlarge`\n\nand `g6.2xlarge`\n\n`google/gemma-4-E2B-it`\n\nThe port lives in `ports/gemma4/`\n\nand is driven by a generation loop behind an\n\nOpenAI-compatible server. No PyTorch, no vLLM, no `torch_xla`\n\n.\n\nThe premise under test is that the same source runs on both cards with nothing\n\nchanged but a config file. It mostly holds. The interesting part is where it does\n\nnot.\n\nAny port has to carry four irregularities, and none of them are optional.\n\n`head_dim=256`\n\n, global layers\nuse That first irregularity is the expensive one. On the vLLM path the heterogeneous\n\nhead dims force the Triton attention backend:\n\n```\nGemma4 model has heterogeneous head dimensions\n(sliding=256, global=512); falling back to the Triton attention backend\n```\n\nOn a Turing GPU that backend then asks for shared memory the hardware does not\n\nhave:\n\n```\ntriton.runtime.errors.OutOfResources: out of resource: shared memory,\nRequired: 147456, Hardware limit: 65536\n```\n\n**JAX never enters that conversation.** Attention is ordinary XLA rather than a\n\nhand-tiled kernel, so there is no per-block shared-memory ceiling in the attention\n\npath at all. The irregular geometry that is a special case everywhere else is just\n\narray shapes here.\n\nThis is the single most expensive lesson in the repository.\n\n**A wrong compute dtype does not raise. It emulates.** `bfloat16`\n\non a pre-Ampere\n\nGPU does not fail — XLA routes it through fp32 and most of decode disappears into\n\nconversion. Nothing in the logs is red.\n\nSo the port does not take the dtype from a config file. It reads the live compute\n\ncapability off the device:\n\n```\nCOMPUTE_DTYPE = float16 if IS_PRE_AMPERE else bfloat16\n```\n\nOn the SM 8.9 Ada card that resolves to `bfloat16`\n\n. On the SM 7.5 Turing card it\n\nresolves to `float16`\n\n— Turing's only real 16-bit datapath, since it has neither\n\nbf16 nor fp8.\n\nThe first line the server emits is the policy, so a misconfiguration is one `grep`\n\naway rather than a mystery in the throughput:\n\n```\nINFO ports.gemma4.jax_e_model: jax_e_model device policy: platform=gpu\ncompute_capability=8.9 compute_dtype=bfloat16 pallas_interpret=False\n```\n\n`pallas_interpret=False`\n\nmatters just as much. It is the difference between\n\nserving 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\n\ndifference wearing a portable API.\n\nThe fused **W4A16 kernel is written in Pallas**, and it was tiled for a device with\n\n16 MB of scratchpad per core. At this model's shapes the tiles want **550 KiB to\n1.1 MiB per block**.\n\nOn a GPU, Pallas lowers through Triton, and those tiles become **shared memory**.\n\nTuring gives you 64 KiB per block. Ada raises the ceiling, but nowhere near a\n\nmegabyte.\n\nSo the fast path runs on **neither card**. The engine computes the requirement at\n\nstartup and refuses with the arithmetic attached, rather than dying as a cryptic\n\n`OutOfResources`\n\nat the first token:\n\n```\ncheck_w4a16_fits_scoped_memory()\n```\n\nThe practical consequence is that both GPU rigs serve the **dense reference\ncheckpoint** at 16-bit.\n\nA padding-eviction bug in the KV ring cache cost a week, and it is the kind only\n\nGemma 4's geometry produces.\n\nThe invariant is that **a cache index is an absolute real position, and padding\nnever occupies an index a real position uses.** A port that right-pads into the\n\n`200`\n\n, `status: \"success\"`\n\n, and output like`The The The The`\n\n.Nothing in the logs is red. Nothing in the metrics is red. The only thing that\n\ncatches it is a degeneracy check on the output itself, which the server now runs on\n\nevery response.\n\nThe scariest bugs in this project all returned success.\n\n`jax[cuda13]`\n\nsupplies CUDA as wheels, so the install needs no CUDA toolkit, no\n\nRust, and no compiler on the box.\n\n```\nInstall: 117 s, with the cache restore included\n```\n\nXLA's persistent compilation cache ports as-is. On the T4G rig it restores **805\nfiles / 12 MB in 6 seconds** onto a fresh instance, from a box that had already\n\n`max_new_tokens`\n\nis a `static_argnames`\n\nentry, so `(bucket, max_tokens)`\n\nis the\n\ncompiled shape on every backend. A harness that does not warm up misreports the\n\nrig badly.\n\nOn the T4G the first request off a fresh engine took **18.06 s against 4.50 s\nwarm** — a 4.0x whole-request ratio, measured in\n\n`2026-08-21-cuda13-py314-g5g`\n\n.That run also notes something worth repeating: the 56x figure from the first-serve\n\nbaseline is **TTFT specifically**, not the same measurement as the whole-request\n\nratio. They are not interchangeable.\n\n64 output tokens, concurrency 1, 3 repeats per cell, median. \"Decode, gauge\" is the\n\nengine's steady-state counter. \"End-to-end\" is wall time over the whole request,\n\nprefill included.\n\n| Input tokens | T4G gauge | T4G end-to-end | L4 gauge | L4 end-to-end |\n|---|---|---|---|---|\n| 41 | 12.9 tok/s | 12.43 tok/s | 48.5 tok/s |\n46.23 tok/s |\n| 521 | 13.0 tok/s | 11.28 tok/s | 48.4 tok/s |\n42.87 tok/s |\n| 2,057 | 12.9 tok/s | 8.22 tok/s | 48.3 tok/s |\n34.57 tok/s |\n| 3,593 | — | — | 48.3 tok/s |\n27.55 tok/s |\n\nDecode moves 0.8% across a 50x context range on the T4G and 0.4% on the L4. End-to-end\n\nfalls hard on both.\n\nThat fall is prefill being linear in the padded bucket, not decode degrading. They\n\nare two different claims, and conflating them makes a benchmark a lie. **Quote the\ngauge.**\n\nA cost proportional to the **weights** rather than the context produces exactly this\n\nshape, which is why KV is not what sets decode speed on either card — despite\n\nGemma 4's whole KV story.\n\nOn context specifically: `MAX_MODEL_LEN=4096`\n\nis the honest number on the T4G.\n\n4,105 prompt tokens serve; 5,120 fails on a prefill transient.\n\nProfiling decode with xprof on the Turing card, 20 decode steps with the service\n\nstopped:\n\n| 🥉 T4G (SM 7.5) | 🥇 L4 (SM 8.9) | |\n|---|---|---|\n| dtype conversion | 54.1% | 0.0% |\nfp32 `gemvx`\n|\n32.8% | absent |\n| Tensor Core | 0.0% | 0.0% |\n| Total kernel time | 1,466.0 ms | 362.8 ms |\n| Decode, gauge | 12.9 tok/s | 48.4 tok/s |\n| Peak HBM bandwidth | 298.083 GiB/s | 279.441 GiB/s |\n| Share of bandwidth roofline | 26% | ~100% |\n\n1,466 ms of kernels across 108 distinct kernels on a Tensor Core GPU, without one\n\nTensor Core firing. More than half of decode went to converting numbers between\n\nformats before any math happened.\n\nThe obvious hypothesis was bf16 weights being converted on a chip with no bf16\n\ndatapath. So the checkpoint was converted to float16 host-side and re-run.\n\nParameter dtypes read `{'float16': 541, 'uint8': 1, 'int8': 1}`\n\n— and **conversion\nstayed at 54.0%**.\n\nThe measurement itself is solid. The same profile on a different instance, a\n\ndifferent AMI and a restored cache landed at 1466.0 ms against 1467.1 ms. **1.1 ms\napart on 1467.**\n\nThe Ada card resolves it. Converting the stored weights changed nothing because\n\nstorage dtype was never the problem: Turing has no native bf16, and the fp32\n\n`gemvx`\n\nline is the tell — XLA was round-tripping through fp32 regardless of what\n\nthe file on disk said.\n\nGive it a card where storage and compute dtype actually match, and the 54%\n\nconversion and the 32.8% fp32 path vanish **together**. An 87% tax gone, for 3.7x\n\nthe throughput, and a rig sitting at its bandwidth roofline instead of 26% of it.\n\nThe `/health`\n\nendpoint on the L4 reports `weights=bfloat16 activations=bfloat16`\n\n— storage dtype and compute dtype matching for\n\nkv_cache=bfloat16 pre_ampere=false\n\nthe first time on this engine.\n\nTensor Core utilization is **0.0% on the Ada card too** — 100 distinct kernels,\n\n362.8 ms of them, and not one Tensor Core firing.\n\nRemoving the dtype pressure made the machine roughly four times faster without\n\nmaking it touch the hardware it was sold for. That is the open question now, and it\n\nis a better one than the question this started with.\n\nBoth rigs run on spot capacity and are terminated after collection. The XLA cache\n\nis pushed to S3 before teardown, which is what makes the 6-second restore on a\n\nfresh instance possible.\n\nThe goal of this article was to find out which parts of \"it's just JAX\" survive a\n\nmove between GPU generations. The key to the solution was reading the compute dtype\n\noff the live device rather than a config file. The measured results were:\n\nScope: two spot instances, one in `us-east-1d`\n\n, each measured once with 3 repeats\n\nper sweep cell and medians reported. The two boxes differ in host architecture\n\n(aarch64 against x86_64) and base image as well as in GPU, so this is not a\n\nsingle-variable experiment; the payload is byte-identical across them — same build\n\n`51bc52c9e2e9`\n\n, same config, same 6,155,450,950 bytes of weights — which is the\n\nbasis for attributing the difference to the chip. The Turing profile was reproduced\n\non a second instance at 1466.0 ms against 1467.1 ms; the Ada profile was measured\n\nonce.\n\nThe strategy for using MCP for Gemma 4 serving across GPU generations was validated\n\nwith an incremental step by step approach.", "url": "https://wpnews.pro/news/gemma-4-in-pure-jax-what-changes-between-turing-and-ada-and-what-doesn-t", "canonical_source": "https://dev.to/xbill/gemma-4-in-pure-jax-what-changes-between-turing-and-ada-and-what-doesnt-4c5e", "published_at": "2026-08-31 01:40:09+00:00", "updated_at": "2026-08-31 01:51:35.139189+00:00", "lang": "en", "topics": ["machine-learning", "large-language-models", "developer-tools", "ai-infrastructure"], "entities": ["Gemma 4", "JAX", "NVIDIA T4G", "NVIDIA L4", "XLA", "Pallas", "Triton", "Google"], "alternates": {"html": "https://wpnews.pro/news/gemma-4-in-pure-jax-what-changes-between-turing-and-ada-and-what-doesn-t", "markdown": "https://wpnews.pro/news/gemma-4-in-pure-jax-what-changes-between-turing-and-ada-and-what-doesn-t.md", "text": "https://wpnews.pro/news/gemma-4-in-pure-jax-what-changes-between-turing-and-ada-and-what-doesn-t.txt", "jsonld": "https://wpnews.pro/news/gemma-4-in-pure-jax-what-changes-between-turing-and-ada-and-what-doesn-t.jsonld"}}