Porting Gemma-4 to AWS Inferentia2: A Field Report

A developer documents the obstacles of running Google's Gemma-4 (2B/4B/12B) on AWS Inferentia2, including mixed attention heads, dead ends with vLLM/optimum-neuron/NxD, and neuronx-cc compiler limits. The solution bypasses the vendor stack using torch-neuronx trace and achieves correct, performant inference.

4 min readJul 14, 2026

Porting Gemma-4 to AWS Inferentia2: A Field Report

The Problem: Gemma-4 on Inferentia2

Google's Gemma-4 model family (2B, 4B, 12B) is designed for TPU/XLA. Running it on AWS Inferentia2 (inf2) breaks the vendor stack in three distinct ways: mixed attention heads, KV-sharing, and compiler limits. The author achieved 44 tok/s (2B), 33-39 tok/s (4B), and 15 tok/s (12B) on inf2 hardware, with token-for-token identical output to CPU reference. The key insight: bypass the vendor inference path entirely and trace the Hugging Face forward pass directly.

Three Faces of Mixed Heads

1. Mixed KV-Sharing Across Layers (E2B, E4B)

On E2B/E4B, some layers don't compute their own Key/Value projections—they reuse a neighbor's. Neuronx-distributed's ModelBuilder expects a static per-layer weight→buffer mapping and cannot express this. The fix: trace the Hugging Face forward pass directly, where KV-sharing becomes a live data dependency. Only 15 layers (out of many) actually own a KV buffer:

NONSHARED = []
for i, lyr in enumerate(lang.layers[:cfg.num_hidden_layers]):
    a = lyr.self_attn
    if not a.is_kv_shared_layer:
        NONSHARED.append(i)
        LINFO[i] = (a.k_proj.out_features // a.head_dim, a.head_dim)

2. Mixed Query vs KV Head Counts (E4B, TP=2)

E4B has 8 query heads and 2 KV heads (GQA). Under tensor parallelism (TP=2), column-sharding the query projection gives 4 q-heads/rank. The KV projection has only 2 heads—you cannot halve them independently. The rule: shard KV only when divisible, and keep num_key_value_groups unchanged:

def _kv_rank_width(nkv):
    return nkv // TP if nkv % TP == 0 else nkv

3. Mixed Attention Types with Different KV-Head Counts (12B)

The 12B model interleaves sliding-window layers (nkv=8, divisible by TP=2) and global layers (nkv=1, indivisible). For global layers, leave K/V replicated and shrink groups to match sharded query-head count:

if nkv % TP == 0:
    a.k_proj = col(a.k_proj); a.v_proj = col(a.v_proj)
else:
    a.num_key_value_groups = (nq // TP) // nkv

Gotcha: read nq before replacing q_proj with ColumnParallelLinear, or you'll compute groups from already-halved numbers.

Why vLLM / optimum-neuron / NxD Failed

optimum-neuron: No Gemma-4 model class (gemma4_text, gemma4_unified).
NxD ModelBuilder: Cannot represent cross-layer KV-sharing.
Vendor endpoint: Served fluent-looking gibberish (worst failure mode).
Auto-port: Claimed 100% PASS using a PLE-stripped checkpoint (both sides wrong the same way).

Option B: Point torch_neuronx.trace() (single-core) or neuronx_distributed.trace.parallel_model_trace (TP) directly at the Hugging Face forward pass. KV-sharing, PLE, GQA trace as ordinary live dependencies.

Neuron Compiler (neuronx-cc) Limits

PLE Table Off-Device (E2B/E4B)

The Per-Layer Embedding table (262144 × hidden) over-commits 16 GB core. Fix: keep on host, gather on CPU, feed as activations.

Logit Soft-Capping Overflow (12B)

tanh over 262144-token vocabulary in fp32 requires 524288 bytes/partition, exceeding 196608 B SBUF limit. Fix: skip soft-capping on device for greedy decode (monotonic, so argmax(softcap(x)) == argmax(x)). Apply only when sampling.

Fused Attention Overflow (12B)

Sliding window 1024 doubles attention tile vs 512 of E4B, overflowing SBUF. Fix: force eager attention (cfg._attn_implementation = "eager").

fp32 Constants Near 16 GB Limit (E4B)

A single Option-B neff for E4B uses ~15.4 GB of fp32 constants. The second neff (prefill+decode) pushes past 16 GB. Solution: use bf16 where possible and rely on tensor parallelism.

Results

Model	Tokens/sec	Hardware
E2B	~44	1 core
E4B	33-39	TP=2
12B	~15	TP=2

Artifacts: HF xbill9/gemma-4-{E2B,E4B,12B}-it-inferentia2, Docker xbill9/gemma4-optb{,-e4b,-12b}.

Next Steps

For developers attempting similar ports: (1) Start by tracing the reference forward pass, not the vendor abstraction. (2) Handle each "mixed heads" variant separately—don't assume uniform sharding. (3) Profile SBUF usage early; compiler limits will hit you. (4) Test with a known-correct reference to avoid false positives from auto-ported oracles.

If you're deploying Gemma-4 on AWS Inferentia2, use the provided artifacts or follow the Option B approach. The vendor stack will likely remain unsupported until Neuron SDK explicitly adds Gemma-4 model classes—which may never happen given the architectural mismatches.

Editor's Take

I've spent months wrestling with Neuron SDK, and this article validates my frustration: the vendor stack is too rigid for modern architectures. The author's decision to trace the Hugging Face forward pass directly is exactly what I'd do—it's more work upfront but avoids the black-box debugging. I'm impressed they got all three model sizes working with identical output to CPU reference. If AWS wants Inferentia to compete with NVIDIA, they need to either open-source more of the compiler or provide better support for non-standard models. For now, this approach is the only viable path.

— DevDigest Editorial

Key Takeaways

•Bypass vendor inference stacks (optimum-neuron, NxD) when they lack model support; trace the Hugging Face forward pass directly using torch-neuronx.
•Handle GQA under tensor parallelism by sharding KV heads only when divisible, and keep num_key_value_groups unchanged to avoid repeat_kv mismatch.
•For mixed attention types with indivisible KV heads, replicate K/V across ranks and adjust group counts manually.

Why It Matters

This is a concrete recipe for running the latest Google Gemma models on cost-effective AWS Inferentia2 hardware, bypassing vendor stack limitations. Developers deploying LLMs on AWS can save significant costs by using inf2 instead of GPUs, but need to know these workarounds. The techniques for handling mixed attention heads and KV-sharing apply to other non-standard transformer architectures.

#gemma-4#LLM inference#transformers#AWS Inferentia2#Neuron SDK

Get the weekly digest

Every Sunday - top tech stories, industry breakthroughs, and developer tools delivered to your inbox.

No spam, unsubscribe anytime.