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/
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.