justinchuby/onnx-genai-example-gemma4-e2b

Gemma-4 E2B-it — ONNX text decoder (fp16, CUDA) with canonical inference metadata

A real-weight ONNX export of the text decoder of google/gemma-4-E2B-it, built with mobius and shipped with a canonical, hashless inference_metadata.yaml. This is one of two paired examples; the drafter is justinchuby/onnx-genai-example-gemma4-e2b-assistant.

• Source: google/gemma-4-E2B-it @ 3e22461f65e89153144f8adb70e3b8c2cc9845a7 (Apache-2.0)
• dtype / EP: float16 / CUDAExecutionProvider (validated on NVIDIA H200)
• mobius: 710d4927 (PR #546 branch, pending review) · onnxruntime 1.27.0 · onnx_ir 1.0.0 · transformers 5.14.1
• Graph also exposes hidden_states.34 — the post-final-norm hidden state (the lm_head input, == HF output_hidden_states hidden_states[-1]) — so the paired borrowed-KV drafter can seed its folded carry from a real target output (onnx-genai #1716 folded_carry_seed). Weights are byte-identical to the prior build (model.onnx.data sha256 e525066e…); only the graph gained the extra output tap. Standalone greedy generation is unchanged.

Truthful graph contracts

The exported decoder faithfully preserves Gemma-4's exotic geometry:

Contract	Value
Layers	35 (GroupQueryAttention × 35)
Hybrid attention	sliding-window(512) + full_attention every 5th layer (indices 4,9,14,19,24,29,34)
Heterogeneous head_dim	local/sliding 256, global/full 512
Heads	num_attention_heads=8, num_key_value_heads=1 (MQA)
Shared-KV	num_kv_shared_layers=20 → only 15 layers own a KV cache; the last 20 borrow
Per-layer input	hidden_size_per_layer_input=256
MLP	double-wide GeLU (use_double_wide_mlp=true)
Logit softcap	final_logit_softcapping=30.0, tie_word_embeddings=true
MoE	disabled in this checkpoint (enable_moe_block=false) — dense MLP; not invented in metadata

inference_metadata.yaml (canonical onnx-genai v1, hashless) encodes the hybrid cache as two state-service groups — full_attention (real owner layers 4/9/14) and sliding_attention (the 12 sliding owners, evictable_prefix: true) — over the real graph ports (past_key_values.N.key / present.N.key), with the 20 shared-KV layers correctly owning no cache. It is validated with the authoritative onnx-genai PR #1716 validate_metadata Rust validator (cross-references + real ONNX graph ports, not just JSON-Schema; see evidence/rust_validation.json). Per-layer head_dim (256/512) lives in the graph IO. inference_metadata.mobius.yaml is the raw mobius --runtime onnx-genai emitter output (its own schema_version: 1.0) and policies/ holds that workflow's sampler / termination / cache-update graphs.

Parity evidence (real weights, ONNX fp16 CUDA vs HF)

• L4 (single forward, 15 tokens): top-1 argmax agreement 1.0 at every position, mean cosine 0.99995, no divergence (evidence/l4_parity.json). The added hidden_states.34 output matches HF hidden_states[-1] at cosine 0.999999 (evidence/l4_hidden_state.json).
• L5 greedy (32 new tokens): 32/32 token-for-token match with HF greedy; ~151 tok/s decode, ~22 GB VRAM on H200 (evidence/l5_generation.json).

Usage (direct ONNX Runtime CUDA)

import numpy as np, onnxruntime as ort
s = ort.InferenceSession("model.onnx", providers=["CUDAExecutionProvider"])
S = ids.shape[1]                          # int64 [1, S] prompt token ids
feeds = {"input_ids": ids, "attention_mask": np.ones((1, S), np.int64)}
head = {4:512,9:512,14:512}               # full layers use head_dim 512, else 256
for i in range(15):
    hd = head.get(i, 256)
    feeds[f"past_key_values.{i}.key"]   = np.zeros((1,1,0,hd), np.float16)
    feeds[f"past_key_values.{i}.value"] = np.zeros((1,1,0,hd), np.float16)
logits = s.run(["logits"], feeds)[0]       # [1, S, 262144]

Runtime notes

inference_metadata.yaml is the canonical onnx-genai v1 decoder workflow, validated against the onnx-genai #1716 schema. The released onnxruntime-genai genai_config.json path cannot represent this model's heterogeneous head_dim (it exposes one flat head_size), so direct ORT CUDA sessions are the supported runtime today. See evidence/runtime_limitations.json.

License

Apache-2.0, inherited from the source model. See SOURCE_LICENSE.md.