justinchuby/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

Contract	Value
Layers	4 (sliding, sliding, sliding, full)
Borrowed read-only KV	reads the target's shared_kv.{sliding,full}_attention.{key,value}; owns no cache, emits no present.*
Heterogeneous head_dim	sliding 256, full 512
Inputs	inputs_embeds [B,q,3072] = concat(target token-embedding, recycled hidden), attention_mask, shared_kv.*
Outputs	logits [B,q,262144], projected_state [B,q,1536] (recurrent carry)
Ordered embeddings	use_ordered_embeddings=true
Centroid-routed sparse pruned LM head	num_centroids=2048, top_k=32 → 4096 active logits / 262144 vocab per position (TopK→Gather→ReduceMin→ScatterElements)

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.