The Problem: Gemma-4 on Inferentia2
Google's Gemma-4 model family (2B, 4B, 12B) is designed for TPU/XLA. Running it on AWS Inferentia2 (inf2) breaks the vendor stack in three distinct ways: mixed attention heads, KV-sharing, and compiler limits. The author achieved 44 tok/s (2B), 33-39 tok/s (4B), and 15 tok/s (12B) on inf2 hardware, with token-for-token identical output to CPU reference. The key insight: bypass the vendor inference path entirely and trace the Hugging Face forward pass directly.
Three Faces of Mixed Heads
1. Mixed KV-Sharing Across Layers (E2B, E4B)
On E2B/E4B, some layers don't compute their own Key/Value projections—they reuse a neighbor's. Neuronx-distributed's ModelBuilder expects a static per-layer weight→buffer mapping and cannot express this. The fix: trace the Hugging Face forward pass directly, where KV-sharing becomes a live data dependency. Only 15 layers (out of many) actually own a KV buffer:
NONSHARED = []
for i, lyr in enumerate(lang.layers[:cfg.num_hidden_layers]):
a = lyr.self_attn
if not a.is_kv_shared_layer:
NONSHARED.append(i)
LINFO[i] = (a.k_proj.out_features // a.head_dim, a.head_dim)
2. Mixed Query vs KV Head Counts (E4B, TP=2)
E4B has 8 query heads and 2 KV heads (GQA). Under tensor parallelism (TP=2), column-sharding the query projection gives 4 q-heads/rank. The KV projection has only 2 heads—you cannot halve them independently. The rule: shard KV only when divisible, and keep num_key_value_groups unchanged:
def _kv_rank_width(nkv):
return nkv // TP if nkv % TP == 0 else nkv
3. Mixed Attention Types with Different KV-Head Counts (12B)
The 12B model interleaves sliding-window layers (nkv=8, divisible by TP=2) and global layers (nkv=1, indivisible). For global layers, leave K/V replicated and shrink groups to match sharded query-head count:
if nkv % TP == 0:
a.k_proj = col(a.k_proj); a.v_proj = col(a.v_proj)
else:
a.num_key_value_groups = (nq // TP) // nkv
Gotcha: read nq before replacing q_proj with ColumnParallelLinear, or you'll compute groups from already-halved numbers.
Why vLLM / optimum-neuron / NxD Failed
- optimum-neuron: No Gemma-4 model class (
gemma4_text,gemma4_unified). - NxD ModelBuilder: Cannot represent cross-layer KV-sharing.
- Vendor endpoint: Served fluent-looking gibberish (worst failure mode).
- Auto-port: Claimed 100% PASS using a PLE-stripped checkpoint (both sides wrong the same way).
Option B: Point torch_neuronx.trace() (single-core) or neuronx_distributed.trace.parallel_model_trace (TP) directly at the Hugging Face forward pass. KV-sharing, PLE, GQA trace as ordinary live dependencies.
Neuron Compiler (neuronx-cc) Limits
PLE Table Off-Device (E2B/E4B)
The Per-Layer Embedding table (262144 × hidden) over-commits 16 GB core. Fix: keep on host, gather on CPU, feed as activations.
Logit Soft-Capping Overflow (12B)
tanh over 262144-token vocabulary in fp32 requires 524288 bytes/partition, exceeding 196608 B SBUF limit. Fix: skip soft-capping on device for greedy decode (monotonic, so argmax(softcap(x)) == argmax(x)). Apply only when sampling.
Fused Attention Overflow (12B)
Sliding window 1024 doubles attention tile vs 512 of E4B, overflowing SBUF. Fix: force eager attention (cfg._attn_implementation = "eager").
fp32 Constants Near 16 GB Limit (E4B)
A single Option-B neff for E4B uses ~15.4 GB of fp32 constants. The second neff (prefill+decode) pushes past 16 GB. Solution: use bf16 where possible and rely on tensor parallelism.
Results
| Model | Tokens/sec | Hardware |
|---|---|---|
| E2B | ~44 | 1 core |
| E4B | 33-39 | TP=2 |
| 12B | ~15 | TP=2 |
Artifacts: HF xbill9/gemma-4-{E2B,E4B,12B}-it-inferentia2, Docker xbill9/gemma4-optb{,-e4b,-12b}.
Next Steps
For developers attempting similar ports: (1) Start by tracing the reference forward pass, not the vendor abstraction. (2) Handle each "mixed heads" variant separately—don't assume uniform sharding. (3) Profile SBUF usage early; compiler limits will hit you. (4) Test with a known-correct reference to avoid false positives from auto-ported oracles.
If you're deploying Gemma-4 on AWS Inferentia2, use the provided artifacts or follow the Option B approach. The vendor stack will likely remain unsupported until Neuron SDK explicitly adds Gemma-4 model classes—which may never happen given the architectural mismatches.
