{"slug": "pure-jax-on-g5g-serving-gemma-4-on-graviton-and-a-t4g", "title": "Pure JAX on G5g: Serving Gemma 4 on Graviton and a T4G", "summary": "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.", "body_md": "This article provides a step by step deployment guide for serving Google's Gemma 4 on an AWS EC2 G5g instance using pure JAX.\n\nThe code is here:\n\nThis 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.\n\nProbably! The T4G is a Turing chip from 2018. It has no bfloat16 and no fp8.\n\nBut 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**.\n\nSo let's give pure JAX a shot on G5g!\n\nG5g instances pair an AWS Graviton2 (64-bit Arm) processor with NVIDIA T4G Tensor Core GPUs. At `g5g.xlarge`\n\nthey are the **cheapest EC2 instance carrying a whole NVIDIA GPU**, and the only Arm-based GPU family AWS offers.\n\nTwo 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):\n\n`g6f.large`\n\nat `g6f`\n\nthat fits is `g6f.4xlarge`\n\nat $0.9500, which is 1.7x this rig's `g5g.2xlarge`\n\n.`g4ad.xlarge`\n\nat Among whole NVIDIA GPUs, G5g is the floor: `g5g.xlarge`\n\nat $0.4200, and the next one up is `g4dn.xlarge`\n\nat $0.5260.\n\nMore information is available here:\n\n[https://aws.amazon.com/ec2/instance-types/g5g/](https://aws.amazon.com/ec2/instance-types/g5g/)\n\nThe default in this rig is `g5g.2xlarge`\n\n— 1 GPU, 8 vCPU, 16 GiB RAM.\n\nNote- the T4G reports **15,360 MiB** of device memory, not the nominal 16 GB. Budget against the measured number.\n\nGemma is Google's family of open models built from the same research as Gemini. This rig serves `google/gemma-4-E2B-it`\n\n, the instruction-tuned reference release.\n\nJAX is Google's array computing library — NumPy semantics, composable transformations, and compilation to XLA. On NVIDIA hardware, `pip`\n\nsupplies the CUDA libraries, so there is nothing to build.\n\nMore information is available here:\n\n\"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.\n\nYou need four things before starting:\n\n`us-east-1`\n\n`pip`\n\nThe instance profile needs `AmazonSSMManagedInstanceCore`\n\nplus 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.\n\nClone the monorepo and install the control plane:\n\n```\ngit clone https://github.com/xbill9/gemma4-dev\ncd gemma4-dev/gpu-jax-g5g-2b\npip install -r requirements.txt\n```\n\nThat installs `boto3`\n\nand FastMCP only. Nothing here needs a GPU — the GPU is on the other end.\n\n```\npython3 -m unittest discover -s tests -v\n```\n\n105 tests, fully offline. Every cloud, subprocess, and network boundary is mocked. If these do not pass, do not launch an instance.\n\nThe whole rig is driven by an MCP server exposing a devops agent:\n\n```\n./project-setup.sh\n```\n\nThis installs the bundled skill and registers `.mcp.json`\n\n:\n\n```\n{\n  \"mcpServers\": {\n    \"gpu-jax-g5g-2b\": {\n      \"command\": \"python3\",\n      \"args\": [\".claude/skills/gpu-jax-g5g-2b-management/mcp/server.py\"],\n      \"env\": {\n        \"AWS_REGION\": \"us-east-1\",\n        \"MODEL_NAME\": \"google/gemma-4-E2B-it\",\n        \"INSTANCE_TYPE\": \"g5g.2xlarge\",\n        \"MCP_SERVER_NAME\": \"gpu-jax-g5g-2b\"\n      }\n    }\n  }\n}\n```\n\nEvery tool is now available as `mcp__gpu-jax-g5g-2b__<tool>`\n\n.\n\n```\nsave_hf_token(token=\"hf_...\")\n```\n\nThis writes to AWS Secrets Manager under `vllm/hf-token`\n\n. The instance fetches it at boot into a root-only `EnvironmentFile`\n\n.\n\nNote- the token never goes in user data. Instance metadata is readable by anything running on the box.\n\n```\ncheck_g5g_quotas()\n```\n\nReports your On-Demand and Spot G instance vCPU limits for the region. `g5g.2xlarge`\n\nis 8 vCPU. Check this before launching, not after the launch fails.\n\n```\ncreate_g5g_instance(\n    subnet_id=\"subnet-...\",\n    security_group_id=\"sg-...\",\n    iam_instance_profile=\"...\",\n    spot=True\n)\n```\n\nThe AMI is resolved at launch time from SSM Parameter Store:\n\n```\n/aws/service/deeplearning/ami/arm64/base-oss-nvidia-driver-gpu-ubuntu-26.04/latest/ami-id\n```\n\n**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/`\n\nparameter also moves — this rig has seen `ami-0bff4343bfd56a20e`\n\nbecome `ami-025a6e5b3b786cf61`\n\novernight as Ubuntu 26.04 became 26.04.1.\n\n```\nget_install_progress(instance_id=\"i-...\")\n```\n\nCloud-init installs `jax[cuda13]`\n\non Python 3.14. The stages are timed:\n\n```\n[stage] jax-wheels      43s   (total  84s)\n[stage] serving-deps    14s   (total  98s)\n[stage] gpu-verify      13s   (total 111s)\n[stage] cache-restore    6s   (total 117s)\n[stage] unit-rewrite     0s   (total 117s)\n```\n\n**117 seconds, and there is no compile step anywhere in it.** That is the entire reason this rig exists — more on that below.\n\nThe `cache-restore`\n\nstage 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.\n\n```\nverify_gpu_arch(instance_id=\"i-...\")\n```\n\nThis 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.\n\n```\nmake skill\ndeploy_jax_server(instance_id=\"i-...\")\n```\n\n**Always make skill first.** The deploy ships the\n\nThe first line the process emits is the device-policy banner:\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\nBoth halves matter. `float16`\n\nis **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`\n\nis the difference between serving and silently running a simulator.\n\nThen the whole resolved configuration lands on one greppable line:\n\n```\nREADY build_id=6852f5680f43 ... compute_dtype=float16 kv_cache_dtype=float16\nkv_cache_requested=auto pre_ampere=True quant_mode=fp16 window_kv=True\n```\n\nLoad is staged, so a hang is attributable:\n\n```\ndownload        87.7s\nread_shards     73.5s   (1 shard, 600 tensors, 0.95 GB of non-text towers skipped)\nconvert_params   3.4s\ndevice_put       0.0s\n                164.7s total, 9.26 GB\n```\n\nThen confirm the served build matches what you shipped:\n\n```\nverify_model_health(instance_id=\"i-...\")\nquery_model(instance_id=\"i-...\", prompt=\"Explain Graviton in one sentence.\")\n```\n\nThe endpoint is OpenAI-compatible on `:8000`\n\n, so anything that speaks that API works:\n\n```\nget_endpoint(instance_id=\"i-...\")\n```\n\nNote- **warm up at the shape you measure.** `max_new_tokens`\n\nis a `static_argnames`\n\nentry, so `(bucket, max_tokens)`\n\nis the compiled shape. The same request measured **18.77 s cold against 4.35 s warm**.\n\n```\nget_metrics(instance_id=\"i-...\")\ntpu_jax_decode_tokens_per_second   13.0\ntpu_jax_prefill_milliseconds      692.6\ntpu_jax_hbm_used_bytes            6296892160\ntpu_jax_weight_bytes              6155450950\ntpu_jax_degenerate_responses_total   0\n```\n\n**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.\n\nBecause I tried, on identical silicon, and it works — at 43 tok/s — but only after this:\n\n`cuda-toolkit`\n\nfrom NVIDIA's `sbsa`\n\nrepo, because the DLAMI ships a driver but no `nvcc`\n\nThat 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:\n\n```\nGemma4 model has heterogeneous head dimensions\n{'sliding_attention': 256, 'full_attention': 512}. FA4 not available,\nforcing TRITON_ATTN backend.\n```\n\nWhose 512-wide tile then asks Turing for memory Turing does not have:\n\n```\ntriton.runtime.errors.OutOfResources: out of resource: shared memory,\nRequired: 98304, Hardware limit: 65536\n```\n\nJAX sidesteps all four. `pip`\n\nsupplies CUDA, so no build, no toolkit, no Rust. The plugin's precompiled cubins already cover `sm_75`\n\n. 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.\n\n**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.\n\nI profiled decode with xprof. The kernel table is the whole story:\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, and not one Tensor Core fired. Over half of decode went to converting numbers between formats before the math could start.\n\nThe 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}`\n\n— and **conversion is still 54.0%**.\n\nThe 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.\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\n| Service | Role |\n|---|---|\nEC2 (`g5g.2xlarge` , Spot) |\nGraviton2 + NVIDIA T4G |\nSystems Manager — Parameter Store |\nResolves the arm64 GPU DLAMI id at launch |\nSystems Manager — Run Command |\nShips the payload, runs every diagnostic. No SSH |\nS3 |\nXLA compilation cache, shared across instances |\nSecrets Manager |\nHugging Face token |\nIAM |\nInstance profile scoping all of the above |\nService Quotas |\nPre-flights G vCPU limits before a doomed launch |\nEBS gp3 |\n100 GB at 500 MiB/s, 6,000 IOPS, for a 9.5 GB checkpoint |\n\nEverything goes through `boto3`\n\n. The rig never shells out to the AWS CLI.\n\n```\nterminate_g5g_instance(instance_id=\"i-...\")\n```\n\nTermination 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.\n\n**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.\n\n**A wrong dtype does not error. It emulates.** bfloat16 on Turing does not fail loudly; it routes through fp32 and quietly eats your decode.\n\n**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`\n\nat the first token.\n\n**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.\n\nThe 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.", "url": "https://wpnews.pro/news/pure-jax-on-g5g-serving-gemma-4-on-graviton-and-a-t4g", "canonical_source": "https://dev.to/aws-builders/pure-jax-on-g5g-serving-gemma-4-on-graviton-and-a-t4g-3glo", "published_at": "2026-08-29 00:38:39+00:00", "updated_at": "2026-08-29 00:48:19.243378+00:00", "lang": "en", "topics": ["artificial-intelligence", "large-language-models", "ai-infrastructure", "developer-tools"], "entities": ["Google", "AWS", "NVIDIA", "Gemma 4", "JAX", "Graviton2", "T4G", "FastAPI"], "alternates": {"html": "https://wpnews.pro/news/pure-jax-on-g5g-serving-gemma-4-on-graviton-and-a-t4g", "markdown": "https://wpnews.pro/news/pure-jax-on-g5g-serving-gemma-4-on-graviton-and-a-t4g.md", "text": "https://wpnews.pro/news/pure-jax-on-g5g-serving-gemma-4-on-graviton-and-a-t4g.txt", "jsonld": "https://wpnews.pro/news/pure-jax-on-g5g-serving-gemma-4-on-graviton-and-a-t4g.jsonld"}}