Pure JAX on G5g: Serving Gemma 4 on Graviton and a T4G

A developer has published a step-by-step guide for serving Google's Gemma 4 model on an AWS EC2 G5g instance using pure JAX, targeting the cheapest whole NVIDIA GPU available on AWS. The guide details deploying on the Arm-based Graviton2 host with a T4G GPU, using an MCP server for management, and emphasizes cost efficiency and honest benchmarking.

This article provides a step by step deployment guide for serving Google's Gemma 4 on an AWS EC2 G5g instance using pure JAX. The code is here: This project aims to serve a modern open model on the cheapest whole CUDA GPU AWS will rent you, and to measure honestly what that costs. Probably The T4G is a Turing chip from 2018. It has no bfloat16 and no fp8. But it is cheap, it is available when nothing else is, and it is attached to a Graviton2 host — which makes G5g the rare hardware axis that almost nothing in the ML ecosystem targets: aarch64 and CUDA together . So let's give pure JAX a shot on G5g G5g instances pair an AWS Graviton2 64-bit Arm processor with NVIDIA T4G Tensor Core GPUs. At g5g.xlarge they are the cheapest EC2 instance carrying a whole NVIDIA GPU , and the only Arm-based GPU family AWS offers. Two GPU instances are cheaper per hour and neither can serve this model us-east-1, Linux, on-demand, checked against the Pricing API on 2026-08-28 : g6f.large at g6f that fits is g6f.4xlarge at $0.9500, which is 1.7x this rig's g5g.2xlarge . g4ad.xlarge at Among whole NVIDIA GPUs, G5g is the floor: g5g.xlarge at $0.4200, and the next one up is g4dn.xlarge at $0.5260. More information is available here: https://aws.amazon.com/ec2/instance-types/g5g/ https://aws.amazon.com/ec2/instance-types/g5g/ The default in this rig is g5g.2xlarge — 1 GPU, 8 vCPU, 16 GiB RAM. Note- the T4G reports 15,360 MiB of device memory, not the nominal 16 GB. Budget against the measured number. Gemma is Google's family of open models built from the same research as Gemini. This rig serves google/gemma-4-E2B-it , the instruction-tuned reference release. JAX is Google's array computing library — NumPy semantics, composable transformations, and compilation to XLA. On NVIDIA hardware, pip supplies the CUDA libraries, so there is nothing to build. More information is available here: "Pure JAX" here is literal. The engine is this repo's own Gemma 4 port driven by a JAX generation loop behind an OpenAI-compatible FastAPI server, running under systemd. You need four things before starting: us-east-1 pip The instance profile needs AmazonSSMManagedInstanceCore plus read access to your Secrets Manager secret and your S3 cache bucket. There is no inbound SSH rule and no private key — all remote administration goes over SSM Run Command. Clone the monorepo and install the control plane: git clone https://github.com/xbill9/gemma4-dev cd gemma4-dev/gpu-jax-g5g-2b pip install -r requirements.txt That installs boto3 and FastMCP only. Nothing here needs a GPU — the GPU is on the other end. python3 -m unittest discover -s tests -v 105 tests, fully offline. Every cloud, subprocess, and network boundary is mocked. If these do not pass, do not launch an instance. The whole rig is driven by an MCP server exposing a devops agent: ./project-setup.sh This installs the bundled skill and registers .mcp.json : { "mcpServers": { "gpu-jax-g5g-2b": { "command": "python3", "args": ".claude/skills/gpu-jax-g5g-2b-management/mcp/server.py" , "env": { "AWS REGION": "us-east-1", "MODEL NAME": "google/gemma-4-E2B-it", "INSTANCE TYPE": "g5g.2xlarge", "MCP SERVER NAME": "gpu-jax-g5g-2b" } } } } Every tool is now available as mcp gpu-jax-g5g-2b <tool . save hf token token="hf ..." This writes to AWS Secrets Manager under vllm/hf-token . The instance fetches it at boot into a root-only EnvironmentFile . Note- the token never goes in user data. Instance metadata is readable by anything running on the box. check g5g quotas Reports your On-Demand and Spot G instance vCPU limits for the region. g5g.2xlarge is 8 vCPU. Check this before launching, not after the launch fails. create g5g instance subnet id="subnet-...", security group id="sg-...", iam instance profile="...", spot=True The AMI is resolved at launch time from SSM Parameter Store: /aws/service/deeplearning/ami/arm64/base-oss-nvidia-driver-gpu-ubuntu-26.04/latest/ami-id Never hardcode an AMI id here. AWS also ships ARM64 DLAMIs built for Graviton CPU inference. They boot perfectly and simply have no GPU. The /latest/ parameter also moves — this rig has seen ami-0bff4343bfd56a20e become ami-025a6e5b3b786cf61 overnight as Ubuntu 26.04 became 26.04.1. get install progress instance id="i-..." Cloud-init installs jax cuda13 on Python 3.14. The stages are timed: stage jax-wheels 43s total 84s stage serving-deps 14s total 98s stage gpu-verify 13s total 111s stage cache-restore 6s total 117s stage unit-rewrite 0s total 117s 117 seconds, and there is no compile step anywhere in it. That is the entire reason this rig exists — more on that below. The cache-restore stage pulled 805 files / 12 MB in 6 seconds from S3, on a fresh instance, compiled by a box that had already been terminated. XLA's cold-compile penalty becomes a rounding error on Spot. verify gpu arch instance id="i-..." This measures whether JAX's CUDA kernels actually cover this GPU, rather than trusting that they do. You want to see SM 7.5 claimed and a real device, not a silent CPU fallback. make skill deploy jax server instance id="i-..." Always make skill first. The deploy ships the The first line the process emits is the device-policy banner: INFO ports.gemma4.jax e model: jax e model device policy: platform=gpu compute capability=7.5 compute dtype=float16 pallas interpret=False Both halves matter. float16 is the device choosing Turing's only real 16-bit datapath — it is read from the live compute capability, not from a config file. And pallas interpret=False is the difference between serving and silently running a simulator. Then the whole resolved configuration lands on one greppable line: READY build id=6852f5680f43 ... compute dtype=float16 kv cache dtype=float16 kv cache requested=auto pre ampere=True quant mode=fp16 window kv=True Load is staged, so a hang is attributable: download 87.7s read shards 73.5s 1 shard, 600 tensors, 0.95 GB of non-text towers skipped convert params 3.4s device put 0.0s 164.7s total, 9.26 GB Then confirm the served build matches what you shipped: verify model health instance id="i-..." query model instance id="i-...", prompt="Explain Graviton in one sentence." The endpoint is OpenAI-compatible on :8000 , so anything that speaks that API works: get endpoint instance id="i-..." Note- warm up at the shape you measure. max new tokens is a static argnames entry, so bucket, max tokens is the compiled shape. The same request measured 18.77 s cold against 4.35 s warm . get metrics instance id="i-..." tpu jax decode tokens per second 13.0 tpu jax prefill milliseconds 692.6 tpu jax hbm used bytes 6296892160 tpu jax weight bytes 6155450950 tpu jax degenerate responses total 0 Quote the gauge, not end-to-end. Decode is flat at 12.9 / 13.0 / 12.9 tok/s across 41 → 2,057 input tokens. End-to-end throughput does fall 12.43 → 8.22 — but that is prefill being linear in the padded bucket, not decode degrading. Two different claims. Because I tried, on identical silicon, and it works — at 43 tok/s — but only after this: cuda-toolkit from NVIDIA's sbsa repo, because the DLAMI ships a driver but no nvcc That last one is the interesting failure. Gemma 4 has heterogeneous attention head dimensions — sliding layers at 256, global layers at 512 . Only two vLLM backends support that, and with FA4 unavailable it force-selects Triton: Gemma4 model has heterogeneous head dimensions {'sliding attention': 256, 'full attention': 512}. FA4 not available, forcing TRITON ATTN backend. Whose 512-wide tile then asks Turing for memory Turing does not have: triton.runtime.errors.OutOfResources: out of resource: shared memory, Required: 98304, Hardware limit: 65536 JAX sidesteps all four. pip supplies CUDA, so no build, no toolkit, no Rust. The plugin's precompiled cubins already cover sm 75 . And attention is ordinary XLA rather than a hand-tiled Triton kernel, so there is no per-block shared-memory ceiling and no patch to carry. The honest trade: 13.10 tok/s against 43, for a 117-second install with nothing to reapply. For a measurement rig I re-provision constantly on Spot, that was right. For a production endpoint it probably is not. I profiled decode with xprof. The kernel table is the whole story: conversion 54.0% <-- dtype conversion fp32 gemv 32.9% fusion 12.2% TensorCore 0.0% Zero. 1,466 ms of kernels across 108 distinct kernels on a Tensor Core GPU, and not one Tensor Core fired. Over half of decode went to converting numbers between formats before the math could start. The obvious hypothesis was bf16 weights being converted on a chip with no bf16. So this weekend I converted the checkpoint to float16 host-side and re-ran. Parameter dtypes now read {'float16': 541, 'uint8': 1, 'int8': 1} — and conversion is still 54.0% . The obvious explanation is wrong and I do not yet know the real one. That is the next thing to profile. I would rather publish the open question than a tidy story. What 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. | Service | Role | |---|---| EC2 g5g.2xlarge , Spot | Graviton2 + NVIDIA T4G | Systems Manager — Parameter Store | Resolves the arm64 GPU DLAMI id at launch | Systems Manager — Run Command | Ships the payload, runs every diagnostic. No SSH | S3 | XLA compilation cache, shared across instances | Secrets Manager | Hugging Face token | IAM | Instance profile scoping all of the above | Service Quotas | Pre-flights G vCPU limits before a doomed launch | EBS gp3 | 100 GB at 500 MiB/s, 6,000 IOPS, for a 9.5 GB checkpoint | Everything goes through boto3 . The rig never shells out to the AWS CLI. terminate g5g instance instance id="i-..." Termination is cheap on this rig — there is no built image to lose with the root volume, only a pip install and a model cache, and the compilation cache is already in S3. Check reachability on paper before spending a provisioning cycle. A nine-step analysis order — does the compute dtype match the chip, is there a fused kernel for that format on that chip — runs in an afternoon. A launch costs a day. It has killed two bad plans before either touched hardware. A wrong dtype does not error. It emulates. bfloat16 on Turing does not fail loudly; it routes through fp32 and quietly eats your decode. Refuse early, with the arithmetic attached. The fused W4A16 kernel wants 550 KiB–1.1 MiB per block and Turing gives you 64 KiB. The rig computes that at startup and refuses with the numbers in the message, rather than dying as a cryptic OutOfResources at the first token. The scariest bugs return status: "success". A padding-eviction bug in the KV ring cache produced a token loop, not a crash. Nothing in the logs was red. It took a week. The AWS G5g instance provides a genuinely cheap environment for serving open models, and pure JAX reaches a served token on it without a single line of compiled code. The throughput is not competitive with a patched vLLM - but the deployment is 117 seconds, reproducible to 1.1 ms, and has nothing to reapply.