Porting Gemma-4 (2B / 4B / 12B) to AWS Inferentia2

*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

/ NxDneuronx-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:

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

def _kv_rank_width(nkv):
    return nkv // TP if nkv % TP == 0 else nkv     # 2 // 2 = 1 KV head / rank on E4B

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)

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) andrepeat_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
    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 classoptimum-neuron

. The Neuron vLLM backend dispatches through optimum-neuron

, so vLLM has nothingNxD 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 attentionsliding_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 timebuf*(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 byneuronx-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/Oinput_output_aliases

), updated in place each step, so the cache never leaves the cores. Each rank's KV is seeded with head ModelBuilder

)

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 <unk>

. The model then faithfully emitted garbage (unused high-id tokens), which looked exactly like a device/reload/precision failure and consumed hours. The decisive diagnostic was a tok("hello world").input_ids

before suspecting the compiler.inf2.xlarge

(16 GB) needs swap.inf2.xlarge

.inf2.8xlarge

(700 GB root); the push only uploads the tiny server layer since the base is already on Hub.The three obstacles above were the entry point. The port also produced a set of findings that

generalize beyond Gemma-4 and were, in aggregate, worth more than any single fix.

Every time this port produced garbage, the NeuronCore was innocent. The causes, in order of how

often they bit: a broken/missing tokenizer (§7), a mis-restored weight reload (§8.2), a wrong

buffer (layer_scalar

, §4.6), and a mismatched driver/SDK version (a precompiled neff can

mis-execute into garbage rather than error on the wrong runtime). All four are cheaper to rule out than a device bug, and all four

The fastest oracle turned out to be a CPU-reference run on the same box: load the model in bf16

with from_pretrained

(≈5 s off the page cache once the weights are downloaded) and run one forward.

If the CPU reference produces the same garbage as the device, the accelerator is exonerated and the

bug is upstream (tokenizer, input, weights). This single technique collapsed a multi-hour "device

bug" into a two-minute tokenizer fix. Reach for the CPU oracle before profiling the neff.

The most dangerous illusion in the whole project was a green SEQ_MATCH

. Correctness was validated

in-process on the freshly traced model (ModelBuilder.trace(initialize_model_weights=True)

) — but

the server loads a saved model in a fresh process and calls

initialize_with_saved_weights()

. Those are different code paths, and the in-process pass never

exercised the one the server actually uses. (torch.jit.load

without the init call fails outright —

"This model is not initialized" — so the init step is load-bearing, not cosmetic.)

Combined with the auto-port's false "100% PASS" against a PLE-stripped golden (§3), the lesson is

blunt: a passing test against the wrong oracle, or on the wrong execution path, is worse than a failing one. Validate (a) the exact artifact you will ship, (b) in the exact way you will load it,

E2B is marketed as "2B" and E4B as "4B" — the MatFormer/PLE effective parameter counts. But the

device footprint is the full parameter count (~5B and ~8B), and that is what has to fit 16 GB. This

is exactly why E4B, the "4B" model, does not fit one core while E2B, the "2B" model, does.

Plan capacity from real parameters, never the effective headline.

The second half of the trap: bf16 halves the on-disk neff but not the on-device constant count.

The intuition "just use lower precision to fit" is wrong here — E4B's ~15.4 GB of resident fp32

constants stay ~15.4 GB of slots regardless of dtype tricks at the boundaries, and a second neff

still tips past 16 GB. Tensor parallelism, not precision, is the lever that actually fits the model on the core (§4.4).

The same cross-layer KV-sharing that NxD can't represent (§2.1, §3) is what makes E2B fit one core in

the first place. Only the 15 non-shared layers allocate a KV buffer; the rest read through. The

device-resident cache is therefore far smaller than a naïve one-buffer-per-layer implementation, which

is a large part of why the memory budget closes. An architectural feature that breaks the vendor

abstraction can still be a net win once you stop fighting it.

Two of the compiler wins came from math, not from the compiler:

argmax

— greedy decode is identical whether or not it runs. Moving it host-side (only needed for sampling) freed the SBUF it was overflowing (§4.2). Generalize: layer_scalar

(§4.6). When a model multiplies by a learned scalar that lives in a buffer, that scalar will be wrong after any param-only load — a bug with no error message, only degraded output.inf2.xlarge

and inf2.8xlarge

carry the identical 2-NeuronCore / 32 GB-HBM accelerator; they

differ only in host vCPU (4 vs 32) and RAM (16 vs 128 GB). Since Gemma-4's transformer runs

entirely on the cores, the only thing standing between the ~4× cheaper box and full performance is

host memory — solved by "slim" servers that keep just the embedding table on the CPU and drop the

transformer layers (they live in the neff). Result: all three models serve on a single inf2.xlarge, and because throughput barely drops, the cheap box is

Two prefill strategies coexist and trade off cleanly, so the "right" one depends on the workload:

host-seed (tp_alias ) | device-prefill (ModelBuilder ) | | |---|---|---| | First token | ~1.4–1.6 s (CPU prefill) | ~0.1–0.16 s | | Decode | up to ~60 tok/s (E4B slim) | ~36–39 tok/s (E4B) | | Best for | long generations | chat / first-token-bound |

Because the KV cache is device-resident in both, per-token latency is essentially flat across context length — there is no host round-trip that grows with the sequence. The knob to turn is

Hugging Face (recipe + compiled neffs + Dockerfiles + model cards):

xbill9/gemma-4-E2B-it-inferentia2

xbill9/gemma-4-E4B-it-inferentia2

xbill9/gemma-4-12B-it-inferentia2

Docker Hub (prebuilt, run with --device /dev/neuron0 --ipc=host

):

xbill9/gemma4-optb

— E2B: latest

/512-128

, slim

, tp2-slim

, tp2-2048

xbill9/gemma4-optb-e4b

— E4B: latest

/512-128

, tp2-devprefill-512

, slim-devprefill

xbill9/gemma4-optb-12b

— 12B: tp2-devprefill-256

, slim-devprefill

Key source (branch gemma4-inf2-nxd-kvshare

): optb_kv.py

(E2B two-graph), tp_alias_trace.py

(E4B TP+alias), tp_mb.py

(E4B/12B device-prefill ModelBuilder

), optb_server_*.py

(full / slim /

device-prefill HTTP servers with OpenAI routes). Fix history: e785f6d

(layer_scalar), d70dc94

(12B global-layer sharding), 328ca88

(drop on-device softcap), a0a2a75

(force eager attention),

78967ac

(bf16 weights).

gemma4_unified

audio (640 samples/40 ms) and image (48×48 patch) projection paths exist but are not yet wired to the device graph.Natural next steps: static batching, a 12B 512-context recompile, a tp_alias

-style 12B decode path

(the E4B data suggests ~2× throughput headroom), and wiring the 12B multimodal projections.

Base models © Google, Apache-2.0. The compiled neffs embed the base weights in bf16 and are redistributed as Apache-2.0 derivatives with attribution. The Option-B recipe, TP sharding rules, device-resident KV design, and servers are Apache-2.0. Not affiliated with or endorsed by Google or AWS; provided for research/testing.

source & further reading

dev.to — original article I built an open-source MCP server so AI agents can read the web - even sites that block bots Our Rust file watcher ate 23.6 GB of RAM, and our ignore rules never had a chance System Design the Agentic Way — Phase 1: The Single Machine Ceiling