Gemma-4 在 AWS Inferentia2 上跑通了

如果你直接拿现成的 vLLM 或者 optimum-neuron 去跑 Gemma-4，大概率会得到一堆逻辑不通的乱码。原因很简单：Gemma-4 根本不是那种规规矩矩的 Decoder 结构，它在 E2B/E4B 版本里搞了层间 KV 共享（Cross-layer KV-sharing），甚至在 12B 版本里玩起了滑动窗口（Sliding-window）和全局注意力（Global attention）混用的“混合头”机制。AWS 的官方工具链目前根本没意识到这种复杂的注意力机制，它的 Graph Builder 压根表达不了这种 KV 共享逻辑。

我实测了三个规模的模型，情况各不相同：

E2B：虽然参数量“有效”只有 2B，但因为有 PLE（Per-Layer Embeddings）和复杂的 KV 共享，单核 bf16 勉强能塞进 16GB 显存，跑出约 44 tok/s。

E4B：由于结构太绕，必须开 TP=2（张量并行）才能跑起来，速度在 33–39 tok/s 左右。

12B：这是最头疼的，混合注意力机制导致 head_dim 在不同层之间变动，必须用 TP=2 跑，实测能到 15 tok/s。

最关键的是，我对比了 CPU 的 Reference 输出，这几个模型在 Greedy Decode 下是 token-for-token 完全一致的，说明模型权重和计算逻辑没跑偏。

如果你也想在 inf2 实例上折腾，别盯着官方默认的 vLLM 路径死磕，得自己处理 neuronx-cc 编译器的限制。

相关的模型权重和 Docker 镜像我已经整理好了，直接拿去用：

Hugging Face: xbill9/gemma-4-{E2B,E4B,12B}-it-inferentia2
Docker Hub: xbill9/gemma4-optb{,-e4b,-12b}

如果你在做高性能推理落地，这种针对特定硬件优化的工程实现比单纯调参重要得多。