Porting Gemma-4 12B (the encoder-free multimodal one) to AWS Inferentia2

Short summary

The author ports Google's Gemma-4 12B (encoder-free multimodal) to AWS Inferentia2 on a single inf2.8xlarge, achieving token-for-token identical output to the CPU fp32 reference. Three non-obvious fixes were required: extracting model.language_model from the unified wrapper, dropping the 30·tanh softcap on-device since it's argmax-invariant, and forcing eager attention when the fused sliding-window kernel overflowed on-chip SBUF. Result: ~101ms prefill, ~12GB/rank in bf16.

• •Gemma-4 12B's 'unified' multimodal class is encoder-free — grab model.language_model and ignore the wrapper
• •Softcap tanh is monotonic and argmax-invariant — skip it on-device over a 262K vocab to save transcendental compute
• •Sliding-window attention with window=1024 overflows SBUF on Neuron — force eager attention as the escape hatch

Generated with AI, which can make mistakes.