Porting Gemma-4 (2B/4B/12B) to AWS Inferentia2: A Field Report
How Gemma-4's mixed attention heads and KV-sharing broke AWS's vendor stack, and the tracing/sharding fix that made Inferentia2 inference work.
Running Google's Gemma-4 family (E2B, E4B, 12B) on AWS Inferentia2 exposes three architectural features—cross-layer KV-sharing, mixed query/KV head counts under GQA, and interleaved sliding-window/global attention with different KV-head counts—that each map cleanly to TPU/XLA but break AWS's vendor stack. optimum-neuron has no Gemma-4 model class, neuronx-distributed's static graph builder can't express KV-sharing, and the public Neuron vLLM endpoint produced fluent-looking but incorrect output.
The fix bypasses NxD's ModelBuilder abstraction and traces the Hugging Face forward pass directly, letting KV-sharing emerge as an ordinary live data dependency—only layers that actually own a KV cache allocate a buffer. Under tensor parallelism (TP=2) for E4B, GQA's 8 query/2 KV heads require sharding KV heads across ranks while keeping num_key_value_groups unchanged, or repeat_kv silently mismatches. For the 12B, indivisible single-KV-head global layers are left replicated across ranks with shrunk group counts to keep attention correct.
The result: all three sizes produce greedy decode output token-for-token identical to the CPU reference, at E2B ~44 tok/s, E4B ~33-39 tok/s (TP=2), and 12B ~15 tok/s (TP=2). The report is a concrete playbook for engineers hitting architecture mismatches when porting non-standard transformer designs to fixed-function AI accelerators.