Migrate task files to new GraphBuilder input/output APIs - #231
Merged
Conversation
Replace manual ir.Value() graph input creation with builder.input() and graph.outputs.append() with builder.add_output() across all 24 task files under src/mobius/tasks/. Key changes: - _make_graph() no longer takes an inputs parameter; graph starts empty and inputs are registered via builder.input() after creation. - _make_kv_cache_inputs() and _make_hybrid_cache_inputs() now take builder as first parameter, use builder.input(), and return only the cache pairs (no flat input list). - _register_kv_cache_outputs() and _register_hybrid_cache_outputs() now take builder instead of graph and use builder.add_output(). - All caller task files updated to the new control flow: create graph first, then register inputs via builder. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
- Rename graph_builder to builder in _gemma4.py and _vision_language_3model.py for consistency with other 22 task files - Change GraphBuilder import from onnxscript._internal.builder (private) to onnxscript (public API) in _base.py, _cache_utils.py, _causal_lm.py, _gemma4.py - Fix import sorting and line length issues flagged by lintrunner - Update _make_graph() docstring to mention builder.add_output() Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Performance Comparison
|
Contributor
There was a problem hiding this comment.
Pull request overview
Migrates src/mobius/tasks/ graph construction to the newer onnxscript.GraphBuilder.input() / GraphBuilder.add_output() APIs, removing manual ir.Value(...) creation and direct graph.outputs.append(...) mutations to simplify and standardize task I/O registration.
Changes:
- Refactors
_make_graph()to create empty graphs and rely onbuilder.input()for input registration. - Updates cache helper utilities to accept
builderand to return only structured cache tuples (no flat “graph_inputs” lists). - Converts all task implementations to use
builder.input(...)andbuilder.add_output(...), with related import cleanup.
Reviewed changes
Copilot reviewed 24 out of 24 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| src/mobius/tasks/_base.py | Updates _make_graph() to start with empty inputs/outputs and documents builder-driven I/O registration. |
| src/mobius/tasks/_cache_utils.py | Refactors KV/hybrid cache input/output helpers to register I/O via builder and simplify return types. |
| src/mobius/tasks/_causal_lm.py | Migrates dynamic/static cache task graph wiring to builder-driven input/output registration. |
| src/mobius/tasks/_gemma4.py | Migrates Gemma4 task graphs and its custom KV-cache input builder to builder.input() / builder.add_output(). |
| src/mobius/tasks/_vision_language.py | Migrates vision-language single-model task I/O to builder APIs and updated cache helper signatures. |
| src/mobius/tasks/_vision_language_3model.py | Migrates 3-model VLM task component graphs (vision + decoder) to builder APIs. |
| src/mobius/tasks/_multimodal.py | Migrates multimodal task inputs/outputs and KV-cache wiring to builder APIs. |
| src/mobius/tasks/_phi4mm_multimodal.py | Migrates Phi4mm multi-component graphs (vision/audio/embedding) to builder APIs. |
| src/mobius/tasks/_speech_to_text.py | Migrates encoder/decoder graphs and KV-cache wiring to builder APIs. |
| src/mobius/tasks/_speech_language.py | Migrates audio encoder graph I/O to builder APIs. |
| src/mobius/tasks/_tts.py | Migrates talker/code-predictor/embedding/speaker-encoder graphs and KV-cache wiring to builder APIs. |
| src/mobius/tasks/_seq2seq.py | Migrates encoder/decoder graphs and rewires self/cross KV-cache inputs/outputs using builder APIs. |
| src/mobius/tasks/_ssm_causal_lm.py | Migrates SSM causal LM state inputs/outputs to builder APIs. |
| src/mobius/tasks/_vae.py | Migrates VAE encoder/decoder graphs to builder APIs. |
| src/mobius/tasks/_qwen_image_vae.py | Migrates Qwen Image VAE encoder/decoder graphs to builder APIs. |
| src/mobius/tasks/_video_denoising.py | Migrates video denoising graph I/O to builder APIs. |
| src/mobius/tasks/_denoising.py | Migrates denoising (UNet2D) graph I/O to builder APIs. |
| src/mobius/tasks/_controlnet.py | Migrates ControlNet graph I/O and multi-output registration to builder APIs. |
| src/mobius/tasks/_codec.py | Migrates codec encoder/decoder graphs to builder APIs. |
| src/mobius/tasks/_feature_extraction.py | Migrates feature extraction graph inputs/outputs to builder APIs. |
| src/mobius/tasks/_audio_feature_extraction.py | Migrates audio feature extraction graph I/O to builder APIs. |
| src/mobius/tasks/_image_classification.py | Migrates image classification graph I/O to builder APIs. |
| src/mobius/tasks/_object_detection.py | Migrates object detection graph I/O (multiple outputs) to builder APIs. |
| src/mobius/tasks/_adapter.py | Migrates adapter task graph I/O (conditional input + variable outputs) to builder APIs. |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
justinchuby
marked this pull request as draft
May 1, 2026 18:53
justinchuby
marked this pull request as ready for review
May 1, 2026 18:55
Member
Author
titaiwangms
approved these changes
May 1, 2026
gramalingam
approved these changes
May 1, 2026
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
Migrate all 24 task files under
src/mobius/tasks/to use the new onnxscriptGraphBuilder.input()andGraphBuilder.add_output()APIs, replacing manualir.Value()construction andgraph.outputs.append()patterns.Changes
_base.py:_make_graph()no longer takes an inputs parameter. Inputs are now created viabuilder.input(name, dtype, shape)which auto-registers them on the graph._cache_utils.py: Cache helper functions (_make_kv_cache_inputs,_make_hybrid_cache_inputs,_register_kv_cache_outputs,_register_hybrid_cache_outputs) now takebuilderas first parameter. Return types simplified — flat input lists removed sincebuilder.input()auto-registers.ir.Value(name=..., shape=ir.Shape(...), type=ir.TensorType(...))→builder.input(name, dtype=..., shape=[...]);value.name = name; graph.outputs.append(value)→builder.add_output(value, name)GraphBuilderimported from publiconnxscriptnamespace instead ofonnxscript._internal.builderbuildervariable name across all filesImpact
Note
op.call()migration for module/function invocations is out of scope —op.call()currently only acceptsir.Function, notnn.Module, and treats kwargs as ONNX attributes rather than tensor inputs. This requires upstream onnxscript changes.