Fix workflow batching capacity declarations - #649
Merged
Conversation
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
Performance Comparison
|
Contributor
There was a problem hiding this comment.
Pull request overview
This PR updates Mobius’s ONNX GenAI workflow metadata producer and fixtures to explicitly declare batching capacity (batch_capacity) for audited multi-request workflow components, while improving how request axes are discovered/stamped (supporting non-leading batch axes) and preserving hierarchical-audio’s fixed internal row-expansion semantics.
Changes:
- Add producer-side helpers to author
batch_capacity(and bumpschema_versiontov1.1when capacity is declared) for decoder, diffusion, video, TTS, and adapter workflows. - Generalize request-axis detection to the unique batch symbol anywhere in the shape (not only leading axis), and update fixtures accordingly (including axis=1 cases).
- Regenerate/adjust workflow fixtures and add stronger contract tests: multi-request components must declare capacity, YAML/schema round-trip must be stable, and hierarchical-audio must retain
request_expandedwith the correct factor.
Reviewed changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| tests/generate_onnx_genai_validation_packages.py | Updates generated adapter-package metadata to schema v1.1 and adds an explicit batch_capacity declaration for the overlay component. |
| tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml | Adds batch_capacity declarations across audited video workflow components and bumps schema version to v1.1. |
| tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml | Adds batch_capacity declarations and updates/stamps request-axis layouts for non-leading batch axes; bumps schema version to v1.1. |
| tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml | Adds batch_capacity for audited components and bumps schema version to v1.1. |
| tests/fixtures/onnx_genai_workflows/shared_state_pixel_flow/inference_metadata.yaml | Updates request-axis batch_layout stamping for shapes where the batch symbol is not on axis 0. |
| tests/fixtures/onnx_genai_workflows/hierarchical_audio/inference_metadata.yaml | Corrects fixed 2-row expansion semantics to request_expanded with factor=2 (without reclassifying as multiple requests). |
| tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml | Adds batch_capacity declarations and specifies append axes for diffusion emits; bumps schema version to v1.1. |
| tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml | Same as diffusion: capacity declarations + emit append axes; bumps schema version to v1.1. |
| tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml | Adds batch_capacity for decoder workflow components and bumps schema version to v1.1. |
| tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml | Adds batch_capacity to the adapter overlay component and bumps schema version to v1.1. |
| tests/canonical_workflow_contract_test.py | Adds mutation-sensitive checks for capacity declarations, YAML/schema round-trip, and hierarchical-audio row-expansion invariants. |
| src/mobius/integrations/onnx_genai/workflow_metadata.py | Implements batch_capacity authoring, preserves hierarchical-audio expansion, adds diffusion emit append axes, and wires capacity stamping into workflow builders. |
| src/mobius/integrations/onnx_genai/inference_metadata.py | Generalizes request-axis discovery/stamping to the unique batch symbol anywhere in shape, and adds unit coverage for the new logic. |
| src/mobius/integrations/onnx_genai/inference_metadata_test.py | Adds direct tests for request_batch_layout behavior (axis 0, axis 1, ambiguous/no-batch cases). |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
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
batch_capacityfor audited row-independent decoder, image-diffusion, video, TTS, and adapter workflow componentsrequest_expandedinstead of misclassifying it as two requestsbatchsymbol, including multi-axis rotary tensorsRoot cause and authority
onnx-genai's runtime admission is intentionally fail-closed: absent
batch_capacitymeans exactly one logical request per component invocation. Mobius had updated to the new schema but still omitted the semantic permission from workflows whose conformance tests intentionally execute multiple rows. Shapes can locate a request axis, but cannot prove row independence, so the runtime check is not relaxed and capacity is not inferred from a dynamic leading dimension.The seven signal failures were producer-side metadata gaps:
decoder_state_initializertensor_scaleoverlaylatent_noisevideo_latent_inittts_state_initializerglobal_initializerusedrequest_alignedfor a fixed 2-row expansion; this is corrected torequest_expandedThis follows onnx-genai #2009, #2010, and runtime admission #2137, while preserving Mobius #636's rule that unproven ESM/ProtBERT encoders remain capacity-absent.
Validation
validate_metadata: all 15 generated packages validencoder_batching: 109 passedundeclared_capacity_never_admits_multiple_requestspassedonnx-genai-metadata,onnx-genai-engine, all targets, warnings denied): passedOriginal failing signal: https://github.com/justinchuby/onnx-genai/actions/runs/32925051072