Gemma-4 E2B-it — self-contained speculative package (target + assistant drafter, fp16/CUDA)

A real-weight, self-contained ONNX speculative-decoding package: the Gemma-4 E2B target decoder and its borrowed-KV drafter co-located under one root, with a single canonical inference_metadata.yaml that spans both as one speculative: workflow. Built with mobius.

Both models are dense (enable_moe_block=false) — MoE is real only in the separate …-gemma4-26b-a4b example.

Why this repo exists (packaging)

The onnx-genai speculative contract is one workflow over components: {target, assistant}, and a package must be self-contained — every artifact resolves within the package root (the loader opens each ONNX and rejects ../absolute/symlink escapes). So the combined metadata references model.onnx (assistant, at root) + target/model.onnx (target) physically present here. The two standalone repos remain for individual ORT execution:

The graphs here are byte-identical to those repos (target model.onnx sha256 c1a83ab6… — now additionally exposing the hidden_states.34 folded-carry seed output; assistant 385e8152…).

Layout

model.onnx (+ .data)          # assistant drafter (4 layers, borrowed read-only KV, no own cache)
target/model.onnx (+ .data)   # target decoder (35 layers, hybrid attn, 20 shared-KV)
inference_metadata.yaml       # ONE speculative workflow: components {target, assistant}
component_contract.yaml       # assistant component contract
tokenizer.json, tokenizer_config.json, chat_template.jinja
source_provenance.json, runtime_versions.json, evidence/

Speculative contract (canonical onnx-genai v1, Rust-validated on #1716)

  • components: {target, assistant}; speculative.proposer: assistant, speculative.target: target.
  • Borrowed read-only KV, no own cache: the assistant's shared_kv.* inputs bind to the target's owner cells via output-less read_only aliases (alias key = a representative target owner cell past_key_values.14/13, input = the drafter's shared_kv.<type>.<role> port, access: read_only, role, no output, no layer). Its resolved decode ABI is kv_ownership: shared with zero KV transitions. shared_state: [full_attention, sliding_attention].
  • Folded output-only carry: proposal_execution: {kind: chained, token_embedding_input: inputs_embeds, logits_output: logits, folded_carry_output: projected_state, folded_carry_seed: {component: target, output: hidden_states.34}, token_embedding: {component: target, table: model.embed_tokens.weight}} — the drafter emits projected_state as output only; it re-enters as the trailing half of the fused inputs_embeds = concat(target_embedding(token), carry). The final #1716 schema makes the carry explicit with real graph ports (no placeholders): folded_carry_seed names the carry_0 source = the target's post-final-norm hidden_states.34 (== HF hidden_states[-1], L4 cos 0.999999; emitted by mobius PR #546), and token_embedding names the tied embedding table. No state cell, not in rollback_state (recomputed on rejection).
  • Vocabulary identical (drafter emits the full 262144 axis; the ordered/ centroid-pruned LM head is graph-internal), distribution_preserving: true (standard speculative rejection sampling corrects any drafter).
  • rollback_state = the 30 target KV cells only; rollback_positions = 6 = max_proposal_width.

This layout validates directly with the onnx-genai #1716 validate_metadata Rust binary (valid, exit 0) — no temp/hardlink layout needed, because both graphs are physically co-located.

Parity evidence (NVIDIA H200) — see evidence/

  • L4 drafter: centroid-route Jaccard 1.0, argmax match ✓, projected_state cosine 0.9998 (using target-produced hidden states + shared KV).
  • L5 assisted target+assistant: 45 proposed / 8 accepted / 37 rejected; final output == target greedy ✓; ONNX drafter teacher-forced matches HF at every step (free-run drift root-caused to fp16-vs-bf16, see evidence/l5_drafter_rootcause.json).

Deterministic: the workflow emits logits; token selection is external greedy argmax (no RNG uint64 sampler), so it runs identically under ORT and native.

Licensed Apache-2.0 (see SOURCE_LICENSE.md); both graphs derive from pinned Apache-2.0 checkpoints.

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for justinchuby/onnx-genai-example-gemma4-e2b-speculative

Quantized
(317)
this model

Collection including justinchuby/onnx-genai-example-gemma4-e2b-speculative