From dbcc48809c910daddacff655573445fa8212e8bc Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 22 Aug 2026 18:44:53 -0700 Subject: [PATCH 1/3] Migrate onnx-genai metadata emitters to the current pipeline/speculative contract mobius's legacy onnx-genai emitters had drifted from two upstream contract redesigns, and the drift was invisible because every schema conformance test skipped in CI. Why it stayed hidden: the conformance tests searched a couple of hard-coded local onnx-genai checkout paths and `pytest.skip`ped when none existed. CI has no such checkout, so six tests never ran there, and on the machines that did have one the result depended on whatever revision that clone sat on. The two contract redesigns: 1. `pipeline` is now `PipelineSpec`, whose only property is `workflow` (a typed SSA graph) with `additionalProperties: false`. The legacy emitters produced `{models, dataflow, strategy, phases}`, which the schema now rejects outright. 2. `speculative` is now `SpeculativeContract`, requiring `{proposer, target, vocabulary, max_proposal_width}` and forbidding every field of the old flat MTP block. The flat `SpeculatorConfig` that block was modelled on still exists upstream, but it describes a HuggingFace `config.json` section, not `InferenceMetadata.speculative`. What changed: * Diffusion (`build_diffusion_pipeline_metadata`). Now emits the denoise loop as an explicit workflow: the solver, schedule, timestep table, guidance combine and output clamp are real ONNX components built from mobius's policy library and shipped with the document. It cannot delegate to `build_diffusion_workflow_metadata`, because the ComfyUI conversion path has no component graphs at all -- it deliberately does not build or export them -- so it builds the workflow directly while reusing that module's `_invoke` and `_publish_workflow_v1`. The schedule is derived from the scheduler's own betas via the same diffusers-compatible helpers the package exporter uses, so a ComfyUI conversion and a package export of one checkpoint describe the same dynamics. img2img's `start_step` lowers to a sliced schedule. Three things now fail closed rather than being silently mis-described: an ancestral sampler (no deterministic solver exists), Karras/exponential sigma spacing (the workflow ships the sigma table as a constant, so a hint field is not enough), and a latent-only graph with no VAE decode. * MTP (`write_mtp_speculator_metadata`). Now emits a `SpeculativeContract` anchored to the backbone's workflow: it registers the head as a workflow component, names the target by its declared `logits` port role, states the hidden handoff as `port_bindings.target_hidden_context` plus a `hidden_states` role on the target output, and completes the rollback capacity its own claim requires. It is `block` rather than `chained` because a chained proposer must expose a `logits_output` and this sidecar emits only `mtp_hidden` -- the runtime decodes it through the shared LM head, which is why that initializer is in `shared_weights`. * VLM. `build_native_vlm_package_metadata` produces mobius's internal structural descriptor, not a publishable document -- `build_vlm_workflow_metadata` already consumes it and republishes only `preprocessing` under a real `pipeline.workflow`. The bug was that `write_native_vlm_package_metadata` wrote the descriptor to `inference_metadata.yaml`; it now writes the workflow document, and the tests validate that published document instead. * CI visibility. The schema is vendored under `_schema/` and is the default, so conformance never skips and drift is a test failure. A local checkout is no longer consulted implicitly, since one that is ahead of or behind `main` reintroduces the same machine-dependent result; set `ONNX_GENAI_SCHEMA` to validate against a specific revision. Three further conformance tests in the codec, speech-to-text and duplex suites were skipping for the same reason and now run. Tests updated in step: the MTP tests assert the new contract (with the banned-name test repointed at the now-legacy field names) and are anchored to a workflow mobius actually emits, and the diffusion/ComfyUI tests assert workflow structure rather than the removed `strategy` block. Full suite: 8 failed / 4613 passed -> 1 failed / 4640 passed. The remaining failure (`qwen_image_test.py::test_deterministic_l4_l5_image_edit_golden`) is pre-existing and handled separately. Signed-off-by: Justin Chu Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- pyproject.toml | 1 + .../integrations/onnx_genai/__init__.py | 5 +- .../integrations/onnx_genai/_schema/README.md | 30 + .../_schema/inference_metadata.schema.json | 5726 +++++++++++++++++ .../codec_workflow_metadata_test.py | 7 +- src/mobius/integrations/onnx_genai/comfyui.py | 11 +- .../integrations/onnx_genai/comfyui_test.py | 146 +- src/mobius/integrations/onnx_genai/convert.py | 30 +- .../integrations/onnx_genai/convert_test.py | 46 +- .../onnx_genai/decoder_metadata_test.py | 25 +- .../duplex_workflow_metadata_test.py | 11 +- .../onnx_genai/inference_metadata.py | 907 ++- .../onnx_genai/inference_metadata_test.py | 554 +- .../speech_to_text_workflow_metadata_test.py | 7 +- 14 files changed, 7096 insertions(+), 410 deletions(-) create mode 100644 src/mobius/integrations/onnx_genai/_schema/README.md create mode 100644 src/mobius/integrations/onnx_genai/_schema/inference_metadata.schema.json diff --git a/pyproject.toml b/pyproject.toml index 82c0a8f72..4fd2e0b87 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,6 +60,7 @@ where = ["src"] [tool.setuptools.package-data] "mobius.upstream_patches" = ["data/*/*/*.patch", "data/*/*/*.json", "data/*/*/*.md"] +"mobius.integrations.onnx_genai" = ["_schema/*.json"] [tool.setuptools.dynamic] version = { attr = "mobius.__version__" } diff --git a/src/mobius/integrations/onnx_genai/__init__.py b/src/mobius/integrations/onnx_genai/__init__.py index 0eec54ade..46f68738e 100644 --- a/src/mobius/integrations/onnx_genai/__init__.py +++ b/src/mobius/integrations/onnx_genai/__init__.py @@ -22,8 +22,9 @@ :func:`write_audio_codec_workflow_metadata` emits typed codec SSA, while :func:`write_tts_workflow_metadata` reports the current nested-loop induction contract blocker precisely. -* **Diffusion pipelines** — :func:`write_diffusion_pipeline_metadata` emits an - iterative pipeline for a denoiser plus optional VAE / text encoder. +* **Diffusion pipelines** — :func:`write_diffusion_pipeline_metadata` emits a + typed SSA denoise loop for a denoiser plus VAE and optional text encoder, + shipping the sampler's solver/schedule components alongside it. :func:`write_onnx_genai_config` is the unified entry point: it inspects the built package and dispatches to the matching writer, so diff --git a/src/mobius/integrations/onnx_genai/_schema/README.md b/src/mobius/integrations/onnx_genai/_schema/README.md new file mode 100644 index 000000000..691905472 --- /dev/null +++ b/src/mobius/integrations/onnx_genai/_schema/README.md @@ -0,0 +1,30 @@ +# Vendored onnx-genai metadata schema + +`inference_metadata.schema.json` is a verbatim copy of onnx-genai's published +`schema/inference_metadata.schema.json`, which that repo generates from its Rust +types (`crates/onnx-genai-metadata/src/schema/`). + +**Why it is vendored.** The schema conformance tests used to look for a local +onnx-genai checkout and `pytest.skip` when they could not find one. CI has no +such checkout, so every one of those tests skipped there and the emitters were +only ever validated on the few developer machines that happened to have +onnx-genai cloned — and then only against whatever revision that clone sat on. +Two upstream contract redesigns (`pipeline` becoming a `PipelineSpec` whose only +property is `workflow`, and `speculative` becoming a `SpeculativeContract`) went +unnoticed for exactly that reason. Pinning the schema here makes the contract +mobius targets an explicit, reviewable file and makes drift a CI failure rather +than a silent skip. + +**Updating it.** Copy the file from onnx-genai `main` and run the onnx-genai +tests in this package: + +```bash +cp /schema/inference_metadata.schema.json \ + src/mobius/integrations/onnx_genai/_schema/ +python -m pytest src/mobius/integrations/onnx_genai/ -q +``` + +Set `ONNX_GENAI_SCHEMA=/path/to/inference_metadata.schema.json` to validate +against a different revision without editing this copy. + +Synced from onnx-genai `10dfcd788` (2026-08-22). diff --git a/src/mobius/integrations/onnx_genai/_schema/inference_metadata.schema.json b/src/mobius/integrations/onnx_genai/_schema/inference_metadata.schema.json new file mode 100644 index 000000000..52b3ef494 --- /dev/null +++ b/src/mobius/integrations/onnx_genai/_schema/inference_metadata.schema.json @@ -0,0 +1,5726 @@ +{ + "$defs": { + "AbsentInputKind": { + "description": "Supported absent-input fallback kinds.", + "oneOf": [ + { + "const": "zeros", + "description": "Materialize a zero-initialized tensor.", + "type": "string" + } + ] + }, + "AbsentInputSpec": { + "description": "Explicit tensor fallback for an absent optional graph input.", + "properties": { + "kind": { + "$ref": "#/$defs/AbsentInputKind", + "description": "Fallback materialization kind." + }, + "shape": { + "description": "Runtime-resolved shape of the fallback tensor.", + "items": { + "$ref": "#/$defs/TensorDimension" + }, + "type": "array" + } + }, + "required": [ + "kind", + "shape" + ], + "type": "object" + }, + "AdapterArtifact": { + "additionalProperties": false, + "properties": { + "alpha": { + "format": "double", + "type": "number" + }, + "bindings": { + "default": [], + "items": { + "$ref": "#/$defs/AdapterTargetBinding" + }, + "type": "array" + }, + "dtype": { + "$ref": "#/$defs/TensorDType" + }, + "identity": { + "type": "string" + }, + "index": { + "description": "Stable non-negative wire ID used by selection.segments.", + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "provenance": { + "anyOf": [ + { + "$ref": "#/$defs/AdapterProvenance" + }, + { + "type": "null" + } + ] + }, + "rank": { + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "version": { + "type": "string" + }, + "weights": { + "default": [], + "items": { + "$ref": "#/$defs/AdapterWeightArtifact" + }, + "type": "array" + } + }, + "required": [ + "index", + "identity", + "version", + "rank", + "alpha", + "dtype" + ], + "type": "object" + }, + "AdapterCacheContract": { + "additionalProperties": false, + "properties": { + "eviction": { + "$ref": "#/$defs/AdapterEvictionPolicy", + "default": "lru" + }, + "max_entries": { + "default": 16, + "format": "uint", + "minimum": 0, + "type": "integer" + } + }, + "type": "object" + }, + "AdapterDiscoveryFallback": { + "oneOf": [ + { + "enum": [ + "disabled" + ], + "type": "string" + }, + { + "const": "tooling_only", + "description": "Tooling/load-time graph discovery may produce a resolved manifest; execution may not guess.", + "type": "string" + } + ] + }, + "AdapterEvictionPolicy": { + "enum": [ + "lru" + ], + "type": "string" + }, + "AdapterPlanningContract": { + "additionalProperties": false, + "properties": { + "bucket_by_adapter_set": { + "default": true, + "type": "boolean" + }, + "invalidate_capture_on_eviction": { + "default": true, + "type": "boolean" + }, + "stable_buffers": { + "default": true, + "type": "boolean" + } + }, + "type": "object" + }, + "AdapterProvenance": { + "additionalProperties": false, + "properties": { + "producer": { + "description": "Producer/importer that resolved the manifest and normalized the artifact.", + "type": "string" + }, + "revision": { + "description": "Immutable source revision, commit, or content identifier.", + "type": [ + "string", + "null" + ] + }, + "source": { + "description": "Source model or adapter URI without credentials.", + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "producer" + ], + "type": "object" + }, + "AdapterScaleEncoding": { + "oneOf": [ + { + "const": "alpha_over_rank", + "description": "Runtime applies the binding/artifact `alpha / rank` factor.", + "type": "string" + }, + { + "const": "baked", + "description": "Source factors already encode the complete static scale (ORT `TORT` convention).", + "type": "string" + } + ] + }, + "AdapterSelectionContract": { + "additionalProperties": false, + "properties": { + "active": { + "description": "Optional bool[batch]; inactive rows never load or apply adapters.", + "type": [ + "string", + "null" + ] + }, + "adapter_counts": { + "description": "Int64[batch] number of valid adapter IDs in each row.", + "type": "string" + }, + "max_adapters": { + "description": "Fixed second dimension of segments and scales.", + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "scales": { + "description": "Float32[batch,max_adapters] effective request scales.", + "type": "string" + }, + "segments": { + "description": "Int64[batch,max_adapters] segment IDs in composition order.", + "type": "string" + } + }, + "required": [ + "segments", + "adapter_counts", + "scales", + "max_adapters" + ], + "type": "object" + }, + "AdapterServiceContract": { + "additionalProperties": false, + "properties": { + "application_capability": { + "description": "Generic application capability required from the runtime or execution provider.", + "type": "string" + }, + "artifacts": { + "additionalProperties": { + "$ref": "#/$defs/AdapterArtifact" + }, + "default": {}, + "type": "object" + }, + "cache": { + "$ref": "#/$defs/AdapterCacheContract", + "default": { + "eviction": "lru", + "max_entries": 16 + } + }, + "discovery_fallback": { + "$ref": "#/$defs/AdapterDiscoveryFallback", + "default": "disabled", + "description": "Explicit load-time tooling fallback. Runtime execution never guesses targets." + }, + "planning": { + "$ref": "#/$defs/AdapterPlanningContract", + "default": { + "bucket_by_adapter_set": true, + "invalidate_capture_on_eviction": true, + "stable_buffers": true + } + }, + "portable_fallback": { + "default": false, + "type": "boolean" + }, + "selection": { + "$ref": "#/$defs/AdapterSelectionContract", + "description": "Request-scoped adapter-set inputs. These are immutable SSA inputs for one request." + }, + "target_manifest": { + "$ref": "#/$defs/LoraTargetManifest", + "description": "Authoritative, architecture-neutral bindings resolved by producer/import tooling." + } + }, + "required": [ + "target_manifest", + "selection", + "application_capability" + ], + "type": "object" + }, + "AdapterTargetBinding": { + "additionalProperties": false, + "properties": { + "alpha": { + "description": "Per-target alpha override; absent uses the artifact alpha.", + "format": "double", + "type": [ + "number", + "null" + ] + }, + "rank": { + "description": "Per-target rank override; absent uses the artifact rank.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "target": { + "description": "Stable target ID declared by the authoritative target manifest.", + "type": "string" + }, + "weight_key": { + "description": "Key used to find this target's A/B tensors in the adapter bundle.", + "type": "string" + } + }, + "required": [ + "target", + "weight_key" + ], + "type": "object" + }, + "AdapterWeightArtifact": { + "additionalProperties": false, + "properties": { + "config_location": { + "description": "PEFT `adapter_config.json` paired with the safetensors file.", + "type": [ + "string", + "null" + ] + }, + "format": { + "$ref": "#/$defs/AdapterWeightFormat", + "default": "json" + }, + "loader_capability": { + "description": "Loader capability required to normalize this source into the canonical artifact.", + "type": "string" + }, + "location": { + "type": "string" + }, + "scale_encoding": { + "$ref": "#/$defs/AdapterScaleEncoding", + "description": "Whether alpha/rank remains to be applied or is already baked into B." + } + }, + "required": [ + "location", + "loader_capability", + "scale_encoding" + ], + "type": "object" + }, + "AdapterWeightFormat": { + "oneOf": [ + { + "const": "json", + "description": "Portable JSON tensor bundle used by the reference fallback.", + "type": "string" + }, + { + "const": "ort_genai", + "description": "Native ONNX Runtime GenAI adapter bundle.", + "type": "string" + }, + { + "const": "hf_peft", + "description": "Hugging Face PEFT `adapter_config.json` plus safetensors.", + "type": "string" + }, + { + "const": "safetensors", + "description": "Manifest-keyed safetensors parameter bundle.", + "type": "string" + } + ] + }, + "AttentionConfig": { + "description": "Build-time attention architecture and dimensions.", + "properties": { + "fallback_behavior": { + "anyOf": [ + { + "$ref": "#/$defs/AttentionType" + }, + { + "type": "null" + } + ], + "description": "Compatible attention behavior for runtimes that do not recognize `type`." + }, + "head_dim": { + "description": "Per-head hidden dimension.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "key_sequence_lengths": { + "anyOf": [ + { + "$ref": "#/$defs/KeySequenceLengthsSpec" + }, + { + "type": "null" + } + ], + "description": "Representation compatibility for the attention key-sequence lengths.\n\nAbsent means the canonical contiguous `int32 [batch_size]` representation\nis required." + }, + "num_attention_heads": { + "description": "Number of query/attention heads.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "num_kv_heads": { + "description": "Number of key/value heads; required by runtimes that need explicit GQA dimensions.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "sink_tokens": { + "description": "Number of leading \"attention sink\" tokens always retained alongside the\nsliding window (StreamingLLM). Only meaningful when `sliding_window` is\nset; `null` or `0` disables sink retention. These first tokens stabilize\nthe attention distribution and are never evicted by the window.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "sliding_window": { + "description": "Sliding-window length in tokens, or null for full-context attention.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "type": { + "$ref": "#/$defs/AttentionType", + "description": "Attention architecture.\n\nCanonical values include `multi_head`, `grouped_query`, and\n`multi_latent`; future values are allowed when paired with a usable\n`fallback_behavior`." + } + }, + "required": [ + "type" + ], + "type": "object" + }, + "AttentionType": { + "description": "Attention architecture vocabulary with an extension branch.", + "oneOf": [ + { + "enum": [ + "multi_head", + "multi_head_attention", + "grouped_query", + "group_query_attention", + "grouped_query_attention", + "gqa", + "multi_latent", + "multi_latent_attention", + "mla" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "multi_head", + "multi_head_attention", + "grouped_query", + "group_query_attention", + "grouped_query_attention", + "gqa", + "multi_latent", + "multi_latent_attention", + "mla" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "AudioOutputBinding": { + "description": "One named tensor output produced by an audio preprocessing program.\n\nThe output binds a processor-local value to a typed workflow SSA name.\nNeither the name nor the content role is inferred from a model identity.", + "properties": { + "content": { + "$ref": "#/$defs/AudioOutputContent", + "description": "Generic content role of this output." + }, + "contract": { + "anyOf": [ + { + "$ref": "#/$defs/TensorContract" + }, + { + "type": "null" + } + ], + "description": "Full workflow tensor contract. Required when `pipeline.workflow` is present." + }, + "dtype": { + "$ref": "#/$defs/TensorDType", + "description": "Element type of the emitted tensor." + }, + "name": { + "description": "Workflow SSA value this output binds to.", + "examples": [ + "audio.input_features" + ], + "minLength": 1, + "type": "string" + }, + "optional": { + "description": "Whether the runtime may omit this output when a model does not need it.", + "type": [ + "boolean", + "null" + ] + }, + "pad_value": { + "description": "Optional sentinel/pad value for padded entries.", + "format": "double", + "type": [ + "number", + "null" + ] + }, + "rank": { + "description": "Tensor rank of the emitted tensor.\n\nRedundant with `contract.rank` and retained only for programs that\ndeclare no full contract; workflow packages declare `contract`.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "source": { + "description": "Program-local value produced by a transform.", + "minLength": 1, + "type": "string" + } + }, + "required": [ + "name", + "source", + "content", + "dtype" + ], + "type": "object" + }, + "AudioOutputContent": { + "description": "Generic audio-output content-role vocabulary.", + "oneOf": [ + { + "enum": [ + "waveform", + "features", + "audio_features", + "valid_frames", + "valid_samples", + "sample_lengths", + "frame_lengths", + "validity_mask" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "waveform", + "features", + "audio_features", + "valid_frames", + "valid_samples", + "sample_lengths", + "frame_lengths", + "validity_mask" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "AudioPreprocessingProgram": { + "description": "Generic audio preprocessing program: an ordered transform pipeline plus the\nnamed workflow SSA tensor outputs it emits.\n\nThe program is expressed entirely as parameterized, architecture-neutral\ndata, mirroring `ImagePreprocessingProgram`. Transform operations are\ngeneric (decode, resample, downmix, rescale, normalize, pad, frame,\nspectrogram, log_mel). In workflow metadata, outputs are materialized by a\nmanifest-pinned preprocessing adapter invocation and bind processor-local\nvalues to typed SSA names. A package may name an output `input_values`,\n`input_features`, `attention_mask`, or anything else without introducing\nruntime model-family dispatch.\n\nOne program type covers every audio family. A CTC acoustic model declares\nresample/downmix/zero_mean_unit_variance over raw samples; an\nencoder-decoder speech model declares resample/pad/log_mel over a fixed\nwindow. The runtime reads the same fields either way.", + "properties": { + "outputs": { + "description": "Named tensor outputs the program emits, each bound to a workflow SSA value.", + "items": { + "$ref": "#/$defs/AudioOutputBinding" + }, + "minItems": 1, + "type": "array" + }, + "transforms": { + "description": "Ordered list of generic transform operations applied to decoded audio.", + "items": { + "$ref": "#/$defs/AudioTransform" + }, + "type": "array" + } + }, + "required": [ + "outputs" + ], + "type": "object" + }, + "AudioTransform": { + "description": "One generic audio transform operation.\n\n`op` selects the operation from a generic vocabulary; the remaining fields\nare the parameters that operation reads (only the relevant ones are set).\nEvery parameter is model DATA — concrete sample rates, channel counts, mel\nbin counts, and so on live in a model's fixture, never as constants baked\ninto this schema.", + "properties": { + "channels": { + "description": "Target channel count for a `downmix` operation.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "epsilon": { + "description": "Numerical stabilizer added to the variance for `zero_mean_unit_variance`.", + "format": "double", + "type": [ + "number", + "null" + ] + }, + "hop_length": { + "description": "Hop length in samples for a spectrogram operation.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "inputs": { + "description": "Named values consumed by this transform.\n\nAbsent means the operation consumes the immediately preceding value.\nExplicit names allow branching programs without tensor-name heuristics.", + "items": { + "minLength": 1, + "type": "string" + }, + "minItems": 1, + "type": [ + "array", + "null" + ] + }, + "mel_scale": { + "description": "Mel scale convention — generic string data (e.g. `slaney`, `htk`).", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "mode": { + "description": "Padding side / normalization mode selector — generic string data\n(e.g. `right`, `left`, `fixed_window`).", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "n_fft": { + "description": "FFT size for a spectrogram operation.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "num_mel_bins": { + "description": "Mel filterbank size for a `log_mel`/`log_mel_spectrogram` operation.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "op": { + "$ref": "#/$defs/AudioTransformOp", + "description": "Generic operation selector (e.g. `resample`, `zero_mean_unit_variance`,\n`log_mel`)." + }, + "outputs": { + "description": "Named values produced by this transform.\n\nThese names are processor-local data. Final graph bindings select them\nthrough `AudioOutputBinding::source`.", + "items": { + "minLength": 1, + "type": "string" + }, + "minItems": 1, + "type": [ + "array", + "null" + ] + }, + "pad_value": { + "description": "Fill value used by a `pad` operation.", + "format": "double", + "type": [ + "number", + "null" + ] + }, + "sample_rate": { + "description": "Target sample rate in Hz for a `resample` operation, and the analysis\nrate a mel filterbank is built for.", + "format": "uint32", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "scale": { + "description": "Scalar multiplier for a `rescale` operation.", + "format": "double", + "type": [ + "number", + "null" + ] + }, + "target_length": { + "description": "Fixed target length in samples or frames for a `pad`/`trim` operation.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "win_length": { + "description": "Analysis window length in samples for a spectrogram operation.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "window": { + "description": "Analysis window function — generic string data (e.g. `hann`).", + "minLength": 1, + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "op" + ], + "type": "object" + }, + "AudioTransformOp": { + "description": "Generic audio transform-operation vocabulary.\n\nOne vocabulary spans every declared audio program. A CTC acoustic\nmodel normalizes raw samples and never builds a spectrogram; a\nspeech-to-text encoder pads to a fixed window and takes a log-mel.\nBoth are the same kind of declaration, so both draw their operation\nnames from here rather than from a per-family list.", + "oneOf": [ + { + "enum": [ + "decode", + "resample", + "downmix", + "rescale", + "zero_mean_unit_variance", + "normalize", + "pad", + "trim", + "frame", + "spectrogram", + "log_mel", + "log_mel_spectrogram", + "emit_valid_frames", + "emit_valid_samples", + "emit_sample_lengths", + "emit_validity_mask" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "decode", + "resample", + "downmix", + "rescale", + "zero_mean_unit_variance", + "normalize", + "pad", + "trim", + "frame", + "spectrogram", + "log_mel", + "log_mel_spectrogram", + "emit_valid_frames", + "emit_valid_samples", + "emit_sample_lengths", + "emit_validity_mask" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "BatchInvariance": { + "description": "Dependence of a row's outputs on the rows batched with it.", + "oneOf": [ + { + "enum": [ + "row_independent", + "padding_sensitive" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "row_independent", + "padding_sensitive" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "BatchLayout": { + "description": "Structural relationship between a typed value and the runtime batch.\n\n`shared` values are invariant across requests. `request_aligned` values carry\nexactly one entry per in-flight request along `axis`, so compaction permutes\nthat axis. `token_packed` values are ragged: `offsets` names the\nrequest-aligned exclusive-prefix offset value and `owner` names the\nper-item owner mapping, which together let a runtime split and regroup the\npacked value without any serialized request ID. `runtime_sequence_state`\nmarks a value whose per-sequence storage the runtime owns outright.", + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "shared", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "axis": { + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "kind": { + "const": "request_aligned", + "type": "string" + } + }, + "required": [ + "kind", + "axis" + ], + "type": "object" + }, + { + "additionalProperties": false, + "description": "Each request owns a fixed-size contiguous group on ``axis`` (for\nexample conditional/unconditional classifier-free-guidance rows).", + "properties": { + "axis": { + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "factor": { + "description": "Number of contiguous physical rows owned by each logical request.\nValidation requires this to be at least one.", + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "kind": { + "const": "request_expanded", + "type": "string" + } + }, + "required": [ + "kind", + "axis", + "factor" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "axis": { + "description": "Packed axis of this value.", + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "kind": { + "const": "token_packed", + "type": "string" + }, + "offsets": { + "description": "Request-aligned value holding the exclusive prefix offset of each request's items.", + "type": "string" + }, + "owner": { + "description": "Item-aligned value mapping each packed item to its owning request row.", + "type": "string" + } + }, + "required": [ + "kind", + "offsets", + "owner", + "axis" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "runtime_sequence_state", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + } + ] + }, + "ChunkedPrefillConfig": { + "description": "Runtime chunked-prefill preference.", + "properties": { + "chunk_size": { + "description": "Preferred number of prompt tokens processed in each prefill chunk.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + } + }, + "type": "object" + }, + "ComponentContract": { + "additionalProperties": false, + "properties": { + "bindings": { + "additionalProperties": { + "type": "string" + }, + "default": {}, + "description": "Semantic role to concrete component port name.", + "type": "object" + }, + "equivalence": { + "$ref": "#/$defs/EquivalenceClass", + "default": "semantic", + "description": "How closely a substituted implementation must reproduce this contract.\n\nA runtime may freely choose any equivalent implementation, but the\ndeclared class bounds what \"equivalent\" means. Only a\n`distribution_preserving` (or `bitwise`) contract may be optimized\nspeculatively without caller opt-in." + }, + "id": { + "description": "Versioned semantic capability identifier. It never selects execution behavior.", + "type": "string" + }, + "parameters": { + "additionalProperties": { + "$ref": "#/$defs/ScalarValue" + }, + "default": {}, + "description": "Contract parameters that are not tensor ports, such as adapter actions.", + "type": "object" + }, + "version": { + "type": "string" + } + }, + "required": [ + "id", + "version" + ], + "type": "object" + }, + "ComponentImplementation": { + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "artifact": { + "type": "string" + }, + "kind": { + "const": "onnx", + "type": "string" + } + }, + "required": [ + "kind", + "artifact" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "abi": { + "type": "string" + }, + "artifact": { + "type": [ + "string", + "null" + ] + }, + "kind": { + "const": "adapter", + "type": "string" + }, + "version": { + "type": "string" + } + }, + "required": [ + "kind", + "abi", + "version" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "binding", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + } + ] + }, + "ComponentPorts": { + "additionalProperties": false, + "description": "Explicit input/output ports of one executable component.", + "properties": { + "inputs": { + "additionalProperties": { + "$ref": "#/$defs/TensorContract" + }, + "default": {}, + "type": "object" + }, + "outputs": { + "additionalProperties": { + "$ref": "#/$defs/TensorContract" + }, + "default": {}, + "type": "object" + }, + "roles": { + "additionalProperties": { + "$ref": "#/$defs/PortRole" + }, + "description": "Semantic role of the ports whose meaning the workflow's own structure\ncannot recover, keyed by port name.\n\nState ports never need an entry: a state group already names its\nper-component `(input, output)` pair, and the fixed-capacity control\nports are named by [`StateUpdate::IndexedScatter`]. What is left is the\nper-step dataflow a workflow binds by SSA value, where the binding\nrecords WHICH value reaches a port but not WHAT the component does with\nit. A runtime that specializes a decode step — packing tokens, reusing a\nlogits buffer, skipping a mask it can prove is causal — needs that\nsecond fact, and inferring it from a port's spelling is exactly the\nname-guessing this schema refuses everywhere else.\n\nRoles are architecture-neutral and describe the port, never the model\nfamily that happens to expose it.", + "type": "object" + } + }, + "type": "object" + }, + "ComponentRowScope": { + "additionalProperties": false, + "description": "Row scope of a component's runtime-private state.", + "properties": { + "axis": { + "description": "Batch axis of the component's row-scoped ports.", + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "stateful": { + "default": false, + "description": "Whether the component retains state between invocations for each row.", + "type": "boolean" + } + }, + "required": [ + "axis" + ], + "type": "object" + }, + "ConstraintLanguageFacts": { + "additionalProperties": false, + "description": "A constraint language the package's parser accepts.\n\nThe parser implementation may be native; only the dialect and version are\nportable facts. A request carries the grammar or JSON Schema itself.", + "properties": { + "component": { + "description": "Workflow component that interprets this dialect.", + "minLength": 1, + "type": "string" + }, + "dialect": { + "description": "Dialect identifier, e.g. `json_schema`, `ebnf`, `regex`.", + "minLength": 1, + "type": "string" + }, + "version": { + "description": "Exact dialect version, e.g. `2020-12` for JSON Schema.", + "minLength": 1, + "type": "string" + } + }, + "required": [ + "dialect", + "version", + "component" + ], + "type": "object" + }, + "DType": { + "description": "Scalar dtype vocabulary with common ONNX and runtime aliases.", + "oneOf": [ + { + "enum": [ + "float32", + "fp32", + "float16", + "fp16", + "half", + "bfloat16", + "bf16", + "float8_e4m3fn", + "fp8_e4m3fn", + "float8_e4m3", + "fp8_e4m3", + "float8_e5m2", + "fp8_e5m2", + "int8", + "uint8", + "int4", + "uint4" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "float32", + "fp32", + "float16", + "fp16", + "half", + "bfloat16", + "bf16", + "float8_e4m3fn", + "fp8_e4m3fn", + "float8_e4m3", + "fp8_e4m3", + "float8_e5m2", + "fp8_e5m2", + "int8", + "uint8", + "int4", + "uint4" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "DecodingVocabulary": { + "additionalProperties": false, + "description": "Source of the class-id -> string mapping used to render a transcript.", + "properties": { + "ignored_tokens": { + "description": "Token strings dropped before rendering (e.g. ``, ``).\n\nRemoval happens after CTC collapsing and before word splitting, so an\nignored token never joins or separates words.\n\nEvery entry must be present in `tokens` when `source` is `inline`.", + "items": { + "type": "string" + }, + "type": "array" + }, + "size": { + "description": "Number of classes in the decoding vocabulary.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "source": { + "$ref": "#/$defs/DecodingVocabularySource", + "description": "Generic source selector (e.g. `tokenizer`, `inline`)." + }, + "tokens": { + "description": "Inline class-id -> string table, ordered by class id.", + "items": { + "type": "string" + }, + "type": "array" + }, + "word_delimiter": { + "description": "Token string that separates words when rendering (e.g. `|`).\n\nA delimiter *separates* words, so it never contributes whitespace of its\nown: a reader splits the decoded token run on this token, discards empty\ngroups, and joins the remaining groups with a single U+0020. Leading,\ntrailing, and repeated delimiters therefore produce no empty words and no\nleading or trailing space. When absent, tokens are concatenated verbatim.\n\nMust be present in `tokens` when `source` is `inline`.", + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "source" + ], + "type": "object" + }, + "DecodingVocabularySource": { + "description": "Class-id -> string mapping source vocabulary.", + "oneOf": [ + { + "enum": [ + "tokenizer", + "inline" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "tokenizer", + "inline" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "EffectContract": { + "additionalProperties": false, + "description": "Declared semantics of one external effect domain.\n\nRetry class and speculation safety are independent axes. A transactional\neffect may still be unsafe to speculate, and an idempotent effect is not\nautomatically rewindable.", + "properties": { + "retry": { + "$ref": "#/$defs/EffectRetryClass", + "description": "Minimum retry-relevant class for runtime/server recovery orchestration." + }, + "speculation_safety": { + "$ref": "#/$defs/SpeculationSafety", + "default": { + "kind": "none" + }, + "description": "Whether and how far this effect can participate in speculative execution." + } + }, + "required": [ + "retry" + ], + "type": "object" + }, + "EffectRetryClass": { + "description": "Retry-relevant classification of an external effect.", + "oneOf": [ + { + "const": "pure", + "description": "No observable external effect; replay is always safe.", + "type": "string" + }, + { + "const": "idempotent", + "description": "Repeating the effect with the same inputs is observationally equivalent.", + "type": "string" + }, + { + "const": "transactional", + "description": "The effect participates in an external transaction that can be aborted.", + "type": "string" + }, + { + "const": "non_retryable", + "description": "The effect must never be repeated or rolled back.", + "type": "string" + } + ] + }, + "EquivalenceClass": { + "description": "Correctness bound on substituting an equivalent component implementation.", + "oneOf": [ + { + "const": "bitwise", + "description": "Every substituted implementation must produce bit-identical outputs.", + "type": "string" + }, + { + "const": "distribution_preserving", + "description": "Outputs may differ numerically but must preserve the output distribution.", + "type": "string" + }, + { + "const": "semantic", + "description": "Only the declared semantics are preserved; the distribution may change.", + "type": "string" + } + ] + }, + "ExpertShardFacts": { + "additionalProperties": false, + "description": "Expert identity and routing facts for expert parallelism.", + "properties": { + "contiguous_groups_only": { + "default": false, + "description": "Whether experts may be split across ranks in arbitrary contiguous groups.", + "type": "boolean" + }, + "expert_count": { + "description": "Total number of experts per sparse layer.", + "format": "uint", + "minimum": 1, + "type": "integer" + }, + "replicated": { + "description": "Values replicated on every expert-parallel rank, such as the router.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": true + } + }, + "required": [ + "expert_count" + ], + "type": "object" + }, + "GenerationContract": { + "additionalProperties": false, + "description": "Authoritative generation defaults and the structural override surface.", + "properties": { + "defaults": { + "anyOf": [ + { + "$ref": "#/$defs/GenerationDefaults" + }, + { + "type": "null" + } + ], + "description": "Authoritative package defaults." + }, + "overrides": { + "additionalProperties": { + "$ref": "#/$defs/GenerationOverride" + }, + "description": "Overridable fields, each bound to a request-sourced workflow input.\n\nA caller may override exactly these fields and no others. An override of\nany unlisted field must fail loudly rather than being silently dropped.", + "type": "object" + } + }, + "type": "object" + }, + "GenerationDefaults": { + "description": "Author-declared text-generation defaults (sampling and beam search).\n\nMirrors the `search` section of an onnxruntime-genai `genai_config.json`.\nEvery field is optional so only values the author declared are carried over.", + "properties": { + "diversity_penalty": { + "description": "Diversity penalty for diverse beam groups.", + "format": "float", + "type": [ + "number", + "null" + ] + }, + "do_sample": { + "description": "Whether to randomize sampling through `top_k`/`top_p` (else greedy).", + "type": [ + "boolean", + "null" + ] + }, + "early_stopping": { + "description": "Whether beam search stops once enough beams have finished.", + "type": [ + "boolean", + "null" + ] + }, + "length_penalty": { + "description": "Exponential length penalty used with beam search.", + "format": "float", + "type": [ + "number", + "null" + ] + }, + "max_length": { + "description": "Maximum final sequence length.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "min_length": { + "description": "Minimum final sequence length.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "no_repeat_ngram_size": { + "description": "Disallow repeating n-grams of this size (`0` = disabled).", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "num_beams": { + "description": "Number of beams for beam search (`1` = no beam search).", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "num_return_sequences": { + "description": "Number of sequences returned after search.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "repetition_penalty": { + "description": "Penalty applied to already-generated tokens (`1.0` = no penalty).", + "format": "float", + "type": [ + "number", + "null" + ] + }, + "temperature": { + "description": "Softmax temperature applied before sampling.", + "format": "float", + "type": [ + "number", + "null" + ] + }, + "top_k": { + "description": "Number of highest-probability tokens kept for top-k filtering.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "top_p": { + "description": "Nucleus (top-p) cumulative-probability threshold.", + "format": "float", + "type": [ + "number", + "null" + ] + } + }, + "type": "object" + }, + "GenerationOverride": { + "additionalProperties": false, + "description": "One caller-overridable generation field.", + "properties": { + "constraint": { + "anyOf": [ + { + "$ref": "#/$defs/GenerationOverrideConstraint" + }, + { + "type": "null" + } + ], + "description": "Declared bounds the runtime enforces before executing the request." + }, + "input": { + "description": "Request-sourced workflow input that carries the override value.", + "minLength": 1, + "type": "string" + } + }, + "required": [ + "input" + ], + "type": "object" + }, + "GenerationOverrideConstraint": { + "additionalProperties": false, + "description": "Declared bounds on one overridable generation field.", + "properties": { + "maximum": { + "format": "double", + "type": [ + "number", + "null" + ] + }, + "minimum": { + "format": "double", + "type": [ + "number", + "null" + ] + } + }, + "type": "object" + }, + "HardwareRequirements": { + "description": "Model-side hardware requirements and distribution-matching hints.", + "properties": { + "beneficial_dtypes": { + "description": "Dtypes that improve performance or memory use but are not mandatory.", + "items": { + "$ref": "#/$defs/DType" + }, + "type": [ + "array", + "null" + ] + }, + "kv_cache_memory_per_1k_tokens_mb": { + "description": "Estimated KV-cache memory in MiB per 1,000 cached tokens.", + "format": "float", + "minimum": 0.0, + "type": [ + "number", + "null" + ] + }, + "min_memory_gb": { + "description": "Minimum aggregate accelerator or system memory in GiB.", + "format": "float", + "minimum": 0.0, + "type": [ + "number", + "null" + ] + }, + "min_tp_degree": { + "description": "Minimum useful tensor-parallel degree when tensor parallelism is selected.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "required_dtypes": { + "description": "Dtypes the selected device or execution provider must support.", + "items": { + "$ref": "#/$defs/DType" + }, + "type": [ + "array", + "null" + ] + }, + "supports_tensor_parallel": { + "description": "Whether the model can be partitioned with tensor parallelism.", + "type": [ + "boolean", + "null" + ] + } + }, + "type": "object" + }, + "ImageOutputBinding": { + "description": "One named tensor output produced by an image preprocessing program.\n\nThe output binds a processor-local value to a typed workflow SSA name.\nNeither the name nor the content role is inferred from a model identity.", + "properties": { + "content": { + "$ref": "#/$defs/ImageOutputContent", + "description": "Generic content role this tensor carries (pixels, coordinates, grid,\noriginal size, or validity mask) — never a model-family label." + }, + "contract": { + "anyOf": [ + { + "$ref": "#/$defs/TensorContract" + }, + { + "type": "null" + } + ], + "description": "Full workflow tensor contract. Required when `pipeline.workflow` is present." + }, + "dtype": { + "$ref": "#/$defs/TensorDType", + "description": "Declared output dtype. Always explicit; never inferred from the model." + }, + "name": { + "description": "Workflow SSA value produced by the preprocessing adapter invocation.", + "examples": [ + "image.pixel_values" + ], + "minLength": 1, + "type": "string" + }, + "optional": { + "description": "Whether the runtime may omit this output when a model does not need it.", + "type": [ + "boolean", + "null" + ] + }, + "pad_value": { + "description": "Optional sentinel/pad value for padded entries (e.g. `-1` coordinates).", + "format": "double", + "type": [ + "number", + "null" + ] + }, + "source": { + "description": "Named processor-local value produced by a transform.", + "minLength": 1, + "type": "string" + } + }, + "required": [ + "source", + "name", + "content", + "dtype" + ], + "type": "object" + }, + "ImageOutputContent": { + "description": "Generic image-output content-role vocabulary.", + "oneOf": [ + { + "enum": [ + "pixels", + "patch_coordinates", + "grid_dimensions", + "original_size", + "transformed_size", + "validity_mask" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "pixels", + "patch_coordinates", + "grid_dimensions", + "original_size", + "transformed_size", + "validity_mask" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "ImageOutputValueRange": { + "description": "Numeric interpretation of pixels emitted by an image workflow output.\n\nThis is an output contract, not a model-family hint: consumers must never\ninfer normalization from observed pixel values.", + "enum": [ + "zero_to_one", + "negative_one_to_one", + "zero_to_255" + ], + "type": "string" + }, + "ImagePreprocessingProgram": { + "description": "Generic image preprocessing program: an ordered transform pipeline plus the\nnamed workflow SSA tensor outputs it emits.\n\nThe program is expressed entirely as parameterized, architecture-neutral\ndata. Transform operations are generic (decode, resize, rescale, normalize,\ntile, patchify, pad). In workflow metadata, outputs are materialized by a\nmanifest-pinned preprocessing adapter invocation and bind processor-local\nvalues to typed SSA names. A package may name an output `pixel_position_ids`,\n`image_grid_thw`, or anything else without introducing runtime model-family\ndispatch.", + "properties": { + "outputs": { + "description": "Named tensor outputs the program emits, each bound to a workflow SSA value.", + "items": { + "$ref": "#/$defs/ImageOutputBinding" + }, + "minItems": 1, + "type": "array" + }, + "transforms": { + "description": "Ordered list of generic transform operations applied to decoded pixels.", + "items": { + "$ref": "#/$defs/ImageTransform" + }, + "type": "array" + } + }, + "required": [ + "outputs" + ], + "type": "object" + }, + "ImageSizeSpec": { + "anyOf": [ + { + "description": "A single edge length applied to both dimensions.", + "format": "uint32", + "minimum": 0, + "type": "integer" + }, + { + "description": "Explicit width and height.", + "properties": { + "height": { + "description": "Target height in pixels.", + "format": "uint32", + "minimum": 0, + "type": "integer" + }, + "width": { + "description": "Target width in pixels.", + "format": "uint32", + "minimum": 0, + "type": "integer" + } + }, + "required": [ + "width", + "height" + ], + "type": "object" + } + ], + "description": "A square size or an explicit width/height for an image transform." + }, + "ImageTransform": { + "description": "One generic image transform operation.\n\n`op` selects the operation from a generic vocabulary; the remaining fields\nare the parameters that operation reads (only the relevant ones are set).\nEvery parameter is model DATA — concrete sizes, patch sizes, means, and so on\nlive in a model's fixture, never as constants baked into this schema.", + "properties": { + "canvas_pad_value": { + "description": "RGB canvas fill value applied before dynamic tiling.", + "format": "double", + "type": [ + "number", + "null" + ] + }, + "channel_order": { + "description": "Flattened patch feature order (`channels_first` or `channels_last`).", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "coordinate_order": { + "description": "Patch-coordinate component order (`yx` or `xy`).", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "flatten": { + "description": "Whether `patchify` flattens each patch into a single feature vector.", + "type": [ + "boolean", + "null" + ] + }, + "include_thumbnail": { + "description": "Whether a `tile` operation also emits a global thumbnail tile.", + "type": [ + "boolean", + "null" + ] + }, + "inputs": { + "description": "Named values consumed by this transform.\n\nAbsent means the operation consumes the immediately preceding value.\nExplicit names allow branching programs without tensor-name heuristics.", + "items": { + "minLength": 1, + "type": "string" + }, + "minItems": 1, + "type": [ + "array", + "null" + ] + }, + "interpolation": { + "description": "Interpolation filter for a `resize` operation — generic string data.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "mask_patch_size": { + "description": "Pixel edge represented by one validity-mask cell.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "max_patches": { + "description": "Maximum number of spatial patches for a patch-budget resize.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "max_pixels": { + "description": "Maximum pixel area for an aspect-preserving `pixel_area` resize.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "max_tiles": { + "description": "Maximum number of local tiles for a `tile` operation.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "mean": { + "description": "Per-channel mean for a `normalize` operation (length is model data).", + "items": { + "format": "float", + "type": "number" + }, + "type": [ + "array", + "null" + ] + }, + "merge_size": { + "description": "Spatial patch-group edge controlling packed patch traversal order.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "min_pixels": { + "description": "Minimum pixel area for an aspect-preserving `pixel_area` resize.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "mode": { + "description": "Resize/crop mode (e.g. `pad`, `crop`, `stretch`) — generic string data.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "op": { + "$ref": "#/$defs/ImageTransformOp", + "description": "Generic operation selector (e.g. `resize`, `normalize`, `patchify`)." + }, + "outputs": { + "description": "Named values produced by this transform.\n\nThese names are processor-local data. Final graph bindings select them\nthrough `ImageOutputBinding::source`.", + "items": { + "minLength": 1, + "type": "string" + }, + "minItems": 1, + "type": [ + "array", + "null" + ] + }, + "pad_value": { + "description": "Fill value for a `pad` operation, or sentinel for padded coordinates.", + "format": "double", + "type": [ + "number", + "null" + ] + }, + "patch_order": { + "description": "Order patches are emitted in (`merge_groups`, the default, or `raster`).\nIndependent of `merge_size`, which only sets how many patches collapse\ninto one image token.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "patch_size": { + "description": "Edge length of a square patch for a `patchify` operation.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "pooling_kernel_size": { + "description": "Spatial pooling edge used when resolving a patch-budget resize.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "scale": { + "description": "Scalar multiplier for a `rescale` operation.", + "format": "double", + "type": [ + "number", + "null" + ] + }, + "size": { + "anyOf": [ + { + "$ref": "#/$defs/ImageSizeSpec" + }, + { + "type": "null" + } + ], + "description": "Target size for a `resize` operation." + }, + "size_multiple": { + "description": "Required divisibility of both resized dimensions.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "std": { + "description": "Per-channel standard deviation for a `normalize` operation.", + "items": { + "format": "float", + "type": "number" + }, + "type": [ + "array", + "null" + ] + }, + "target_length": { + "description": "Exact first-axis length produced by a `pad` operation.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "temporal_order": { + "description": "Relative nesting of the temporal and channel axes inside a flattened\n`channels_first` patch (`channel_major` or `temporal_major`).", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "temporal_patch_size": { + "description": "Number of identical temporal frames packed into each spatial patch.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "thumbnail_interpolation": { + "description": "Interpolation filter used specifically for a global thumbnail.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "thumbnail_order": { + "anyOf": [ + { + "$ref": "#/$defs/ThumbnailOrder" + }, + { + "type": "null" + } + ], + "description": "Ordering of a global thumbnail relative to local tiles." + }, + "tile_size": { + "description": "Edge length of a square tile for a `tile` operation.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + } + }, + "required": [ + "op" + ], + "type": "object" + }, + "ImageTransformOp": { + "description": "Generic image transform-operation vocabulary.", + "oneOf": [ + { + "enum": [ + "decode", + "decode_rgb", + "convert_rgb", + "resize", + "rescale", + "normalize", + "tile", + "flatten", + "patchify", + "pad", + "emit_original_size", + "emit_transformed_size", + "emit_validity_mask", + "emit_patch_coordinates", + "emit_grid_coordinates" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "decode", + "decode_rgb", + "convert_rgb", + "resize", + "rescale", + "normalize", + "tile", + "flatten", + "patchify", + "pad", + "emit_original_size", + "emit_transformed_size", + "emit_validity_mask", + "emit_patch_coordinates", + "emit_grid_coordinates" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "KeySequenceLengthsSpec": { + "description": "Explicit compatibility rules for attention key-sequence-length metadata.", + "properties": { + "scalar_broadcast": { + "anyOf": [ + { + "$ref": "#/$defs/SequenceLengthScalarBroadcast" + }, + { + "type": "null" + } + ], + "description": "Optional scalar compatibility. `unit_batch` authorizes a contiguous\nrank-0 one-element `int32` tensor only when the attention batch is one." + } + }, + "type": "object" + }, + "KvAxisStrides": { + "description": "Symbolic element stride of each of the four logical KV axes.\n\nThe stride of an axis is the product of the runtime dimensions in its factor\nlist; an **empty** list means unit stride (the innermost, contiguous axis).\nThe innermost axis of every layout the converted kernels honor is\n`head_dim`, whose stride is `1` (empty), because the fp16 read vectorizes\n`head_dim` as `half2` and the fused write addresses it as `dst + d`.\n\nThe two historical layouts map onto this as:\n\n| axis | head-major BNSH | seq-major BSNH |\n|----------|------------------------|-----------------------|\n| batch | `kv_heads·seq·head_dim`| `seq·kv_heads·head_dim`|\n| head | `seq·head_dim` | `head_dim` |\n| seq | `head_dim` | `kv_heads·head_dim` |\n| head_dim | `1` | `1` |", + "properties": { + "batch": { + "default": [], + "description": "Factors of the batch-axis stride.", + "items": { + "$ref": "#/$defs/KvStrideDim" + }, + "type": "array" + }, + "head": { + "default": [], + "description": "Factors of the KV-head-axis stride.", + "items": { + "$ref": "#/$defs/KvStrideDim" + }, + "type": "array" + }, + "head_dim": { + "default": [], + "description": "Factors of the head-dim-axis stride. Unit (empty) for every honored\nlayout.", + "items": { + "$ref": "#/$defs/KvStrideDim" + }, + "type": "array" + }, + "seq": { + "default": [], + "description": "Factors of the sequence-axis (per-token) stride.", + "items": { + "$ref": "#/$defs/KvStrideDim" + }, + "type": "array" + } + }, + "type": "object" + }, + "KvCacheLayout": { + "anyOf": [ + { + "$ref": "#/$defs/KvNamedLayout", + "description": "A readable shorthand for a standard layout. Deserializes from the strings\n`\"head_major_bnsh\"` and `\"seq_major_bsnh\"`; expands to explicit strides\nvia [`KvCacheLayout::resolve_strides`]." + }, + { + "$ref": "#/$defs/KvStrideDescriptor", + "description": "A fully explicit stride descriptor for layouts the named forms cannot\nexpress." + } + ], + "description": "Physical memory layout of a backend's KV cache tensors, as a stride\ndescriptor.\n\nThis is a **per-backend capability**, not a cross-backend constant: the two\nbackends own their KV buffers independently and never read each other's KV\nbytes, so they may store the cache differently. The ONNX Runtime backend\nrequires head-major BNSH (`[batch, kv_heads, seq, head_dim]`) because ORT's\nGroupQueryAttention past/present is BNSH on every dispatch path (Flash,\ncuDNN SDPA, memory-efficient, XQA). The native backend additionally supports\nseq-major BSNH (`[batch, seq, kv_heads, head_dim]`), which makes each token's\nlive prefix contiguous across heads — shrinking the VMM granule floor by the\n`kv_heads` factor, removing growth-triggered graph re-capture (the append\nstride is sequence-length independent), and making page-level prefix sharing\n(#777) practical. Absent preserves the historical head-major behavior.\n\nLayout preference is per-EP and per-platform rather than a global constant,\nand a JIT backend compiles a specialized kernel per descriptor, so this is a\ndescriptor rather than a closed enum: a raw stride tuple is unreadable, so\nthe common cases are still nameable (`head_major_bnsh`, `seq_major_bsnh`)\nwhile an explicit [`KvStrideDescriptor`] expresses anything the named forms\ncannot (e.g. a token-major view).\n\nOn-device, the native backend selects the layout by stamping the `kv_layout`\nattribute (`0` = BNSH, `1` = BSNH) on its GroupQueryAttention nodes; the\nCUDA EP honors it on the fused fp16 single-token decode pair. Seq-major is\nonly enabled end-to-end once the prefill (flash) read is also converted, so\nthe two never disagree about how a shared cache is physically laid out." + }, + "KvNamedLayout": { + "description": "The named, human-readable KV cache layouts.", + "oneOf": [ + { + "const": "head_major_bnsh", + "description": "Head-major BNSH `[batch, kv_heads, seq, head_dim]`. ORT-compatible; the\ndefault for both backends.", + "type": "string" + }, + { + "const": "seq_major_bsnh", + "description": "Seq-major BSNH `[batch, seq, kv_heads, head_dim]`. Native backend only.", + "type": "string" + } + ] + }, + "KvOwnership": { + "description": "Ownership model for a graph's KV cache inputs.", + "oneOf": [ + { + "const": "owned", + "description": "The graph consumes past KV and emits replacement/extended present KV.", + "type": "string" + }, + { + "const": "shared", + "description": "The graph reads references to KV owned and advanced by another decoder.", + "type": "string" + } + ] + }, + "KvStrideDescriptor": { + "description": "A fully explicit KV-cache stride descriptor.\n\nThis is the general form the two named layouts expand into, and the shape a\nfuture layout (e.g. token-major) is expressed in without adding an enum\nvariant. The `reservation_*` fields describe a binding that is a **view into\na larger reservation** rather than the owner of its whole buffer:\ntoken-major stores every layer's tokens in one reservation and hands each\n`(layer, side)` a sub-view, so its per-token (seq) stride is taken over the\nreservation's total token count and its data starts at a non-zero offset.\nBoth historical layouts are whole-buffer bindings: `offset == 0` and no\nreservation override.", + "properties": { + "reservation_offset_elements": { + "description": "Element offset of this binding's first element within the reservation it\nviews. `0` for a binding that owns its whole buffer — the only case the\nconverted kernels honor today.", + "format": "uint64", + "minimum": 0, + "type": "integer" + }, + "reservation_seq_slots": { + "description": "Sequence-axis extent, in token slots, of the reservation this binding\nviews when the reservation is larger than the binding's own\n`cache_capacity`. Absent means the binding spans its own capacity (a\nwhole-buffer binding). Present expresses a token-major view whose seq\nstride collapses the per-`(layer, side)` buffer boundary; not honored by\nthe converted path yet.", + "format": "uint64", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "strides": { + "$ref": "#/$defs/KvAxisStrides", + "description": "Symbolic stride of each logical axis." + } + }, + "required": [ + "strides" + ], + "type": "object" + }, + "KvStrideDim": { + "description": "A runtime KV-cache dimension that an axis stride can be a multiple of.\n\nAbsolute element strides are a serving-time property — they depend on the\n`cache_capacity` a runtime picks — so metadata cannot store them as numbers.\nA stride is therefore stored **symbolically**, as the (unordered) set of\nruntime dimensions it multiplies. The concrete element stride of an axis is\nthe product of the sizes of the dimensions in its factor list.", + "oneOf": [ + { + "const": "kv_heads", + "description": "Number of KV heads (`kv_heads` / `N`).", + "type": "string" + }, + { + "const": "seq_capacity", + "description": "Sequence capacity of the growing axis (`cache_capacity` / `S`).", + "type": "string" + }, + { + "const": "head_dim", + "description": "Per-token head width (`head_dim` / `H`).", + "type": "string" + } + ] + }, + "LiteralValue": { + "anyOf": [ + { + "$ref": "#/$defs/ScalarValue" + }, + { + "items": { + "$ref": "#/$defs/ScalarValue" + }, + "type": "array" + } + ], + "description": "A literal tensor initializer.\n\nA single scalar broadcasts to every element of the declared contract, which\ncovers flags, counters, and zero-filled buffers. Workflows whose constants\nare genuinely per-position -- interleaved stream delay patterns, per-stream\ninitial tokens, fixed schedule tables -- declare the elements explicitly in\nrow-major order instead, so the value stays inside the metadata document and\ndoes not become an out-of-band artifact." + }, + "LoopStatePair": { + "description": "One fixed-shape loop-carried recurrent-state port pair.\n\nGeneric and architecture-neutral: the runtime zero/other-initializes `input`\non the first step, runs the graph, and copies `output` back into `input` for\nthe next step (`replace` update). This models any fixed recurrent tensor\n(convolution state, linear-attention recurrent state, and so on) without\nreferencing a model family. It is intentionally distinct from growing KV\nstate, whose logical cells and storage service are declared by a workflow.", + "properties": { + "init": { + "$ref": "#/$defs/StateInitKind", + "description": "How `input` is initialized before the first step (e.g. `zeros`)." + }, + "input": { + "description": "Graph input port that receives the carried state for this step.", + "minLength": 1, + "type": "string" + }, + "output": { + "description": "Graph output port that produces the next-step state.", + "minLength": 1, + "type": "string" + }, + "update": { + "$ref": "#/$defs/StateUpdateKind", + "description": "How `output` becomes the next step's `input` (fixed state uses `replace`)." + } + }, + "required": [ + "input", + "output", + "init", + "update" + ], + "type": "object" + }, + "LoraGraphInputBinding": { + "additionalProperties": false, + "properties": { + "a": { + "type": "string" + }, + "b": { + "type": "string" + }, + "scale": { + "description": "Optional graph input for request scale; otherwise scale is folded into stable factors.", + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "a", + "b" + ], + "type": "object" + }, + "LoraTargetDescriptor": { + "additionalProperties": false, + "description": "One resolved base projection. Fused-QKV knowledge is lowered to an optional slice.", + "properties": { + "activation_dtype": { + "$ref": "#/$defs/TensorDType", + "description": "Projection activation dtype used by graph-native delta application." + }, + "alpha": { + "description": "Optional alpha policy for artifacts binding this target.", + "format": "double", + "type": [ + "number", + "null" + ] + }, + "component": { + "description": "Workflow component name, or `model` for a bare decoder package.", + "type": "string" + }, + "graph_inputs": { + "anyOf": [ + { + "$ref": "#/$defs/LoraGraphInputBinding" + }, + { + "type": "null" + } + ], + "description": "Phase-1 graph-native optional A/B inputs. Base-only omits both and is bit-identical." + }, + "id": { + "type": "string" + }, + "initializer": { + "description": "Exact immutable base initializer name.", + "type": "string" + }, + "input_features": { + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "layer_index": { + "description": "Optional producer/importer layer identity retained from Phase-2 manifests.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "node_name": { + "description": "Exact ONNX projection node name used for load-time manifest validation.", + "type": "string" + }, + "output_features": { + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "output_name": { + "description": "Exact graph value produced by the projection.", + "type": "string" + }, + "output_slice": { + "anyOf": [ + { + "$ref": "#/$defs/LoraTargetSlice" + }, + { + "type": "null" + } + ], + "description": "Resolved child range within a fused output; producer/import tooling owns discovery." + }, + "rank": { + "description": "Optional rank policy for artifacts binding this target.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + } + }, + "required": [ + "id", + "component", + "initializer", + "node_name", + "output_name", + "activation_dtype", + "input_features", + "output_features" + ], + "type": "object" + }, + "LoraTargetManifest": { + "additionalProperties": false, + "description": "Authoritative generic target map migrated from Phase-2 `LoraTargetManifest`.", + "properties": { + "targets": { + "items": { + "$ref": "#/$defs/LoraTargetDescriptor" + }, + "type": "array" + } + }, + "required": [ + "targets" + ], + "type": "object" + }, + "LoraTargetSlice": { + "additionalProperties": false, + "properties": { + "alpha": { + "format": "double", + "type": [ + "number", + "null" + ] + }, + "offset": { + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "rank": { + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "role": { + "description": "Producer-defined semantic label; runtime execution uses only the resolved range.", + "type": "string" + }, + "width": { + "format": "uint", + "minimum": 0, + "type": "integer" + } + }, + "required": [ + "role", + "offset", + "width" + ], + "type": "object" + }, + "MediaContainer": { + "enum": [ + "raw", + "wav" + ], + "type": "string" + }, + "MediaDelivery": { + "enum": [ + "buffered", + "streaming" + ], + "type": "string" + }, + "MediaEncoding": { + "enum": [ + "pcm_s16_le", + "pcm_f32_le" + ], + "type": "string" + }, + "MediaOutputContract": { + "additionalProperties": false, + "properties": { + "channels": { + "format": "uint16", + "maximum": 65535, + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "container": { + "$ref": "#/$defs/MediaContainer" + }, + "delivery": { + "$ref": "#/$defs/MediaDelivery", + "default": "buffered" + }, + "encoding": { + "$ref": "#/$defs/MediaEncoding" + }, + "sample_rate_hz": { + "description": "Sample rate of the encoded response.", + "format": "uint32", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "source_sample_rate_hz": { + "description": "Sample rate of a pre-adapter waveform. When it differs from\n`sample_rate_hz`, the API boundary resamples before encoding.", + "format": "uint32", + "minimum": 0, + "type": [ + "integer", + "null" + ] + } + }, + "required": [ + "container", + "encoding" + ], + "type": "object" + }, + "MixtureOfExpertsSpec": { + "description": "Explicit sparse mixture-of-experts structure and graph representation.", + "properties": { + "activation": { + "description": "Expert FFN activation name, such as `silu`.", + "minLength": 1, + "type": "string" + }, + "expert_intermediate_size": { + "description": "Intermediate width of each routed expert FFN.", + "format": "uint", + "minimum": 1, + "type": "integer" + }, + "experts_per_token": { + "description": "Number of routed experts selected for each token.", + "format": "uint", + "minimum": 1, + "type": "integer" + }, + "representation": { + "$ref": "#/$defs/MoERepresentation", + "description": "Expert graph representation: `dense_fallback`, `moe`, or `qmoe`." + }, + "routed_expert_count": { + "description": "Number of independently routed experts.", + "format": "uint", + "minimum": 1, + "type": "integer" + }, + "router": { + "$ref": "#/$defs/MoERouterSpec", + "description": "Router scoring, selection, normalization, and scaling semantics." + }, + "shared_expert_count": { + "description": "Number of dense shared experts evaluated for every token.", + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "shared_expert_intermediate_size": { + "description": "Total intermediate width of the always-on shared-expert FFN.", + "format": "uint", + "minimum": 0, + "type": "integer" + } + }, + "required": [ + "representation", + "routed_expert_count", + "shared_expert_count", + "experts_per_token", + "expert_intermediate_size", + "shared_expert_intermediate_size", + "activation", + "router" + ], + "type": "object" + }, + "MoEGroupScore": { + "description": "Group-scoring reduction vocabulary.", + "oneOf": [ + { + "enum": [ + "maximum", + "top_2_sum" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "maximum", + "top_2_sum" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "MoERepresentation": { + "description": "Sparse expert graph representation vocabulary.", + "oneOf": [ + { + "enum": [ + "dense_fallback", + "moe", + "qmoe" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "dense_fallback", + "moe", + "qmoe" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "MoERouterScoreFunction": { + "description": "Router score-operation vocabulary.", + "oneOf": [ + { + "enum": [ + "softmax", + "sigmoid" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "softmax", + "sigmoid" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "MoERouterSelectionMethod": { + "description": "Router expert-selection vocabulary.", + "oneOf": [ + { + "enum": [ + "top_k", + "grouped_top_k", + "sparse_mixer" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "top_k", + "grouped_top_k", + "sparse_mixer" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "MoERouterSpec": { + "allOf": [ + { + "if": { + "properties": { + "selection_method": { + "const": "grouped_top_k" + } + }, + "required": [ + "selection_method" + ] + }, + "then": { + "required": [ + "group_count", + "groups_per_token", + "group_score" + ] + } + } + ], + "description": "Explicit router semantics, kept separate from expert FFN execution.", + "properties": { + "group_count": { + "description": "Number of expert groups considered by grouped selection.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "group_score": { + "anyOf": [ + { + "$ref": "#/$defs/MoEGroupScore" + }, + { + "type": "null" + } + ], + "description": "Reduction used to score a group before group TopK." + }, + "groups_per_token": { + "description": "Number of groups retained per token by grouped selection.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "normalize_weights": { + "description": "Whether selected aggregation weights are normalized to sum to one.", + "type": "boolean" + }, + "scaling_factor": { + "description": "Multiplicative scale applied to final aggregation weights.", + "format": "float", + "minimum": 0.0, + "type": "number" + }, + "score_function": { + "$ref": "#/$defs/MoERouterScoreFunction", + "description": "Elementwise score operation applied to router logits." + }, + "selection_method": { + "$ref": "#/$defs/MoERouterSelectionMethod", + "description": "Expert selection operation applied to the scores." + } + }, + "required": [ + "score_function", + "selection_method", + "normalize_weights", + "scaling_factor" + ], + "type": "object" + }, + "ModelCapabilities": { + "description": "Model properties that are baked into the graph or advertised as configurable.", + "properties": { + "attention": { + "anyOf": [ + { + "$ref": "#/$defs/AttentionConfig" + }, + { + "type": "null" + } + ], + "description": "Attention architecture and dimensions." + }, + "io": { + "anyOf": [ + { + "$ref": "#/$defs/ModelIoSpec" + }, + { + "type": "null" + } + ], + "default": null, + "deprecated": true, + "description": "DEPRECATED, import-only: legacy explicit graph I/O for a bare\nsingle-decoder package.\n\nThe canonical, and only authoritative, expression of a package's\nexecutable graph ABI is the workflow:\n`pipeline.workflow.components..ports` (with `ports.roles`),\nthe invoke bindings that connect them, and the `state_service` groups\nthat declare model state. A composite package and a bare one-file\ndecoder use that same representation.\n\nThis block remains only so packages written before the workflow existed\nstill load. It is never read directly: [`ModelCapabilities::io`] resolves\nthe workflow first and falls back here, and a document carrying both is\nrejected so the two can never disagree. New producers must not emit it." + }, + "max_sequence_length": { + "description": "Maximum total sequence length, in tokens.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "mixture_of_experts": { + "anyOf": [ + { + "$ref": "#/$defs/MixtureOfExpertsSpec" + }, + { + "type": "null" + } + ], + "description": "Explicit sparse mixture-of-experts graph and routing contract.\n\nThis describes graph structure, never a model family. Runtimes use the\ndeclared representation and dimensions instead of inferring them from\nnode names, initializer shapes, or architecture strings." + }, + "runtime_configurable": { + "anyOf": [ + { + "$ref": "#/$defs/RuntimeConfigurable" + }, + { + "type": "null" + } + ], + "description": "Features that a serving runtime may configure at load time." + }, + "sharding": { + "anyOf": [ + { + "$ref": "#/$defs/ShardingContract" + }, + { + "type": "null" + } + ], + "description": "Legal tensor, pipeline, and expert sharding facts.\n\nThe caller and runtime choose degree, device mapping, and collective\nbackend. Portable metadata never standardizes a cross-runtime KV or\ncache wire format." + }, + "vocab_size": { + "description": "Vocabulary size (rows of the token-embedding / logits table).", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + } + }, + "type": "object" + }, + "ModelIoSpec": { + "additionalProperties": false, + "description": "Explicit binding of the graph ports the decode step reads and writes.\n\nEvery field is optional so a model package can declare only the ports its\ngraph exposes. A port left unset is resolved ONLY from an unambiguous\ndtype/shape signal; when the shape cannot disambiguate the port, the runtime\nfails with an actionable error naming the key to declare rather than\ninterpreting a tensor name. A declared port is always authoritative.", + "properties": { + "aliasing": { + "anyOf": [ + { + "$ref": "#/$defs/StateAliasing" + }, + { + "type": "null" + } + ], + "description": "Whether the graph permits, requires, or forbids the runtime aliasing a\n`present` output onto its paired `past` input.\n\nThis is the graph ABI fact that replaced the old `shared_buffer` policy\nflag: the package states what aliasing its graph is CORRECT under, and\nthe runtime alone decides whether to exploit it (execution provider\ncapability, buffer capacity, and batching are runtime concerns). A graph\nthat reads a `past` region after the paired `present` write would touch\nit must declare `forbidden`, which is the default when the package is\nsilent — silence never grants an optimization." + }, + "attention_mask_input": { + "description": "Attention-mask input, if the graph takes one.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "audio_features_input": { + "description": "Raw audio-feature prompt input for an encoder-decoder encoder graph\n(e.g. Whisper `audio_features`, a log-mel `[batch, mels, frames]`\ntensor). Declared on the encoder component; a text encoder-decoder uses\n`token_input` instead.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "cross_kv_inputs": { + "description": "Cross-attention past-KV cache inputs for an encoder-decoder decoder, in\nthe SAME order as `cross_kv_outputs`. These are the encoder-derived KV\ntensors, distinct from the self-attention `kv_inputs`.", + "items": { + "minLength": 1, + "type": "string" + }, + "type": [ + "array", + "null" + ] + }, + "cross_kv_outputs": { + "description": "Cross-attention present-KV cache outputs (produced by the encoder for an\nencoder-decoder model), paired positionally with `cross_kv_inputs`.", + "items": { + "minLength": 1, + "type": "string" + }, + "type": [ + "array", + "null" + ] + }, + "encoder_hidden_states_input": { + "description": "Encoder-hidden-states input for an encoder-decoder (cross-attention)\ndecoder graph (e.g. `encoder_hidden_states`).", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "hidden_output": { + "description": "Per-token hidden-state output for embedding / VLM hidden extraction, if\nthe graph exposes a distinct hidden output (e.g. `last_hidden_state`).", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "inputs_embeds_input": { + "description": "Pre-embedded / routed sequence input (e.g. `inputs_embeds`).\n\nMay be declared alongside `token_input` (see its documentation): a graph\nthat consumes both a raw token input and one or more routed sequence\ninputs is explicitly permitted.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "kv_inputs": { + "description": "Past-KV cache inputs, in the SAME order as `kv_outputs` (positional\npairing). Length must match `kv_outputs`.", + "items": { + "minLength": 1, + "type": "string" + }, + "type": [ + "array", + "null" + ] + }, + "kv_layout": { + "anyOf": [ + { + "$ref": "#/$defs/KvCacheLayout" + }, + { + "type": "null" + } + ], + "description": "Physical layout of this backend's KV cache tensors, as a stride\ndescriptor. Accepts a readable named layout (`head_major_bnsh` or\n`seq_major_bsnh`) or a fully explicit [`KvStrideDescriptor`]. This is a\nper-backend capability — each backend owns its KV buffers and never reads\nthe other's KV bytes — so the ORT backend stays head-major while the\nnative backend may declare seq-major. Absent preserves the historical\nhead-major (BNSH) behavior. See [`KvCacheLayout`]." + }, + "kv_outputs": { + "description": "Present-KV cache outputs, paired positionally with `kv_inputs`.", + "items": { + "minLength": 1, + "type": "string" + }, + "type": [ + "array", + "null" + ] + }, + "kv_ownership": { + "anyOf": [ + { + "$ref": "#/$defs/KvOwnership" + }, + { + "type": "null" + } + ], + "description": "Whether this graph owns past/present KV state or reads target-owned KV.\n\nAbsent preserves the historical `owned` behavior." + }, + "logits_output": { + "description": "Logits output.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "optional_inputs": { + "additionalProperties": { + "$ref": "#/$defs/OptionalInputSpec" + }, + "description": "Optional graph inputs and their explicit absent-value contracts, keyed by\nthe real ONNX input port name.", + "type": "object" + }, + "position_ids_input": { + "description": "Position-ids input, if the graph takes one.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "sequence_source": { + "anyOf": [ + { + "$ref": "#/$defs/SequenceInputKind" + }, + { + "type": "null" + } + ], + "description": "Which declared sequence port drives autoregressive execution.\n\nAbsent preserves the historical `token_ids` behavior. Declaring\n`inputs_embeds` requires `inputs_embeds_input`; declaring `token_ids`\nrequires `token_input`." + }, + "state_pairs": { + "description": "Fixed-shape loop-carried recurrent state ports, distinct from KV cache.\n\nEach pair binds an input port to its matching output port and declares\nhow the input is initialized and how the output feeds the next step\n(`replace` semantics for fixed recurrent tensors). These are neither KV\ncache nor fixed conditioning; the sparse set of state ports comes from\nthis declared list, never expanded from a layer count.", + "items": { + "$ref": "#/$defs/LoopStatePair" + }, + "minItems": 1, + "type": [ + "array", + "null" + ] + }, + "static_cache": { + "anyOf": [ + { + "$ref": "#/$defs/StaticCacheIoSpec" + }, + { + "type": "null" + } + ], + "description": "Explicit port binding for a fixed-buffer TensorScatter static KV cache.\n\nA static-cache decoder scatters each step's K/V into pre-allocated,\nfixed-length buffers via an integer write-index vector and a non-pad\nsequence-length vector, rather than growing/appending a cache. These\ncontrol ports are integer vectors and are therefore SHAPE-indistinguish-\nable from one another, so shape cannot disambiguate them: the ABI must be\ndeclared explicitly. When present, this spec is authoritative and the\nruntime binds exactly these ports. When absent, a graph that exposes the\nscatter ABI is REJECTED with an actionable error naming this key rather\nthan having its integer control ports guessed by name." + }, + "token_input": { + "description": "Token-id input (e.g. `input_ids`).\n\nA graph MAY declare this together with `inputs_embeds_input`: some fused\ndecoders consume a raw token stream AND a routed pre-embedded sequence in\nthe same forward pass. The two are not mutually exclusive; declaring both\nis a valid, explicit contract.", + "minLength": 1, + "type": [ + "string", + "null" + ] + } + }, + "type": "object" + }, + "OptionalInputSpec": { + "description": "Presence and absent-value contract for one optional graph input.", + "properties": { + "absent": { + "$ref": "#/$defs/AbsentInputSpec", + "description": "Tensor value supplied when the presence key is absent." + }, + "presence": { + "description": "Opaque, non-empty request presence key; not a port or model name.", + "minLength": 1, + "type": "string" + } + }, + "required": [ + "presence", + "absent" + ], + "type": "object" + }, + "OutputStage": { + "enum": [ + "pre_adapter", + "post_adapter" + ], + "type": "string" + }, + "PackageFacts": { + "additionalProperties": false, + "description": "Exact package facts required to interpret request data correctly.", + "properties": { + "constraint_languages": { + "description": "Constraint/grammar dialects this package's parser can interpret.", + "items": { + "$ref": "#/$defs/ConstraintLanguageFacts" + }, + "type": "array" + }, + "tokenizer": { + "anyOf": [ + { + "$ref": "#/$defs/TokenizerFacts" + }, + { + "type": "null" + } + ], + "description": "Exact tokenizer, vocabulary, and special-token facts." + } + }, + "type": "object" + }, + "PipelineSpec": { + "additionalProperties": false, + "description": "Executable package described by the universal typed workflow IR.", + "properties": { + "workflow": { + "$ref": "#/$defs/WorkflowSpec", + "description": "Required component-centric SSA workflow." + } + }, + "required": [ + "workflow" + ], + "type": "object" + }, + "PipelineStageFacts": { + "additionalProperties": false, + "description": "One legal pipeline stage boundary.", + "properties": { + "components": { + "description": "Workflow components executed by this stage, in order.", + "items": { + "type": "string" + }, + "minItems": 1, + "type": "array" + }, + "outputs": { + "description": "Typed values crossing into the next stage.", + "items": { + "type": "string" + }, + "type": "array" + }, + "state_groups": { + "description": "State groups this stage owns.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": true + } + }, + "required": [ + "components" + ], + "type": "object" + }, + "PoolingKind": { + "description": "Pooling reduction kind.", + "enum": [ + "mean", + "max", + "cls", + "last_token" + ], + "type": "string" + }, + "PoolingSpec": { + "additionalProperties": false, + "description": "Pooling reduction applied to a sequence-valued task output.", + "properties": { + "axis": { + "description": "Axis reduced by the pooling operation.", + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "kind": { + "$ref": "#/$defs/PoolingKind" + }, + "normalize": { + "default": false, + "description": "Whether the pooled vector is L2-normalized.", + "type": "boolean" + } + }, + "required": [ + "kind", + "axis" + ], + "type": "object" + }, + "PortRole": { + "description": "Architecture-neutral semantic role of one component port.\n\nThis vocabulary names what a value MEANS to the component that consumes or\nproduces it. It deliberately excludes anything a state group already\ndeclares, so a role and a state binding can never disagree about the same\nport.", + "oneOf": [ + { + "const": "token_ids", + "description": "Discrete token identifiers driving autoregressive execution.", + "type": "string" + }, + { + "const": "inputs_embeds", + "description": "Pre-embedded sequence, used when another component owns the embedding.", + "type": "string" + }, + { + "const": "attention_mask", + "description": "Attention mask over the sequence.", + "type": "string" + }, + { + "const": "position_ids", + "description": "Per-position indices used by position embedding.", + "type": "string" + }, + { + "const": "logits", + "description": "Unnormalized next-token scores.", + "type": "string" + }, + { + "const": "hidden_states", + "description": "Per-token hidden states exposed as a distinct output.", + "type": "string" + }, + { + "const": "encoder_hidden_states", + "description": "Encoder result consumed by a cross-attending decoder.", + "type": "string" + }, + { + "const": "audio_features", + "description": "Encoded audio features consumed by a speech decoder.", + "type": "string" + } + ] + }, + "Precision": { + "description": "Weight precision and quantization-recipe vocabulary.", + "oneOf": [ + { + "enum": [ + "float32", + "fp32", + "float16", + "fp16", + "bfloat16", + "bf16", + "float8_e4m3fn", + "float8_e5m2", + "int8", + "int4", + "int4_group128" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "float32", + "fp32", + "float16", + "fp16", + "bfloat16", + "bf16", + "float8_e4m3fn", + "float8_e5m2", + "int8", + "int4", + "int4_group128" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "PreprocessingSpec": { + "description": "Declared, architecture-neutral input preprocessing programs.", + "properties": { + "audio": { + "anyOf": [ + { + "$ref": "#/$defs/AudioPreprocessingProgram" + }, + { + "type": "null" + } + ], + "description": "Typed audio preprocessing transform program and its named tensor outputs." + }, + "image": { + "anyOf": [ + { + "$ref": "#/$defs/ImagePreprocessingProgram" + }, + { + "type": "null" + } + ], + "description": "Typed image preprocessing transform program and its named tensor outputs." + } + }, + "type": "object" + }, + "ProfileRequirement": { + "description": "Whether a reader may ignore a profile it does not understand.", + "oneOf": [ + { + "const": "required", + "description": "A reader that cannot execute this profile must refuse to load the package.", + "type": "string" + }, + { + "const": "ignorable", + "description": "A reader that does not understand this profile may skip it.", + "type": "string" + } + ] + }, + "QuantizationIntent": { + "description": "Runtime-independent model-weight quantization intent.", + "properties": { + "default": { + "anyOf": [ + { + "$ref": "#/$defs/Precision" + }, + { + "type": "null" + } + ], + "description": "Default precision or quantization recipe for model weights." + }, + "overrides": { + "description": "Layer- or component-specific precision overrides.", + "items": { + "$ref": "#/$defs/QuantizationOverride" + }, + "type": [ + "array", + "null" + ] + } + }, + "type": "object" + }, + "QuantizationOverride": { + "description": "Precision override for selected layers or a named graph component.", + "properties": { + "component": { + "description": "Logical component path, for example `attention.qk` or `lm_head`.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "layers": { + "description": "Layer indices to which the override applies; negative indices count from the end.", + "items": { + "format": "int32", + "type": "integer" + }, + "type": [ + "array", + "null" + ] + }, + "precision": { + "$ref": "#/$defs/Precision", + "description": "Required precision or quantization recipe." + } + }, + "required": [ + "precision" + ], + "type": "object" + }, + "RuntimeConfigurable": { + "additionalProperties": false, + "description": "Features whose concrete settings may be selected by the runtime.", + "properties": { + "chunked_prefill": { + "anyOf": [ + { + "$ref": "#/$defs/ChunkedPrefillConfig" + }, + { + "type": "null" + } + ], + "description": "Chunked-prefill support and preferred chunk size." + }, + "continuous_batching": { + "description": "Whether continuous batching may be enabled.", + "type": [ + "boolean", + "null" + ] + }, + "prefix_cache": { + "description": "Whether prefix caching may be enabled.", + "type": [ + "boolean", + "null" + ] + } + }, + "type": "object" + }, + "RuntimeInputRole": { + "oneOf": [ + { + "enum": [ + "prompt_text", + "prompt_tokens", + "negative_prompt_text", + "negative_prompt_tokens", + "media", + "max_iterations", + "max_output_tokens", + "seed", + "guidance_scale", + "width", + "height", + "denoising_strength", + "sampling_temperature", + "sampling_top_k", + "sampling_top_p", + "sampling_min_p", + "constraint", + "session_id", + "adapter_segments", + "adapter_counts", + "adapter_scales", + "adapter_active" + ], + "type": "string" + }, + { + "const": "row_selection", + "description": "Runtime-minted gather of source batch positions for beam or speculative\nrow expansion. Values are positions inside the current batch, never\nscheduler slots, request IDs, or epochs.", + "type": "string" + } + ] + }, + "ScalarValue": { + "anyOf": [ + { + "type": "boolean" + }, + { + "format": "int64", + "type": "integer" + }, + { + "format": "double", + "type": "number" + }, + { + "type": "string" + } + ] + }, + "SemanticInputRole": { + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "runtime", + "type": "string" + }, + "role": { + "$ref": "#/$defs/RuntimeInputRole" + }, + "version": { + "type": "string" + } + }, + "required": [ + "kind", + "version", + "role" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "opaque", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + } + ] + }, + "SequenceDecodingKind": { + "description": "Frame-synchronous sequence-decoding algorithm vocabulary.", + "oneOf": [ + { + "enum": [ + "ctc", + "greedy_argmax" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "ctc", + "greedy_argmax" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "SequenceDecodingSpec": { + "additionalProperties": false, + "description": "Frame-synchronous decoding contract for non-autoregressive sequence models.\n\nA CTC acoustic model emits one class distribution per encoder frame. The\ntranscript is recovered by taking the per-frame argmax, collapsing runs of\nrepeated ids, and dropping the blank id — no generation loop and no\nautoregressive state are involved.", + "properties": { + "blank_id": { + "description": "Class id reserved for the CTC blank symbol.", + "format": "uint32", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "class_axis": { + "description": "Axis of the logits tensor that enumerates classes.", + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "collapse_repeats": { + "default": false, + "description": "Whether runs of identical consecutive ids collapse to one id.", + "type": "boolean" + }, + "kind": { + "$ref": "#/$defs/SequenceDecodingKind", + "description": "Generic decoding algorithm selector (e.g. `ctc`)." + }, + "lengths": { + "description": "Profile output role naming the per-row count of valid frames.\n\nPresent when the package is batched with padding: rows are decoded only\nover their own valid frame prefix so a padded batch produces the same\ntranscript per row as an unpadded single-row run.", + "type": [ + "string", + "null" + ] + }, + "time_axis": { + "description": "Axis of the logits tensor that enumerates frames.", + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "vocabulary": { + "anyOf": [ + { + "$ref": "#/$defs/DecodingVocabulary" + }, + { + "type": "null" + } + ], + "description": "Where the class-id -> string mapping comes from." + } + }, + "required": [ + "kind", + "time_axis", + "class_axis" + ], + "type": "object" + }, + "SequenceInputKind": { + "description": "Primary autoregressive sequence source for a decoder or proposer graph.", + "oneOf": [ + { + "const": "token_ids", + "description": "Integer token ids supplied through `token_input`.", + "type": "string" + }, + { + "const": "inputs_embeds", + "description": "Precomputed floating-point embeddings supplied through\n`inputs_embeds_input`.", + "type": "string" + } + ] + }, + "SequenceLengthScalarBroadcast": { + "description": "Permitted scalar compatibility for attention key-sequence lengths.", + "oneOf": [ + { + "const": "unit_batch", + "description": "Interpret one rank-0 value as the canonical one-element vector only for\nan attention batch of exactly one.", + "type": "string" + } + ] + }, + "ServingServiceContract": { + "additionalProperties": false, + "properties": { + "accepted_len": { + "type": [ + "string", + "null" + ] + }, + "active": { + "type": "string" + }, + "done": { + "type": "string" + }, + "state_service": { + "$ref": "#/$defs/StateServiceContract", + "description": "Semantic state groups whose graph ABI the runtime must honor." + } + }, + "required": [ + "active", + "done", + "state_service" + ], + "type": "object" + }, + "SessionLeaseContract": { + "additionalProperties": false, + "properties": { + "optimistic_metadata_version": { + "default": false, + "type": "boolean" + }, + "policy": { + "$ref": "#/$defs/SessionMutationPolicy", + "default": "exclusive" + }, + "ttl_seconds": { + "format": "uint64", + "minimum": 0, + "type": [ + "integer", + "null" + ] + } + }, + "type": "object" + }, + "SessionMutationPolicy": { + "enum": [ + "exclusive", + "copy_on_write" + ], + "type": "string" + }, + "ShapeRecurrence": { + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "invariant", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "axis": { + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "increment": { + "type": "string" + }, + "kind": { + "const": "growing", + "type": "string" + }, + "max": { + "type": "string" + } + }, + "required": [ + "kind", + "axis", + "increment", + "max" + ], + "type": "object" + }, + { + "additionalProperties": false, + "description": "The selected axis may grow or shrink between iterations, but never exceed `max`.", + "properties": { + "axis": { + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "kind": { + "const": "bounded", + "type": "string" + }, + "max": { + "type": "string" + } + }, + "required": [ + "kind", + "axis", + "max" + ], + "type": "object" + } + ] + }, + "ShardingContract": { + "additionalProperties": false, + "description": "Legal sharding and replication facts for distributed execution.\n\nThe caller and runtime choose degree, device mapping, placement, and\ncollective backend. Metadata only declares what is legal.", + "properties": { + "expert_parallel": { + "anyOf": [ + { + "$ref": "#/$defs/ExpertShardFacts" + }, + { + "type": "null" + } + ], + "description": "Legal expert-parallel facts for sparse mixture-of-experts layers." + }, + "pipeline_parallel": { + "description": "Legal pipeline-parallel stage boundaries and their cross-stage state.", + "items": { + "$ref": "#/$defs/PipelineStageFacts" + }, + "type": "array" + }, + "tensor_parallel": { + "additionalProperties": { + "$ref": "#/$defs/TensorShardFacts" + }, + "description": "Legal tensor-parallel shard axes by logical parameter group.", + "type": "object" + } + }, + "type": "object" + }, + "SpecialTokenFact": { + "additionalProperties": false, + "description": "One special token, pinned by id and exact surface bytes.", + "properties": { + "content": { + "description": "Exact UTF-8 surface form of the token.", + "minLength": 1, + "type": "string" + }, + "id": { + "description": "Vocabulary id of the token.", + "format": "uint32", + "minimum": 0, + "type": "integer" + } + }, + "required": [ + "id", + "content" + ], + "type": "object" + }, + "SpeculationSafety": { + "description": "Whether an effect may be executed inside a speculative region.", + "oneOf": [ + { + "additionalProperties": false, + "description": "The effect must not run speculatively.", + "properties": { + "kind": { + "const": "none", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + { + "additionalProperties": false, + "description": "The effect's observable state can be cloned before a speculative region.", + "properties": { + "kind": { + "const": "clonable", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + { + "additionalProperties": false, + "description": "The effect can be rewound by at most `max_depth` proposed positions.", + "properties": { + "kind": { + "const": "rewindable", + "type": "string" + }, + "max_depth": { + "format": "uint", + "minimum": 0, + "type": "integer" + } + }, + "required": [ + "kind", + "max_depth" + ], + "type": "object" + } + ] + }, + "SpeculativeContract": { + "additionalProperties": false, + "description": "Portable compatibility facts for speculative decoding.\n\nProposal width, tree shape, scheduling, kernels, and whether speculation is\nenabled at all are runtime decisions.", + "properties": { + "distribution_preserving": { + "default": false, + "description": "Whether accepting proposals preserves the target's output distribution.\n\nA runtime may auto-enable speculation only when this is true; otherwise\nthe caller must opt in explicitly.", + "type": "boolean" + }, + "max_proposal_width": { + "description": "Maximum number of proposed positions this package can undo.\n\nThis is a rollback bound, not a choice of proposal width: the runtime\npicks any K up to this bound. A validator rejects a package whose state\nor speculative effects cannot be undone this far.", + "format": "uint", + "minimum": 1, + "type": "integer" + }, + "port_bindings": { + "additionalProperties": { + "type": "string" + }, + "description": "Proposer ports bound to target-owned values, by semantic role.", + "type": "object" + }, + "proposal_execution": { + "$ref": "#/$defs/SpeculativeProposalExecution", + "default": { + "kind": "block" + }, + "description": "How the proposer materializes one candidate block.\n\n`block` components emit the complete proposal in one invocation.\n`chained` components emit one distribution and recurrence update per\ninvocation; the runtime repeatedly invokes the same typed component up\nto `max_proposal_width`." + }, + "proposer": { + "description": "Workflow component that proposes tokens.", + "minLength": 1, + "type": "string" + }, + "rollback_state": { + "description": "State groups that must roll back when a proposal is rejected.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": true + }, + "shared_state": { + "description": "State groups the proposer shares with the target.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": true + }, + "shared_weights": { + "description": "Target initializers the proposer borrows.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": true + }, + "target": { + "description": "Workflow component that verifies proposals.", + "minLength": 1, + "type": "string" + }, + "vocabulary": { + "$ref": "#/$defs/SpeculativeVocabulary", + "description": "How the proposer's vocabulary relates to the target's." + } + }, + "required": [ + "proposer", + "target", + "vocabulary", + "max_proposal_width" + ], + "type": "object" + }, + "SpeculativeProposalExecution": { + "description": "Execution shape of a speculative proposer.", + "oneOf": [ + { + "additionalProperties": false, + "description": "One proposer invocation returns the complete token block.", + "properties": { + "kind": { + "const": "block", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + { + "additionalProperties": false, + "description": "Repeated proposer invocations form an autoregressive proposal chain.", + "properties": { + "folded_carry_output": { + "description": "Proposer OUTPUT port producing the folded carry for the NEXT step\n(carry_k, k>=1). The carry has no dedicated input port: it re-enters\nas the trailing segment of the proposer's fused\n`token_embedding_input` (`concat(embed(last_token), carry)`), so it\nowns no workflow state cell. Three ports pin the fold EXPLICITLY, so\na runtime never infers by convention:\n\n* DESTINATION: `port_bindings.target_hidden_context` names the\n proposer input port the carry lands in. For a folded carry it must\n equal `token_embedding_input`: the carry occupies the fused input's\n trailing half, so the destination is that fused input, never a\n separate port.\n* carry_0 SOURCE: `folded_carry_seed` names the target output read\n as the carry on the first step.\n* carry_k SOURCE: this field, the proposer output on every later\n step.\n\nBecause a folded carry is recomputed from committed tokens on\nrejection rather than restored, it does not appear in\n`rollback_state`. A chained proposer declares at least one of\n`recurrent` or `folded_carry_output`.", + "type": [ + "string", + "null" + ] + }, + "folded_carry_seed": { + "anyOf": [ + { + "$ref": "#/$defs/SpeculativeValueRef" + }, + { + "type": "null" + } + ], + "description": "carry_0 seed: the target component OUTPUT read as the folded carry on\nthe FIRST step, before the proposer has produced a carry. Named\nexplicitly (`component` + `output`) so a runtime reads it rather than\ninferring \"the target hidden output\" by convention. Its `component`\nmust be the speculative target, since carry_0 is the target's own\nper-token hidden output. Required whenever `folded_carry_output` is\npresent." + }, + "kind": { + "const": "chained", + "type": "string" + }, + "logits_output": { + "description": "Proposer output port carrying the next-token distribution.", + "minLength": 1, + "type": "string" + }, + "recurrent": { + "default": [], + "description": "Loop-carried hidden/cache state updated by every proposer invocation\nthrough its own input port.", + "items": { + "$ref": "#/$defs/SpeculativeRecurrenceBinding" + }, + "type": "array" + }, + "token_embedding": { + "anyOf": [ + { + "$ref": "#/$defs/TokenEmbeddingSource" + }, + { + "type": "null" + } + ], + "description": "Where a runtime obtains `embed(last_token)` for the LEADING half of\nthe fused `token_embedding_input`. A folded-carry proposer graph\nconsumes only the fused input, so it reads no embedding initializer of\nits own and `speculative.shared_weights` stays empty; this names the\nmodel-agnostic embedding table the runtime gathers the leading half\nfrom (never extracted heuristically from a graph). Its `component`\nmust be the speculative target — an ONNX model that owns the named\n`table` initializer — so the table resolves to a real initializer in\nthe target model/artifact. Required whenever `folded_carry_output` is\npresent." + }, + "token_embedding_input": { + "description": "Proposer input port receiving the previous selected token embedding.", + "minLength": 1, + "type": "string" + } + }, + "required": [ + "kind", + "token_embedding_input", + "logits_output" + ], + "type": "object" + } + ] + }, + "SpeculativeRecurrenceBinding": { + "additionalProperties": false, + "description": "One loop-carried proposer value in a chained proposal.", + "properties": { + "input": { + "description": "Proposer input port receiving the current value.", + "minLength": 1, + "type": "string" + }, + "output": { + "description": "Proposer output port producing the next value.", + "minLength": 1, + "type": "string" + }, + "state": { + "description": "Workflow state cell checkpointed before proposal and restored on rejection.", + "minLength": 1, + "type": "string" + } + }, + "required": [ + "state", + "input", + "output" + ], + "type": "object" + }, + "SpeculativeValueRef": { + "additionalProperties": false, + "description": "An explicit reference to a value a workflow component produces.\n\nBoth halves are named so a speculative runtime resolves the value from the\ndeclared graph I/O, never by string convention or shape guessing.", + "properties": { + "component": { + "description": "The workflow component that produces the value.", + "minLength": 1, + "type": "string" + }, + "output": { + "description": "The declared output port of that component carrying the value.", + "minLength": 1, + "type": "string" + } + }, + "required": [ + "component", + "output" + ], + "type": "object" + }, + "SpeculativeVocabulary": { + "description": "Vocabulary relationship between a proposer and its target.", + "oneOf": [ + { + "additionalProperties": false, + "description": "Proposer and target share one identical vocabulary.", + "properties": { + "kind": { + "const": "identical", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + { + "additionalProperties": false, + "description": "The proposer's vocabulary is a prefix-compatible subset of the target's.", + "properties": { + "kind": { + "const": "subset", + "type": "string" + }, + "proposer_vocab_size": { + "format": "uint", + "minimum": 0, + "type": "integer" + } + }, + "required": [ + "kind", + "proposer_vocab_size" + ], + "type": "object" + }, + { + "additionalProperties": false, + "description": "A declared mapping artifact translates proposer ids into target ids.", + "properties": { + "artifact": { + "type": "string" + }, + "kind": { + "const": "mapped", + "type": "string" + } + }, + "required": [ + "kind", + "artifact" + ], + "type": "object" + } + ] + }, + "StateAliasing": { + "description": "Legality of aliasing a component's `present` output onto its `past` input.", + "oneOf": [ + { + "const": "permitted", + "description": "The runtime may alias input and output buffers but need not.", + "type": "string" + }, + { + "const": "required", + "description": "The component only works correctly when input and output alias.", + "type": "string" + }, + { + "const": "forbidden", + "description": "Aliasing input and output buffers is incorrect for this component.", + "type": "string" + } + ] + }, + "StateCheckpointContract": { + "additionalProperties": false, + "description": "The only portable, cross-build path for a state group's contents.\n\nThis is deliberately not a wire format: metadata names the adapter and its\nversion, and the adapter owns the encoding. A portable checkpoint is slow and\nsurvives a version change; a private transfer is fast and does not.", + "properties": { + "adapter": { + "description": "Versioned adapter identifier, for example `onnx-genai.kv-checkpoint`.", + "minLength": 1, + "type": "string" + }, + "version": { + "description": "Adapter version this group's checkpoints are written against.", + "minLength": 1, + "type": "string" + } + }, + "required": [ + "adapter", + "version" + ], + "type": "object" + }, + "StateGroupCapabilities": { + "additionalProperties": false, + "description": "Rollback, snapshot, and fork bounds declared for one state group.", + "properties": { + "cascade": { + "description": "Other groups that must be rolled back, snapshotted, or forked together.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": true + }, + "fork": { + "default": false, + "description": "Whether the runtime can fork this group into an independent row.", + "type": "boolean" + }, + "rollback_positions": { + "description": "Maximum number of trailing positions that can be discarded correctly.\n\nAbsent means the group cannot be rolled back at all.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "snapshot": { + "default": false, + "description": "Whether the runtime can snapshot and restore this group.", + "type": "boolean" + } + }, + "type": "object" + }, + "StateGroupContract": { + "additionalProperties": false, + "description": "One semantic group of model state sharing a kind, geometry, and graph ABI.", + "properties": { + "aliasing": { + "$ref": "#/$defs/StateAliasing", + "default": "forbidden", + "description": "Whether the component may write `present` into the `past` buffer." + }, + "capabilities": { + "$ref": "#/$defs/StateGroupCapabilities", + "default": { + "fork": false, + "snapshot": false + }, + "description": "Rollback, snapshot, and fork bounds usable by the runtime." + }, + "checkpoint": { + "anyOf": [ + { + "$ref": "#/$defs/StateCheckpointContract" + }, + { + "type": "null" + } + ], + "description": "The versioned checkpoint adapter through which this group's state may\nleave the process portably.\n\nAbsent means the group's state is private: it may still move between\nprocesses, but only through a private runtime protocol that requires a\nmatching protocol and build on both ends (prefill/decode disaggregation,\nencoder/decoder interchange). Those transfers are fast precisely because\nthey are not portable, and treating one as a portable export is how a\ncluster silently corrupts state across a rolling upgrade." + }, + "kind": { + "$ref": "#/$defs/StateKind", + "description": "Semantic kind of the state in this group." + }, + "layout": { + "description": "Graph-visible element layout of the state tensors.", + "type": "string" + }, + "logical_lengths": { + "description": "Semantic state cell holding the current logical length of each row.\n\nRequired only when the logical length is graph-visible. Absent means the\nruntime derives length from its private sequence table.", + "type": [ + "string", + "null" + ] + }, + "ports": { + "additionalProperties": { + "additionalProperties": { + "$ref": "#/$defs/StatePortAlias" + }, + "type": "object" + }, + "default": {}, + "type": "object" + }, + "reuse": { + "$ref": "#/$defs/StateReuse", + "default": { + "evictable_prefix": false, + "prefix_reusable": false + }, + "description": "Prefix reuse and eviction semantics of this group." + }, + "sequence_axis": { + "description": "Axis whose extent represents logical sequence positions.\n\nRequired for sequence-growing and indexed-scatter state. Fixed-size\nrecurrent state updated by replacement has no logical sequence extent\nand therefore omits this field.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "total_length": { + "description": "Graph-visible total-length input, when the component reads one.", + "type": [ + "string", + "null" + ] + }, + "update": { + "anyOf": [ + { + "$ref": "#/$defs/StateUpdate" + }, + { + "type": "null" + } + ], + "description": "How the graph writes each step's positions into this group's buffers.\n\nAbsent means the buffer extends along `sequence_axis` as the sequence\ngrows, which is the historical behavior. `indexed_scatter` declares a\nbuffer of fixed capacity whose new positions are written at destinations\nthe graph reads from a declared value. That distinction is what makes\nrewind, row replacement, and inactive-row compaction expressible: in a\ngrowing buffer the valid region is the whole tensor, while in a\nfixed-capacity buffer it is a declared prefix that the shape cannot\nreveal.\n\nThis describes the GRAPH's update discipline, not an allocator. The\nphysical buffer and where it lives remain runtime-owned." + } + }, + "required": [ + "kind", + "layout" + ], + "type": "object" + }, + "StateInitKind": { + "description": "Loop-carried state initialization vocabulary.", + "oneOf": [ + { + "enum": [ + "zeros" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "zeros" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "StateKind": { + "description": "Semantic kind of a model state group.", + "oneOf": [ + { + "const": "full_attention", + "description": "Dense causal attention over the full sequence.", + "type": "string" + }, + { + "const": "sliding_attention", + "description": "Causal attention restricted to a sliding window.", + "type": "string" + }, + { + "const": "multi_latent_attention", + "description": "Compressed latent attention state (MLA).", + "type": "string" + }, + { + "const": "recurrent", + "description": "Fixed-size recurrent or state-space carry.", + "type": "string" + }, + { + "const": "cross_attention", + "description": "Cross-attention state keyed by an encoder result.", + "type": "string" + }, + { + "const": "encoder", + "description": "Encoder output retained across decoder steps.", + "type": "string" + } + ] + }, + "StateManagement": { + "description": "Storage ownership of one workflow state cell.", + "oneOf": [ + { + "const": "workflow", + "description": "Ordinary SSA-managed value; liveness is derived from the workflow graph.", + "type": "string" + }, + { + "const": "runtime", + "description": "The runtime owns the physical storage and its allocation policy.", + "type": "string" + }, + { + "const": "external", + "description": "An external service owns the storage; only the typed handle is portable.", + "type": "string" + } + ] + }, + "StatePortAccess": { + "description": "State access performed by one component binding.", + "oneOf": [ + { + "const": "read_write", + "description": "The component consumes the current value and produces its successor.", + "type": "string" + }, + { + "const": "read_only", + "description": "The component consumes a frozen value; any graph output is discarded.", + "type": "string" + } + ] + }, + "StatePortAlias": { + "additionalProperties": false, + "properties": { + "access": { + "$ref": "#/$defs/StatePortAccess", + "description": "Whether this component advances the state or only observes a frozen\nvalue produced by another component in the same service group.\n\nA read-only binding still names the graph's present output when the\nartifact exposes one for kernel ABI reasons, but that output is not a\nstate transition and must not be aliased back onto the input." + }, + "input": { + "type": "string" + }, + "layer": { + "description": "Zero-based layer index of this port pair within its group.\n\nRequired when a group binds more than one alias of the same\n[`StatePortRole`], because the map key is a producer-chosen label and\nits lexicographic order is not the layer order (`layer.10` sorts before\n`layer.2`). A runtime that pairs per-layer buffers positionally would\notherwise silently transpose two layers' caches.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "output": { + "description": "Graph output port carrying this pair's next-step value.\n\nRequired for a read-write transition. A `read_only` binding MAY omit it:\na pure borrowed-state reader — e.g. a shared-KV drafter that consumes\nanother decoder's cache and advances nothing — exposes no present output\nat all, so there is nothing to name. A read-only reader whose artifact\nstill emits a discarded present output for kernel-ABI reasons may name it\nhere, but that value is never a state transition.", + "type": [ + "string", + "null" + ] + }, + "role": { + "anyOf": [ + { + "$ref": "#/$defs/StatePortRole" + }, + { + "type": "null" + } + ], + "description": "Which half of an attention cache this port pair carries.\n\nA graph that splits keys and values into separate buffers exposes two\naliases per layer that are shape-identical and therefore\nindistinguishable; a graph that packs them exposes one. Only the\nproducer knows which it built, and recovering it from a port's spelling\nwould be the name-guessing this schema refuses. Absent means the group\ndoes not distinguish halves, which is correct for recurrent and latent\nstate." + } + }, + "required": [ + "input" + ], + "type": "object" + }, + "StatePortRole": { + "description": "Which half of a split attention cache a state port pair carries.", + "oneOf": [ + { + "const": "key", + "description": "Keys of a split key/value cache.", + "type": "string" + }, + { + "const": "value", + "description": "Values of a split key/value cache.", + "type": "string" + }, + { + "const": "combined", + "description": "A single buffer holding keys and values together.", + "type": "string" + } + ] + }, + "StateReleaseBoundary": { + "description": "Logical boundary after which a managed or external state cell is releasable.", + "oneOf": [ + { + "const": "invocation", + "description": "Releasable when the invocation that created it completes.", + "type": "string" + }, + { + "const": "session", + "description": "Releasable when the session that owns it ends.", + "type": "string" + }, + { + "const": "row", + "description": "Releasable when its owning batch row is released.", + "type": "string" + } + ] + }, + "StateReuse": { + "additionalProperties": false, + "description": "Semantic reuse and eviction legality of a state group.", + "properties": { + "evictable_prefix": { + "default": false, + "description": "Whether dropping the oldest positions preserves declared semantics.", + "type": "boolean" + }, + "prefix_reusable": { + "default": false, + "description": "Whether a shared token prefix produces identical state that may be reused.", + "type": "boolean" + } + }, + "type": "object" + }, + "StateServiceContract": { + "additionalProperties": false, + "description": "Semantic model-state contract.\n\nThis declares what the state *means* and which graph ABI facts constrain the\nruntime. It never selects paged, shared-buffer, or separate storage, a slot\nallocation algorithm, a compaction algorithm, or a device.", + "properties": { + "groups": { + "additionalProperties": { + "$ref": "#/$defs/StateGroupContract" + }, + "default": {}, + "type": "object" + } + }, + "type": "object" + }, + "StateUpdate": { + "description": "How a state group's buffers absorb each step's new positions.\n\nBoth variants describe what the GRAPH does. Neither selects a storage\nstrategy, a slot allocator, or a device: a runtime is free to back an\n`append` group with a fixed arena or an `indexed_scatter` group with paged\nstorage, so long as the graph sees what it declared.", + "oneOf": [ + { + "additionalProperties": false, + "description": "Each step's positions extend the buffer along `sequence_axis`.\n\nThe valid region is the whole tensor, so no write cursor is graph-visible\nand the buffer's shape carries the length.", + "properties": { + "kind": { + "const": "append", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + { + "additionalProperties": false, + "description": "Each step replaces the complete fixed-size state tensor.\n\nThis is the common discipline for recurrent accumulators, state-space\ncarries, and causal-convolution history. The algorithm does not need a\ndistinct state kind: separate groups already declare each tensor's\nshape, ports, lifetime, and rollback behavior.", + "properties": { + "kind": { + "const": "replace", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + { + "additionalProperties": false, + "description": "Each step's positions are scattered into a buffer of FIXED capacity at\ndestinations the graph reads from `write_indices`.\n\nThe tensor's extent along `sequence_axis` is the capacity, not the\nlength: the valid region is the prefix named by the group's\n`logical_lengths`. Because destinations are data rather than position,\nrewinding a row is a cursor move, replacing a row reuses its slots, and\nrows of unequal length share one rectangular buffer.", + "properties": { + "capacity": { + "description": "Integer-scalar workflow value giving the fixed extent of\n`sequence_axis` that the graph was built against.\n\nThis is a graph fact, not a deployment budget. It bounds legal write\ndestinations; it does not say where the buffer lives, when it is\nallocated, or how many rows a deployment admits.", + "minLength": 1, + "type": "string" + }, + "kind": { + "const": "indexed_scatter", + "type": "string" + }, + "kv_length_ports": { + "additionalProperties": { + "type": "string" + }, + "description": "Per-component input port that receives the graph-visible valid\nlength, keyed by component name.\n\nExactly the same problem as `write_indices_ports`, and unsolvable the\nsame way: the length is a rank-1 integer vector, so it is\nshape-indistinguishable from the destinations sitting next to it. A\ngraph that reads no length port declares none.", + "type": "object" + }, + "write_indices": { + "description": "Semantic state cell carrying this step's per-row destination\npositions along `sequence_axis`.\n\nA cell rather than a step output because the write cursor is part of\nthe group's state: it must be checkpointed, forked, and rewound with\nthe buffer it indexes, or a restored row would overwrite live\npositions.", + "minLength": 1, + "type": "string" + }, + "write_indices_ports": { + "additionalProperties": { + "type": "string" + }, + "default": {}, + "description": "Per-component input port that receives the destinations, keyed by\ncomponent name.\n\nThe same class of fact as [`StateGroupContract::ports`]: which port\nof which component carries this group's ABI. A runtime cannot\nrecover it from the step graph, because destinations are an ordinary\ninteger vector and are shape-indistinguishable from every other\ninteger control input.", + "type": "object" + } + }, + "required": [ + "kind", + "write_indices", + "capacity" + ], + "type": "object" + } + ] + }, + "StateUpdateKind": { + "description": "Loop-carried state update-semantics vocabulary.", + "oneOf": [ + { + "enum": [ + "replace" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "replace" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "StaticCacheIoSpec": { + "description": "Explicit port ABI for a fixed-buffer TensorScatter static KV cache.\n\nDescribes GRAPH STRUCTURE, never a model family. The four per-layer cache\nlists pair positionally per layer and must all have the same length: index\n`i` in each list is layer `i`'s key/value input and updated key/value output.", + "properties": { + "key_cache_inputs": { + "description": "Per-layer static key-cache buffer inputs, positional per layer. Length\nmust equal `value_cache_inputs`, `key_cache_outputs`, and\n`value_cache_outputs`.", + "items": { + "minLength": 1, + "type": "string" + }, + "minItems": 1, + "type": "array" + }, + "key_cache_outputs": { + "description": "Per-layer updated key-cache outputs, paired positionally with\n`key_cache_inputs`.", + "items": { + "minLength": 1, + "type": "string" + }, + "minItems": 1, + "type": "array" + }, + "kv_sequence_length_input": { + "description": "Input port carrying the non-pad KV sequence length (`int` vector).\nShape-indistinguishable from `write_indices_input`, so it too must be\nnamed explicitly.", + "minLength": 1, + "type": "string" + }, + "value_cache_inputs": { + "description": "Per-layer static value-cache buffer inputs, paired positionally with\n`key_cache_inputs`.", + "items": { + "minLength": 1, + "type": "string" + }, + "minItems": 1, + "type": "array" + }, + "value_cache_outputs": { + "description": "Per-layer updated value-cache outputs, paired positionally with\n`value_cache_inputs`.", + "items": { + "minLength": 1, + "type": "string" + }, + "minItems": 1, + "type": "array" + }, + "write_indices_input": { + "description": "Input port carrying the per-token scatter write positions\n(`int` vector). Shape-indistinguishable from other integer control\ninputs, so it must be named explicitly.", + "minLength": 1, + "type": "string" + } + }, + "required": [ + "write_indices_input", + "kv_sequence_length_input", + "key_cache_inputs", + "value_cache_inputs", + "key_cache_outputs", + "value_cache_outputs" + ], + "type": "object" + }, + "TaskProfile": { + "additionalProperties": false, + "description": "One executable task profile that shares the package's common facts.\n\nGenerative and non-generative tasks live in one document. Each profile\ncarries its own version and a requirement class so a strict reader can skip\nan optional profile it does not understand while still rejecting unknown\ncore fields.", + "properties": { + "batch_invariance": { + "anyOf": [ + { + "$ref": "#/$defs/BatchInvariance" + }, + { + "type": "null" + } + ], + "description": "Whether a row's outputs depend on the other rows batched with it.\n\n`row_independent` means a row produces identical values whether it is\nrun alone or co-batched with rows of any other length, so a runtime may\nbatch freely. `padding_sensitive` means padding a row to the batch width\nchanges its values — for example when a normalization reduces over the\npadded time axis — so a runtime that batches trades accuracy for\nthroughput and must not treat batched results as equal to solo results.\n\nAbsent means unstated, not `row_independent`." + }, + "decoding": { + "anyOf": [ + { + "$ref": "#/$defs/SequenceDecodingSpec" + }, + { + "type": "null" + } + ], + "description": "How a non-generative sequence output is decoded into discrete tokens." + }, + "generation_affecting": { + "default": false, + "description": "Whether this profile changes generated output and therefore participates\nin cache correctness dependencies.", + "type": "boolean" + }, + "kind": { + "description": "Task kind identifier, e.g. `generation`, `embedding`, `reranking`.", + "minLength": 1, + "type": "string" + }, + "outputs": { + "additionalProperties": { + "type": "string" + }, + "description": "Workflow outputs this profile consumes, by semantic role.", + "type": "object" + }, + "pooling": { + "anyOf": [ + { + "$ref": "#/$defs/PoolingSpec" + }, + { + "type": "null" + } + ], + "description": "Pooling applied to a sequence-valued output, when the task needs it." + }, + "requirement": { + "$ref": "#/$defs/ProfileRequirement", + "default": "required", + "description": "Whether a reader that does not understand this profile may skip it." + }, + "version": { + "description": "Version of this profile's own contract.", + "minLength": 1, + "type": "string" + } + }, + "required": [ + "kind", + "version" + ], + "type": "object" + }, + "TensorContract": { + "additionalProperties": false, + "description": "Typed tensor contract used at package and component boundaries.", + "properties": { + "batch_layout": { + "$ref": "#/$defs/BatchLayout", + "description": "How this value relates to the runtime's private request/sequence table.\n\nThis is a structural batching fact, never a row identity. It is the only\ninformation a runtime needs to move, split, or drop this value during\ncompaction; scheduler slots, epochs, block tables, and sequence handles\nstay runtime-private." + }, + "dtype": { + "$ref": "#/$defs/TensorDType" + }, + "optional": { + "default": false, + "type": "boolean" + }, + "rank": { + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "shape": { + "items": { + "$ref": "#/$defs/TensorDimension" + }, + "type": [ + "array", + "null" + ] + } + }, + "required": [ + "dtype", + "rank" + ], + "type": "object" + }, + "TensorDType": { + "description": "Tensor-boundary dtype vocabulary, including non-numeric pipeline values.", + "oneOf": [ + { + "enum": [ + "float32", + "fp32", + "float16", + "fp16", + "bfloat16", + "bf16", + "float8_e4m3fn", + "float8_e5m2", + "int64", + "int32", + "int8", + "uint8", + "bool", + "string" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "float32", + "fp32", + "float16", + "fp16", + "bfloat16", + "bf16", + "float8_e4m3fn", + "float8_e5m2", + "int64", + "int32", + "int8", + "uint8", + "bool", + "string" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "TensorDimension": { + "anyOf": [ + { + "description": "A fixed, non-negative dimension.", + "format": "int64", + "minimum": 0, + "type": "integer" + }, + { + "description": "A runtime shape symbol.", + "minLength": 1, + "type": "string" + } + ], + "description": "One fixed or runtime-resolved tensor-shape dimension." + }, + "TensorShardFacts": { + "additionalProperties": false, + "description": "Shard axis and replication facts of one logical parameter group.", + "properties": { + "max_shards": { + "description": "Largest legal shard count; the caller may choose any divisor.", + "format": "uint", + "minimum": 1, + "type": "integer" + }, + "replicated": { + "description": "Values that must be replicated identically on every rank.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": true + }, + "shard_axis": { + "description": "Axis of the parameter that may be split across ranks.", + "format": "uint", + "minimum": 0, + "type": "integer" + } + }, + "required": [ + "shard_axis", + "max_shards" + ], + "type": "object" + }, + "ThumbnailOrder": { + "description": "Optional-thumbnail ordering vocabulary.", + "oneOf": [ + { + "enum": [ + "none", + "prepend", + "append" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "none", + "prepend", + "append" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "TokenEmbeddingSource": { + "additionalProperties": false, + "description": "Explicit, model-agnostic source of the token embedding a chained proposer\nfolds into the leading half of its fused input.\n\nA folded-carry proposer never reads an embedding initializer inside its own\ngraph, so the table cannot be recovered from the proposer graph. This names\nthe component whose embedding the runtime reuses and the initializer that\nholds it, so gathering `embed(last_token)` is a declared contract rather than\na per-model heuristic.", + "properties": { + "component": { + "description": "The workflow component whose token-embedding table is reused. For a\nfolded carry this must be the speculative target, whose vocabulary the\nproposer shares and whose ONNX model owns the `table` initializer.", + "minLength": 1, + "type": "string" + }, + "table": { + "description": "The named embedding table (graph initializer) on that component, e.g.\n`model.embed_tokens.weight`. A `[vocab, hidden]` row-major matrix. It\nmust name a real initializer in the target model/artifact.", + "minLength": 1, + "type": "string" + } + }, + "required": [ + "component", + "table" + ], + "type": "object" + }, + "TokenizerArtifact": { + "additionalProperties": false, + "description": "One package-relative tokenizer artifact.", + "properties": { + "location": { + "description": "Package-relative path of the artifact.", + "minLength": 1, + "type": "string" + } + }, + "required": [ + "location" + ], + "type": "object" + }, + "TokenizerFacts": { + "additionalProperties": false, + "description": "Tokenizer facts and package-relative artifacts.\n\nA request carries text, tokens, grammars, and JSON Schemas. Interpreting any\nof them requires the vocabulary contract the package was built against.", + "properties": { + "algorithm": { + "description": "Tokenizer algorithm identifier, e.g. `bpe`, `unigram`, `wordpiece`.", + "minLength": 1, + "type": "string" + }, + "artifacts": { + "description": "Package-relative tokenizer artifacts.", + "items": { + "$ref": "#/$defs/TokenizerArtifact" + }, + "minItems": 1, + "type": "array" + }, + "byte_level": { + "default": false, + "description": "Whether the tokenizer operates on raw bytes rather than Unicode scalars.", + "type": "boolean" + }, + "special_tokens": { + "additionalProperties": { + "$ref": "#/$defs/SpecialTokenFact" + }, + "description": "Special tokens by semantic role, e.g. `bos`, `eos`, `pad`.", + "type": "object" + }, + "vocab_size": { + "description": "Number of entries in the vocabulary, including added tokens.", + "format": "uint", + "minimum": 1, + "type": "integer" + } + }, + "required": [ + "algorithm", + "vocab_size" + ], + "type": "object" + }, + "WorkflowBranchOutput": { + "additionalProperties": false, + "properties": { + "cases": { + "additionalProperties": { + "type": "string" + }, + "type": "object" + }, + "default": { + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "cases" + ], + "type": "object" + }, + "WorkflowCarry": { + "additionalProperties": false, + "properties": { + "cell": { + "type": "string" + }, + "initial": { + "type": [ + "string", + "null" + ] + }, + "next": { + "type": "string" + } + }, + "required": [ + "cell", + "next" + ], + "type": "object" + }, + "WorkflowComponent": { + "additionalProperties": false, + "properties": { + "application_overridable": { + "default": false, + "description": "Allow an application to select another declared component with the same\nversioned contract ABI for this invocation.", + "type": "boolean" + }, + "cache_affects_state": { + "description": "Non-dataflow facts that change this component's observable state.\n\nCache correctness dependencies of ONNX components are derived from the\nworkflow SSA graph. Native and external components must declare any\nadditional state they read that is not visible as a typed input.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": true + }, + "contract": { + "anyOf": [ + { + "$ref": "#/$defs/ComponentContract" + }, + { + "type": "null" + } + ] + }, + "effects": { + "default": [], + "items": { + "type": "string" + }, + "type": "array" + }, + "implementation": { + "$ref": "#/$defs/ComponentImplementation" + }, + "ports": { + "$ref": "#/$defs/ComponentPorts", + "default": { + "inputs": {}, + "outputs": {} + } + }, + "row_scope": { + "anyOf": [ + { + "$ref": "#/$defs/ComponentRowScope" + }, + { + "type": "null" + } + ], + "description": "Declared per-request row scope of this component's private state.\n\nA component with row scope must implement the mandatory row ABI\n(`compact(selection)` and `release(row)`). This is an ABI invariant, not\na negotiated capability: a runtime may not load a package whose\nrow-scoped component cannot be compacted." + } + }, + "required": [ + "implementation" + ], + "type": "object" + }, + "WorkflowEmitMode": { + "enum": [ + "replace", + "append", + "event" + ], + "type": "string" + }, + "WorkflowInput": { + "additionalProperties": false, + "properties": { + "contract": { + "$ref": "#/$defs/TensorContract" + }, + "default": { + "anyOf": [ + { + "$ref": "#/$defs/LiteralValue" + }, + { + "type": "null" + } + ] + }, + "externally_suppliable": { + "description": "Whether an application may supply a previously computed typed value in\nplace of recomputing it, such as a cached encoder result.\n\nTransport, remote caching, and identity of the supplied value remain\nruntime-owned; metadata only declares that the substitution is legal.", + "type": "boolean" + }, + "present_as": { + "description": "Initial scalar bool SSA value indicating whether the caller supplied this input.", + "type": [ + "string", + "null" + ] + }, + "required": { + "default": true, + "type": "boolean" + }, + "role": { + "$ref": "#/$defs/SemanticInputRole" + }, + "source": { + "$ref": "#/$defs/WorkflowInputSource" + } + }, + "required": [ + "contract", + "role", + "source" + ], + "type": "object" + }, + "WorkflowInputSource": { + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "request", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "application", + "type": "string" + }, + "name": { + "type": "string" + } + }, + "required": [ + "kind", + "name" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "literal", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "artifact", + "type": "string" + }, + "path": { + "type": "string" + } + }, + "required": [ + "kind", + "path" + ], + "type": "object" + } + ] + }, + "WorkflowLoopIteration": { + "additionalProperties": false, + "properties": { + "contract": { + "$ref": "#/$defs/TensorContract", + "description": "`int64` scalar or rank-one broadcast contract." + }, + "value": { + "description": "SSA value containing the current zero-based iteration.", + "type": "string" + } + }, + "required": [ + "value", + "contract" + ], + "type": "object" + }, + "WorkflowLoopTermination": { + "enum": [ + "predicate", + "generation_eos" + ], + "type": "string" + }, + "WorkflowManifest": { + "additionalProperties": false, + "properties": { + "adapter_abis": { + "additionalProperties": { + "type": "string" + }, + "default": {}, + "type": "object" + }, + "capabilities": { + "default": [], + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": true + } + }, + "type": "object" + }, + "WorkflowOutput": { + "additionalProperties": false, + "properties": { + "contract": { + "$ref": "#/$defs/TensorContract" + }, + "media": { + "anyOf": [ + { + "$ref": "#/$defs/MediaOutputContract" + }, + { + "type": "null" + } + ], + "description": "Concrete media delivery contract for a post-processing output.\n\nTensor shape alone cannot distinguish PCM samples from encoded WAV bytes,\nnor can it carry the sample rate and channel count required by an audio\nserving API. This remains architecture-neutral and intentionally contains\nno model-family identifiers or artifact fingerprints." + }, + "role": { + "$ref": "#/$defs/WorkflowOutputRole" + }, + "stage": { + "$ref": "#/$defs/OutputStage" + }, + "value_range": { + "anyOf": [ + { + "$ref": "#/$defs/ImageOutputValueRange" + }, + { + "type": "null" + } + ] + } + }, + "required": [ + "contract", + "role", + "stage" + ], + "type": "object" + }, + "WorkflowOutputRole": { + "oneOf": [ + { + "enum": [ + "tokens", + "text", + "image", + "audio", + "tensor", + "event" + ], + "type": "string" + }, + { + "const": "video", + "description": "A frame sequence. Distinct from `Image` because a consumer has to know\nthe value carries a temporal axis and may be published incrementally.", + "type": "string" + } + ] + }, + "WorkflowSpec": { + "additionalProperties": false, + "description": "Sound, component-centric workflow IR. Tensor math lives in invoked components.", + "properties": { + "components": { + "additionalProperties": { + "$ref": "#/$defs/WorkflowComponent" + }, + "type": "object" + }, + "effects": { + "additionalProperties": { + "$ref": "#/$defs/EffectContract" + }, + "default": {}, + "description": "Retry and speculation semantics of every declared external effect domain.", + "type": "object" + }, + "inputs": { + "additionalProperties": { + "$ref": "#/$defs/WorkflowInput" + }, + "default": {}, + "type": "object" + }, + "manifest": { + "$ref": "#/$defs/WorkflowManifest" + }, + "outputs": { + "additionalProperties": { + "$ref": "#/$defs/WorkflowOutput" + }, + "default": {}, + "type": "object" + }, + "serving": { + "anyOf": [ + { + "$ref": "#/$defs/ServingServiceContract" + }, + { + "type": "null" + } + ] + }, + "state": { + "additionalProperties": { + "$ref": "#/$defs/WorkflowStateCell" + }, + "default": {}, + "type": "object" + }, + "steps": { + "items": { + "$ref": "#/$defs/WorkflowStep" + }, + "type": "array" + } + }, + "required": [ + "manifest", + "components", + "steps" + ], + "type": "object" + }, + "WorkflowStateCell": { + "additionalProperties": false, + "properties": { + "class": { + "$ref": "#/$defs/WorkflowStateClass", + "default": "semantic" + }, + "contract": { + "$ref": "#/$defs/TensorContract" + }, + "initializer": { + "type": "string" + }, + "management": { + "$ref": "#/$defs/StateManagement", + "default": "workflow", + "description": "Who owns the physical storage of this cell.\n\nWorkflow-managed cells follow ordinary SSA liveness. Runtime-managed and\nexternal cells must also declare a logical release boundary so a runtime\nknows when the semantic value stops being reachable." + }, + "recurrence": { + "$ref": "#/$defs/ShapeRecurrence" + }, + "release_boundary": { + "anyOf": [ + { + "$ref": "#/$defs/StateReleaseBoundary" + }, + { + "type": "null" + } + ], + "description": "Logical point at which the runtime may release this cell's storage." + }, + "scope": { + "$ref": "#/$defs/WorkflowStateScope" + }, + "service_group": { + "type": [ + "string", + "null" + ] + }, + "session": { + "anyOf": [ + { + "$ref": "#/$defs/SessionLeaseContract" + }, + { + "type": "null" + } + ] + } + }, + "required": [ + "contract", + "scope", + "initializer", + "recurrence" + ], + "type": "object" + }, + "WorkflowStateClass": { + "enum": [ + "semantic", + "advisory" + ], + "type": "string" + }, + "WorkflowStateScope": { + "enum": [ + "invocation", + "session" + ], + "type": "string" + }, + "WorkflowStep": { + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "sequence", + "type": "string" + }, + "steps": { + "items": { + "$ref": "#/$defs/WorkflowStep" + }, + "type": "array" + } + }, + "required": [ + "kind", + "steps" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "component": { + "type": "string" + }, + "inputs": { + "additionalProperties": { + "type": "string" + }, + "default": {}, + "type": "object" + }, + "kind": { + "const": "invoke", + "type": "string" + }, + "outputs": { + "additionalProperties": { + "type": "string" + }, + "default": {}, + "type": "object" + } + }, + "required": [ + "kind", + "component" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "carried": { + "default": [], + "items": { + "$ref": "#/$defs/WorkflowCarry" + }, + "type": "array" + }, + "continue_when": { + "type": "string" + }, + "iteration": { + "anyOf": [ + { + "$ref": "#/$defs/WorkflowLoopIteration" + }, + { + "type": "null" + } + ] + }, + "kind": { + "const": "loop", + "type": "string" + }, + "max_iterations": { + "type": "string" + }, + "setup": { + "default": [], + "items": { + "$ref": "#/$defs/WorkflowStep" + }, + "type": "array" + }, + "steps": { + "items": { + "$ref": "#/$defs/WorkflowStep" + }, + "type": "array" + }, + "termination": { + "$ref": "#/$defs/WorkflowLoopTermination", + "default": "predicate" + } + }, + "required": [ + "kind", + "steps", + "continue_when", + "max_iterations" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "cases": { + "additionalProperties": { + "$ref": "#/$defs/WorkflowStep" + }, + "type": "object" + }, + "default": { + "anyOf": [ + { + "$ref": "#/$defs/WorkflowStep" + }, + { + "type": "null" + } + ] + }, + "kind": { + "const": "branch", + "type": "string" + }, + "outputs": { + "additionalProperties": { + "$ref": "#/$defs/WorkflowBranchOutput" + }, + "default": {}, + "type": "object" + }, + "predicate": { + "type": "string" + } + }, + "required": [ + "kind", + "predicate", + "cases" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "axis": { + "description": "Axis along which the output grows; defaults to the final axis.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "kind": { + "const": "emit", + "type": "string" + }, + "mode": { + "$ref": "#/$defs/WorkflowEmitMode" + }, + "output": { + "type": "string" + }, + "valid_length": { + "type": [ + "string", + "null" + ] + }, + "value": { + "type": "string" + }, + "when": { + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "kind", + "value", + "output", + "mode" + ], + "type": "object" + } + ] + } + }, + "$id": "https://github.com/onnx/onnx/issues/8184", + "$schema": "https://json-schema.org/draft/2020-12/schema", + "additionalProperties": false, + "allOf": [ + { + "not": { + "properties": { + "model": { + "required": [ + "io" + ] + } + }, + "required": [ + "pipeline", + "model" + ] + } + } + ], + "description": "Portable, runtime-agnostic inference metadata for ONNX generative models. The v1 top-level surface is closed; executable composite packages use pipeline.workflow.", + "properties": { + "adapters": { + "anyOf": [ + { + "$ref": "#/$defs/AdapterServiceContract" + }, + { + "type": "null" + } + ], + "description": "Runtime-managed LoRA adapters for bare or composite model packages.\n\nThis is the migrated `InferenceMetadata.adapters` contract from native\nLoRA phases 1 and 2. Composite execution references workflow SSA inputs,\nbut artifact identity, target resolution, and lifecycle remain package\nmetadata rather than workflow control-flow nodes." + }, + "generation": { + "anyOf": [ + { + "$ref": "#/$defs/GenerationContract" + }, + { + "type": "null" + } + ], + "description": "Authoritative generation defaults and the structural override surface." + }, + "hardware_requirements": { + "anyOf": [ + { + "$ref": "#/$defs/HardwareRequirements" + }, + { + "type": "null" + } + ], + "description": "Minimum and beneficial hardware capabilities used for distribution matching." + }, + "model": { + "anyOf": [ + { + "$ref": "#/$defs/ModelCapabilities" + }, + { + "type": "null" + } + ], + "description": "Build-time model properties and runtime-configurable capabilities." + }, + "package": { + "anyOf": [ + { + "$ref": "#/$defs/PackageFacts" + }, + { + "type": "null" + } + ], + "description": "Exact package facts needed to interpret request data correctly.\n\nTokenizer artifacts, vocabulary size, special tokens, and the constraint\ndialects the package's parser accepts. Grammars and JSON Schemas\nthemselves are request data, not package metadata." + }, + "pipeline": { + "anyOf": [ + { + "$ref": "#/$defs/PipelineSpec" + }, + { + "type": "null" + } + ], + "description": "Declarative multi-model pipeline and its dataflow graph." + }, + "preprocessing": { + "anyOf": [ + { + "$ref": "#/$defs/PreprocessingSpec" + }, + { + "type": "null" + } + ], + "description": "Declared, architecture-neutral input preprocessing programs.\n\nCarries the typed multimodal preprocessing contract (currently the image\ntransform program and its named tensor outputs). Every operation and\noutput is generic, parameterized data — never a model family, vendor\nstring, or baked-in shape. Absent means the model declares no native\npreprocessing program and a runtime must obtain it elsewhere or fail." + }, + "profiles": { + "additionalProperties": { + "$ref": "#/$defs/TaskProfile" + }, + "description": "Executable task profiles sharing this package's common facts.\n\nEvery profile carries its own version and requirement class. A strict\nreader may skip an `ignorable` profile it does not understand; unknown\ncore fields still fail.", + "type": "object" + }, + "quantization": { + "anyOf": [ + { + "$ref": "#/$defs/QuantizationIntent" + }, + { + "type": "null" + } + ], + "description": "Model weight quantization intent, independent of the packed representation." + }, + "required_capabilities": { + "default": [], + "description": "Capability identifiers that a runtime MUST support or refuse to load the model.", + "examples": [ + [ + "kv_cache", + "grouped_query_attention" + ] + ], + "items": { + "minLength": 1, + "type": "string" + }, + "type": "array" + }, + "schema_version": { + "description": "Schema version of this inference-metadata document, e.g. `\"v1\"`.\n\nAbsent means the initial `\"v1\"` contract (readers default to `v1`).\nBump this only for breaking schema changes; additive fields keep the\nsame major version and rely on the forward-compatible \"ignore unknown\nfields\" rule.", + "type": [ + "string", + "null" + ] + }, + "speculative": { + "anyOf": [ + { + "$ref": "#/$defs/SpeculativeContract" + }, + { + "type": "null" + } + ], + "description": "Portable speculative-decoding compatibility facts." + } + }, + "title": "ONNX Inference Metadata", + "type": "object" +} diff --git a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py index ac52e262c..065d1c481 100644 --- a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py @@ -5,7 +5,6 @@ import dataclasses import json -import os import jsonschema import onnx_ir as ir @@ -15,6 +14,7 @@ from mobius._model_package import ModelPackage from mobius.integrations.onnx_genai.inference_metadata_test import ( _model, + _onnx_genai_schema_path, _value, ) from mobius.integrations.onnx_genai.workflow_metadata import ( @@ -93,10 +93,7 @@ def test_codec_workflow_roundtrips_yaml(tmp_path): def test_codec_workflow_matches_producer_schema(): - schema_path = os.environ.get("ONNX_GENAI_SCHEMA") - if not schema_path: - pytest.skip("set ONNX_GENAI_SCHEMA to the producer-contract schema") - with open(schema_path, encoding="utf-8") as handle: + with open(_onnx_genai_schema_path(), encoding="utf-8") as handle: schema = json.load(handle) jsonschema.validate(build_audio_codec_workflow_metadata(_codec_package()), schema) diff --git a/src/mobius/integrations/onnx_genai/comfyui.py b/src/mobius/integrations/onnx_genai/comfyui.py index 8e810e41b..386e7dc0b 100644 --- a/src/mobius/integrations/onnx_genai/comfyui.py +++ b/src/mobius/integrations/onnx_genai/comfyui.py @@ -318,13 +318,22 @@ def parse_comfyui_workflow( sched.kind, ) + if not has_vae: + raise ValueError( + "ComfyUI workflow has no VAE decode node, so it produces latents rather " + "than an image. Why: onnx-genai's pipeline workflow terminates in a decoded " + f"image output, and there is nothing to decode with; supported decode nodes: " + f"{', '.join(sorted(_VAE_DECODE_NODES))}. How to fix: add a VAEDecode node to " + "the workflow before converting it." + ) + metadata = build_diffusion_pipeline_metadata( num_inference_steps=steps, scheduler=sched, guidance_scale=guidance, start_step=start_step or None, denoiser_filename=denoiser_filename, - vae_filename=vae_filename if has_vae else None, + vae_filename=vae_filename, text_encoder_filename=text_encoder_filename if has_text_encoder else None, ) return ComfyUIWorkflow( diff --git a/src/mobius/integrations/onnx_genai/comfyui_test.py b/src/mobius/integrations/onnx_genai/comfyui_test.py index 0c9d7f7f1..af31bdf1b 100644 --- a/src/mobius/integrations/onnx_genai/comfyui_test.py +++ b/src/mobius/integrations/onnx_genai/comfyui_test.py @@ -17,6 +17,9 @@ ) from mobius.integrations.onnx_genai.convert import convert_comfyui_workflow from mobius.integrations.onnx_genai.inference_metadata import SchedulerConfig +from mobius.integrations.onnx_genai.inference_metadata_test import ( + _onnx_genai_schema_path, +) # ComfyUI's canonical "default" text-to-image API-format workflow (trimmed). _DEFAULT_TXT2IMG = { @@ -52,15 +55,15 @@ def test_translate_default_txt2img(): meta = translate_comfyui_workflow(_DEFAULT_TXT2IMG) - strat = meta["pipeline"]["strategy"] - assert strat["kind"] == "iterative" - assert strat["num_steps"] == 20 - assert strat["scheduler_config"]["kind"] == "euler" - assert strat["guidance_scale"] == pytest.approx(8.0) - # CFG conditioning + text encoder + VAE all present. - assert strat["cfg_conditioning_input"] == "encoder_hidden_states" - models = meta["pipeline"]["models"] - assert "denoiser" in models and "vae" in models and "text_encoder" in models + # onnx-genai's PipelineSpec carries a typed SSA workflow and nothing else; + # the denoise loop is executable steps, not an "iterative" strategy block. + assert set(meta["pipeline"]) == {"workflow"} + workflow = meta["pipeline"]["workflow"] + assert {"denoiser", "vae", "text_encoder", "solver_step"} <= set(workflow["components"]) + assert workflow["inputs"]["request.max_iterations"]["default"] == 20 + # CFG is two denoiser invocations plus a combine component. + assert "guidance_combine" in workflow["components"] + assert workflow["inputs"]["request.guidance_scale"]["default"] == pytest.approx(8.0) def test_parse_recovers_full_run_params(): @@ -186,79 +189,84 @@ def test_parse_batch_size(): assert parse_comfyui_workflow(wf).batch_size == 4 -def test_ddim_sampler_maps_to_ddim(): +def test_ddim_sampler_selects_an_unscaled_solver(): wf = json.loads(json.dumps(_DEFAULT_TXT2IMG)) wf["3"]["inputs"]["sampler_name"] = "ddim" - meta = translate_comfyui_workflow(wf) - assert meta["pipeline"]["strategy"]["scheduler_config"]["kind"] == "ddim" - - -def test_dpmpp_2m_sampler_maps_to_dpmpp_2m(): - wf = json.loads(json.dumps(_DEFAULT_TXT2IMG)) - wf["3"]["inputs"]["sampler_name"] = "dpmpp_2m" - meta = translate_comfyui_workflow(wf) - assert meta["pipeline"]["strategy"]["scheduler_config"]["kind"] == "dpmpp_2m" + parsed = parse_comfyui_workflow(wf) + assert parsed.scheduler_kind == "ddim" + # DDIM consumes the latent directly; only Euler pre-scales it by sigma. + components = parsed.metadata["pipeline"]["workflow"]["components"] + assert "solver_step" in components + assert "model_input_scale" not in components -def test_euler_ancestral_sampler_maps(): - wf = json.loads(json.dumps(_DEFAULT_TXT2IMG)) - wf["3"]["inputs"]["sampler_name"] = "euler_ancestral" - meta = translate_comfyui_workflow(wf) - assert meta["pipeline"]["strategy"]["scheduler_config"]["kind"] == "euler_ancestral" +def test_euler_sampler_scales_the_model_input(): + components = translate_comfyui_workflow(_DEFAULT_TXT2IMG)["pipeline"]["workflow"][ + "components" + ] + assert "model_input_scale" in components -def test_karras_scheduler_enables_karras_sigmas(): +def test_dpmpp_2m_sampler_carries_solver_history(): wf = json.loads(json.dumps(_DEFAULT_TXT2IMG)) wf["3"]["inputs"]["sampler_name"] = "dpmpp_2m" - wf["3"]["inputs"]["scheduler"] = "karras" parsed = parse_comfyui_workflow(wf) - assert parsed.scheduler_spacing == "karras" - assert ( - parsed.metadata["pipeline"]["strategy"]["scheduler_config"]["use_karras_sigmas"] - is True - ) + assert parsed.scheduler_kind == "dpmpp_2m" + # A multistep solver's previous data estimate is a declared state cell, not + # a hidden scheduler attribute. + assert "history" in parsed.metadata["pipeline"]["workflow"]["state"] -def test_normal_scheduler_omits_karras_sigmas(): - meta = translate_comfyui_workflow(_DEFAULT_TXT2IMG) - assert "use_karras_sigmas" not in meta["pipeline"]["strategy"]["scheduler_config"] +def test_ancestral_sampler_is_rejected(): + wf = json.loads(json.dumps(_DEFAULT_TXT2IMG)) + wf["3"]["inputs"]["sampler_name"] = "euler_ancestral" + # The runtime executes a declared solver component; mobius ships no + # stochastic solver, so lowering it to deterministic Euler would silently + # run the wrong dynamics. + with pytest.raises(ValueError, match="euler_ancestral"): + translate_comfyui_workflow(wf) -def test_exponential_scheduler_enables_exponential_sigmas(): +@pytest.mark.parametrize("spacing", ["karras", "exponential"]) +def test_unmaterializable_sigma_spacing_is_rejected(spacing): wf = json.loads(json.dumps(_DEFAULT_TXT2IMG)) - wf["3"]["inputs"]["scheduler"] = "exponential" - parsed = parse_comfyui_workflow(wf) - assert parsed.scheduler_spacing == "exponential" - assert ( - parsed.metadata["pipeline"]["strategy"]["scheduler_config"]["use_exponential_sigmas"] - is True - ) + wf["3"]["inputs"]["scheduler"] = spacing + # The published workflow ships the sigma table as a constant component, so a + # spacing mobius cannot materialize has to fail rather than be dropped: the + # old document merely carried ``use_karras_sigmas`` as a hint for a runtime + # scheduler that no longer exists. + with pytest.raises(ValueError, match="Karras or exponential sigmas"): + parse_comfyui_workflow(wf) -def test_denoise_less_than_one_sets_start_step(): +def test_denoise_less_than_one_shortens_the_schedule(): wf = json.loads(json.dumps(_DEFAULT_TXT2IMG)) wf["3"]["inputs"]["denoise"] = 0.5 # steps=20 -> start_step = 20 - round(10) = 10 parsed = parse_comfyui_workflow(wf) assert parsed.denoise == pytest.approx(0.5) assert parsed.start_step == 10 - assert parsed.metadata["pipeline"]["strategy"]["start_step"] == 10 + # img2img skips the noisiest steps, which lowers to a sliced schedule. + inputs = parsed.metadata["pipeline"]["workflow"]["inputs"] + assert inputs["request.max_iterations"]["default"] == 10 -def test_denoise_one_omits_start_step(): - meta = translate_comfyui_workflow(_DEFAULT_TXT2IMG) - assert "start_step" not in meta["pipeline"]["strategy"] +def test_denoise_one_runs_every_step(): + workflow = translate_comfyui_workflow(_DEFAULT_TXT2IMG)["pipeline"]["workflow"] + assert workflow["inputs"]["request.max_iterations"]["default"] == 20 def test_cfg_one_disables_guidance(): wf = json.loads(json.dumps(_DEFAULT_TXT2IMG)) wf["3"]["inputs"]["cfg"] = 1.0 - meta = translate_comfyui_workflow(wf) - assert "guidance_scale" not in meta["pipeline"]["strategy"] + workflow = translate_comfyui_workflow(wf)["pipeline"]["workflow"] + assert "guidance_combine" not in workflow["components"] + assert "request.guidance_scale" not in workflow["inputs"] def test_prompt_wrapper_is_accepted(): meta = translate_comfyui_workflow({"prompt": _DEFAULT_TXT2IMG}) - assert meta["pipeline"]["strategy"]["num_steps"] == 20 + workflow = meta["pipeline"]["workflow"] + assert workflow["inputs"]["request.max_iterations"]["default"] == 20 def test_unsupported_sampler_rejected(): @@ -280,54 +288,32 @@ def test_multiple_samplers_rejected(): translate_comfyui_workflow(wf) -def test_no_vae_no_text_encoder_denoiser_only(): - # A latent-only graph (no VAEDecode / no CLIPTextEncode) yields a denoiser-only - # pipeline. +def test_latent_only_graph_is_rejected(): + # A graph with no VAEDecode produces latents, and the published workflow + # terminates in a decoded image output, so there is nothing to declare. wf = { "3": { "class_type": "KSampler", "inputs": {"steps": 5, "cfg": 1.0, "sampler_name": "euler", "scheduler": "normal"}, }, } - meta = translate_comfyui_workflow(wf) - models = meta["pipeline"]["models"] - assert "denoiser" in models and "vae" not in models and "text_encoder" not in models + with pytest.raises(ValueError, match="no VAE decode node"): + translate_comfyui_workflow(wf) def test_translate_from_file(tmp_path): path = tmp_path / "workflow.json" path.write_text(json.dumps(_DEFAULT_TXT2IMG)) meta = translate_comfyui_workflow_file(str(path)) - assert meta["pipeline"]["strategy"]["num_steps"] == 20 - - -def _onnx_genai_schema_path() -> str | None: - import os - - candidates = [ - os.environ.get("ONNX_GENAI_SCHEMA"), - os.path.join( - os.path.dirname(__file__), - "../../../../../onnx-genai/schema/inference_metadata.schema.json", - ), - os.path.expanduser( - "~/Documents/GitHub/onnx-genai/schema/inference_metadata.schema.json" - ), - ] - for candidate in candidates: - if candidate and os.path.exists(candidate): - return candidate - return None + workflow = meta["pipeline"]["workflow"] + assert workflow["inputs"]["request.max_iterations"]["default"] == 20 def test_translated_metadata_matches_onnx_genai_schema(): """A ComfyUI-translated pipeline validates against onnx-genai's real schema.""" - schema_path = _onnx_genai_schema_path() - if schema_path is None: - pytest.skip("onnx-genai schema not found (set ONNX_GENAI_SCHEMA)") import jsonschema - with open(schema_path) as handle: + with open(_onnx_genai_schema_path()) as handle: schema = json.load(handle) meta = translate_comfyui_workflow(_DEFAULT_TXT2IMG) jsonschema.validate(instance=meta, schema=schema) diff --git a/src/mobius/integrations/onnx_genai/convert.py b/src/mobius/integrations/onnx_genai/convert.py index fa72671e7..a5a0074f6 100644 --- a/src/mobius/integrations/onnx_genai/convert.py +++ b/src/mobius/integrations/onnx_genai/convert.py @@ -38,6 +38,7 @@ SchedulerConfig, build_diffusion_pipeline_metadata, load_diffusers_scheduler_config, + load_diffusers_vae_scaling_factor, ) _LOGGER = logging.getLogger(__name__) @@ -85,6 +86,8 @@ def build_pipeline_metadata_for_workflow( *, sdxl: bool = False, timesteps: list[float] | None = None, + vae_scaling_factor: float | None = None, + package: Any | None = None, ) -> dict[str, Any]: """Reconcile a parsed workflow with a scheduler config into pipeline metadata. @@ -92,8 +95,9 @@ def build_pipeline_metadata_for_workflow( *schedule* comes from ``scheduler`` (read from the checkpoint's diffusers config — the ComfyUI JSON never carries betas). """ - has_vae = "vae" in workflow.metadata["pipeline"]["models"] - has_text = "text_encoder" in workflow.metadata["pipeline"]["models"] + components = workflow.metadata["pipeline"]["workflow"]["components"] + has_vae = "vae" in components + has_text = "text_encoder" in components guidance = workflow.cfg if not math.isclose(workflow.cfg, 1.0) else None # SDXL routes two conditioning edges (concatenated hidden states + pooled # text_embeds); its time_ids is an external denoiser input the driver supplies. @@ -114,6 +118,8 @@ def build_pipeline_metadata_for_workflow( vae_latent_input="latent", text_encoder_filename="text_encoder.onnx" if has_text else None, text_encoder_edges=text_encoder_edges, + vae_scaling_factor=vae_scaling_factor, + package=package, ) @@ -209,6 +215,9 @@ def convert_comfyui_workflow( """ parsed_workflow = parse_comfyui_workflow(workflow) os.makedirs(output_dir, exist_ok=True) + from mobius._model_package import ModelPackage + + package = ModelPackage({}) use_karras = parsed_workflow.scheduler_spacing == "karras" use_exponential = parsed_workflow.scheduler_spacing == "exponential" scheduler = _scheduler_for_workflow( @@ -227,9 +236,24 @@ def convert_comfyui_workflow( use_exponential, ) metadata = build_pipeline_metadata_for_workflow( - parsed_workflow, scheduler, sdxl=sdxl, timesteps=timesteps + parsed_workflow, + scheduler, + sdxl=sdxl, + timesteps=timesteps, + # The VAE normalizes its latents, so the decoder input has to be scaled + # back before decoding; the factor lives in the checkpoint, never in the + # ComfyUI JSON. + vae_scaling_factor=( + load_diffusers_vae_scaling_factor(checkpoint_source, revision=revision) + if checkpoint_source + else None + ), + package=package, ) + # The emitted workflow references the sampler policy components as ONNX + # artifacts, so they ship next to the document that declares them. + package.save_policy_components(output_dir) metadata_path = os.path.join(output_dir, "inference_metadata.yaml") with open(metadata_path, "w", encoding="utf-8") as handle: yaml.safe_dump(metadata, handle, sort_keys=False) diff --git a/src/mobius/integrations/onnx_genai/convert_test.py b/src/mobius/integrations/onnx_genai/convert_test.py index 8d291f802..5f4e89522 100644 --- a/src/mobius/integrations/onnx_genai/convert_test.py +++ b/src/mobius/integrations/onnx_genai/convert_test.py @@ -48,17 +48,17 @@ def test_reconciles_workflow_sampler_with_checkpoint_schedule(): wf = parse_comfyui_workflow(_WF) fake_ts = [float(x) for x in range(25)] meta = build_pipeline_metadata_for_workflow(wf, _SCHEDULER, timesteps=fake_ts) - strat = meta["pipeline"]["strategy"] - # sampler kind/steps/cfg from the ComfyUI graph ... - assert strat["num_steps"] == 25 - assert strat["scheduler_config"]["kind"] == "ddim" - assert strat["guidance_scale"] == pytest.approx(6.5) - assert strat["timesteps"] == fake_ts - # ... betas from the checkpoint (never present in the ComfyUI JSON). - assert strat["scheduler_config"]["beta_start"] == pytest.approx(0.00085) - assert strat["scheduler_config"]["beta_schedule"] == "scaled_linear" - models = meta["pipeline"]["models"] - assert {"denoiser", "vae", "text_encoder"} <= set(models) + workflow = meta["pipeline"]["workflow"] + # sampler steps/cfg from the ComfyUI graph ... + assert workflow["inputs"]["request.max_iterations"]["default"] == 25 + assert workflow["inputs"]["request.guidance_scale"]["default"] == pytest.approx(6.5) + assert {"denoiser", "vae", "text_encoder"} <= set(workflow["components"]) + # ... and the checkpoint-derived timestep table (never present in the + # ComfyUI JSON) is materialized as a constant component rather than a + # scheduler_config block the runtime would have to interpret. + assert workflow["components"]["diffusion_timesteps"]["implementation"]["kind"] == "onnx" + # DDIM consumes the latent unscaled; only Euler pre-scales it. + assert "model_input_scale" not in workflow["components"] def test_run_params_capture_prompt_and_dims(): @@ -74,15 +74,25 @@ def test_cfg_one_reconciles_without_guidance(): wf_json["3"]["inputs"]["cfg"] = 1.0 wf = parse_comfyui_workflow(wf_json) meta = build_pipeline_metadata_for_workflow(wf, _SCHEDULER) - assert "guidance_scale" not in meta["pipeline"]["strategy"] + workflow = meta["pipeline"]["workflow"] + assert "guidance_combine" not in workflow["components"] + assert "request.guidance_scale" not in workflow["inputs"] def test_sdxl_exported_routes_dual_conditioning(): wf = parse_comfyui_workflow(_WF) meta = build_pipeline_metadata_for_workflow(wf, _SCHEDULER, sdxl=True) - flow = meta["pipeline"]["dataflow"] - assert { - "from": "text_encoder.encoder_hidden_states", - "to": "denoiser.encoder_hidden_states", - } in flow - assert {"from": "text_encoder.text_embeds", "to": "denoiser.text_embeds"} in flow + loop = next( + step for step in meta["pipeline"]["workflow"]["steps"] if step["kind"] == "loop" + ) + encoder = next(step for step in loop["setup"] if step.get("component") == "text_encoder") + assert encoder["outputs"] == { + "encoder_hidden_states": "conditioning.encoder_hidden_states", + "text_embeds": "conditioning.text_embeds", + } + # CFG runs the denoiser twice; the second call carries the positive prompt. + denoiser = [step for step in loop["steps"] if step.get("component") == "denoiser"][-1] + assert denoiser["inputs"]["encoder_hidden_states"] == ( + "conditioning.encoder_hidden_states" + ) + assert denoiser["inputs"]["text_embeds"] == "conditioning.text_embeds" diff --git a/src/mobius/integrations/onnx_genai/decoder_metadata_test.py b/src/mobius/integrations/onnx_genai/decoder_metadata_test.py index 287cbc367..39c9eb1eb 100644 --- a/src/mobius/integrations/onnx_genai/decoder_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/decoder_metadata_test.py @@ -6,7 +6,6 @@ from __future__ import annotations import dataclasses -import os import onnx_ir as ir import pytest @@ -19,22 +18,9 @@ moe_metadata_from_config, write_decoder_metadata, ) - - -def _schema_path() -> str | None: - for c in [ - os.environ.get("ONNX_GENAI_SCHEMA"), - os.path.join( - os.path.dirname(__file__), - "../../../../../onnx-genai/schema/inference_metadata.schema.json", - ), - os.path.expanduser( - "~/Documents/GitHub/onnx-genai/schema/inference_metadata.schema.json" - ), - ]: - if c and os.path.exists(c): - return c - return None +from mobius.integrations.onnx_genai.inference_metadata_test import ( + _onnx_genai_schema_path, +) @dataclasses.dataclass @@ -140,14 +126,11 @@ def test_matches_onnx_genai_schema(self): cfg.topk_group = 2 meta = decoder_metadata_from_config(cfg) - schema_path = _schema_path() - if schema_path is None: - pytest.skip("onnx-genai schema not found") import json import jsonschema - with open(schema_path) as handle: + with open(_onnx_genai_schema_path()) as handle: schema = json.load(handle) jsonschema.validate(instance=meta, schema=schema) diff --git a/src/mobius/integrations/onnx_genai/duplex_workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/duplex_workflow_metadata_test.py index 5613940fa..82ea1a5bc 100644 --- a/src/mobius/integrations/onnx_genai/duplex_workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/duplex_workflow_metadata_test.py @@ -20,7 +20,11 @@ import yaml from mobius._model_package import ModelPackage -from mobius.integrations.onnx_genai.inference_metadata_test import _model, _value +from mobius.integrations.onnx_genai.inference_metadata_test import ( + _model, + _onnx_genai_schema_path, + _value, +) from mobius.integrations.onnx_genai.workflow_metadata import ( build_full_duplex_workflow_metadata, write_full_duplex_workflow_metadata, @@ -325,10 +329,7 @@ def test_duplex_workflow_writes_policy_artifacts(tmp_path) -> None: def test_duplex_workflow_matches_producer_schema() -> None: - schema_path = os.environ.get("ONNX_GENAI_SCHEMA") - if not schema_path or not os.path.isfile(schema_path): - pytest.skip("set ONNX_GENAI_SCHEMA to the producer-contract schema") - with open(schema_path, encoding="utf-8") as handle: + with open(_onnx_genai_schema_path(), encoding="utf-8") as handle: schema = json.load(handle) metadata = build_full_duplex_workflow_metadata(_duplex_package(), _DuplexConfig()) jsonschema.validate(metadata, schema) diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index d7570928e..112983f46 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -5,30 +5,38 @@ Mobius builds the neural components of a diffusion model (denoiser transformer, VAE, and — externally — a text encoder) as separate ONNX graphs, but does not -itself carry a scheduler loop. onnx-genai's *iterative* pipeline supplies that -loop declaratively: given an ``inference_metadata`` document describing the -components, the loop-carried dataflow, a timestep input, a scheduler, and -(optionally) classifier-free guidance, it drives the denoise loop and returns -the decoded output. - -This module produces that document from the component filenames + a scheduler -config. It reads no torch/diffusers state — only plain values — so it is cheap -to unit-test and safe to call anywhere. - -The emitted contract matches onnx-genai's pipeline schema: -``schema/inference_metadata.schema.json`` (kind ``iterative`` with -``denoiser`` / ``num_steps`` / ``timestep_input`` / ``scheduler_config`` / -``cfg_conditioning_input`` and denoiser self-edge loop-carried dataflow). +itself carry a scheduler loop. onnx-genai's pipeline contract supplies that loop +declaratively: given an ``inference_metadata`` document, the runtime executes +the denoise loop and returns the decoded output. + +That contract is a *typed SSA workflow*. onnx-genai's ``PipelineSpec`` has a +single property, ``workflow`` (``crates/onnx-genai-metadata/src/schema/``), so +the sampler is not a ``strategy`` block the runtime interprets but an ordinary +executable component the package ships: the sigma schedule and timestep table +are constant components, the step index is the loop induction value, and +classifier-free guidance is two denoiser invocations plus a combine component. +This module builds that document from the component filenames plus a scheduler +config, materializing the sampler components from mobius's policy library. It +reads no torch/diffusers state — only plain values — so it is cheap to unit-test +and safe to call anywhere. + +Not everything here is publishable. :func:`build_native_vlm_package_metadata` +returns mobius's *internal* structural descriptor of a VLM package (a +``models``/``dataflow``/``strategy`` view used to reason about wiring and +validate the executable closure); the published contract for such a package is +the workflow that +:func:`~mobius.integrations.onnx_genai.workflow_metadata.build_vlm_workflow_metadata` +derives from it. Autoregressive decoder-only LLM metadata (``model.attention`` + ``kv_cache``) lives in the sibling :mod:`mobius.integrations.onnx_genai.decoder_metadata` -module. Composite multimodal pipelines retain those decoder properties while -declaring their encoder, fusion, and decoder execution stages here. +module. """ from __future__ import annotations import dataclasses +import inspect import json import logging import math @@ -1736,7 +1744,17 @@ def build_native_vlm_package_metadata( revision: str | None = None, decoder_metadata: dict[str, Any] | None = None, ) -> dict[str, Any]: - """Emit a native VLM contract by inspecting every component graph. + """Describe a native VLM package by inspecting every component graph. + + This returns mobius's **internal** structural descriptor, not the published + onnx-genai document. Its ``pipeline`` key is a ``models``/``dataflow``/ + ``strategy`` view used to reason about component wiring and to validate the + executable closure; onnx-genai's ``PipelineSpec`` accepts only a typed SSA + ``workflow`` and rejects every other property, so this view is consumed by + :func:`~mobius.integrations.onnx_genai.workflow_metadata.build_vlm_workflow_metadata` + (which publishes the workflow and the ``preprocessing`` program) rather than + written to disk. Use :func:`write_native_vlm_package_metadata` to emit the + package contract. Processor selection is registry-driven from graph rank/dtype/shape signatures. No model type, architecture name, or model-name branch @@ -2025,23 +2043,22 @@ def write_native_vlm_package_metadata( config: Any, source: str | None = None, revision: str | None = None, - filename: str = "inference_metadata.yaml", ) -> dict[str, str]: - """Write native VLM metadata and the runtime's tokenizer/processor assets.""" - from mobius.integrations.onnx_genai.decoder_metadata import ( - decoder_metadata_from_config, + """Write the published VLM package contract and its tokenizer/processor assets. + + What lands in ``inference_metadata.yaml`` is the typed SSA workflow built by + :func:`~mobius.integrations.onnx_genai.workflow_metadata.build_vlm_workflow_metadata`, + not the structural descriptor :func:`build_native_vlm_package_metadata` + returns. onnx-genai's ``PipelineSpec`` has exactly one property + (``workflow``) and forbids anything else, so the descriptor's + ``models``/``dataflow``/``strategy`` view is mobius-internal and is never + published. + """ + from mobius.integrations.onnx_genai.workflow_metadata import ( + write_vlm_workflow_metadata, ) - metadata = build_native_vlm_package_metadata( - pkg, - config=config, - source=source, - decoder_metadata=decoder_metadata_from_config(config), - ) - os.makedirs(directory, exist_ok=True) - path = os.path.join(directory, filename) - with open(path, "w", encoding="utf-8") as handle: - yaml.safe_dump(metadata, handle, sort_keys=False) + path = write_vlm_workflow_metadata(pkg, directory, config, source=source) artifacts = {"inference_metadata": path} artifacts.update(_copy_runtime_assets(directory, source, revision=revision)) return artifacts @@ -2309,6 +2326,99 @@ def load_diffusers_vae_scaling_factor( return float(factor) if factor else None +#: Diffusion solvers mobius can materialize as an ONNX policy component, keyed +#: by the scheduler ``kind`` a diffusers config or a ComfyUI sampler resolves +#: to, with whether the solver consumes the latent pre-scaled by the current +#: sigma. A stochastic sampler injects fresh noise every step and has no +#: deterministic solver here, so it is rejected rather than silently lowered to +#: the closest deterministic one. +_SCHEDULER_SOLVERS: dict[str, tuple[str, bool]] = { + "ddim": ("ddim", False), + "euler": ("euler", True), + "dpmpp_2m": ("multistep", False), +} + +#: Latent axis names shared by the solver, guidance and clamp policy components +#: (``mobius.generation._policy_components._IMAGE_LATENT_DIMS``). The workflow's +#: latent contract uses the same symbols so a declared value and the component +#: port it binds to can never disagree about an axis. +_LATENT_DIMS: tuple[str, ...] = ("batch", "channels", "height", "width") + +_WORKFLOW_DTYPES: dict[str, str] = { + "fp32": "float32", + "fp16": "float16", + "bf16": "bfloat16", +} + + +def _diffusion_solver(scheduler: SchedulerConfig) -> tuple[str, bool]: + """Resolve the solver component and whether it pre-scales the model input.""" + resolved = _SCHEDULER_SOLVERS.get(scheduler.kind) + if resolved is None: + raise ValueError( + f"Cannot publish a diffusion workflow for scheduler kind {scheduler.kind!r}. " + "Why: the runtime executes the sampler as a declared solver component, and " + f"only {sorted(_SCHEDULER_SOLVERS)} have one; a stochastic sampler injects " + "fresh noise per step and has no deterministic equivalent. How to fix: export " + "with a deterministic scheduler (DDIM, Euler or DPMSolverMultistep)." + ) + return resolved + + +def _diffusion_schedule_values( + scheduler: SchedulerConfig, + num_inference_steps: int, + timesteps: list[float] | None, + schedule: list[float] | None, + start_step: int | None, +) -> tuple[list[float], list[float], int]: + """Resolve the solver schedule, the timestep table and the executed step count. + + The schedule is what ``solver_step`` integrates, so it is derived from the + scheduler's own noise schedule rather than invented: a variance-preserving + DDIM solver reads cumulative alphas, a sigma-space Euler / DPM-Solver++ one + reads sigmas. Both come from the same diffusers-compatible derivations the + package exporter uses, so a ComfyUI conversion and a package export of the + same checkpoint describe the same dynamics. + + ``start_step`` is img2img's "skip the noisiest steps": diffusers implements + it by starting from a later entry of the same table, so it lowers to a + sliced schedule plus a smaller loop bound rather than to a separate control + knob the typed workflow has no room for. + """ + from mobius.integrations.onnx_genai.auto_export import ( + _ddim_alpha_schedule, + _diffusion_schedule, + ) + + if timesteps is not None and len(timesteps) != num_inference_steps: + raise ValueError( + f"timesteps has {len(timesteps)} entries but num_inference_steps is " + f"{num_inference_steps}" + ) + if schedule is not None and len(schedule) != num_inference_steps + 1: + raise ValueError( + f"schedule has {len(schedule)} entries but num_inference_steps + 1 is " + f"{num_inference_steps + 1}" + ) + derive = _ddim_alpha_schedule if scheduler.kind == "ddim" else _diffusion_schedule + derived_timesteps, derived_schedule = derive(scheduler, num_inference_steps) + schedule = ( + [float(value) for value in schedule] if schedule is not None else derived_schedule + ) + table = ( + [float(value) for value in timesteps] if timesteps is not None else derived_timesteps + ) + if start_step: + if not 0 < start_step < num_inference_steps: + raise ValueError( + f"start_step ({start_step}) must be in 1..{num_inference_steps - 1}" + ) + schedule = schedule[start_step:] + table = table[start_step:] + return schedule, table, len(table) + + def build_diffusion_pipeline_metadata( *, num_inference_steps: int, @@ -2319,115 +2429,467 @@ def build_diffusion_pipeline_metadata( denoiser_output: str = "noise_pred", scheduler: SchedulerConfig | None = None, timesteps: list[float] | None = None, + schedule: list[float] | None = None, guidance_scale: float | None = None, start_step: int | None = None, vae_filename: str | None = None, vae_latent_input: str = "latent", + vae_output: str = "sample", text_encoder_filename: str | None = None, + text_encoder_input: str = "input_ids", text_encoder_output: str = "last_hidden_state", text_encoder_edges: list[tuple[str, str]] | None = None, + vae_scaling_factor: float | None = None, + activation_dtype: str = "fp32", + package: Any | None = None, ) -> dict[str, Any]: - """Build the onnx-genai ``inference_metadata`` dict for a diffusion pipeline. + """Build the onnx-genai ``inference_metadata`` document for a diffusion pipeline. + + onnx-genai's ``PipelineSpec`` has a single property, ``workflow``: a typed + SSA graph in which the sampler is an ordinary executable component rather + than a ``strategy`` block the runtime interprets. So the denoise loop is + emitted explicitly — the sigma schedule and timestep table are constant + components, the step index is the loop induction value, and classifier-free + guidance is two denoiser invocations plus a combine component. - The denoiser runs an iterative loop: its ``denoiser_output`` (a noise - prediction) is fed back to ``denoiser_sample_input`` each step (a - loop-carried self-edge), the scheduler combines it with the current latent, - the per-step timestep is injected into ``denoiser_timestep_input``, and the - conditioning is supplied on ``denoiser_conditioning_input``. + Only the ONNX components a caller *names* are described by artifact; their + graphs stay authoritative for which ports exist. The solver, schedule, + lookup, guidance and clamp components are built here from mobius's policy + library and are attached to ``package`` so a writer can save them next to + the metadata. Args: - num_inference_steps: Number of denoise steps (``strategy.num_steps``). + num_inference_steps: Number of denoise steps. denoiser_*: Denoiser component filename and I/O port names. - scheduler: Noise-schedule config (defaults to DDIM defaults). - guidance_scale: When set and != 1.0, enables classifier-free guidance - (the conditioning input is zeroed on the unconditional pass). - vae_filename: Optional VAE decoder; runs ``final_only`` on the final - latent (``denoiser_sample_input``). - vae_latent_input: VAE latent input port name. - text_encoder_filename: Optional text encoder; runs ``prompt_only`` and - feeds ``denoiser_conditioning_input``. - text_encoder_output: Text encoder output port name. + scheduler: Noise-schedule config (defaults to DDIM defaults). Its + ``kind`` selects the solver component. + timesteps: Explicit per-step timestep table; a linear ramp otherwise. + schedule: Explicit sigma schedule (``num_inference_steps + 1`` values); + a linear ramp otherwise. + guidance_scale: When set and != 1.0, enables classifier-free guidance. + start_step: img2img skip count; lowered to a sliced schedule. + vae_filename: VAE decoder producing the image output. + text_encoder_filename: Optional text encoder feeding the conditioning. + text_encoder_edges: Every ``(encoder_output, denoiser_input)`` edge to + route; defaults to the single primary conditioning edge. + vae_scaling_factor: The diffusers ``scaling_factor`` the VAE latents are + normalized by; the decoder input is divided by it before decoding. + activation_dtype: Declared dtype of the latent and image values. + package: Optional :class:`~mobius._model_package.ModelPackage` the + generated policy components are attached to. Returns: - A dict with a top-level ``pipeline`` key, ready to serialize to - ``inference_metadata.yaml``. + A dict with ``schema_version`` and a top-level ``pipeline.workflow``. """ + from mobius._model_package import ModelPackage + from mobius.generation import ( + SOLVER_BUILDERS, + build_boolean_not, + build_euler_model_input, + build_guidance_combine, + build_scalar_constant, + build_schedule_constant, + build_schedule_lookup, + build_tensor_clamp, + build_tensor_scale, + build_zeros_like, + ) + from mobius.integrations.onnx_genai.workflow_metadata import ( + _invoke, + _publish_workflow_v1, + _request_aligned, + ) + if num_inference_steps < 1: raise ValueError("num_inference_steps must be >= 1") + if vae_filename is None: + raise ValueError( + "Cannot publish a diffusion workflow without a VAE decoder. Why: the " + "workflow terminates in a decoded image output, so a latent-only pipeline " + "has no executable result to declare. How to fix: pass vae_filename for the " + "decoder the package ships." + ) + if activation_dtype not in _WORKFLOW_DTYPES: + raise ValueError( + f"unsupported diffusion activation dtype {activation_dtype!r}; " + f"expected one of {sorted(_WORKFLOW_DTYPES)}" + ) scheduler = scheduler or SchedulerConfig() + solver, scales_model_input = _diffusion_solver(scheduler) + schedule_values, timestep_values, executed_steps = _diffusion_schedule_values( + scheduler, num_inference_steps, timesteps, schedule, start_step + ) + conditioned = text_encoder_filename is not None + guided = guidance_scale is not None and not math.isclose(guidance_scale, 1.0) + if guided and not conditioned: + raise ValueError("classifier-free guidance requires a text encoder to condition on") + + dtype = _ir_dtype_for(activation_dtype) + pkg = package if package is not None else ModelPackage({}) + solver_builder = SOLVER_BUILDERS[solver] + # A solver that fixes its own latent axis names takes only a dtype. + solver_component = ( + solver_builder(dtype, _LATENT_DIMS) + if "latent_dims" in inspect.signature(solver_builder).parameters + else solver_builder(dtype) + ) + # A multistep solver keeps the previous data estimate; a single-step one + # does not, so only then is a history cell part of the loop. + carries_history = "history" in { + value.name for value in solver_component.model.graph.inputs + } + pkg.add_policy_component("solver_step", solver_component) + pkg.add_policy_component("diffusion_schedule", build_schedule_constant(schedule_values)) + pkg.add_policy_component("diffusion_timesteps", build_schedule_constant(timestep_values)) + pkg.add_policy_component("schedule_lookup", build_schedule_lookup(dtype)) + pkg.add_policy_component("continue_predicate", build_boolean_not()) + if scales_model_input: + pkg.add_policy_component( + "model_input_scale", build_euler_model_input(dtype, _LATENT_DIMS) + ) + if carries_history: + pkg.add_policy_component("history_initializer", build_zeros_like(dtype)) + if guided: + pkg.add_policy_component( + "guidance_combine", build_guidance_combine(dtype, _LATENT_DIMS) + ) + pkg.add_policy_component( + "image_output_clamp", + build_tensor_clamp(dtype, _LATENT_DIMS, minimum=-1.0, maximum=1.0), + ) + # A sigma-space sampler starts from noise scaled by the largest sigma; a + # variance-preserving one starts from the unit-variance draw itself. And a + # VAE whose latents are normalized needs them un-normalized before decoding. + # Only emit the constant and the multiply the pipeline actually performs. + initial_state_scale = schedule_values[0] if scales_model_input else 1.0 + decoder_input_scale = 1.0 / vae_scaling_factor if vae_scaling_factor else 1.0 + scales_initial_state = not math.isclose(initial_state_scale, 1.0) + scales_decoder_input = not math.isclose(decoder_input_scale, 1.0) + if scales_initial_state or scales_decoder_input: + pkg.add_policy_component("tensor_scale", build_tensor_scale(dtype)) + if scales_initial_state: + pkg.add_policy_component( + "initial_state_scale", build_scalar_constant(initial_state_scale) + ) + if scales_decoder_input: + pkg.add_policy_component( + "decoder_input_scale", build_scalar_constant(decoder_input_scale) + ) + + workflow_dtype = _WORKFLOW_DTYPES[activation_dtype] + latent_contract = _request_aligned( + {"dtype": workflow_dtype, "rank": 4, "shape": list(_LATENT_DIMS)} + ) + row_float = _request_aligned({"dtype": "float32", "rank": 1, "shape": ["batch"]}) + batch_bool = _request_aligned({"dtype": "bool", "rank": 1, "shape": ["batch"]}) + prompt_contract = _request_aligned( + {"dtype": "int64", "rank": 2, "shape": ["batch", "prompt_sequence"]} + ) - models: dict[str, Any] = { - "denoiser": {"filename": denoiser_filename, "type": "denoiser"}, + inputs: dict[str, Any] = { + "request.max_iterations": { + "contract": {"dtype": "int64", "rank": 1, "shape": [1]}, + "role": {"kind": "runtime", "version": "1.0", "role": "max_iterations"}, + "source": {"kind": "request", "field": "max_iterations"}, + "required": False, + "default": executed_steps, + }, + "package.false": { + "contract": batch_bool, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": False, + }, + "request.noise": { + "contract": latent_contract, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "noise"}, + "required": True, + "externally_suppliable": True, + }, } - dataflow: list[dict[str, Any]] = [ - # Loop-carried self-edge: previous step's prediction seeds the next. - { - "from": f"denoiser.{denoiser_output}", - "to": f"denoiser.{denoiser_sample_input}", + outputs: dict[str, Any] = { + "image": { + "contract": latent_contract, + "role": "image", + "value_range": "negative_one_to_one", + "stage": "pre_adapter", + }, + "latent": { + "contract": latent_contract, + "role": "tensor", + "stage": "pre_adapter", }, + } + + components: dict[str, Any] = { + "denoiser": {"implementation": {"kind": "onnx", "artifact": denoiser_filename}}, + "vae": {"implementation": {"kind": "onnx", "artifact": vae_filename}}, + } + + setup_nodes: list[dict[str, Any]] = [ + _invoke("diffusion_schedule", {}, {"schedule": "diffusion.schedule"}), + _invoke("diffusion_timesteps", {}, {"schedule": "diffusion.timesteps"}), ] - phases: dict[str, Any] = {} + if scales_initial_state: + setup_nodes.append( + _invoke("initial_state_scale", {}, {"value": "diffusion.initial_scale"}) + ) + if scales_decoder_input: + setup_nodes.append( + _invoke("decoder_input_scale", {}, {"value": "diffusion.decoder_scale"}) + ) + initial_state_value = "request.noise" + if scales_initial_state: + initial_state_value = "diffusion.initial_state" + setup_nodes.append( + _invoke( + "tensor_scale", + {"tensor": "request.noise", "scale": "diffusion.initial_scale"}, + {"scaled": initial_state_value}, + ) + ) - if text_encoder_filename is not None: - models["text_encoder"] = { - "filename": text_encoder_filename, - "type": "encoder", + # Each (encoder_output, denoiser_input) edge becomes one SSA value routed + # from the text encoder into the denoiser. SDXL routes two (concatenated + # hidden states + pooled text_embeds); SD routes one. + edges = list(text_encoder_edges or [(text_encoder_output, denoiser_conditioning_input)]) + conditional_values: dict[str, str] = {} + unconditional_values: dict[str, str] = {} + if conditioned: + components["text_encoder"] = { + "implementation": {"kind": "onnx", "artifact": text_encoder_filename} + } + inputs["request.prompt_tokens"] = { + "contract": prompt_contract, + "role": {"kind": "runtime", "version": "1.0", "role": "prompt_tokens"}, + "source": {"kind": "request", "field": "prompt_tokens"}, + "required": True, + "externally_suppliable": True, + } + conditional_values = { + denoiser_in: f"conditioning.{denoiser_in}" for _, denoiser_in in edges } - # Route each text-encoder output to its denoiser conditioning input. SD - # has one edge (hidden states -> encoder_hidden_states); SDXL has two - # (concatenated hidden states + pooled text_embeds). `time_ids` is not - # routed here — it is an external denoiser input the caller supplies. - edges = text_encoder_edges or [(text_encoder_output, denoiser_conditioning_input)] - for enc_out, denoiser_in in edges: - dataflow.append( - {"from": f"text_encoder.{enc_out}", "to": f"denoiser.{denoiser_in}"} + setup_nodes.append( + _invoke( + "text_encoder", + {text_encoder_input: "request.prompt_tokens"}, + { + encoder_out: conditional_values[denoiser_in] + for encoder_out, denoiser_in in edges + }, + ) + ) + if guided: + assert guidance_scale is not None + inputs["request.negative_prompt_tokens"] = { + "contract": prompt_contract, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "negative_prompt_tokens", + }, + "source": {"kind": "request", "field": "negative_prompt_tokens"}, + "required": True, + "externally_suppliable": True, + } + inputs["request.guidance_scale"] = { + "contract": row_float, + "role": {"kind": "runtime", "version": "1.0", "role": "guidance_scale"}, + "source": {"kind": "request", "field": "guidance_scale"}, + "required": False, + "default": float(guidance_scale), + } + unconditional_values = { + denoiser_in: f"conditioning.unconditional.{denoiser_in}" + for _, denoiser_in in edges + } + setup_nodes.append( + _invoke( + "text_encoder", + {text_encoder_input: "request.negative_prompt_tokens"}, + { + encoder_out: unconditional_values[denoiser_in] + for encoder_out, denoiser_in in edges + }, + ) ) - phases["text_encoder"] = {"run_on": "prompt_only"} - if vae_filename is not None: - models["vae"] = {"filename": vae_filename, "type": "vae"} - # The VAE decodes the final post-scheduler latent (the sample port). - dataflow.append( + state: dict[str, Any] = { + "latent": { + "contract": latent_contract, + "scope": "invocation", + "initializer": initial_state_value, + "recurrence": {"kind": "invariant"}, + } + } + carried: list[dict[str, Any]] = [ + { + "cell": "latent", + "current": initial_state_value, + "body_input": "state.latent.body", + "body_output": "latent.body", + "next": "latent.final", + } + ] + if carries_history: + setup_nodes.append( + _invoke( + "history_initializer", + {"reference": initial_state_value}, + {"zeros": "diffusion.initial_history"}, + ) + ) + state["history"] = { + "contract": latent_contract, + "scope": "invocation", + "initializer": "diffusion.initial_history", + "recurrence": {"kind": "invariant"}, + } + carried.append( { - "from": f"denoiser.{denoiser_sample_input}", - "to": f"vae.{vae_latent_input}", + "cell": "history", + "current": "diffusion.initial_history", + "body_input": "state.history.body", + "body_output": "history.body", + "next": "history.final", } ) - phases["vae"] = {"run_on": "final_only"} - - strategy: dict[str, Any] = { - "kind": "iterative", - "denoiser": "denoiser", - "num_steps": num_inference_steps, - "timestep_input": denoiser_timestep_input, - "scheduler_config": scheduler.to_metadata(), - } - if timesteps is not None: - if len(timesteps) != num_inference_steps: - raise ValueError( - f"timesteps has {len(timesteps)} entries but num_inference_steps is " - f"{num_inference_steps}" + setup_nodes.append( + _invoke( + "continue_predicate", {"done": "package.false"}, {"continue": "setup.continue"} + ) + ) + + body_nodes: list[dict[str, Any]] = [ + _invoke( + "schedule_lookup", + {"schedule": "diffusion.timesteps", "step": "loop.iteration"}, + {"timestep": "diffusion.timestep"}, + ) + ] + model_input_value = "state.latent.body" + if scales_model_input: + model_input_value = "diffusion.model_input" + body_nodes.append( + _invoke( + "model_input_scale", + { + "sample": "state.latent.body", + "step": "loop.iteration", + "schedule": "diffusion.schedule", + }, + {"model_input": model_input_value}, ) - strategy["timesteps"] = [float(t) for t in timesteps] - if guidance_scale is not None: - strategy["guidance_scale"] = guidance_scale - if not math.isclose(guidance_scale, 1.0): - strategy["cfg_conditioning_input"] = denoiser_conditioning_input - if start_step: - if not 0 < start_step < num_inference_steps: - raise ValueError( - f"start_step ({start_step}) must be in 1..{num_inference_steps - 1}" + ) + + def denoiser_call(conditioning: dict[str, str], estimate: str) -> dict[str, Any]: + call_inputs = { + denoiser_sample_input: model_input_value, + denoiser_timestep_input: "diffusion.timestep", + **conditioning, + } + return _invoke("denoiser", call_inputs, {denoiser_output: estimate}) + + if guided: + body_nodes.append(denoiser_call(unconditional_values, "denoiser.unconditional")) + body_nodes.append(denoiser_call(conditional_values, "denoiser.conditional")) + body_nodes.append( + _invoke( + "guidance_combine", + { + "unconditional": "denoiser.unconditional", + "conditional": "denoiser.conditional", + "scale": "request.guidance_scale", + }, + {"estimate": "denoiser.estimate"}, ) - strategy["start_step"] = start_step + ) + else: + body_nodes.append(denoiser_call(conditional_values, "denoiser.estimate")) - pipeline: dict[str, Any] = { - "models": models, - "dataflow": dataflow, - "strategy": strategy, + solver_inputs = { + "sample": "state.latent.body", + "step": "loop.iteration", + "schedule": "diffusion.schedule", + "estimate" if carries_history else "derivative": "denoiser.estimate", + } + solver_outputs = {"next_state": "latent.body"} + if carries_history: + solver_inputs["history"] = "state.history.body" + solver_outputs["next_history"] = "history.body" + body_nodes.append(_invoke("solver_step", solver_inputs, solver_outputs)) + body_nodes.append( + _invoke("continue_predicate", {"done": "package.false"}, {"continue": "loop.continue"}) + ) + + decoder_input_value = "latent.final" + tail_nodes: list[dict[str, Any]] = [] + if scales_decoder_input: + decoder_input_value = "diffusion.decoder_input" + tail_nodes.append( + _invoke( + "tensor_scale", + {"tensor": "latent.final", "scale": "diffusion.decoder_scale"}, + {"scaled": decoder_input_value}, + ) + ) + tail_nodes += [ + _invoke("vae", {vae_latent_input: decoder_input_value}, {vae_output: "vae.raw_image"}), + _invoke("image_output_clamp", {"tensor": "vae.raw_image"}, {"clamped": "vae.image"}), + {"kind": "emit", "value": "latent.final", "output": "latent", "mode": "replace"}, + {"kind": "emit", "value": "vae.image", "output": "image", "mode": "replace"}, + ] + + workflow = { + "manifest": { + "capabilities": [ + "workflow_ssa", + "nested_control_flow", + "loop_induction_values", + "typed_emit", + ] + }, + "inputs": inputs, + "outputs": outputs, + "components": components, + "state": state, + "graph": { + "kind": "sequence", + "nodes": [ + { + "kind": "loop", + "setup": {"kind": "sequence", "nodes": setup_nodes}, + "body": {"kind": "sequence", "nodes": body_nodes}, + "condition": "loop.continue", + "max_iterations": "request.max_iterations", + "iteration": { + "value": "loop.iteration", + "contract": {"dtype": "int64", "rank": 1, "shape": ["batch"]}, + }, + "carried": carried, + }, + *tail_nodes, + ], + }, + } + metadata = { + "schema_version": "v1", + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, } - if phases: - pipeline["phases"] = phases - return {"pipeline": pipeline} + add_policy_components_to_workflow(metadata, pkg) + return metadata + + +def _ir_dtype_for(activation_dtype: str) -> Any: + import onnx_ir as ir + + return { + "fp32": ir.DataType.FLOAT, + "fp16": ir.DataType.FLOAT16, + "bf16": ir.DataType.BFLOAT16, + }[activation_dtype] def build_multimodal_pipeline_metadata( @@ -2660,17 +3122,63 @@ def write_diffusion_pipeline_metadata( ) -> str: """Build and write ``inference_metadata.yaml`` into ``directory``. + The generated sampler policy components (solver, schedule constants, + lookup, guidance, clamp) are saved alongside it, because the emitted + workflow references them as ONNX artifacts. + Extra keyword arguments are forwarded to :func:`build_diffusion_pipeline_metadata`. Returns the written path. """ - metadata = build_diffusion_pipeline_metadata(**kwargs) + from mobius._model_package import ModelPackage + + package = ModelPackage({}) + metadata = build_diffusion_pipeline_metadata(package=package, **kwargs) os.makedirs(directory, exist_ok=True) + package.save_policy_components(directory) path = os.path.join(directory, filename) with open(path, "w", encoding="utf-8") as handle: yaml.safe_dump(metadata, handle, sort_keys=False) return path +#: Proposer input port that receives the target's per-token hidden state. +#: +#: ``SpeculativeContract.port_bindings`` is keyed by semantic role; onnx-genai +#: names this role ``target_hidden_context`` +#: (``crates/onnx-genai-metadata/src/schema/package.rs``). +_TARGET_HIDDEN_CONTEXT_ROLE = "target_hidden_context" + +#: Port roles the Qwen3.5/3.8 MTP sidecar exposes. The graph is authoritative +#: for which ports exist and their contracts; only what a port *means* is +#: declared here (see ``workflow_metadata._component``). +_MTP_PROPOSER_PORT_ROLES: dict[str, str] = { + "inputs_embeds": "inputs_embeds", + "hidden_states": "hidden_states", + "attention_mask": "attention_mask", + "position_ids": "position_ids", + "mtp_hidden": "hidden_states", +} + + +def _speculative_target_candidates(workflow: dict[str, Any], exclude: str) -> list[str]: + """Workflow components that can verify a speculative proposal. + + A verifier is the component that scores the next token, so it is selected by + its declared ``logits`` output role. Selecting on ``implementation.kind == + "onnx"`` would not work: the generated policy graphs a workflow ships (the + sampler, the termination predicate, the state updates) are ONNX components + too, and a real decoder package declares about a dozen of them. + """ + return [ + name + for name, declaration in (workflow.get("components") or {}).items() + if name != exclude + and isinstance(declaration, dict) + and declaration.get("implementation", {}).get("kind") == "onnx" + and "logits" in (declaration.get("ports", {}).get("roles") or {}).values() + ] + + def write_mtp_speculator_metadata( directory: str, *, @@ -2680,78 +3188,151 @@ def write_mtp_speculator_metadata( num_speculative_tokens: int = 1, embedding_weights: str = "model.embed_tokens.weight", lm_head_weights: str = "lm_head.weight", + proposer_name: str = "mtp", ) -> str | None: - """Add a ``speculative`` block for the exported MTP head to the backbone metadata. + """Declare the exported MTP head as the backbone's speculative proposer. The Qwen3.5/3.8 MTP head is a self-speculative drafter saved next to the backbone (``mtp/model.onnx``). It borrows the target's shared embedding / - LM head and is seeded by the backbone's final-layer hidden state - (``hidden_states.``). - - The emitted block conforms exactly to the authoritative onnx-genai runtime - schema (``crates/onnx-genai-metadata/src/schema/generation.rs`` - ``SpeculatorConfig`` + ``parser.rs`` ``resolve_mtp`` + - ``config.rs`` ``validate_resolved_mtp_config``): - - - Published under the top-level ``speculative`` key (the runtime - deserializes ``InferenceMetadata.speculative``; a bare ``speculator`` - key is unknown and silently dropped). - - ``model`` (not ``model_path``), ``target_hidden_size`` (not - ``hidden_size``). - - ``kv_mode: proposal_local`` — the only valid ``MtpKvMode`` for this k=1 - head (``hidden_threaded`` is an engine-internal enum, not a metadata one). - - ``embedding`` / ``lm_head`` as ``MtpTargetInitializer`` objects - (``source: target_initializer`` + ``name``), not flat name strings. - - ``target_hidden_layout: BSH`` because ``hidden_states.`` is rank-3 - ``[batch, seq, hidden]``; ``hc_mult: 1`` is required by ``resolve_mtp`` - (``hc_mult > 0``) and pinned to 1 for the BSH layout by - ``validate_resolved_mtp_config``. - - ``mtp_state_output`` is omitted: an ``hc_mult == 1`` head has no recurrent - Hyper-Connection state (the sidecar emits only ``mtp_hidden`` + - ``present.0.{key,value}``); the runtime looks it up optionally. + LM head and is seeded by the backbone's final-layer hidden state. + + The emitted block is onnx-genai's ``SpeculativeContract`` + (``crates/onnx-genai-metadata/src/schema/package.rs``), which replaced the + older flat ``SpeculatorConfig`` block. That contract is expressed entirely + in terms of the *workflow*: ``proposer`` and ``target`` name workflow + components, ``rollback_state`` names workflow state cells, and the hidden + handoff is a ``port_bindings`` role rather than a pair of free-form port + names. So this writer also registers the head as a workflow component and + completes the rollback capabilities the claim depends on: + + - ``proposer`` / ``target`` are workflow component names. The backbone + metadata must already declare a ``pipeline.workflow`` with exactly one + ONNX component (the decoder) to anchor against; anything else fails + closed rather than publishing an unanchored contract. + - ``proposal_execution: block``. The head returns its complete proposal in + one invocation. It is deliberately *not* declared ``chained``: a chained + proposer must expose a ``logits_output`` carrying the next-token + distribution, and this sidecar emits only ``mtp_hidden`` — the runtime + obtains draft logits by decoding it through the target's LM head, which + is why that initializer is listed in ``shared_weights``. + - ``port_bindings.target_hidden_context`` names the proposer input port the + target's per-token hidden state lands in. The *source* of that value is + declared structurally, as a ``hidden_states`` port role on the target + component, instead of being spelled out as a layer-indexed port name. + - ``vocabulary: identical`` — the head shares the target's LM head, so it + scores the target's own vocabulary axis. + - ``rollback_state`` lists the target's service-group-backed state cells, + and each reached group is given a ``rollback_positions`` bound of at + least ``max_proposal_width``: a rejected proposal must be undoable that + far or the runtime rejects the package. The backbone ``inference_metadata.yaml`` must already exist. Returns the metadata path, or ``None`` when it is missing. """ + if num_speculative_tokens < 1: + raise ValueError("num_speculative_tokens must be >= 1") path = os.path.join(directory, filename) if not os.path.isfile(path): return None with open(path, encoding="utf-8") as handle: metadata = yaml.safe_load(handle) or {} + workflow = metadata.get("pipeline", {}).get("workflow") + if not isinstance(workflow, dict) or not workflow.get("components"): + raise ValueError( + "Cannot declare an MTP speculator: the backbone " + f"{filename!r} has no pipeline.workflow components. Why: a " + "SpeculativeContract names its proposer and target as workflow " + "components, so there is nothing to anchor the claim to. How to " + "fix: write the backbone workflow metadata (write_onnx_genai_config) " + "before attaching the MTP head." + ) + candidates = _speculative_target_candidates(workflow, exclude=proposer_name) + if len(candidates) != 1: + raise ValueError( + "Cannot declare an MTP speculator: expected exactly one workflow " + "component declaring a 'logits' output role to verify proposals, found " + f"{sorted(candidates)}. Why: the speculative target must be unambiguous. " + "How to fix: pass a backbone package whose workflow declares a single " + "decoder component." + ) + target_name = candidates[0] + + # The head borrows the target's embedding + LM head, which live as + # initializers in the backbone ``model.onnx`` (borrowed, not duplicated). + shared_weights = sorted({embedding_weights, lm_head_weights}) + + components = workflow["components"] + components[proposer_name] = { + "implementation": {"kind": "onnx", "artifact": model_path}, + "ports": {"roles": dict(_MTP_PROPOSER_PORT_ROLES)}, + } + # Name the target output the head is seeded from. The backbone exposes its + # final-layer hidden state as ``hidden_states.``; the role, not the + # spelling, is what the runtime resolves. num_layers = getattr(backbone_config, "num_hidden_layers", None) - hidden_size = getattr(backbone_config, "hidden_size", None) - vocab_size = getattr(backbone_config, "vocab_size", None) - target_hidden_output = ( - f"hidden_states.{int(num_layers) - 1}" if num_layers is not None else None + if num_layers is not None: + target_roles = components[target_name].setdefault("ports", {}).setdefault("roles", {}) + target_roles[f"hidden_states.{int(num_layers) - 1}"] = "hidden_states" + + # A rejected proposal must rewind every state cell the target advances. + rollback_cells = sorted( + cell + for cell, declaration in (workflow.get("state") or {}).items() + if isinstance(declaration, dict) and declaration.get("service_group") ) - - speculator: dict[str, Any] = { - "proposal_type": "mtp", - "num_speculative_tokens": int(num_speculative_tokens), - "model": model_path, - # rank-3 [batch, seq, hidden] seed => BSH layout with a single - # Hyper-Connection lane (hc_mult == 1). - "target_hidden_layout": "BSH", - "hc_mult": 1, - # The head threads its final hidden state forward and shares the - # target's embedding + LM head; k=1 head resets KV each verify step. - "mtp_hidden_output": "mtp_hidden", - "kv_mode": "proposal_local", - # The MTP head reuses the backbone's shared embedding + LM head, which - # live as initializers in the main ``model.onnx`` (borrowed, not - # duplicated). - "embedding": {"source": "target_initializer", "name": embedding_weights}, - "lm_head": {"source": "target_initializer", "name": lm_head_weights}, + _declare_rollback_capacity(workflow, rollback_cells, int(num_speculative_tokens)) + + metadata["speculative"] = { + "proposer": proposer_name, + "target": target_name, + # One invocation yields the whole k-token proposal. + "proposal_execution": {"kind": "block"}, + "port_bindings": {_TARGET_HIDDEN_CONTEXT_ROLE: "hidden_states"}, + "shared_weights": shared_weights, + # The head scores the target's own vocabulary through the shared LM head. + "vocabulary": {"kind": "identical"}, + "max_proposal_width": int(num_speculative_tokens), + # Standard rejection sampling against the target keeps the target's + # output distribution exact. + "distribution_preserving": True, + **({"rollback_state": rollback_cells} if rollback_cells else {}), } - if target_hidden_output is not None: - speculator["target_hidden_output"] = target_hidden_output - if hidden_size is not None: - speculator["target_hidden_size"] = int(hidden_size) - if vocab_size is not None: - speculator["vocab_size"] = int(vocab_size) - - metadata["speculative"] = speculator with open(path, "w", encoding="utf-8") as handle: yaml.safe_dump(metadata, handle, sort_keys=False) return path + + +def _declare_rollback_capacity( + workflow: dict[str, Any], + cells: Sequence[str], + positions: int, +) -> None: + """Guarantee every group reached by ``cells`` can rewind ``positions``. + + A speculative package is rejected when a rolled-back cell resolves to a + state group that declares no ``rollback_positions``, or fewer than the + declared maximum proposal width. Attaching the speculator is what creates + that requirement, so it is also what states the bound. + """ + groups = (workflow.get("serving") or {}).get("state_service", {}).get("groups", {}) + state = workflow.get("state") or {} + pending = [ + state[cell]["service_group"] + for cell in cells + if isinstance(state.get(cell), dict) and state[cell].get("service_group") + ] + seen: set[str] = set() + while pending: + name = pending.pop() + if name in seen: + continue + seen.add(name) + contract = groups.get(name) + if not isinstance(contract, dict): + continue + capabilities = contract.setdefault("capabilities", {}) + declared = capabilities.get("rollback_positions") + if not isinstance(declared, int) or declared < positions: + capabilities["rollback_positions"] = positions + pending.extend(capabilities.get("cascade") or []) diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index 8b20d1e64..20dde82fa 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -9,6 +9,7 @@ import json import math import os +import re from pathlib import Path import jsonschema @@ -41,6 +42,9 @@ write_mtp_speculator_metadata, write_native_vlm_package_metadata, ) +from mobius.integrations.onnx_genai.workflow_metadata import ( + build_vlm_workflow_metadata, +) def test_ort_extensions_processor_config_supplies_structural_values(tmp_path): @@ -179,23 +183,22 @@ def test_workflow_policy_components_reference_saved_onnx_artifacts(tmp_path): assert (tmp_path / component["implementation"]["artifact"]).is_file() -def _onnx_genai_schema_path() -> str | None: - """Locate onnx-genai's committed pipeline JSON schema, if available.""" - candidates = [ - os.environ.get("ONNX_GENAI_SCHEMA"), - os.path.join( - os.path.dirname(__file__), - "../../../../../onnx-genai/schema/inference_metadata.schema.json", - ), - "/home/justinchu/onnx-genai/schema/inference_metadata.schema.json", - os.path.expanduser( - "~/Documents/GitHub/onnx-genai/schema/inference_metadata.schema.json" - ), - ] - for candidate in candidates: - if candidate and os.path.exists(candidate): - return candidate - return None +def _onnx_genai_schema_path() -> str: + """Locate onnx-genai's published pipeline JSON schema. + + The vendored copy under ``_schema/`` is the default so this never skips: + the previous behaviour — search a few local onnx-genai checkouts and + ``pytest.skip`` when none is found — meant these tests were silently + inactive in CI, which is how two upstream contract redesigns went unnoticed. + A developer checkout is deliberately *not* consulted implicitly either, + because a clone that is ahead of or behind ``main`` would make the result + machine-dependent in exactly the same way. Set ``ONNX_GENAI_SCHEMA`` to + validate against a specific revision. + """ + override = os.environ.get("ONNX_GENAI_SCHEMA") + if override: + return override + return os.path.join(os.path.dirname(__file__), "_schema", "inference_metadata.schema.json") def _value( @@ -575,12 +578,21 @@ def _assert_all_graph_ports_declared( class TestNativeVlmPackageMetadata: - def _validate(self, metadata: dict) -> None: - schema_path = _onnx_genai_schema_path() - if schema_path is None: - pytest.skip("onnx-genai schema not found (set ONNX_GENAI_SCHEMA)") - with open(schema_path, encoding="utf-8") as handle: - jsonschema.validate(instance=metadata, schema=json.load(handle)) + def _validate(self, package, config, source=None) -> None: + """Validate the *published* package contract against onnx-genai's schema. + + ``build_native_vlm_package_metadata`` returns mobius's internal + structural descriptor; what a package publishes is the typed SSA + workflow ``build_vlm_workflow_metadata`` derives from it. onnx-genai's + ``PipelineSpec`` has ``workflow`` as its only property, so the + descriptor's ``models``/``dataflow``/``strategy`` view is not a + publishable document and validating it against the schema would assert + the wrong contract. + """ + published = build_vlm_workflow_metadata(package, config, source=source) + assert set(published["pipeline"]) == {"workflow"} + with open(_onnx_genai_schema_path(), encoding="utf-8") as handle: + jsonschema.validate(instance=published, schema=json.load(handle)) def test_gemma4_routes_all_embedding_outputs(self, tmp_path): config = _VlmConfig( @@ -706,7 +718,7 @@ def test_gemma4_routes_all_embedding_outputs(self, tmp_path): ) emitted_yaml = yaml.safe_load(yaml.safe_dump(metadata, sort_keys=False)) validate_executable_closure(package, metadata) - self._validate(metadata) + self._validate(package, config, source=str(source)) _assert_all_graph_ports_declared(package, metadata) assert metadata["schema_version"] == "v1" assert { @@ -907,7 +919,7 @@ def test_qwen_packed_grid_rank3_positions_sparse_and_fixed_state(self, tmp_path) metadata = build_native_vlm_package_metadata( package, config=config, source=str(source) ) - self._validate(metadata) + self._validate(package, config, source=str(source)) _assert_all_graph_ports_declared(package, metadata) assert { "position_program", @@ -1051,7 +1063,7 @@ def test_phi_routes_both_modality_gates_and_mask_processor(self, tmp_path): metadata = build_native_vlm_package_metadata( package, config=config, source=str(source) ) - self._validate(metadata) + self._validate(package, config, source=str(source)) _assert_all_graph_ports_declared(package, metadata) flow = metadata["pipeline"]["dataflow"] for gate in ("vision_gate", "speech_gate"): @@ -1532,47 +1544,131 @@ def test_writer_copies_local_runtime_assets(self, tmp_path): ) +def _constant_values(component) -> list[float]: + """Read the constant a schedule policy component materializes.""" + initializer = next(node for node in component.model.graph if node.op_type == "Constant") + return initializer.attributes["value"].as_tensor().numpy().tolist() + + +def _invocations(workflow: dict, component: str) -> list[dict]: + """Every ``invoke`` step of ``component``, at any nesting depth.""" + found: list[dict] = [] + + def walk(step: dict) -> None: + kind = step["kind"] + if kind == "invoke": + if step["component"] == component: + found.append(step) + elif kind == "loop": + for child in (*step.get("setup", []), *step["steps"]): + walk(child) + elif kind == "sequence": + for child in step["steps"]: + walk(child) + elif kind == "branch": + for case in step["cases"].values(): + walk(case) + if step.get("default"): + walk(step["default"]) + + for step in workflow["steps"]: + walk(step) + return found + + class TestBuildDiffusionPipelineMetadata: - def test_denoiser_only_minimal(self): - meta = build_diffusion_pipeline_metadata(num_inference_steps=20) - pipe = meta["pipeline"] - assert set(pipe["models"]) == {"denoiser"} - assert pipe["strategy"]["kind"] == "iterative" - assert pipe["strategy"]["denoiser"] == "denoiser" - assert pipe["strategy"]["num_steps"] == 20 - assert pipe["strategy"]["timestep_input"] == "timestep" - # Loop-carried self-edge is present. - assert {"from": "denoiser.noise_pred", "to": "denoiser.sample"} in pipe["dataflow"] - # Default DDIM scheduler config. - sched = pipe["strategy"]["scheduler_config"] - assert sched["kind"] == "ddim" - assert sched["num_train_timesteps"] == 1000 + """The published document is a typed SSA workflow, not an ``iterative`` strategy. - def test_full_pipeline_with_vae_and_text_encoder(self): + onnx-genai's ``PipelineSpec`` declares ``workflow`` as its only property + (``crates/onnx-genai-metadata/src/schema/pipeline.rs``), so the sampler is + an executable component the package ships rather than a ``strategy`` block + the runtime interprets. + """ + + def _workflow(self, **kwargs) -> dict: meta = build_diffusion_pipeline_metadata( - num_inference_steps=4, vae_filename="vae.onnx", + **kwargs, + ) + assert set(meta["pipeline"]) == {"workflow"} + return meta["pipeline"]["workflow"] + + def test_a_latent_only_pipeline_is_not_publishable(self): + # The workflow terminates in a decoded image, so a package with no + # decoder has no executable result to declare. + with pytest.raises(ValueError, match="without a VAE decoder"): + build_diffusion_pipeline_metadata(num_inference_steps=20) + + def test_denoise_loop_carries_the_latent_through_the_solver(self): + workflow = self._workflow(num_inference_steps=20) + assert set(workflow["components"]) >= { + "denoiser", + "vae", + "solver_step", + "diffusion_schedule", + "diffusion_timesteps", + "schedule_lookup", + } + loop = next(step for step in workflow["steps"] if step["kind"] == "loop") + assert loop["max_iterations"] == "request.max_iterations" + assert workflow["inputs"]["request.max_iterations"]["default"] == 20 + # The loop-carried latent replaces the old denoiser output self-edge. + # ``latent`` is also a workflow output, so the published cell is + # disambiguated to ``latent_state``. + latent_carry = next( + carry for carry in loop["carried"] if carry["cell"] == "latent_state" + ) + assert latent_carry["next"] == "latent.body" + solver = _invocations(workflow, "solver_step")[0] + assert solver["inputs"]["sample"] == "latent_state" + assert solver["inputs"]["derivative"] == "denoiser.estimate" + assert solver["outputs"]["next_state"] == "latent.body" + # The step index is the loop induction value, not a timestep port name. + assert loop["iteration"]["value"] == "loop.iteration" + lookup = _invocations(workflow, "schedule_lookup")[0] + assert lookup["inputs"] == { + "schedule": "diffusion.timesteps", + "step": "loop.iteration", + } + + def test_full_pipeline_with_vae_and_text_encoder(self): + workflow = self._workflow( + num_inference_steps=4, text_encoder_filename="text_encoder.onnx", guidance_scale=7.5, ) - pipe = meta["pipeline"] - assert set(pipe["models"]) == {"denoiser", "vae", "text_encoder"} - # Text encoder feeds conditioning (prompt-phase); VAE decodes final latent. - assert { - "from": "text_encoder.last_hidden_state", - "to": "denoiser.encoder_hidden_states", - } in pipe["dataflow"] - assert {"from": "denoiser.sample", "to": "vae.latent"} in pipe["dataflow"] - assert pipe["phases"]["text_encoder"] == {"run_on": "prompt_only"} - assert pipe["phases"]["vae"] == {"run_on": "final_only"} - # CFG enabled -> conditioning input declared for the unconditional pass. - assert pipe["strategy"]["guidance_scale"] == pytest.approx(7.5) - assert pipe["strategy"]["cfg_conditioning_input"] == "encoder_hidden_states" + assert set(workflow["components"]) >= { + "denoiser", + "vae", + "text_encoder", + "guidance_combine", + } + # The text encoder runs once in the loop setup (the old "prompt_only" + # phase); the VAE decodes the final latent after the loop. + loop = next(step for step in workflow["steps"] if step["kind"] == "loop") + encoder_calls = [ + step for step in loop["setup"] if step.get("component") == "text_encoder" + ] + assert len(encoder_calls) == 2 # conditional + unconditional + assert encoder_calls[0]["outputs"]["last_hidden_state"] == ( + "conditioning.encoder_hidden_states" + ) + vae = _invocations(workflow, "vae")[0] + assert vae["inputs"]["latent"] == "latent_state" + # CFG is two denoiser invocations plus a combine, not a hidden flag. + denoiser_calls = _invocations(workflow, "denoiser") + assert len(denoiser_calls) == 2 + assert denoiser_calls[0]["inputs"]["encoder_hidden_states"] == ( + "conditioning.unconditional.encoder_hidden_states" + ) + assert denoiser_calls[1]["inputs"]["encoder_hidden_states"] == ( + "conditioning.encoder_hidden_states" + ) + assert workflow["inputs"]["request.guidance_scale"]["default"] == pytest.approx(7.5) def test_sdxl_dual_conditioning_edges(self): - meta = build_diffusion_pipeline_metadata( + workflow = self._workflow( num_inference_steps=4, - vae_filename="vae.onnx", text_encoder_filename="text_encoder.onnx", guidance_scale=7.5, text_encoder_edges=[ @@ -1580,16 +1676,86 @@ def test_sdxl_dual_conditioning_edges(self): ("text_embeds", "text_embeds"), ], ) - flow = meta["pipeline"]["dataflow"] - assert { - "from": "text_encoder.encoder_hidden_states", - "to": "denoiser.encoder_hidden_states", - } in flow - assert {"from": "text_encoder.text_embeds", "to": "denoiser.text_embeds"} in flow + loop = next(step for step in workflow["steps"] if step["kind"] == "loop") + encoder = next( + step for step in loop["setup"] if step.get("component") == "text_encoder" + ) + assert encoder["outputs"] == { + "encoder_hidden_states": "conditioning.encoder_hidden_states", + "text_embeds": "conditioning.text_embeds", + } + denoiser = _invocations(workflow, "denoiser")[-1] + assert denoiser["inputs"]["encoder_hidden_states"] == ( + "conditioning.encoder_hidden_states" + ) + assert denoiser["inputs"]["text_embeds"] == "conditioning.text_embeds" def test_guidance_scale_one_does_not_enable_cfg(self): - meta = build_diffusion_pipeline_metadata(num_inference_steps=2, guidance_scale=1.0) - assert "cfg_conditioning_input" not in meta["pipeline"]["strategy"] + workflow = self._workflow( + num_inference_steps=2, + text_encoder_filename="text_encoder.onnx", + guidance_scale=1.0, + ) + assert "guidance_combine" not in workflow["components"] + assert "request.guidance_scale" not in workflow["inputs"] + assert len(_invocations(workflow, "denoiser")) == 1 + + def test_start_step_slices_the_schedule(self): + # img2img "skip the noisiest steps" is a shorter schedule, not a knob. + workflow = self._workflow(num_inference_steps=10, start_step=4) + assert workflow["inputs"]["request.max_iterations"]["default"] == 6 + + def test_schedule_is_derived_from_the_scheduler_not_invented(self): + """``solver_step`` integrates the checkpoint's own noise schedule. + + The schedule component is what the solver reads as alpha_cumprod (DDIM) + or sigma (Euler / DPM-Solver++), so a placeholder ramp would silently + denoise along the wrong trajectory. + """ + from mobius._model_package import ModelPackage + from mobius.integrations.onnx_genai.auto_export import _ddim_alpha_schedule + + scheduler = SchedulerConfig(kind="ddim", beta_start=0.0001, beta_end=0.02) + package = ModelPackage({}) + build_diffusion_pipeline_metadata( + num_inference_steps=5, + vae_filename="vae.onnx", + scheduler=scheduler, + package=package, + ) + _, expected = _ddim_alpha_schedule(scheduler, 5) + emitted = _constant_values(package.policy_components["diffusion_schedule"]) + assert emitted == pytest.approx(expected, rel=1e-6) + # Guard the specific regression: a 1 - i/n ramp is not a beta schedule. + assert emitted != pytest.approx([1.0 - index / 5 for index in range(6)]) + + def test_vae_scaling_factor_unnormalizes_the_decoder_input(self): + workflow = self._workflow(num_inference_steps=3, vae_scaling_factor=0.18215) + assert {"tensor_scale", "decoder_input_scale"} <= set(workflow["components"]) + vae = _invocations(workflow, "vae")[0] + assert vae["inputs"]["latent"] == "diffusion.decoder_input" + + def test_variance_preserving_solver_does_not_rescale_the_initial_latent(self): + workflow = self._workflow(num_inference_steps=3) + # DDIM starts from the unit-variance draw itself; only a sigma-space + # sampler scales it by the largest sigma. + assert "initial_state_scale" not in workflow["components"] + assert workflow["state"]["latent_state"]["initializer"] == "request.noise" + + def test_sigma_space_solver_scales_the_initial_latent(self): + workflow = self._workflow( + num_inference_steps=3, scheduler=SchedulerConfig(kind="euler") + ) + assert "initial_state_scale" in workflow["components"] + assert workflow["state"]["latent_state"]["initializer"] == "diffusion.initial_state" + + def test_stochastic_scheduler_is_rejected(self): + with pytest.raises(ValueError, match="euler_ancestral"): + build_diffusion_pipeline_metadata( + num_inference_steps=4, + vae_filename="vae.onnx", + scheduler=SchedulerConfig(kind="euler_ancestral"), + ) def test_scheduler_from_diffusers_config(self): sched = SchedulerConfig.from_diffusers( @@ -1732,19 +1898,21 @@ def test_write_roundtrips_yaml(self, tmp_path): ) with open(path) as handle: loaded = yaml.safe_load(handle) - assert loaded["pipeline"]["strategy"]["num_steps"] == 3 - assert "vae" in loaded["pipeline"]["models"] + workflow = loaded["pipeline"]["workflow"] + assert workflow["inputs"]["request.max_iterations"]["default"] == 3 + assert "vae" in workflow["components"] + # The workflow declares the sampler components as ONNX artifacts, so + # the writer has to ship them next to the document. + solver = workflow["components"]["solver_step"]["implementation"]["artifact"] + assert (tmp_path / solver).is_file() def test_matches_onnx_genai_json_schema(self): """The emitted metadata validates against onnx-genai's published schema.""" - schema_path = _onnx_genai_schema_path() - if schema_path is None: - pytest.skip("onnx-genai schema not found (set ONNX_GENAI_SCHEMA)") import json import jsonschema - with open(schema_path) as handle: + with open(_onnx_genai_schema_path()) as handle: schema = json.load(handle) meta = build_diffusion_pipeline_metadata( num_inference_steps=4, @@ -1890,15 +2058,121 @@ class _MtpBackboneConfig: def _seed_backbone_metadata(directory: Path) -> str: - """Write a minimal backbone inference_metadata.yaml for the MTP writer. - - Deliberately empty: the writer only *appends* a ``speculative`` block, and - the runtime schema rejects unknown top-level properties, so seeding a - convenience key such as ``model_type`` would fail schema validation for a - reason that has nothing to do with the speculator block under test. + """Write a backbone inference_metadata.yaml for the MTP writer to extend. + + ``SpeculativeContract`` is expressed against the workflow — ``proposer`` and + ``target`` are component names and ``rollback_state`` names state cells — so + the backbone must already publish a ``pipeline.workflow`` for the writer to + anchor against. This is the shape ``write_onnx_genai_config`` emits for a + single-component decoder package, reduced to what the speculator claim + touches: one ONNX decoder component and one service-group-backed KV cell. """ + kv_contract = { + "dtype": "float16", + "rank": 4, + "shape": ["batch", "kv_heads", "sequence", "head_dim"], + "batch_layout": {"kind": "request_aligned", "axis": 0}, + } + metadata = { + "schema_version": "v1", + "pipeline": { + "workflow": { + "manifest": {"capabilities": ["workflow_ssa", "serving_service_contract"]}, + "inputs": { + "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.decoder_cache": { + "contract": kv_contract, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "decoder_cache"}, + "required": True, + }, + "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, + }, + }, + "components": { + "decoder": { + "implementation": {"kind": "onnx", "artifact": "model.onnx"}, + "ports": {"roles": {"input_ids": "token_ids", "logits": "logits"}}, + } + }, + "state": { + "decoder_cache.0": { + "contract": kv_contract, + "scope": "invocation", + "initializer": "request.decoder_cache", + "recurrence": {"kind": "invariant"}, + "service_group": "decoder_cache", + } + }, + "steps": [ + { + "kind": "invoke", + "component": "decoder", + "inputs": { + "input_ids": "request.input_ids", + "past_key_values.0.key": "decoder_cache.0", + }, + "outputs": {"logits": "decoder.logits"}, + } + ], + "serving": { + "active": "request.active", + "done": "request.done", + "state_service": { + "groups": { + "decoder_cache": { + "kind": "full_attention", + "sequence_axis": 2, + "layout": "bnsh", + "update": {"kind": "append"}, + "capabilities": {"snapshot": True, "fork": True}, + "ports": { + "decoder": { + "decoder_cache.0": { + "input": "past_key_values.0.key", + "output": "present.0.key", + } + } + }, + } + } + }, + }, + } + }, + } path = directory / "inference_metadata.yaml" - path.write_text(yaml.safe_dump({}), encoding="utf-8") + path.write_text(yaml.safe_dump(metadata, sort_keys=False), encoding="utf-8") return str(path) @@ -1906,10 +2180,12 @@ class TestMtpSpeculatorMetadata: """The emitted ``speculative`` block conforms to the onnx-genai runtime schema. Authoritative source: onnx-genai - ``crates/onnx-genai-metadata/src/schema/generation.rs`` (``SpeculatorConfig``, - ``MtpKvMode``, ``MtpHiddenLayout``, ``MtpTargetInitializer``) + - ``parser.rs`` (``resolve_mtp``) + ``config.rs`` - (``validate_resolved_mtp_config``). + ``crates/onnx-genai-metadata/src/schema/package.rs`` + (``SpeculativeContract``, ``SpeculativeProposalExecution``, + ``SpeculativeVocabulary``) + ``validation.rs`` + (``validate_speculative_rollback``). The older flat ``SpeculatorConfig`` + block in ``schema/generation.rs`` describes a HuggingFace ``config.json`` + speculator section, not ``InferenceMetadata.speculative``. """ def _write(self, tmp_path: Path) -> dict: @@ -1931,39 +2207,103 @@ def test_top_level_key_is_speculative(self, tmp_path): def test_exact_schema_keys_and_values(self, tmp_path): spec = self._write(tmp_path)["speculative"] assert spec == { - "proposal_type": "mtp", - "num_speculative_tokens": 1, - "model": "mtp/model.onnx", - "target_hidden_layout": "BSH", - "hc_mult": 1, - "mtp_hidden_output": "mtp_hidden", - "kv_mode": "proposal_local", - "embedding": { - "source": "target_initializer", - "name": "model.embed_tokens.weight", - }, - "lm_head": { - "source": "target_initializer", - "name": "lm_head.weight", - }, - "target_hidden_output": "hidden_states.63", - "target_hidden_size": 5120, - "vocab_size": 248320, + "proposer": "mtp", + "target": "decoder", + "proposal_execution": {"kind": "block"}, + "port_bindings": {"target_hidden_context": "hidden_states"}, + "shared_weights": ["lm_head.weight", "model.embed_tokens.weight"], + "vocabulary": {"kind": "identical"}, + "max_proposal_width": 1, + "distribution_preserving": True, + "rollback_state": ["decoder_cache.0"], } def test_no_legacy_field_names(self, tmp_path): spec = self._write(tmp_path)["speculative"] - # Field names the runtime cannot parse must not appear. - for banned in ("model_path", "hidden_size", "embedding_weights", "lm_head_weights"): + # Fields of the superseded flat MTP block. onnx-genai's + # SpeculativeContract sets ``deny_unknown_fields``, so any of these + # makes the whole package unparseable rather than being ignored. + for banned in ( + "proposal_type", + "num_speculative_tokens", + "model", + "model_path", + "target_hidden_layout", + "hc_mult", + "mtp_hidden_output", + "kv_mode", + "embedding", + "lm_head", + "embedding_weights", + "lm_head_weights", + "target_hidden_output", + "target_hidden_size", + "hidden_size", + "vocab_size", + ): assert banned not in spec - assert spec["kv_mode"] != "hidden_threaded" + + def test_proposer_is_registered_as_a_workflow_component(self, tmp_path): + workflow = self._write(tmp_path)["pipeline"]["workflow"] + # proposer/target are workflow component names, so the head has to be + # declared before it can be referenced. + assert workflow["components"]["mtp"]["implementation"] == { + "kind": "onnx", + "artifact": "mtp/model.onnx", + } + roles = workflow["components"]["mtp"]["ports"]["roles"] + assert roles["hidden_states"] == "hidden_states" + assert roles["inputs_embeds"] == "inputs_embeds" + # The target output the head is seeded from is named by role, not by a + # free-form ``target_hidden_output`` string. + assert ( + workflow["components"]["decoder"]["ports"]["roles"]["hidden_states.63"] + == "hidden_states" + ) + + def test_rollback_capacity_covers_the_proposal_width(self, tmp_path): + workflow = self._write(tmp_path)["pipeline"]["workflow"] + group = workflow["serving"]["state_service"]["groups"]["decoder_cache"] + # A rolled-back cell whose group declares no rollback_positions makes + # the package unloadable, so attaching the speculator states the bound. + assert group["capabilities"]["rollback_positions"] >= 1 + + def test_anchors_to_a_real_mobius_decoder_workflow(self, tmp_path): + """The target is picked out of a workflow mobius actually emits. + + A real decoder workflow ships a dozen generated policy graphs, every one + of which is an ONNX component, so the verifier can only be identified by + its declared ``logits`` role. + """ + from mobius.integrations.onnx_genai.auto_export_test import _decoder_package + from mobius.integrations.onnx_genai.workflow_metadata import ( + write_decoder_workflow_metadata, + ) + + package = _decoder_package() + write_decoder_workflow_metadata(package, str(tmp_path), package.config) + out = write_mtp_speculator_metadata( + str(tmp_path), backbone_config=_MtpBackboneConfig() + ) + assert out is not None + with open(out, encoding="utf-8") as handle: + metadata = yaml.safe_load(handle) + workflow = metadata["pipeline"]["workflow"] + components = set(workflow["components"]) + assert len(components) > 2, "a real decoder workflow ships policy components" + assert metadata["speculative"]["target"] == "model" + assert metadata["speculative"]["proposer"] == "mtp" + with open(_onnx_genai_schema_path()) as handle: + jsonschema.validate(instance=metadata, schema=json.load(handle)) + + def test_requires_a_workflow_to_anchor_against(self, tmp_path): + (tmp_path / "inference_metadata.yaml").write_text(yaml.safe_dump({}), encoding="utf-8") + with pytest.raises(ValueError, match=re.escape("pipeline.workflow")): + write_mtp_speculator_metadata(str(tmp_path), backbone_config=_MtpBackboneConfig()) def test_matches_onnx_genai_json_schema(self, tmp_path): """Emitted metadata validates against onnx-genai's published schema.""" - schema_path = _onnx_genai_schema_path() - if schema_path is None: - pytest.skip("onnx-genai schema not found (set ONNX_GENAI_SCHEMA)") - with open(schema_path) as handle: + with open(_onnx_genai_schema_path()) as handle: schema = json.load(handle) meta = self._write(tmp_path) jsonschema.validate(instance=meta, schema=schema) diff --git a/src/mobius/integrations/onnx_genai/speech_to_text_workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/speech_to_text_workflow_metadata_test.py index 76fd6d13a..c70991b95 100644 --- a/src/mobius/integrations/onnx_genai/speech_to_text_workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/speech_to_text_workflow_metadata_test.py @@ -6,7 +6,6 @@ from __future__ import annotations import json -import os from types import SimpleNamespace import jsonschema @@ -18,6 +17,7 @@ from mobius.integrations.onnx_genai.auto_export import _audio_preprocessing_program from mobius.integrations.onnx_genai.inference_metadata_test import ( _model, + _onnx_genai_schema_path, _value, ) from mobius.integrations.onnx_genai.workflow_metadata import ( @@ -235,10 +235,7 @@ def test_write_round_trips_the_built_metadata(tmp_path): def test_speech_workflow_matches_producer_schema(tmp_path): - schema_path = os.environ.get("ONNX_GENAI_SCHEMA") - if not schema_path: - pytest.skip("set ONNX_GENAI_SCHEMA to the producer-contract schema") - with open(schema_path, encoding="utf-8") as handle: + with open(_onnx_genai_schema_path(), encoding="utf-8") as handle: schema = json.load(handle) processor = tmp_path / "audio_processor.json" processor.write_text(json.dumps(_WHISPER_EXTRACTOR), encoding="utf-8") From e5a00bf8a119cf75dd664d2c740d4e7fa43df1a6 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 22 Aug 2026 18:44:53 -0700 Subject: [PATCH 2/3] Condense the new onnx-genai metadata docstrings Drop the before/after narration comparing against the superseded pipeline and speculative shapes; the docstrings describe the contract the code emits today. The banned-name test keeps its list of now-legacy fields, since rejecting them is the point of that test. Signed-off-by: Justin Chu Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../integrations/onnx_genai/comfyui_test.py | 6 +- .../onnx_genai/inference_metadata.py | 172 +++++++----------- .../onnx_genai/inference_metadata_test.py | 52 ++---- 3 files changed, 91 insertions(+), 139 deletions(-) diff --git a/src/mobius/integrations/onnx_genai/comfyui_test.py b/src/mobius/integrations/onnx_genai/comfyui_test.py index af31bdf1b..a58c71255 100644 --- a/src/mobius/integrations/onnx_genai/comfyui_test.py +++ b/src/mobius/integrations/onnx_genai/comfyui_test.py @@ -231,10 +231,8 @@ def test_ancestral_sampler_is_rejected(): def test_unmaterializable_sigma_spacing_is_rejected(spacing): wf = json.loads(json.dumps(_DEFAULT_TXT2IMG)) wf["3"]["inputs"]["scheduler"] = spacing - # The published workflow ships the sigma table as a constant component, so a - # spacing mobius cannot materialize has to fail rather than be dropped: the - # old document merely carried ``use_karras_sigmas`` as a hint for a runtime - # scheduler that no longer exists. + # The workflow ships the sigma table as a constant component, so a spacing + # mobius cannot materialize has to fail rather than be silently dropped. with pytest.raises(ValueError, match="Karras or exponential sigmas"): parse_comfyui_workflow(wf) diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 112983f46..068721998 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -4,29 +4,22 @@ """Emit onnx-genai ``inference_metadata`` for multi-model pipelines. Mobius builds the neural components of a diffusion model (denoiser transformer, -VAE, and — externally — a text encoder) as separate ONNX graphs, but does not -itself carry a scheduler loop. onnx-genai's pipeline contract supplies that loop -declaratively: given an ``inference_metadata`` document, the runtime executes -the denoise loop and returns the decoded output. - -That contract is a *typed SSA workflow*. onnx-genai's ``PipelineSpec`` has a -single property, ``workflow`` (``crates/onnx-genai-metadata/src/schema/``), so -the sampler is not a ``strategy`` block the runtime interprets but an ordinary -executable component the package ships: the sigma schedule and timestep table -are constant components, the step index is the loop induction value, and -classifier-free guidance is two denoiser invocations plus a combine component. -This module builds that document from the component filenames plus a scheduler -config, materializing the sampler components from mobius's policy library. It -reads no torch/diffusers state — only plain values — so it is cheap to unit-test -and safe to call anywhere. - -Not everything here is publishable. :func:`build_native_vlm_package_metadata` -returns mobius's *internal* structural descriptor of a VLM package (a -``models``/``dataflow``/``strategy`` view used to reason about wiring and -validate the executable closure); the published contract for such a package is -the workflow that +VAE, and — externally — a text encoder) as separate ONNX graphs. The document +this module produces wires them into onnx-genai's ``pipeline.workflow``: a typed +SSA graph in which the sampler is an executable component the package ships, so +the sigma schedule and timestep table are constant components, the step index is +the loop induction value, and classifier-free guidance is two denoiser +invocations plus a combine component. + +It builds that document from the component filenames plus a scheduler config, +materializing the sampler components from mobius's policy library. It reads no +torch/diffusers state — only plain values — so it is cheap to unit-test and safe +to call anywhere. + +Not everything here is publishable: :func:`build_native_vlm_package_metadata` +returns mobius's *internal* structural descriptor of a VLM package, from which :func:`~mobius.integrations.onnx_genai.workflow_metadata.build_vlm_workflow_metadata` -derives from it. +derives the published contract. Autoregressive decoder-only LLM metadata (``model.attention`` + ``kv_cache``) lives in the sibling :mod:`mobius.integrations.onnx_genai.decoder_metadata` @@ -1746,15 +1739,13 @@ def build_native_vlm_package_metadata( ) -> dict[str, Any]: """Describe a native VLM package by inspecting every component graph. - This returns mobius's **internal** structural descriptor, not the published - onnx-genai document. Its ``pipeline`` key is a ``models``/``dataflow``/ - ``strategy`` view used to reason about component wiring and to validate the - executable closure; onnx-genai's ``PipelineSpec`` accepts only a typed SSA - ``workflow`` and rejects every other property, so this view is consumed by - :func:`~mobius.integrations.onnx_genai.workflow_metadata.build_vlm_workflow_metadata` - (which publishes the workflow and the ``preprocessing`` program) rather than - written to disk. Use :func:`write_native_vlm_package_metadata` to emit the - package contract. + This is mobius's **internal** structural descriptor, not a publishable + document: its ``pipeline`` key is a ``models``/``dataflow``/``strategy`` + view used to reason about component wiring and validate the executable + closure. It is consumed by + :func:`~mobius.integrations.onnx_genai.workflow_metadata.build_vlm_workflow_metadata`, + which publishes the workflow and the ``preprocessing`` program. Use + :func:`write_native_vlm_package_metadata` to emit the package contract. Processor selection is registry-driven from graph rank/dtype/shape signatures. No model type, architecture name, or model-name branch @@ -2049,10 +2040,7 @@ def write_native_vlm_package_metadata( What lands in ``inference_metadata.yaml`` is the typed SSA workflow built by :func:`~mobius.integrations.onnx_genai.workflow_metadata.build_vlm_workflow_metadata`, not the structural descriptor :func:`build_native_vlm_package_metadata` - returns. onnx-genai's ``PipelineSpec`` has exactly one property - (``workflow``) and forbids anything else, so the descriptor's - ``models``/``dataflow``/``strategy`` view is mobius-internal and is never - published. + returns. """ from mobius.integrations.onnx_genai.workflow_metadata import ( write_vlm_workflow_metadata, @@ -2340,8 +2328,8 @@ def load_diffusers_vae_scaling_factor( #: Latent axis names shared by the solver, guidance and clamp policy components #: (``mobius.generation._policy_components._IMAGE_LATENT_DIMS``). The workflow's -#: latent contract uses the same symbols so a declared value and the component -#: port it binds to can never disagree about an axis. +#: latent contract reuses them so a declared value and the component port it +#: binds to can never disagree about an axis. _LATENT_DIMS: tuple[str, ...] = ("batch", "channels", "height", "width") _WORKFLOW_DTYPES: dict[str, str] = { @@ -2374,17 +2362,15 @@ def _diffusion_schedule_values( ) -> tuple[list[float], list[float], int]: """Resolve the solver schedule, the timestep table and the executed step count. - The schedule is what ``solver_step`` integrates, so it is derived from the - scheduler's own noise schedule rather than invented: a variance-preserving - DDIM solver reads cumulative alphas, a sigma-space Euler / DPM-Solver++ one - reads sigmas. Both come from the same diffusers-compatible derivations the - package exporter uses, so a ComfyUI conversion and a package export of the - same checkpoint describe the same dynamics. - - ``start_step`` is img2img's "skip the noisiest steps": diffusers implements - it by starting from a later entry of the same table, so it lowers to a - sliced schedule plus a smaller loop bound rather than to a separate control - knob the typed workflow has no room for. + The schedule is what ``solver_step`` integrates: a variance-preserving DDIM + solver reads cumulative alphas, a sigma-space Euler / DPM-Solver++ one reads + sigmas. Both use the same diffusers-compatible derivations the package + exporter does, so a ComfyUI conversion and a package export of one + checkpoint describe the same dynamics. + + ``start_step`` is img2img's "skip the noisiest steps", which diffusers + implements by starting from a later entry of the same table, so it lowers to + a sliced schedule plus a smaller loop bound. """ from mobius.integrations.onnx_genai.auto_export import ( _ddim_alpha_schedule, @@ -2445,12 +2431,10 @@ def build_diffusion_pipeline_metadata( ) -> dict[str, Any]: """Build the onnx-genai ``inference_metadata`` document for a diffusion pipeline. - onnx-genai's ``PipelineSpec`` has a single property, ``workflow``: a typed - SSA graph in which the sampler is an ordinary executable component rather - than a ``strategy`` block the runtime interprets. So the denoise loop is - emitted explicitly — the sigma schedule and timestep table are constant - components, the step index is the loop induction value, and classifier-free - guidance is two denoiser invocations plus a combine component. + The denoise loop is emitted as an explicit ``pipeline.workflow``: the sigma + schedule and timestep table are constant components, the step index is the + loop induction value, and classifier-free guidance is two denoiser + invocations plus a combine component. Only the ONNX components a caller *names* are described by artifact; their graphs stay authoritative for which ports exist. The solver, schedule, @@ -2463,9 +2447,10 @@ def build_diffusion_pipeline_metadata( denoiser_*: Denoiser component filename and I/O port names. scheduler: Noise-schedule config (defaults to DDIM defaults). Its ``kind`` selects the solver component. - timesteps: Explicit per-step timestep table; a linear ramp otherwise. - schedule: Explicit sigma schedule (``num_inference_steps + 1`` values); - a linear ramp otherwise. + timesteps: Explicit per-step timestep table; derived from ``scheduler`` + otherwise. + schedule: Explicit solver schedule (``num_inference_steps + 1`` values); + derived from ``scheduler`` otherwise. guidance_scale: When set and != 1.0, enables classifier-free guidance. start_step: img2img skip count; lowered to a sliced schedule. vae_filename: VAE decoder producing the image output. @@ -3141,11 +3126,8 @@ def write_diffusion_pipeline_metadata( return path -#: Proposer input port that receives the target's per-token hidden state. -#: -#: ``SpeculativeContract.port_bindings`` is keyed by semantic role; onnx-genai -#: names this role ``target_hidden_context`` -#: (``crates/onnx-genai-metadata/src/schema/package.rs``). +#: ``SpeculativeContract.port_bindings`` role for the proposer input port that +#: receives the target's per-token hidden state. _TARGET_HIDDEN_CONTEXT_ROLE = "target_hidden_context" #: Port roles the Qwen3.5/3.8 MTP sidecar exposes. The graph is authoritative @@ -3163,11 +3145,11 @@ def write_diffusion_pipeline_metadata( def _speculative_target_candidates(workflow: dict[str, Any], exclude: str) -> list[str]: """Workflow components that can verify a speculative proposal. - A verifier is the component that scores the next token, so it is selected by - its declared ``logits`` output role. Selecting on ``implementation.kind == - "onnx"`` would not work: the generated policy graphs a workflow ships (the - sampler, the termination predicate, the state updates) are ONNX components - too, and a real decoder package declares about a dozen of them. + A verifier scores the next token, so it is selected by its declared + ``logits`` output role. Selecting on ``implementation.kind == "onnx"`` would + not narrow anything: the generated policy graphs a workflow ships (sampler, + termination predicate, state updates) are ONNX components too, and a real + decoder package declares about a dozen of them. """ return [ name @@ -3196,38 +3178,24 @@ def write_mtp_speculator_metadata( backbone (``mtp/model.onnx``). It borrows the target's shared embedding / LM head and is seeded by the backbone's final-layer hidden state. - The emitted block is onnx-genai's ``SpeculativeContract`` - (``crates/onnx-genai-metadata/src/schema/package.rs``), which replaced the - older flat ``SpeculatorConfig`` block. That contract is expressed entirely - in terms of the *workflow*: ``proposer`` and ``target`` name workflow - components, ``rollback_state`` names workflow state cells, and the hidden - handoff is a ``port_bindings`` role rather than a pair of free-form port - names. So this writer also registers the head as a workflow component and - completes the rollback capabilities the claim depends on: - - - ``proposer`` / ``target`` are workflow component names. The backbone - metadata must already declare a ``pipeline.workflow`` with exactly one - ONNX component (the decoder) to anchor against; anything else fails - closed rather than publishing an unanchored contract. - - ``proposal_execution: block``. The head returns its complete proposal in - one invocation. It is deliberately *not* declared ``chained``: a chained - proposer must expose a ``logits_output`` carrying the next-token - distribution, and this sidecar emits only ``mtp_hidden`` — the runtime - obtains draft logits by decoding it through the target's LM head, which - is why that initializer is listed in ``shared_weights``. + ``SpeculativeContract`` is expressed in terms of the workflow — ``proposer`` + and ``target`` are component names, ``rollback_state`` names state cells — + so this writer also registers the head as a workflow component and completes + the rollback capabilities the claim depends on. Notably: + + - ``proposal_execution: block``, not ``chained``: a chained proposer must + expose a ``logits_output``, and this sidecar emits only ``mtp_hidden``. + The runtime obtains draft logits by decoding it through the target's LM + head, which is why that initializer is listed in ``shared_weights``. - ``port_bindings.target_hidden_context`` names the proposer input port the - target's per-token hidden state lands in. The *source* of that value is - declared structurally, as a ``hidden_states`` port role on the target - component, instead of being spelled out as a layer-indexed port name. - - ``vocabulary: identical`` — the head shares the target's LM head, so it - scores the target's own vocabulary axis. - - ``rollback_state`` lists the target's service-group-backed state cells, - and each reached group is given a ``rollback_positions`` bound of at - least ``max_proposal_width``: a rejected proposal must be undoable that - far or the runtime rejects the package. - - The backbone ``inference_metadata.yaml`` must already exist. Returns the - metadata path, or ``None`` when it is missing. + target's hidden state lands in; its source is declared as a + ``hidden_states`` port role on the target component. + - ``vocabulary: identical`` — sharing the target's LM head means scoring the + target's own vocabulary axis. + + The backbone ``inference_metadata.yaml`` must already exist and declare a + ``pipeline.workflow`` to anchor against. Returns the metadata path, or + ``None`` when the file is missing. """ if num_speculative_tokens < 1: raise ValueError("num_speculative_tokens must be >= 1") @@ -3310,10 +3278,10 @@ def _declare_rollback_capacity( ) -> None: """Guarantee every group reached by ``cells`` can rewind ``positions``. - A speculative package is rejected when a rolled-back cell resolves to a - state group that declares no ``rollback_positions``, or fewer than the - declared maximum proposal width. Attaching the speculator is what creates - that requirement, so it is also what states the bound. + A package is rejected when a rolled-back cell resolves to a group declaring + no ``rollback_positions``, or fewer than the maximum proposal width. + Attaching the speculator creates that requirement, so it also states the + bound. """ groups = (workflow.get("serving") or {}).get("state_service", {}).get("groups", {}) state = workflow.get("state") or {} diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index 20dde82fa..cc8cd2cf3 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -186,14 +186,10 @@ def test_workflow_policy_components_reference_saved_onnx_artifacts(tmp_path): def _onnx_genai_schema_path() -> str: """Locate onnx-genai's published pipeline JSON schema. - The vendored copy under ``_schema/`` is the default so this never skips: - the previous behaviour — search a few local onnx-genai checkouts and - ``pytest.skip`` when none is found — meant these tests were silently - inactive in CI, which is how two upstream contract redesigns went unnoticed. - A developer checkout is deliberately *not* consulted implicitly either, - because a clone that is ahead of or behind ``main`` would make the result - machine-dependent in exactly the same way. Set ``ONNX_GENAI_SCHEMA`` to - validate against a specific revision. + The vendored copy under ``_schema/`` is the default so conformance never + skips. A developer checkout is deliberately not consulted implicitly: one + that is ahead of or behind ``main`` makes the result machine-dependent. Set + ``ONNX_GENAI_SCHEMA`` to validate against a specific revision. """ override = os.environ.get("ONNX_GENAI_SCHEMA") if override: @@ -583,11 +579,7 @@ def _validate(self, package, config, source=None) -> None: ``build_native_vlm_package_metadata`` returns mobius's internal structural descriptor; what a package publishes is the typed SSA - workflow ``build_vlm_workflow_metadata`` derives from it. onnx-genai's - ``PipelineSpec`` has ``workflow`` as its only property, so the - descriptor's ``models``/``dataflow``/``strategy`` view is not a - publishable document and validating it against the schema would assert - the wrong contract. + workflow ``build_vlm_workflow_metadata`` derives from it. """ published = build_vlm_workflow_metadata(package, config, source=source) assert set(published["pipeline"]) == {"workflow"} @@ -1577,12 +1569,11 @@ def walk(step: dict) -> None: class TestBuildDiffusionPipelineMetadata: - """The published document is a typed SSA workflow, not an ``iterative`` strategy. + """The published document is a typed SSA workflow. - onnx-genai's ``PipelineSpec`` declares ``workflow`` as its only property + ``PipelineSpec`` declares ``workflow`` as its only property (``crates/onnx-genai-metadata/src/schema/pipeline.rs``), so the sampler is - an executable component the package ships rather than a ``strategy`` block - the runtime interprets. + an executable component the package ships. """ def _workflow(self, **kwargs) -> dict: @@ -1612,7 +1603,7 @@ def test_denoise_loop_carries_the_latent_through_the_solver(self): loop = next(step for step in workflow["steps"] if step["kind"] == "loop") assert loop["max_iterations"] == "request.max_iterations" assert workflow["inputs"]["request.max_iterations"]["default"] == 20 - # The loop-carried latent replaces the old denoiser output self-edge. + # The latent is loop-carried state, advanced by the solver each step. # ``latent`` is also a workflow output, so the published cell is # disambiguated to ``latent_state``. latent_carry = next( @@ -1643,8 +1634,8 @@ def test_full_pipeline_with_vae_and_text_encoder(self): "text_encoder", "guidance_combine", } - # The text encoder runs once in the loop setup (the old "prompt_only" - # phase); the VAE decodes the final latent after the loop. + # The text encoder runs once in the loop setup; the VAE decodes the + # final latent after the loop. loop = next(step for step in workflow["steps"] if step["kind"] == "loop") encoder_calls = [ step for step in loop["setup"] if step.get("component") == "text_encoder" @@ -2060,12 +2051,11 @@ class _MtpBackboneConfig: def _seed_backbone_metadata(directory: Path) -> str: """Write a backbone inference_metadata.yaml for the MTP writer to extend. - ``SpeculativeContract`` is expressed against the workflow — ``proposer`` and - ``target`` are component names and ``rollback_state`` names state cells — so - the backbone must already publish a ``pipeline.workflow`` for the writer to - anchor against. This is the shape ``write_onnx_genai_config`` emits for a - single-component decoder package, reduced to what the speculator claim - touches: one ONNX decoder component and one service-group-backed KV cell. + ``SpeculativeContract`` names workflow components and state cells, so the + backbone must already publish a ``pipeline.workflow``. This is the shape + ``write_onnx_genai_config`` emits for a single-component decoder package, + reduced to what the speculator claim touches: one decoder component and one + service-group-backed KV cell. """ kv_contract = { "dtype": "float16", @@ -2180,12 +2170,9 @@ class TestMtpSpeculatorMetadata: """The emitted ``speculative`` block conforms to the onnx-genai runtime schema. Authoritative source: onnx-genai - ``crates/onnx-genai-metadata/src/schema/package.rs`` - (``SpeculativeContract``, ``SpeculativeProposalExecution``, - ``SpeculativeVocabulary``) + ``validation.rs`` - (``validate_speculative_rollback``). The older flat ``SpeculatorConfig`` - block in ``schema/generation.rs`` describes a HuggingFace ``config.json`` - speculator section, not ``InferenceMetadata.speculative``. + ``crates/onnx-genai-metadata/src/schema/package.rs`` (``SpeculativeContract``, + ``SpeculativeProposalExecution``, ``SpeculativeVocabulary``) plus + ``validation.rs`` (``validate_speculative_rollback``). """ def _write(self, tmp_path: Path) -> dict: @@ -2220,7 +2207,6 @@ def test_exact_schema_keys_and_values(self, tmp_path): def test_no_legacy_field_names(self, tmp_path): spec = self._write(tmp_path)["speculative"] - # Fields of the superseded flat MTP block. onnx-genai's # SpeculativeContract sets ``deny_unknown_fields``, so any of these # makes the whole package unparseable rather than being ignored. for banned in ( From a52e72cfd25515bbe459ea6c2e4c4a5c31ddeb30 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 22 Aug 2026 18:54:23 -0700 Subject: [PATCH 3/3] Carry the generated sampler graphs with the parsed ComfyUI workflow The translated document declares the sampler as executable components under ``policies/*.onnx``, but ``parse_comfyui_workflow`` built those graphs and threw them away, so a caller holding only the metadata could not produce a loadable package. ``ComfyUIWorkflow`` now carries them and exposes ``save_policy_components``; the dict-only wrappers point at it. ``convert_comfyui_workflow`` keeps building its own package rather than reusing the parse-time one, because the checkpoint's scheduler config can reconcile to a different solver than the ComfyUI sampler implied -- reusing it would leave that run's components behind for a document that never references them. That is now stated where someone would otherwise "simplify" it, and covered by a test asserting the written set equals the referenced set exactly. Signed-off-by: Justin Chu Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/mobius/integrations/onnx_genai/comfyui.py | 38 ++++++++++++++++- .../integrations/onnx_genai/comfyui_test.py | 42 +++++++++++++++++++ src/mobius/integrations/onnx_genai/convert.py | 5 +++ 3 files changed, 84 insertions(+), 1 deletion(-) diff --git a/src/mobius/integrations/onnx_genai/comfyui.py b/src/mobius/integrations/onnx_genai/comfyui.py index 386e7dc0b..2bf0b65a3 100644 --- a/src/mobius/integrations/onnx_genai/comfyui.py +++ b/src/mobius/integrations/onnx_genai/comfyui.py @@ -41,6 +41,7 @@ import json import logging import math +import os from typing import Any from mobius.integrations.onnx_genai.inference_metadata import ( @@ -80,6 +81,12 @@ class ComfyUIWorkflow: ``metadata`` is the onnx-genai ``inference_metadata`` document (topology + scheduler + guidance). The remaining fields are the per-run inputs recovered from the graph so a caller can actually drive the pipeline. + + The document declares the sampler as executable components under + ``policies/*.onnx``. Those graphs are generated during translation rather + than read off disk, so they are carried here in ``policy_components``; + :meth:`save_policy_components` writes them next to the metadata. A package + whose metadata references them without shipping them is not loadable. """ metadata: dict[str, Any] @@ -99,6 +106,18 @@ class ComfyUIWorkflow: start_step: int = 0 loras: tuple[tuple[str, float], ...] = () controlnet: tuple[str, float] | None = None + policy_components: Any | None = None + + def save_policy_components(self, directory: str) -> dict[str, str]: + """Write the generated sampler graphs ``metadata`` references. + + Returns the package-relative artifact paths, or an empty mapping when + the workflow carries no generated components. + """ + if self.policy_components is None: + return {} + os.makedirs(directory, exist_ok=True) + return self.policy_components.save_policy_components(directory) def _nodes(workflow: dict[str, Any]) -> dict[str, Any]: @@ -327,6 +346,11 @@ def parse_comfyui_workflow( "the workflow before converting it." ) + from mobius._model_package import ModelPackage + + # The sampler components the document declares are generated here, so they + # travel with the parsed workflow instead of being discarded. + policy_components = ModelPackage({}) metadata = build_diffusion_pipeline_metadata( num_inference_steps=steps, scheduler=sched, @@ -335,9 +359,11 @@ def parse_comfyui_workflow( denoiser_filename=denoiser_filename, vae_filename=vae_filename, text_encoder_filename=text_encoder_filename if has_text_encoder else None, + package=policy_components, ) return ComfyUIWorkflow( metadata=metadata, + policy_components=policy_components, prompt=prompt, negative_prompt=negative_prompt, width=width, @@ -362,6 +388,12 @@ def translate_comfyui_workflow(workflow: dict[str, Any], **kwargs: Any) -> dict[ Convenience wrapper over :func:`parse_comfyui_workflow` for callers that only need the pipeline document (see that function for the full run parameters). + + The document references generated sampler graphs under ``policies/*.onnx``, + which this wrapper drops. To write a loadable package use + :func:`convert_comfyui_workflow`, or keep the + :class:`ComfyUIWorkflow` and call + :meth:`~ComfyUIWorkflow.save_policy_components`. """ return parse_comfyui_workflow(workflow, **kwargs).metadata @@ -374,5 +406,9 @@ def parse_comfyui_workflow_file(path: str, **kwargs: Any) -> ComfyUIWorkflow: def translate_comfyui_workflow_file(path: str, **kwargs: Any) -> dict[str, Any]: - """Load a ComfyUI API-format JSON file and translate it to metadata.""" + """Load a ComfyUI API-format JSON file and translate it to metadata. + + Returns the document only; see :func:`translate_comfyui_workflow` for how to + materialize the sampler graphs it references. + """ return parse_comfyui_workflow_file(path, **kwargs).metadata diff --git a/src/mobius/integrations/onnx_genai/comfyui_test.py b/src/mobius/integrations/onnx_genai/comfyui_test.py index a58c71255..aec8eac8d 100644 --- a/src/mobius/integrations/onnx_genai/comfyui_test.py +++ b/src/mobius/integrations/onnx_genai/comfyui_test.py @@ -8,6 +8,7 @@ import json import pytest +import yaml from mobius.integrations.onnx_genai.comfyui import ( ComfyUIWorkflow, @@ -307,6 +308,47 @@ def test_translate_from_file(tmp_path): assert workflow["inputs"]["request.max_iterations"]["default"] == 20 +def _referenced_artifacts(metadata: dict) -> set[str]: + """Every ONNX artifact the published workflow declares.""" + return { + declaration["implementation"]["artifact"] + for declaration in metadata["pipeline"]["workflow"]["components"].values() + if declaration["implementation"]["kind"] == "onnx" + } + + +def test_parsed_workflow_can_materialize_the_policies_it_references(tmp_path): + """The generated sampler graphs travel with the parse result. + + The document declares them as ``policies/*.onnx`` artifacts, so a caller + that only had the metadata dict could not produce a loadable package. + """ + parsed = parse_comfyui_workflow(_DEFAULT_TXT2IMG) + written = parsed.save_policy_components(str(tmp_path)) + assert written, "a diffusion workflow always generates sampler components" + for artifact in _referenced_artifacts(parsed.metadata): + if artifact.startswith("policies/"): + assert (tmp_path / artifact).is_file(), artifact + + +def test_conversion_writes_every_referenced_policy_artifact(tmp_path): + convert_comfyui_workflow(_DEFAULT_TXT2IMG, None, str(tmp_path), compute_timesteps=False) + with open(tmp_path / "inference_metadata.yaml", encoding="utf-8") as handle: + metadata = yaml.safe_load(handle) + policies = { + artifact + for artifact in _referenced_artifacts(metadata) + if artifact.startswith("policies/") + } + assert policies + for artifact in policies: + assert (tmp_path / artifact).is_file(), artifact + # And nothing stale: the reconciled solver's components are the only ones + # written, so a package never ships a graph its document does not declare. + written = {f"policies/{path.name}" for path in (tmp_path / "policies").iterdir()} + assert written == policies + + def test_translated_metadata_matches_onnx_genai_schema(): """A ComfyUI-translated pipeline validates against onnx-genai's real schema.""" import jsonschema diff --git a/src/mobius/integrations/onnx_genai/convert.py b/src/mobius/integrations/onnx_genai/convert.py index a5a0074f6..b4e97d07a 100644 --- a/src/mobius/integrations/onnx_genai/convert.py +++ b/src/mobius/integrations/onnx_genai/convert.py @@ -217,6 +217,11 @@ def convert_comfyui_workflow( os.makedirs(output_dir, exist_ok=True) from mobius._model_package import ModelPackage + # Deliberately a fresh package rather than ``parsed_workflow.policy_components``: + # the checkpoint's scheduler config can reconcile to a different solver than + # the ComfyUI sampler implied, and reusing the parse-time package would leave + # that run's components (say Euler's ``model_input_scale``) behind for a DDIM + # document that never references them. package = ModelPackage({}) use_karras = parsed_workflow.scheduler_spacing == "karras" use_exponential = parsed_workflow.scheduler_spacing == "exponential"