Emit explicit ONNX GenAI inference metadata - #434
Merged
Conversation
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Performance Comparison
|
Contributor
There was a problem hiding this comment.
Pull request overview
This PR augments Mobius’s ONNX GenAI sidecar emission to include explicit, name-agnostic inference I/O metadata derived directly from exported ONNX graph ports, enabling runtimes to bind decoder/encoder inputs/outputs (including KV/cache/state) without relying on model-family-specific assumptions.
Changes:
- Extend decoder I/O introspection to emit explicit roles for token/embedding inputs, logits output, owned KV ports, cross-attention KV ports, optional hidden output, and TensorScatter static-cache ABI.
- Add
add_explicit_package_io()to attach port-derived I/O roles onto emitted metadata for decoder-only, composite pipelines (encoder/decoder), and nested autoregressive (TTS) strategies. - Update auto-export to post-process emitted YAML metadata files to include explicit I/O roles; expand tests with dynamic KV, static cache, encoder prompt role, and nested autoregressive coverage.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| src/mobius/integrations/onnx_genai/inference_metadata.py | Adds regex/logic to introspect exported ONNX ports and emit explicit decoder/encoder/static-cache/nested-autoregressive I/O metadata. |
| src/mobius/integrations/onnx_genai/inference_metadata_test.py | Adds unit tests validating explicit decoder roles (dynamic/static KV) and encoder prompt role emission. |
| src/mobius/integrations/onnx_genai/auto_export.py | Post-processes written inference metadata YAML files to inject explicit port-derived I/O metadata. |
| src/mobius/integrations/onnx_genai/auto_export_test.py | Updates fixtures/tests to include graph outputs and assert explicit I/O fields in emitted metadata. |
Comment on lines
1013
to
1027
| @@ -898,6 +1025,25 @@ def _input_source_map( | |||
| "from": f"{decoder_name}.{pair['output']}", | |||
| "update": pair["update"], | |||
| } | |||
Member
Author
There was a problem hiding this comment.
Addressed in 50eb36b. Cross-attention KV input/output pairs are now classified through the same stateful loop mapping as self-attention KV pairs. The new encoder-decoder regression test verifies that both cross key and value inputs point to their matching present outputs with append updates.
Comment on lines
+1283
to
+1292
| if strategy.get("kind") == "nested_autoregressive": | ||
| inner_name = strategy.get("inner") | ||
| inner_io = component_ios.get(inner_name, {}) | ||
| hidden_output = inner_io.get("hidden_output") | ||
| if not hidden_output: | ||
| raise ValueError( | ||
| "Cannot emit pipeline.strategy.inner_embedding_output: the inner " | ||
| f"decoder {inner_name!r} has no unique non-logits float output" | ||
| ) | ||
| strategy["inner_embedding_output"] = hidden_output |
Member
Author
There was a problem hiding this comment.
Addressed in 50eb36b. Explicit inner_embedding_output values are now preserved without derivation or validation against the unique-output heuristic. Added coverage for preserving an explicit value, deriving it when absent, and raising an actionable error only when absent and ambiguous.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (1)
src/mobius/integrations/onnx_genai/auto_export.py:53
yaml.safe_load()returnsNonefor an empty file, which will crash inadd_explicit_package_io(metadata, ...)(expects a dict and usesmetadata.get(...)). Since this helper runs post-write on all generated sidecars, a truncated/empty metadata file would fail the whole export.
Default to {} and guard against non-dict YAML payloads before mutating.
with open(path, encoding="utf-8") as handle:
metadata = yaml.safe_load(handle)
add_explicit_package_io(metadata, pkg, config)
with open(path, "w", encoding="utf-8") as handle:
yaml.safe_dump(metadata, handle, sort_keys=False)
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (1)
src/mobius/integrations/onnx_genai/auto_export.py:51
yaml.safe_load()can returnNonefor an empty document (or a non-mapping type for malformed YAML)._add_explicit_io_to_file()passes that value intoadd_explicit_package_io(), which will then crash with an unhelpfulAttributeError(e.g.,'NoneType' object has no attribute 'get'). Add a small type/None guard here so failures are actionable and don’t surface as a generic attribute error.
with open(path, encoding="utf-8") as handle:
metadata = yaml.safe_load(handle)
add_explicit_package_io(metadata, pkg, config)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Emits the explicit, name-agnostic inference metadata required by
justinchuby/onnx-genai#377from the actual exported ONNX graph ports.model.io.token_input,inputs_embeds_input,logits_output,sequence_source,kv_ownership, positionalkv_inputs/kv_outputs, cross-attention KV ports, and hidden outputmodel.io.static_cachewith the control ports and positionally paired key/value cache input/output listspipeline.strategy.inner_embedding_outputplus explicit talker/code-predictor component I/Oaudio_features_inputortoken_input, and explicit decoderencoder_hidden_states_inputExample:
Validation
81 passed, 1 skippedfor ONNX GenAI auto-export/inference/decoder metadata tests67 passedwhile validating against the regenerated schema fromonnx-genaibranchsquad/377-explicit-metadataA real large-model export smoke test was not run because weights were not available; synthetic graph fixtures cover dynamic KV, static cache, encoder roles, and nested autoregressive metadata.
Closes justinchuby/onnx-genai#377