File size: 7,202 Bytes
b5cc9f2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | """Author a FAITHFUL onnx-genai v1 inference_metadata.yaml for the real
gemma4-e2b target decoder, using the ACTUAL graph port names
(`past_key_values.N.key` / `present.N.key`), grouped into full/sliding state
services. Structure mirrors onnx-genai #1716 example 23; validated against that
PR's schema. Owner layers are read from the graph's per-port head_dim
(512=global/full, 256=local/sliding)."""
from __future__ import annotations
import json, sys
import onnx_ir as ir
import yaml
MODEL = "/datadisks/disk1/justinchu/inference-metadata-catalogue/gemma4/target/package/model.onnx"
SCHEMA = "/datadisks/disk1/justinchu/inference-metadata-catalogue/gemma4/ref-1716/schema.json"
OUT = "/datadisks/disk1/justinchu/inference-metadata-catalogue/gemma4/target/package/inference_metadata.yaml"
g = ir.load(MODEL).graph
ins = {v.name: v for v in g.inputs}
outs = {v.name: v for v in g.outputs}
# discover owner KV layers + their head_dim
owners = {} # layer_idx -> head_dim
for name, v in ins.items():
if name.startswith("past_key_values.") and name.endswith(".key"):
idx = int(name.split(".")[1])
hd = (v.shape[3].value if hasattr(v.shape[3], "value") else int(v.shape[3]))
owners[idx] = hd
full = sorted(i for i, hd in owners.items() if hd == 512)
slide = sorted(i for i, hd in owners.items() if hd == 256)
def kv_contract(kind):
hh, hd = (f"{kind}_kv_heads", f"{kind}_head_dim")
return {"dtype": "float16", "rank": 4, "shape": ["batch", hh, "sequence", hd],
"batch_layout": {"kind": "request_aligned", "axis": 0}}
def opaque_app(name):
return {"role": {"kind": "opaque"}, "source": {"kind": "application", "name": name}, "required": True}
inputs = {
"request.active": {"contract": {"dtype": "bool", "rank": 1, "shape": ["batch"], "batch_layout": {"kind": "request_aligned", "axis": 0}}, "role": {"kind": "opaque"}, "source": {"kind": "application", "name": "active"}, "required": True},
"request.done": {"contract": {"dtype": "bool", "rank": 1, "shape": ["batch"], "batch_layout": {"kind": "request_aligned", "axis": 0}}, "role": {"kind": "opaque"}, "source": {"kind": "application", "name": "done"}, "required": True},
"request.accepted_len": {"contract": {"dtype": "int64", "rank": 1, "shape": ["batch"], "batch_layout": {"kind": "request_aligned", "axis": 0}}, "role": {"kind": "opaque"}, "source": {"kind": "application", "name": "accepted_len"}, "required": True},
"request.input_ids": {"contract": {"dtype": "int64", "rank": 2, "shape": ["batch", "sequence"], "batch_layout": {"kind": "request_aligned", "axis": 0}}, "role": {"kind": "runtime", "version": "1.0", "role": "prompt_tokens"}, "source": {"kind": "request"}, "required": True},
"request.attention_mask": {"contract": {"dtype": "int64", "rank": 2, "shape": ["batch", "sequence"], "batch_layout": {"kind": "request_aligned", "axis": 0}}, **opaque_app("attention_mask")},
}
comp_in = {
"input_ids": {"dtype": "int64", "rank": 2, "shape": ["batch", "sequence"], "batch_layout": {"kind": "request_aligned", "axis": 0}},
"attention_mask": {"dtype": "int64", "rank": 2, "shape": ["batch", "sequence"], "batch_layout": {"kind": "request_aligned", "axis": 0}},
}
comp_out = {"logits": {"dtype": "float16", "rank": 3, "shape": ["batch", "sequence", "vocab"], "batch_layout": {"kind": "request_aligned", "axis": 0}}}
state = {}
step_in = {"input_ids": "request.input_ids", "attention_mask": "request.attention_mask"}
step_out = {"logits": "decoder.logits"}
groups = {"full_attention": {"kind": "full_attention", "sequence_axis": 2, "layout": "bnsh", "update": {"kind": "append"}, "reuse": {"prefix_reusable": True, "evictable_prefix": False}, "ports": {"decoder": {}}},
"sliding_attention": {"kind": "sliding_attention", "sequence_axis": 2, "layout": "bnsh", "update": {"kind": "append"}, "reuse": {"prefix_reusable": True, "evictable_prefix": True}, "ports": {"decoder": {}}}}
for idx in sorted(owners):
kind = "full" if owners[idx] == 512 else "sliding"
grp = f"{kind}_attention"
for role in ("key", "value"):
pin = f"past_key_values.{idx}.{role}"
pout = f"present.{idx}.{role}"
inputs[f"request.{pin}"] = {"contract": kv_contract(kind), **opaque_app(pin)}
comp_in[pin] = kv_contract(kind)
comp_out[pout] = kv_contract(kind)
state[pin] = {"contract": kv_contract(kind), "scope": "invocation", "initializer": f"request.{pin}",
"recurrence": {"kind": "invariant"}, "management": "runtime", "release_boundary": "invocation",
"service_group": grp}
step_in[pin] = f"request.{pin}"
step_out[pout] = f"decoder.{pin}"
groups[grp]["ports"]["decoder"][pin] = {"input": pin, "output": pout, "role": role, "layer": idx}
doc = {
"schema_version": "v1",
"pipeline": {"workflow": {
"manifest": {"capabilities": ["workflow_ssa", "typed_emit", "serving_service_contract"]},
"inputs": inputs,
"outputs": {"logits": {"contract": {"dtype": "float16", "rank": 3, "shape": ["batch", "sequence", "vocab"], "batch_layout": {"kind": "request_aligned", "axis": 0}}, "role": "tensor", "stage": "pre_adapter"}},
"components": {"decoder": {"implementation": {"kind": "onnx", "artifact": "model.onnx"},
"ports": {"inputs": comp_in, "outputs": comp_out, "roles": {"input_ids": "token_ids", "logits": "logits"}}}},
"state": state,
"steps": [{"kind": "invoke", "component": "decoder", "inputs": step_in, "outputs": step_out},
{"kind": "emit", "value": "decoder.logits", "output": "logits", "mode": "replace"}],
"serving": {"active": "request.active", "done": "request.done", "accepted_len": "request.accepted_len",
"state_service": {"groups": groups}},
}},
}
# validate
from jsonschema import Draft202012Validator
schema = json.load(open(SCHEMA))
errs = sorted(Draft202012Validator(schema).iter_errors(doc), key=lambda e: list(e.path))
if errs:
print(f"INVALID ({len(errs)} errors):")
for e in errs[:12]:
print(" -", list(e.path), "->", e.message[:160])
sys.exit(1)
header = ("# Faithful onnx-genai v1 inference_metadata for google/gemma-4-E2B-it text decoder.\n"
"# Authored to the onnx-genai #1716 schema (example 23 shape) using the REAL exported\n"
"# graph port names (past_key_values.N.key / present.N.key) and REAL owner layers:\n"
f"# full_attention owners (head_dim 512): layers {full}\n"
f"# sliding_attention owners (head_dim 256): layers {slide}\n"
"# The 20 shared-KV layers borrow an owner's buffer inside the graph and expose no ports.\n"
"# MoE is DISABLED in this checkpoint (enable_moe_block=false) -> dense MLP, not invented.\n"
"# final_logit_softcapping / tie_word_embeddings are graph-internal. Validates against\n"
"# onnx-genai PR #1716 schema/inference_metadata.schema.json.\n")
with open(OUT, "w") as f:
f.write(header)
yaml.safe_dump(doc, f, sort_keys=False, default_flow_style=None, width=1000)
print(f"VALID. full owners={full} sliding owners={slide}. Wrote {OUT}")
|