Porting Gemma-4 (2B / 4B / 12B) to AWS Inferentia2 A developer ported Google's Gemma-4 family (2B, 4B, 12B) to AWS Inferentia2, overcoming three architectural obstacles that break the vendor stack. The port achieved 44 tok/s for the 2B model, 33-39 tok/s for the 4B model, and 15 tok/s for the 12B model with token-for-token identical output to CPU. The key challenges involved mixed attention heads, cross-layer KV-sharing, and GQA sharding, which were resolved by tracing the Hugging Face forward pass directly. A field report on running Google's Gemma-4 family on AWS Inferentia2 inf2 , covering the three architectural obstacles that break the vendor stack — mixed attention heads optimum-neuron / NxD neuronx-cc limits —| Models | google/gemma-4-E2B-it , google/gemma-4-E4B-it , google/gemma-4-12B-it | | Hardware | AWS Inferentia2 — inf2.xlarge 1 chip / 2 cores / 32 GB HBM and inf2.8xlarge same accelerator, 128 GB host RAM | | Software | Neuron SDK 2.23 · torch-neuronx 2.8.0 · neuronx-cc 2.23.6484 · neuronx-distributed 0.17 · transformers 5.13.0 | | Result | E2B ~44 tok/s 1 core , E4B ~33–39 tok/s TP=2 , 12B ~15 tok/s TP=2 — greedy decode token-for-token identical to the CPU reference on all three | | Artifacts | HF: xbill9/gemma-4-{E2B,E4B,12B}-it-inferentia2 · Docker Hub: xbill9/gemma4-optb{,-e4b,-12b} | Gemma-4 is not a vanilla decoder. Across the family it combines several features that each map cleanly onto TPU/XLA where the model was designed but individually break the AWS inference path: tanh cap at 30 over a The AWS vendor stack optimum-neuron + neuronx-distributed + the Neuron vLLM backend has no Gemma-4 model class at all , and its graph builder cannot express KV-sharing. The public Neuron E2B | E4B | 12B | | |---|---|---|---| | HF class | Gemma4ForConditionalGeneration text | same | Gemma4UnifiedForConditionalGeneration | | model type | gemma4 text | gemma4 text | gemma4 unified encoder-free multimodal | | Params effective | ~5B 2B | ~8B 4B | 12B | | Per-Layer Embeddings | yes | yes | no hidden size per layer input=0 | | Cross-layer KV-sharing | yes 15 non-shared layers own KV | yes | no num kv shared layers=0 , every layer owns KV | | Query / KV heads | GQA | 8 q / 2 kv | 16 q / 8 kv, global layers nkv=1 | | Attention | sliding + global, sliding window=512 | sliding + global, sliding window=512 | sliding + global, sliding window=1024 , attention k eq v=true | | head dim | — | 256 | 256 global 512 | | Logit softcap | 30 | 30 | 30 | | Vocab | 262 144 | 262 144 | 262 144 tied embeddings | | Fits one 16 GB core? | yes bf16 | no → TP=2 | no → TP=2 | Three model sizes, three different reasons the vendor path fails, and — as it turns out — three distinct meanings of "mixed heads." "Mixed heads" was the single most expensive class of bug in this port. It shows up in three completely different forms as you move up the size ladder. On E2B/E4B a layer's attention is flagged self attn.is kv shared layer . Shared layers do not run their own k proj / v proj ; they attend against a KV tensor produced by an earlier "owner" layer. On TPU this is a free graph view. In AWS's neuronx-distributed NxD ModelBuilder , which wants a static, per-layer weight→buffer mapping, it cannot be represented — there is no place to say "this layer's K/V is that layer's K/V, computed once, live." The fix "Option B" , see §3 is to stop fighting the framework and trace the Hugging Face forward pass directly. KV-sharing then falls out as an ordinary live data dependency in the traced graph. The port only has to enumerate which layers actually own a cache: discover the layers that write KV; shared layers never touch the cache NONSHARED = for i, lyr in enumerate lang.layers :cfg.num hidden layers : a = lyr.self attn if not a.is kv shared layer: <- the crux NONSHARED.append i LINFO i = a.k proj.out features // a.head dim, a.head dim n kv, head dim Only the NONSHARED layers 15 of them on E2B allocate a KV buffer; the shared layers read through. This is why the KV cache is far smaller than a naïve "one buffer per layer" implementation would allocate — and why it fits a single core in bf16. E4B does not fit one 16 GB core §4.3 , so it must be sharded across both NeuronCores TP=2 . GQA now bites: E4B has 8 query heads and 2 KV heads . Column-sharding the query projection across 2 ranks is trivial 4 q-heads/rank . The KV projection has only 2 heads, and the subtle failure is that repeat kv inside attention broadcasts each KV head to a group of query heads — so you cannot independently halve the KV heads and the query heads without scrambling which query attends which KV . The rule that makes it correct: shard KV to nkv // TP heads per rank only when it divides , and crucially keep num key value groups unchanged so repeat kv still maps each rank's 4 query python def kv rank width nkv : return nkv // TP if nkv % TP == 0 else nkv 2 // 2 = 1 KV head / rank on E4B ... per layer: nkv = a.k proj.out features // a.head dim if nkv % TP == 0: divisible: shard k/v, keep groups a.k proj = col a.k proj a.v proj = col a.v proj else: see 2.3 Naïvely also dividing num key value groups produces the classic "4-vs-8 repeat kv mismatch": plausible-looking but wrong output. The KV cache is then kept device-resident and sliced per rank head r → rank r , because NxD's forward only returns rank 0's KV §5.2 . The 12B is where "mixed heads" becomes structural. gemma4 unified interleaves two attention types with different KV-head counts : nkv = 8 → divisible by TP=2 → shard to 4 heads/rank the §2.2 path . num global key value heads = 1 → You cannot column-shard a single KV head across two ranks. The working answer is to leave the global layers' K/V replicated a plain nn.Linear , which NxD loads identically on every rank and repeat kv matches the a = lyr.self attn hd = a.head dim nq = a.q proj.out features // hd capture BEFORE replacing q proj see gotcha a.q proj = col a.q proj ; a.o proj = row a.o proj nkv = a.k proj.out features // hd if nkv % TP == 0: sliding layers: nkv=8 - 4/rank, keep groups a.k proj = col a.k proj ; a.v proj = col a.v proj else: global layers: nkv=1, indivisible leave k/v REPLICATED; shrink groups to the per-rank sharded q-head count a.num key value groups = nq // TP // nkv Two gotchas made this a multi-hour fight commit d70dc94 : nq = q proj.out features // head dim q proj with a ColumnParallelLinear , whose .out features reports the AttributeError /wrong groups. else branch had literally never been exercised by E4B all E4B layers are nkv=2 , divisible , so it was dead, untested code until the 12B walked into it. Takeaway: "mixed heads" is not one problem. It is a mixed KV sharing across layers, b mixed query/KV counts within a layer GQA , and c mixed attention types with different KV counts per type. Each needs a different sharding rule, and the naïve "just divide everything by TP" is wrong in all three. optimum-neuron / NxD couldn't do it The port deliberately does not use the AWS vendor inference path. Here is why each vendor layer failed, and what replaced it. optimum-neuron has no Gemma-4. There is simply no gemma4 text / gemma4 unified model class optimum-neuron . The Neuron vLLM backend dispatches through optimum-neuron , so vLLM has nothing NxD ModelBuilder cannot represent KV-sharing. NxD wants a static graph where each layer maps The vendor endpoint served gibberish. The public Neuron vLLM/NxD deployment we began with came up "healthy" and produced fluent-looking but semantically empty text — the worst failure mode, because it looks like it works. An automated port falsely reported success. A framework auto-port pass claimed "100% PASS" . On inspection its golden reference had been built from a PLE-stripped checkpoint, so both sides of the comparison were wrong in the same way. A green check on a broken oracle is worse than a red one. The replacement — "Option B." Instead of teaching NxD about Gemma-4, torch neuronx.trace for single-core E2B and neuronx distributed.trace.parallel model trace / ModelBuilder for TP E4B/12B are pointed straight at the Hugging Face transformers 5.13 text forward pass . KV-sharing, There is still a TP-launch wrinkle: parallel model trace spawns rank workers, and transformers 5.13's FX utilities plus the spawn need a transformers.utils.fx shim and a TP CHILD environment sentinel to guard against infinite re-spawn. neuronx-cc wall Once the graph is traceable, neuronx-cc imposes its own hard limits. The recurring one is SBUF , the on-chip state buffer SRAM , capped at 196 608 B per partition . Several Gemma-4 ops blow straight through it, each with a distinct fix. The Per-Layer Embedding table is 262144 × hidden . Materialised on device it trips the compiler and by itself over-commits the 16 GB core. Fix: keep the PLE and word embedding on the host , gather on CPU, and feed the result in as activations . The device graph never sees the table. Gemma-4 applies softcap tanh logits / softcap over the full 262 144-token vocabulary. In fp32 that tanh is a custom-call whose working set is 128 × 524288 bytes — 524 288 196 608 B/partition : NCC INLA001 Allocated memory out of bound 128x524288 524288 vs 196608 B/partition Fix commit 328ca88 : argmax softcap x == argmax x — greedy decode is completely unaffected. The device returns temperature 0 . Correct and free. tanh cap does fit and is kept in-graph. sliding window=1024 12B Dropping softcap was necessary but not sufficient — the real overflow on the 12B was the fused attention custom-call , because the 12B's sliding window = 1024 is a0a2a75 : cfg. attn implementation = "eager" on both the config and its text config . Eager tiles attention sliding window=512 the fused pathA single Option-B neff for E4B materialises ~15.4 GB of fp32 model constants ; a second neff prefill + decode on the same core tips past the 16 GB NeuronCore budget → NRT RESOURCE / status=4 Allocation Failure . Two things are needed: MB WDTYPE=bf16 : NxD's shard children casts the checkpoint to the layer dtype, halving the neff. E2B is the exception: in bf16 it fits one core, which is why it stays single-core at ~44 tok/s. neuronx-cc is happiest with plain arithmetic. Several things are written out by hand rather than left to library helpers or dynamic control flow: 0.5 x 1+tanh 0.7978845608 x + 0.044715 x^3 . DynamicCache at trace time buf 1-oh + k oh — pure arithmetic, trace-safe, no scatter op or data-dependent indexing. layer scalar is a Each Gemma-4 layer does hidden states = self.layer scalar . layer scalar is a registered buffer default 1.0 , real value ≈ 0.06 . NxD's weight-sharding shard children / get sharded checkpoint loads parameters only, never buffers — so without intervention every layer over-scales ~16×, compounding across all layers into cos ≈ 0 garbage. Fix commit e785f6d : read the per-layer layer scalar from the checkpoint and copy it into the module by neuronx-cc runs on CPU — no NeuronCore needed to compile, so you can build neffs on any box. But a TP=2 compile runs neuronx-cc for both ranks concurrently and peaks past 128 GB host RAM ; add ≥55 GB of swap before compiling the resulting neff runs fine without swap . Correctness and SBUF are only half the story; throughput comes from keeping the KV cache on the device and never round-tripping it through the host. Three designs evolved: Two neffs share one static KV buffer: prefill padded prompt ≤ KV BUCKET , returns the 15 non-shared layers' K/V and decode single-token forward against a fixed KV MAX buffer, one-hot masked write . KV tensors are graph inputs/outputs. ~44 tok/s, ~0.06 s prefill after a one-time ~100 s neff load. Under TP a neff spans both cores, so you cannot park prefill on core 0 and decode on core 1. Two sub-designs handle prefill: tp alias nn.Parameter s aliased as graph I/O input output aliases , updated in place each step, so the cache never leaves the cores. Each rank's KV is seeded with head ModelBuilder KV parameters aliased as graph I/O so the cache is never round-tripped through the host aliases = {} for j in range NK : aliases w.kbuf j = 1 + j for j in range NK : aliases w.vbuf j = 1 + NK + j Per-token cost is essentially flat across context length because the cache is device-resident. The traced executor is saved with torch.jit.save model, path and reloaded with torch.jit.load + nxd model.initialize with saved weights torch.tensor start rank . The executor's own .save is TorchScript's, not NxD's — a subtle trap. All three ports match the CPU float reference token-for-token on greedy decode SEQ MATCH True , e.g. "The capital of France is Paris ." | Model | Build | Hardware | TP | Context | First token | Decode | Host RAM | |---|---|---|---|---|---|---|---| E2B | two-graph static KV | inf2.8xlarge / inf2.xlarge slim | 1 core | 512 / 128 | ~0.06 s | ~44 tok/s | ~6 GB slim | E4B | tp alias | inf2.8xlarge | 2 | 512 / 128 · 2048 / 512 | ~1.4–1.6 s | ~33 tok/s | — | E4B | device-prefill bf16 | inf2.8xlarge / inf2.xlarge slim | 2 | 512 / 128 | ~0.11–0.16 s | ~36–39 tok/s | ~8 GB slim | 12B | device-prefill bf16 | inf2.8xlarge / inf2.xlarge slim | 2 | 256 / 64 | ~0.1 s | ~15 tok/s | ~8 GB slim | The 12B's lower decode rate is inherent dense 12B, all weights active per token , not a port defect; per-token cost matches the 8xlarge full server. Notably, all three run on a single inf2.xlarge tokenizer.json masquerades as a device bug. hf download once left tokenizer.json absent; GemmaTokenizer loaded with no vocab and mapped