Five Gemma-4 models, one accelerator: what porting E2B 31B to AWS Inferentia2 taught me
Scored daily by a customisable AI persona to surface the most relevant engineering leadership news.
Hands-on porting experience with detailed insights, very actionable and relevant.
The author hand-ported all five Gemma-4 models (E2B, E4B, 12B, 31B, 26B-A4B MoE) to AWS Inferentia2 because vendor stacks (optimum-neuron, Neuron vLLM) lack support for architectures like cross-layer KV-sharing and 128-expert MoE. The ports share a skeleton of softcap 30, 262K vocabulary, GQA, and interleaved sliding+global attention, but differ in approach: E2B/E4B use torch_neuronx.trace with per-layer embeddings, 12B forces eager attention, 31B requires NxD ModelBuilder with TP=8, and the MoE model uses ModelBuilder with top-8 routing. All five decode token-for-token identically to CPU fp32 reference, with softcap dropped from the device graph.