Five Gemma-4 models, one accelerator: what porting E2B 31B to AWS Inferentia2 taught me

Short summary

The author ported all five Gemma-4 models (E2B through 31B and the 26B-A4B MoE) to AWS Inferentia2, achieving token-for-token correctness against CPU fp32 references. The post details three progressive recipes — torch_neuronx.trace, tp_alias, and NxD ModelBuilder — and explains architecture-specific challenges like KV-sharing, per-layer embeddings, and MoE routing that the vendor stack couldn't handle natively. Practical guidance includes environment setup quirks and attention implementation workarounds.

• •All five Gemma-4 models ported to AWS Inferentia2 with exact correctness vs CPU fp32
• •Three progressive deployment recipes: torch_neuronx.trace, tp_alias, NxD ModelBuilder
• •Vendor stack (optimum-neuron, Neuron vLLM) couldn't express KV-sharing or MoE — custom tracing required

Generated with AI, which can make mistakes.