Porting Gemma-4 (2B/4B/12B) to AWS Inferentia2: A Field Report
How Gemma-4 was ported to AWS Inferentia2 by bypassing vLLM/NxD, solving mixed attention heads and KV-sharing via direct model tracing.
An engineering team documents how they got Google's Gemma-4 family (E2B, E4B, 12B) running correctly on AWS Inferentia2 by bypassing the vendor's vLLM/optimum-neuron/neuronx-distributed stack entirely and tracing the Hugging Face forward pass directly. The core problem: Gemma-4's TPU/XLA-native features — cross-layer KV-sharing, grouped-query attention with mismatched head counts, and interleaved sliding-window/global attention — cannot be expressed in AWS's static graph builder, causing the public Neuron vLLM endpoint to serve fluent-looking but incorrect output.
The report breaks 'mixed heads' into three distinct engineering problems, each requiring a different fix: on E2B/E4B, identifying which layers actually own a KV cache versus which merely read a neighbor's; on E4B, sharding GQA's KV heads across tensor-parallel ranks only when divisible while preserving num_key_value_groups so repeat_kv stays correct; and on 12B, replicating indivisible single-KV-head global attention layers while shrinking group counts to match sharded query heads. Two subtle bugs — an ordering issue with post-sharding tensor shapes and a dead code path never exercised until the 12B model — consumed hours of debugging.
The payoff: greedy decoding matches the CPU reference token-for-token across all three model sizes, with throughput of ~44 tok/s (E2B, single core), ~33-39 tok/s (E4B, TP=2), and ~15 tok/s (12B, TP=2). The write-up offers a concrete blueprint for engineers who need to bring newer, architecturally unconventional models to production inference hardware ahead of official vendor support.