Running Gemma-4 on AWS Inferentia2: A Performance Field Report

The core issue? Gemma-4 is architecturally "messy" compared to vanilla decoders. It uses Per-Layer Embeddings (PLE), MatFormer nesting, and—most disruptively—cross-layer KV-sharing. On the smaller E2B and E4B models, many layers don't even compute their own Key/Value projections; they just reuse a neighbor's. When I tried the public Neuron vLLM endpoint, it served nothing but fluent-looking gibberish because the graph builder couldn't express this KV-sharing logic.

Here is how the performance actually shook out once I bypassed the broken vendor abstractions:

| Model | Hardware | Speed (Greedy Decode) | Notes |
|---|---|---|---|
| E2B | inf2.xlarge | ~44 tok/s | 1 core |
| E4B | inf2.8xlarge | ~33–39 tok/s | TP=2 required |
| 12B | inf2.8xlarge | ~15 tok/s | TP=2 required |

The most important takeaway for anyone worried about accuracy: the greedy decoding is token-for-token identical to the CPU reference across all three sizes.

The technical headache scales with the model size. While the E2B/E4B models fail due to that KV-sharing logic, the 12B model introduces a "mixed attention" nightmare. It interleaves sliding-window and global attention layers with different KV-head counts. The neuronx-cc compiler simply wasn't ready for this level of complexity out of the box.

If you're looking for the weights to try this yourself, I've pushed the converted versions here:
xbill9/gemma-4-{E2B,E4B,12B}-it-inferentia2 on Hugging Face.

For Docker users, check out:
xbill9/gemma4-optb{,-e4b,-12b}

Is the cost-to-performance ratio worth it? On the E2B, getting 44 tokens per second on an inf2.xlarge is actually quite impressive for the hardware, provided you can get past the initial compilation hurdles. For the larger 12B, the throughput drops significantly, but it's still a functional path if you need to stay within the AWS ecosystem.