---
license: apache-2.0
base_model:
- google/gemma-4-E2B-it
tags:
- onnx
- onnxruntime
- gemma4
- inference-metadata
- text-generation
pipeline_tag: text-generation
library_name: onnxruntime
---

# 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`](https://huggingface.co/google/gemma-4-E2B-it),
built with [mobius](https://github.com/onnxruntime/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`](https://huggingface.co/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)

```python
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`.

<!-- inference-metadata-annotation:start -->
## Annotated inference metadata

Review [`inference_metadata.annotated.yaml`](./inference_metadata.annotated.yaml) for inline explanations of this package's workflow, tensor/state/cache contracts, and fail-closed omissions. [`inference_metadata.yaml`](./inference_metadata.yaml) remains the canonical machine-authored contract; automated validation confirms both files parse to the same metadata object.
<!-- inference-metadata-annotation:end -->