Views
No views yet
google/gemma-4-26B-A4B-it —
the Mixture-of-Experts member of the Gemma-4 family — built with
mobius and shipped with a canonical,
hashless inference_metadata.yaml.This is the MoE example. It is deliberately distinct from the dense pairjustinchuby/onnx-genai-example-gemma4-e2b(target) and…-assistant(speculative drafter), which are dense (enable_moe_block=false). MoE is not invented anywhere — it is real only in this 26B-A4B checkpoint.
google/gemma-4-26B-A4B-it @ 4d7ae4984b7db7de8f8457170b3f1a419ee76d52 (Apache-2.0)float16 / CUDAExecutionProvider (validated on NVIDIA H200)272bce2b · onnxruntime-gpu 1.29.0 · onnx_ir 1.0.0 · transformers 5.14.1| Contract | Value |
|---|---|
| Layers | 30 |
| Mixture of Experts | 128 experts, top-8 per token, moe_intermediate_size=704, gelu_pytorch_tanh |
| MoE topology | hybrid parallel dense MLP (intermediate_size=2112) + MoE block each layer → post_ff_norm(dense + moe) + residual |
| Router | scale-free RMSNorm → learned scale × hidden_size^-0.5 → linear → softmax(fp32) → top-8 → renormalize → per_expert_scale |
| Hybrid attention | sliding-window(1024) + full_attention at layers 5, 11, 17, 23, 29 (5 full / 25 sliding) |
| Heterogeneous head geometry | sliding/local: 8 KV heads × 256 dim (GQA 2×); global/full: 2 KV heads × 512 dim (GQA 8×); 16 query heads |
| Shared-KV | num_kv_shared_layers=0 → all 30 layers own a KV cache (contrast E2B: 20 borrow) |
| Per-layer input | hidden_size_per_layer_input=0 (none; contrast E2B: 256) |
| Logit softcap | final_logit_softcapping=30.0, tie_word_embeddings=true |
inference_metadata.yaml (canonical onnx-genai v1, hashless) encodes the
model's MoE capability (model.mixture_of_experts: 128 routed experts, 8 per
token, expert intermediate 704, softmax top-k router) and the hybrid cache as
two state-service groups — full_attention (the 5 global owner layers, head_dim
512) and sliding_attention (the 25 local owner layers, head_dim 256,
evictable_prefix: true) — over the real graph ports
(past_key_values.N.key / present.N.key). It is validated with the
authoritative onnx-genai PR #1716 validate_metadata Rust validator
(cross-references + real ONNX graph ports, not just JSON-Schema).inference_metadata.mobius.yaml is the raw mobius emitter output (schema 1.0)
for reference.transformers (fp16) for the prompt
"Explain gravity in one sentence.":| Level | Metric | Result |
|---|---|---|
| L4 last-position logits | top-1 argmax / cosine | 1.0 / 0.99999 |
| L4 all 19 positions | top-1 argmax / cosine | 0.947 / 0.9996 |
| L5 greedy, 24 new tokens | exact token match vs HF | 24 / 24 ✓ |
"Gravity is the fundamental force of attraction that pulls objects with mass toward one another, governing everything from the fall of an apple"
performance.json.com.microsoft::MoE op for CUDA. However, the
released-onnxruntime fused MoE CUDA kernel fp16-accumulates the expert GEMMs,
which loses accuracy at this checkpoint's hidden size (2816): per-MoE-layer
cosine drops to ~0.97 and all-position top-1 parity collapses to 0.684. This is
a runtime kernel limitation, not an export error — verified because the
exact same expert weights, evaluated with standard MatMul/Einsum (which
accumulate in fp32 on CUDA even for fp16 inputs), match HuggingFace to
cosine 0.998.supports_fused_moe=False): batched Einsum gate/up + down GEMMs
masked by a ScatterElements'd top-K routing vector — O(1) ONNX nodes per layer
(~2k total) rather than a per-expert unroll (~183k nodes, impractical to load).
Once a future ORT ships an fp32-accumulate fused MoE kernel, the smaller fused
graph becomes usable with identical numerics.["CUDAExecutionProvider","CPUExecutionProvider"]. onnxruntime-gpu 1.29 silently
ignores the ("CUDAExecutionProvider", {"device_id": N}) tuple form (new
plugin-EP API) and falls back to CPU.inference_metadata.yaml workflow is a single deterministic step —
invoke decoder → emit logits — so token selection is external greedy argmax
(the L5 parity above is exact-token). The package ships no RNG token sampler
(no policies/; a counter-based PRNG sampler casts to uint64, which the native
CPU EP cannot execute), so it runs identically under ORT and native backends.model.onnx + model.onnx.data — fp16 MoE text decoder (real weights, ~52 GB)inference_metadata.yaml — canonical onnx-genai v1 (hashless, Rust-validated)inference_metadata.mobius.yaml — raw mobius emitter (schema 1.0)policies/ — generation policy ONNX helperstokenizer.json, tokenizer_config.json, chat_template.jinjasource_provenance.json, runtime_versions.json, performance.jsonevidence/ — L4/L5 harness + JSON results + MoE root-cause scripts1import onnxruntime as ort, numpy as np
2sess = ort.InferenceSession("model.onnx",
3 providers=["CUDAExecutionProvider", "CPUExecutionProvider"]) # plain strings
4# feed input_ids, attention_mask, and empty past_key_values.{0..29}.{key,value}
5# (KV geometry per layer: full layers 2×512, sliding layers 8×256 — see metadata)SOURCE_LICENSE.md); derived from a pinned Apache-2.0
checkpoint.inference_metadata.annotated.yaml for inline explanations of this package's workflow, tensor/state/cache contracts, and fail-closed omissions. inference_metadata.yaml remains the canonical machine-authored contract; automated validation confirms both files parse to the same metadata object.