onnx-genai-example-gemma4-e2b-assistant

Gemma-4 E2B-it assistant — ONNX speculative-decoding drafter (fp16, CUDA)

A real-weight ONNX export of google/gemma-4-E2B-it-assistant, the speculative-decoding drafter for google/gemma-4-E2B-it. Built with mobius. Pairs with the target justinchuby/onnx-genai-example-gemma4-e2b.

• Source: google/gemma-4-E2B-it-assistant @ 2d874ef7d29f9a30599a1e4b3c1cbc9595f005df (Apache-2.0)
• dtype / EP: float16 / CUDAExecutionProvider (validated on NVIDIA H200)
• mobius: 0776f562 · onnxruntime 1.27.0 · onnx_ir 1.0.0 · transformers 5.14.1

Truthful graph contracts

The drafter is driven by the released SinglePositionMultiTokenCandidateGenerator: each step feeds inputs_embeds = concat(target_embed(last_token), recycled_hidden), emits a draft token (argmax) and projected_state, which is recycled as the next hidden; the target's KV is held read-only across the draft window.

Metadata — canonical v1 speculative workflow SHIPPED (Rust-validated vs onnx-genai #1716)

This repo ships inference_metadata.yaml: a faithful onnx-genai v1 speculative workflow spanning components: {target, assistant}, authored to onnx-genai PR #1716 and validated with the authoritative validate_metadata Rust validator (cross-references + real ONNX graph ports, not just JSON-Schema). The assistant reads shared_kv.{full,sliding}_attention.{key,value} via access: read_only aliases keyed on the borrowed target owner cell (no output); the output-only projected_state carry is folded into inputs_embeds via proposal_execution.chained.folded_carry_output. Per the final #1716 schema the chained block also names, with real graph ports (no placeholders): folded_carry_seed: {component: target, output: hidden_states.34} (the carry_0 seed — the target's post-final-norm hidden, == HF hidden_states[-1], L4 cos 0.999999) and token_embedding: {component: target, table: model.embed_tokens.weight} (the tied table the drafter gathers embed(last_token) from). vocabulary: {kind: identical}, distribution_preserving: true, and rollback_state over the 30 target KV cells. Uses a combined-package layout (target artifact = target/model.onnx; the two hosting repos keep model.onnx at their roots for standalone ORT use). The target must be built with mobius PR #546 (which emits hidden_states.34); the drafter graph itself is unchanged. See METADATA_STATUS.md + evidence/rust_validation.json. Also ships the faithful, graph-derived component_contract.yaml + full parity evidence.

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

• L4 drafter step (target-produced hidden + shared KV): the centroid router selects the identical 4096 active vocab positions as HF (Jaccard 1.0), argmax matches, projected_state cosine 0.9998 (evidence/l4_parity.json).
• L5 assisted (target + assistant): 8 rounds, 45 proposed draft tokens, 8 accepted / 37 rejected, output identical to greedy (distribution- preserving) (evidence/l5_assisted_detail.json).
• ONNX-drafter bridge: teacher-forced with HF's per-step inputs, the ONNX drafter matches HF argmax at every step (evidence/l5_drafter_rootcause.json). Free-running, it matches HF for the first ~4 draft tokens then drifts — pure fp16-vs-bf16 accumulation in a 4-layer draft, not a graph bug.

Usage (one drafter step, direct ONNX Runtime CUDA)

import numpy as np, onnxruntime as ort
s = ort.InferenceSession("model.onnx", providers=["CUDAExecutionProvider"])
out = s.run(["logits", "projected_state"], {
    "inputs_embeds": inputs_embeds,                          # f16 [1,q,3072]
    "attention_mask": np.ones((1, kv_len), np.int64),
    "shared_kv.sliding_attention.key":   skv_sliding_k,      # f16 [1,1,kv,256]
    "shared_kv.sliding_attention.value": skv_sliding_v,
    "shared_kv.full_attention.key":      skv_full_k,         # f16 [1,1,kv,512]
    "shared_kv.full_attention.value":    skv_full_v,
})

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

<!-- inference-metadata-annotation:end -->