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:
- Target (standalone):
justinchuby/onnx-genai-example-gemma4-e2b - Assistant (standalone):
justinchuby/onnx-genai-example-gemma4-e2b-assistant
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-lessread_onlyaliases (alias key = a representative target owner cellpast_key_values.14/13,input= the drafter'sshared_kv.<type>.<role>port,access: read_only,role, nooutput, nolayer). Its resolved decode ABI iskv_ownership: sharedwith 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 emitsprojected_stateas output only; it re-enters as the trailing half of the fusedinputs_embeds = concat(target_embedding(token), carry). The final #1716 schema makes the carry explicit with real graph ports (no placeholders):folded_carry_seednames the carry_0 source = the target's post-final-normhidden_states.34(== HFhidden_states[-1], L4 cos 0.999999; emitted by mobius PR #546), andtoken_embeddingnames the tied embedding table. No state cell, not inrollback_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.