From 1e1cbe5ae3f63b94061fb691a7a06267b86a87dc Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 1 May 2026 18:34:23 +0000 Subject: [PATCH 1/2] Migrate all task files to GraphBuilder input/output APIs 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 --- src/mobius/tasks/_adapter.py | 24 +- src/mobius/tasks/_audio_feature_extraction.py | 13 +- src/mobius/tasks/_base.py | 68 +++--- src/mobius/tasks/_cache_utils.py | 167 +++++++------- src/mobius/tasks/_causal_lm.py | 132 +++++------ src/mobius/tasks/_codec.py | 22 +- src/mobius/tasks/_controlnet.py | 38 +--- src/mobius/tasks/_denoising.py | 25 +-- src/mobius/tasks/_feature_extraction.py | 25 +-- src/mobius/tasks/_gemma4.py | 205 ++++++++---------- src/mobius/tasks/_image_classification.py | 13 +- src/mobius/tasks/_multimodal.py | 50 ++--- src/mobius/tasks/_object_detection.py | 16 +- src/mobius/tasks/_phi4mm_multimodal.py | 80 +++---- src/mobius/tasks/_qwen_image_vae.py | 26 +-- src/mobius/tasks/_seq2seq.py | 120 ++++------ src/mobius/tasks/_speech_language.py | 13 +- src/mobius/tasks/_speech_to_text.py | 52 ++--- src/mobius/tasks/_ssm_causal_lm.py | 34 +-- src/mobius/tasks/_tts.py | 119 ++++------ src/mobius/tasks/_vae.py | 26 +-- src/mobius/tasks/_video_denoising.py | 33 +-- src/mobius/tasks/_vision_language.py | 54 ++--- src/mobius/tasks/_vision_language_3model.py | 119 +++++----- 24 files changed, 571 insertions(+), 903 deletions(-) diff --git a/src/mobius/tasks/_adapter.py b/src/mobius/tasks/_adapter.py index ccf5e5dea..17730dbfc 100644 --- a/src/mobius/tasks/_adapter.py +++ b/src/mobius/tasks/_adapter.py @@ -23,33 +23,23 @@ def build( module, config, ) -> ModelPackage: + graph, builder = _make_graph() + op = builder.op + # Determine input shape based on adapter type if hasattr(config, "in_channels"): # T2I-Adapter: conditioning image input - condition = ir.Value( - name="condition", - type=ir.TensorType(ir.DataType.FLOAT), - shape=ir.Shape(("batch", config.in_channels, "height", "width")), - ) + condition = builder.input("condition", dtype=ir.DataType.FLOAT, shape=["batch", config.in_channels, "height", "width"]) else: # IP-Adapter: image embedding input - condition = ir.Value( - name="image_embeds", - type=ir.TensorType(ir.DataType.FLOAT), - shape=ir.Shape(("batch", config.image_embed_dim)), - ) - - graph, builder = _make_graph([condition]) - op = builder.op + condition = builder.input("image_embeds", dtype=ir.DataType.FLOAT, shape=["batch", config.image_embed_dim]) outputs = module(op, condition) if isinstance(outputs, list): for i, out in enumerate(outputs): - out.name = f"feature_{i}" - graph.outputs.append(out) + builder.add_output(out, f"feature_{i}") else: - outputs.name = "adapter_output" - graph.outputs.append(outputs) + builder.add_output(outputs, "adapter_output") return ModelPackage({"model": _make_model(graph)}, config=config) diff --git a/src/mobius/tasks/_audio_feature_extraction.py b/src/mobius/tasks/_audio_feature_extraction.py index 647a7d8af..493183c2f 100644 --- a/src/mobius/tasks/_audio_feature_extraction.py +++ b/src/mobius/tasks/_audio_feature_extraction.py @@ -28,18 +28,13 @@ def build( module, config: ArchitectureConfig, ) -> ModelPackage: - input_values = ir.Value( - name="input_values", - type=ir.TensorType(ir.DataType.FLOAT), - shape=ir.Shape(("batch", "time")), - ) - - graph, builder = _make_graph([input_values]) + graph, builder = _make_graph() op = builder.op + input_values = builder.input("input_values", dtype=ir.DataType.FLOAT, shape=["batch", "time"]) + last_hidden_state = module(op, input_values=input_values) - last_hidden_state.name = "last_hidden_state" - graph.outputs.append(last_hidden_state) + builder.add_output(last_hidden_state, "last_hidden_state") return ModelPackage({"model": _make_model(graph)}, config=config) diff --git a/src/mobius/tasks/_base.py b/src/mobius/tasks/_base.py index ff29723c6..563a88a75 100644 --- a/src/mobius/tasks/_base.py +++ b/src/mobius/tasks/_base.py @@ -103,16 +103,17 @@ def __repr__(self) -> str: def _make_graph( - inputs: list[ir.Value], name: str = "main_graph", ) -> tuple[ir.Graph, GraphBuilder]: """Create an empty graph and its builder. + Inputs should be added after creation via ``builder.input()``. + Returns: ``(graph, builder)`` — call ``builder.op`` to get the op handle. """ graph = ir.Graph( - inputs, + [], [], nodes=[], name=name, @@ -237,35 +238,36 @@ def build_decoder_from_embeds( seq_len = ir.SymbolicDim("sequence_len") past_seq_len = ir.SymbolicDim("past_sequence_len") - inputs_embeds = ir.Value( - name="inputs_embeds", - shape=ir.Shape([batch, seq_len, config.hidden_size]), - type=ir.TensorType(config.dtype), + graph, builder = _make_graph() + inputs_embeds = builder.input( + "inputs_embeds", + dtype=config.dtype, + shape=[batch, seq_len, config.hidden_size], ) - attention_mask = ir.Value( - name="attention_mask", - shape=ir.Shape([batch, "past_seq_len + seq_len"]), - type=ir.TensorType(ir.DataType.INT64), + attention_mask = builder.input( + "attention_mask", + dtype=ir.DataType.INT64, + shape=[batch, "past_seq_len + seq_len"], ) # MRoPE: 3D position IDs (temporal, height, width) — shape [3, batch, seq_len] # Standard: shape [batch, seq_len] - position_ids = ir.Value( - name="position_ids", - shape=ir.Shape([3, batch, seq_len] if mrope else [batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), + position_ids = builder.input( + "position_ids", + dtype=ir.DataType.INT64, + shape=[3, batch, seq_len] if mrope else [batch, seq_len], ) - graph_inputs = [inputs_embeds, attention_mask, position_ids] - if hybrid: - cache_inputs, past_key_values = _make_hybrid_cache_inputs( + past_key_values = _make_hybrid_cache_inputs( + builder, config, config.dtype, batch, past_seq_len, ) else: - cache_inputs, past_key_values = _make_kv_cache_inputs( + past_key_values = _make_kv_cache_inputs( + builder, config.num_hidden_layers, config.num_key_value_heads, config.head_dim, @@ -273,9 +275,7 @@ def build_decoder_from_embeds( batch, past_seq_len, ) - graph_inputs.extend(cache_inputs) - graph, builder = _make_graph(graph_inputs) logits, present_key_values = decoder( builder.op, inputs_embeds=inputs_embeds, @@ -284,12 +284,11 @@ def build_decoder_from_embeds( past_key_values=past_key_values, ) - logits.name = "logits" - graph.outputs.append(logits) + builder.add_output(logits, "logits") if hybrid: _register_hybrid_cache_outputs( - graph, + builder, present_key_values, config.layer_types or [], ) @@ -297,7 +296,7 @@ def build_decoder_from_embeds( _register_linear_attention_functions(model, config) return model else: - _register_kv_cache_outputs(graph, present_key_values) + _register_kv_cache_outputs(builder, present_key_values) return _make_model(graph) @@ -328,24 +327,23 @@ def build_embedding_from_features( seq_len = ir.SymbolicDim("sequence_len") num_feature_tokens = ir.SymbolicDim("num_feature_tokens") - input_ids = ir.Value( - name="input_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), + graph, builder = _make_graph(name="embedding") + input_ids = builder.input( + "input_ids", + dtype=ir.DataType.INT64, + shape=[batch, seq_len], ) - features = ir.Value( - name=feature_name, - shape=ir.Shape([num_feature_tokens, feature_dim]), - type=ir.TensorType(config.dtype), + features = builder.input( + feature_name, + dtype=config.dtype, + shape=[num_feature_tokens, feature_dim], ) - graph, builder = _make_graph([input_ids, features], name="embedding") inputs_embeds = embedding( builder.op, input_ids=input_ids, **{feature_name: features}, ) - inputs_embeds.name = "inputs_embeds" - graph.outputs.append(inputs_embeds) + builder.add_output(inputs_embeds, "inputs_embeds") return _make_model(graph) diff --git a/src/mobius/tasks/_cache_utils.py b/src/mobius/tasks/_cache_utils.py index 6a2d6875d..5f843be13 100644 --- a/src/mobius/tasks/_cache_utils.py +++ b/src/mobius/tasks/_cache_utils.py @@ -14,6 +14,7 @@ from typing import NamedTuple import onnx_ir as ir +from onnxscript._internal.builder import GraphBuilder from mobius._configs import BaseModelConfig @@ -64,6 +65,7 @@ def linear_attention_dims(config: BaseModelConfig) -> LinearAttentionDims: def _make_kv_cache_inputs( + builder: GraphBuilder, num_layers: int, num_kv_heads: int, head_dim: int, @@ -74,42 +76,41 @@ def _make_kv_cache_inputs( prefix: str = "past_key_values", key_head_dim: int | None = None, value_head_dim: int | None = None, -) -> tuple[list[ir.Value], list[tuple[ir.Value, ir.Value]]]: +) -> list[tuple[ir.Value, ir.Value]]: """Create KV cache input values for ``num_layers`` layers. + Uses ``builder.input()`` to create and register graph inputs directly. + Args: + builder: The graph builder to register inputs on. key_head_dim: Head dim for keys. Defaults to ``head_dim``. For MLA attention, this is ``qk_nope_head_dim + qk_rope_head_dim``. value_head_dim: Head dim for values. Defaults to ``head_dim``. For MLA attention, this is ``v_head_dim``. Returns: - ``(flat_inputs, kv_pairs)`` where *flat_inputs* is a flat list - suitable for extending ``graph_inputs`` and *kv_pairs* is a list - of ``(key, value)`` tuples for passing to the module. + A list of ``(key, value)`` tuples for passing to the module. """ k_dim = key_head_dim if key_head_dim is not None else head_dim v_dim = value_head_dim if value_head_dim is not None else head_dim - flat: list[ir.Value] = [] pairs: list[tuple[ir.Value, ir.Value]] = [] for i in range(num_layers): - past_key = ir.Value( - name=f"{prefix}.{i}.key", - shape=ir.Shape([batch, num_kv_heads, past_seq_len, k_dim]), - type=ir.TensorType(dtype), + past_key = builder.input( + f"{prefix}.{i}.key", + dtype=dtype, + shape=[batch, num_kv_heads, past_seq_len, k_dim], ) - past_value = ir.Value( - name=f"{prefix}.{i}.value", - shape=ir.Shape([batch, num_kv_heads, past_seq_len, v_dim]), - type=ir.TensorType(dtype), + past_value = builder.input( + f"{prefix}.{i}.value", + dtype=dtype, + shape=[batch, num_kv_heads, past_seq_len, v_dim], ) - flat.extend([past_key, past_value]) pairs.append((past_key, past_value)) - return flat, pairs + return pairs def _register_kv_cache_outputs( - graph: ir.Graph, + builder: GraphBuilder, present_key_values: list[tuple[ir.Value, ir.Value]], *, prefix: str = "present", @@ -120,23 +121,23 @@ def _register_kv_cache_outputs( that runs during model optimization. """ for i, (present_key, present_value) in enumerate(present_key_values): - present_key.name = f"{prefix}.{i}.key" - present_value.name = f"{prefix}.{i}.value" - - graph.outputs.append(present_key) - graph.outputs.append(present_value) + builder.add_output(present_key, f"{prefix}.{i}.key") + builder.add_output(present_value, f"{prefix}.{i}.value") def _make_hybrid_cache_inputs( + builder: GraphBuilder, config: BaseModelConfig, dtype: ir.DataType, batch: ir.SymbolicDim, past_seq_len: ir.SymbolicDim, *, prefix: str = "past_key_values", -) -> tuple[list[ir.Value], list[StatePair]]: +) -> list[StatePair]: """Create cache inputs for hybrid models with mixed layer types. + Uses ``builder.input()`` to create and register graph inputs directly. + Supported layer types: ``"full_attention"`` — standard KV cache (key + value). ``"lightning_attention"`` — single recurrent state only; no conv_state. @@ -146,13 +147,10 @@ def _make_hybrid_cache_inputs( ``"mlp"`` — stateless, produces ``(None, None)`` pair. Returns: - ``(flat_inputs, state_pairs)`` — *flat_inputs* contains only - the ``ir.Value`` entries (no graph inputs for MLP layers); - *state_pairs* has one entry per layer, with ``(None, None)`` + A list of state pairs, one per layer, with ``(None, None)`` for stateless MLP layers. """ layer_types = config.layer_types or [] - flat: list[ir.Value] = [] pairs: list[StatePair] = [] # DeltaNet dimensions from config (computed once via shared helper) @@ -187,91 +185,79 @@ def _make_hybrid_cache_inputs( if ltype == "lightning_attention": # Lightning Attention: single recurrent state only (no conv_state) # State: (B, num_heads, head_dim, head_dim) — square matrix accumulator - rec_state = ir.Value( - name=f"{prefix}.{i}.recurrent_state", - shape=ir.Shape( - [batch, config.num_attention_heads, config.head_dim, config.head_dim] - ), - type=ir.TensorType(dtype), + rec_state = builder.input( + f"{prefix}.{i}.recurrent_state", + dtype=dtype, + shape=[batch, config.num_attention_heads, config.head_dim, config.head_dim], ) - flat.append(rec_state) pairs.append((rec_state,)) # 1-tuple: lightning has no conv_state elif ltype == "linear_attention": - conv_state = ir.Value( - name=f"{prefix}.{i}.conv_state", - shape=ir.Shape([batch, dims.conv_dim, dims.conv_kernel - 1]), - type=ir.TensorType(dtype), + conv_state = builder.input( + f"{prefix}.{i}.conv_state", + dtype=dtype, + shape=[batch, dims.conv_dim, dims.conv_kernel - 1], ) - rec_state = ir.Value( - name=f"{prefix}.{i}.recurrent_state", - shape=ir.Shape([batch, dims.num_v_heads, dims.head_k_dim, dims.head_v_dim]), - type=ir.TensorType(dtype), + rec_state = builder.input( + f"{prefix}.{i}.recurrent_state", + dtype=dtype, + shape=[batch, dims.num_v_heads, dims.head_k_dim, dims.head_v_dim], ) - flat.extend([conv_state, rec_state]) pairs.append((conv_state, rec_state)) elif ltype == "conv": # ShortConv layers: conv_state only (no SSM state) # State: (batch, hidden_size, short_conv_kernel - 1) short_conv_kernel = getattr(config, "short_conv_kernel", 3) - conv_state = ir.Value( - name=f"{prefix}.{i}.conv_state", - shape=ir.Shape([batch, config.hidden_size, short_conv_kernel - 1]), - type=ir.TensorType(dtype), + conv_state = builder.input( + f"{prefix}.{i}.conv_state", + dtype=dtype, + shape=[batch, config.hidden_size, short_conv_kernel - 1], ) - flat.append(conv_state) pairs.append((conv_state,)) # 1-tuple: conv has no second state elif ltype in ("mlp", "moe"): # MLP and MoE layers are stateless — no cache inputs needed pairs.append((None, None)) elif ltype == "mamba": - conv_state = ir.Value( - name=f"{prefix}.{i}.conv_state", - shape=ir.Shape([batch, mamba_d_inner, mamba_d_conv - 1]), - type=ir.TensorType(dtype), + conv_state = builder.input( + f"{prefix}.{i}.conv_state", + dtype=dtype, + shape=[batch, mamba_d_inner, mamba_d_conv - 1], ) - ssm_state = ir.Value( - name=f"{prefix}.{i}.ssm_state", - shape=ir.Shape([batch, mamba_d_inner, mamba_d_state]), - type=ir.TensorType(dtype), + ssm_state = builder.input( + f"{prefix}.{i}.ssm_state", + dtype=dtype, + shape=[batch, mamba_d_inner, mamba_d_state], ) - flat.extend([conv_state, ssm_state]) pairs.append((conv_state, ssm_state)) elif ltype == "mamba2": - conv_state = ir.Value( - name=f"{prefix}.{i}.conv_state", - shape=ir.Shape([batch, mamba2_conv_dim, mamba_d_conv - 1]), - type=ir.TensorType(dtype), + conv_state = builder.input( + f"{prefix}.{i}.conv_state", + dtype=dtype, + shape=[batch, mamba2_conv_dim, mamba_d_conv - 1], ) - ssm_state = ir.Value( - name=f"{prefix}.{i}.ssm_state", - shape=ir.Shape([batch, mamba2_n_heads, mamba2_d_state, mamba2_d_head]), - type=ir.TensorType(dtype), + ssm_state = builder.input( + f"{prefix}.{i}.ssm_state", + dtype=dtype, + shape=[batch, mamba2_n_heads, mamba2_d_state, mamba2_d_head], ) - flat.extend([conv_state, ssm_state]) pairs.append((conv_state, ssm_state)) else: - past_key = ir.Value( - name=f"{prefix}.{i}.key", - shape=ir.Shape( - [batch, config.num_key_value_heads, past_seq_len, config.head_dim] - ), - type=ir.TensorType(dtype), + past_key = builder.input( + f"{prefix}.{i}.key", + dtype=dtype, + shape=[batch, config.num_key_value_heads, past_seq_len, config.head_dim], ) - past_value = ir.Value( - name=f"{prefix}.{i}.value", - shape=ir.Shape( - [batch, config.num_key_value_heads, past_seq_len, config.head_dim] - ), - type=ir.TensorType(dtype), + past_value = builder.input( + f"{prefix}.{i}.value", + dtype=dtype, + shape=[batch, config.num_key_value_heads, past_seq_len, config.head_dim], ) - flat.extend([past_key, past_value]) pairs.append((past_key, past_value)) - return flat, pairs + return pairs def _register_hybrid_cache_outputs( - graph: ir.Graph, + builder: GraphBuilder, present_key_values: list[tuple[ir.Value, ...]], layer_types: list[str], *, @@ -295,27 +281,22 @@ def _register_hybrid_cache_outputs( if ltype == "lightning_attention": # Single recurrent state only (no conv_state for lightning) (state_a,) = states - state_a.name = f"{prefix}.{i}.recurrent_state" - graph.outputs.append(state_a) + builder.add_output(state_a, f"{prefix}.{i}.recurrent_state") elif ltype == "conv": # ShortConv: single conv_state only (state_a,) = states - state_a.name = f"{prefix}.{i}.conv_state" - graph.outputs.append(state_a) + builder.add_output(state_a, f"{prefix}.{i}.conv_state") else: state_a, state_b = states if ltype == "linear_attention": - state_a.name = f"{prefix}.{i}.conv_state" - state_b.name = f"{prefix}.{i}.recurrent_state" + builder.add_output(state_a, f"{prefix}.{i}.conv_state") + builder.add_output(state_b, f"{prefix}.{i}.recurrent_state") elif ltype in ("mamba", "mamba2"): - state_a.name = f"{prefix}.{i}.conv_state" - state_b.name = f"{prefix}.{i}.ssm_state" + builder.add_output(state_a, f"{prefix}.{i}.conv_state") + builder.add_output(state_b, f"{prefix}.{i}.ssm_state") else: - state_a.name = f"{prefix}.{i}.key" - state_b.name = f"{prefix}.{i}.value" - - graph.outputs.append(state_a) - graph.outputs.append(state_b) + builder.add_output(state_a, f"{prefix}.{i}.key") + builder.add_output(state_b, f"{prefix}.{i}.value") def _register_linear_attention_functions( diff --git a/src/mobius/tasks/_causal_lm.py b/src/mobius/tasks/_causal_lm.py index e6456c6b2..9fb968002 100644 --- a/src/mobius/tasks/_causal_lm.py +++ b/src/mobius/tasks/_causal_lm.py @@ -7,6 +7,7 @@ import onnx_ir as ir from onnxscript import nn +from onnxscript._internal.builder import GraphBuilder from mobius._configs import ArchitectureConfig from mobius._model_package import ModelPackage @@ -110,23 +111,19 @@ def build( batch = ir.SymbolicDim("batch") seq_len = ir.SymbolicDim("sequence_len") + # --- Build graph first, then create inputs via builder --- + graph, builder = _make_graph() + op = builder.op + # --- Inputs common to both modes --- - input_ids = ir.Value( - name="input_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), - ) - position_ids = ir.Value( - name="position_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), - ) + input_ids = builder.input("input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) # --- Cache setup (static vs dynamic) --- if static: attention_mask = None - graph_inputs = [input_ids, position_ids] - cache_inputs, past_key_values = _make_static_cache_inputs( + position_ids = builder.input("position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) + past_key_values = _make_static_cache_inputs( + builder, config.num_hidden_layers, config.num_key_value_heads, config.head_dim, @@ -136,12 +133,10 @@ def build( ) else: past_seq_len = ir.SymbolicDim("past_sequence_len") - attention_mask = ir.Value( - name="attention_mask", - shape=ir.Shape([batch, "past_seq_len + seq_len"]), - type=ir.TensorType(ir.DataType.INT64), + attention_mask = builder.input( + "attention_mask", dtype=ir.DataType.INT64, shape=[batch, "past_seq_len + seq_len"], ) - graph_inputs = [input_ids, attention_mask, position_ids] + position_ids = builder.input("position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) # MLA attention: K/V heads equal q heads (no GQA reduction in # latent space). The ONNX Attention op is called with @@ -154,7 +149,8 @@ def build( config.num_attention_heads if use_mla else config.num_key_value_heads ) - cache_inputs, past_key_values = _make_kv_cache_inputs( + past_key_values = _make_kv_cache_inputs( + builder, config.num_hidden_layers, num_kv_cache_heads, config.head_dim, @@ -166,12 +162,6 @@ def build( value_head_dim=config.v_head_dim or None, ) - graph_inputs.extend(cache_inputs) - - # --- Build graph, invoke module, collect outputs --- - graph, builder = _make_graph(graph_inputs) - op = builder.op - logits, present_key_values = module( op, input_ids=input_ids, @@ -180,18 +170,17 @@ def build( past_key_values=past_key_values, ) - logits.name = "logits" - graph.outputs.append(logits) + builder.add_output(logits, "logits") # --- Output registration (static vs dynamic) --- if static: _register_static_cache_outputs( - graph, + builder, present_key_values, ) else: _register_kv_cache_outputs( - graph, + builder, present_key_values, ) @@ -228,34 +217,22 @@ def build( seq_len = ir.SymbolicDim("sequence_len") past_seq_len = ir.SymbolicDim("past_sequence_len") - input_ids = ir.Value( - name="input_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), - ) - attention_mask = ir.Value( - name="attention_mask", - shape=ir.Shape([batch, "past_seq_len + seq_len"]), - type=ir.TensorType(ir.DataType.INT64), - ) - position_ids = ir.Value( - name="position_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), - ) + graph, builder = _make_graph() + op = builder.op - graph_inputs = [input_ids, attention_mask, position_ids] + input_ids = builder.input("input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) + attention_mask = builder.input( + "attention_mask", dtype=ir.DataType.INT64, shape=[batch, "past_seq_len + seq_len"], + ) + position_ids = builder.input("position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) - cache_inputs, past_key_values = _make_hybrid_cache_inputs( + past_key_values = _make_hybrid_cache_inputs( + builder, config, config.dtype, batch, past_seq_len, ) - graph_inputs.extend(cache_inputs) - - graph, builder = _make_graph(graph_inputs) - op = builder.op logits, present_key_values = module( op, @@ -265,10 +242,9 @@ def build( past_key_values=past_key_values, ) - logits.name = "logits" - graph.outputs.append(logits) + builder.add_output(logits, "logits") _register_hybrid_cache_outputs( - graph, + builder, present_key_values, config.layer_types or [], ) @@ -279,51 +255,45 @@ def build( def _make_static_cache_inputs( + builder: GraphBuilder, num_layers: int, num_key_value_heads: int, head_dim: int, dtype: ir.DataType, batch: ir.SymbolicDim, max_seq_len: int, -) -> tuple[list[ir.Value], list[StaticCacheState]]: +) -> list[StaticCacheState]: """Create static KV cache inputs for ``num_layers`` layers. + Uses ``builder.input()`` to create and register graph inputs directly. + Returns: - ``(flat_inputs, static_caches)`` where *flat_inputs* is a flat - list suitable for extending ``graph_inputs``, and - *static_caches* is a list of :class:`StaticCacheState` tuples - for passing to the module via ``past_key_values``. + A list of :class:`StaticCacheState` tuples for passing to the + module via ``past_key_values``. """ kv_hidden = num_key_value_heads * head_dim - flat: list[ir.Value] = [] cache_pairs: list[tuple[ir.Value, ir.Value]] = [] for i in range(num_layers): - key_cache = ir.Value( - name=f"key_cache.{i}", - shape=ir.Shape([batch, max_seq_len, kv_hidden]), - type=ir.TensorType(dtype), + key_cache = builder.input( + f"key_cache.{i}", + dtype=dtype, + shape=[batch, max_seq_len, kv_hidden], ) - value_cache = ir.Value( - name=f"value_cache.{i}", - shape=ir.Shape([batch, max_seq_len, kv_hidden]), - type=ir.TensorType(dtype), + value_cache = builder.input( + f"value_cache.{i}", + dtype=dtype, + shape=[batch, max_seq_len, kv_hidden], ) - flat.extend([key_cache, value_cache]) cache_pairs.append((key_cache, value_cache)) # Shared inputs across all layers - write_indices = ir.Value( - name="write_indices", - shape=ir.Shape([batch]), - type=ir.TensorType(ir.DataType.INT64), + write_indices = builder.input( + "write_indices", dtype=ir.DataType.INT64, shape=[batch], ) - nonpad_kv_seqlen = ir.Value( - name="nonpad_kv_seqlen", - shape=ir.Shape([batch]), - type=ir.TensorType(ir.DataType.INT64), + nonpad_kv_seqlen = builder.input( + "nonpad_kv_seqlen", dtype=ir.DataType.INT64, shape=[batch], ) - flat.extend([write_indices, nonpad_kv_seqlen]) # Build StaticCacheState for each layer (shared indices) static_caches: list[StaticCacheState] = [] @@ -337,11 +307,11 @@ def _make_static_cache_inputs( ) ) - return flat, static_caches + return static_caches def _register_static_cache_outputs( - graph: ir.Graph, + builder: GraphBuilder, present_key_values: list[tuple[ir.Value, ir.Value]], ) -> None: """Name and register static cache outputs on the graph. @@ -350,10 +320,8 @@ def _register_static_cache_outputs( that runs during model optimization. """ for i, (updated_key, updated_value) in enumerate(present_key_values): - updated_key.name = f"updated_key_cache.{i}" - updated_value.name = f"updated_value_cache.{i}" - graph.outputs.append(updated_key) - graph.outputs.append(updated_value) + builder.add_output(updated_key, f"updated_key_cache.{i}") + builder.add_output(updated_value, f"updated_value_cache.{i}") def _validate_static_cache_support(module: nn.Module) -> None: diff --git a/src/mobius/tasks/_codec.py b/src/mobius/tasks/_codec.py index 604574297..95b535d86 100644 --- a/src/mobius/tasks/_codec.py +++ b/src/mobius/tasks/_codec.py @@ -65,17 +65,12 @@ def _build_decoder( num_q = ir.SymbolicDim("num_quantizers") seq_len = ir.SymbolicDim("sequence_len") - codes = ir.Value( - name="codes", - shape=ir.Shape([batch, num_q, seq_len]), - type=ir.TensorType(ir.DataType.INT64), - ) + graph, builder = _make_graph() + codes = builder.input("codes", dtype=ir.DataType.INT64, shape=[batch, num_q, seq_len]) - graph, builder = _make_graph([codes]) waveform = decoder(builder.op, codes) - waveform.name = "waveform" - graph.outputs.append(waveform) + builder.add_output(waveform, "waveform") return _make_model(graph) def _build_encoder( @@ -93,15 +88,10 @@ def _build_encoder( batch = ir.SymbolicDim("batch") audio_len = ir.SymbolicDim("audio_length") - waveform = ir.Value( - name="waveform", - shape=ir.Shape([batch, 1, audio_len]), - type=ir.TensorType(ir.DataType.FLOAT), - ) + graph, builder = _make_graph(name="encoder") + waveform = builder.input("waveform", dtype=ir.DataType.FLOAT, shape=[batch, 1, audio_len]) - graph, builder = _make_graph([waveform], name="encoder") codes = encoder(builder.op, waveform) - codes.name = "codes" - graph.outputs.append(codes) + builder.add_output(codes, "codes") return _make_model(graph) diff --git a/src/mobius/tasks/_controlnet.py b/src/mobius/tasks/_controlnet.py index c0c82bdd5..5315a2e7f 100644 --- a/src/mobius/tasks/_controlnet.py +++ b/src/mobius/tasks/_controlnet.py @@ -28,34 +28,14 @@ def build( module, config: ControlNetConfig, ) -> ModelPackage: - sample = ir.Value( - name="sample", - type=ir.TensorType(ir.DataType.FLOAT), - shape=ir.Shape(("batch", config.in_channels, "height", "width")), - ) - timestep = ir.Value( - name="timestep", - type=ir.TensorType(ir.DataType.INT64), - shape=ir.Shape(("batch",)), - ) - encoder_hidden_states = ir.Value( - name="encoder_hidden_states", - type=ir.TensorType(ir.DataType.FLOAT), - shape=ir.Shape(("batch", "sequence_length", config.cross_attention_dim)), - ) - controlnet_cond = ir.Value( - name="controlnet_cond", - type=ir.TensorType(ir.DataType.FLOAT), - shape=ir.Shape( - ("batch", config.conditioning_channels, "cond_height", "cond_width") - ), - ) - - graph, builder = _make_graph( - [sample, timestep, encoder_hidden_states, controlnet_cond] - ) + graph, builder = _make_graph() op = builder.op + sample = builder.input("sample", dtype=ir.DataType.FLOAT, shape=["batch", config.in_channels, "height", "width"]) + timestep = builder.input("timestep", dtype=ir.DataType.INT64, shape=["batch"]) + encoder_hidden_states = builder.input("encoder_hidden_states", dtype=ir.DataType.FLOAT, shape=["batch", "sequence_length", config.cross_attention_dim]) + controlnet_cond = builder.input("controlnet_cond", dtype=ir.DataType.FLOAT, shape=["batch", config.conditioning_channels, "cond_height", "cond_width"]) + down_outputs, mid_output = module( op, sample=sample, @@ -66,9 +46,7 @@ def build( # Register outputs for i, out in enumerate(down_outputs): - out.name = f"down_block_res_{i}" - graph.outputs.append(out) - mid_output.name = "mid_block_res" - graph.outputs.append(mid_output) + builder.add_output(out, f"down_block_res_{i}") + builder.add_output(mid_output, "mid_block_res") return ModelPackage({"model": _make_model(graph)}, config=config) diff --git a/src/mobius/tasks/_denoising.py b/src/mobius/tasks/_denoising.py index 96823f214..257ecef69 100644 --- a/src/mobius/tasks/_denoising.py +++ b/src/mobius/tasks/_denoising.py @@ -28,25 +28,13 @@ def build( module, config: UNet2DConfig, ) -> ModelPackage: - sample = ir.Value( - name="sample", - type=ir.TensorType(ir.DataType.FLOAT), - shape=ir.Shape(("batch", config.in_channels, "height", "width")), - ) - timestep = ir.Value( - name="timestep", - type=ir.TensorType(ir.DataType.INT64), - shape=ir.Shape(("batch",)), - ) - encoder_hidden_states = ir.Value( - name="encoder_hidden_states", - type=ir.TensorType(ir.DataType.FLOAT), - shape=ir.Shape(("batch", "sequence_length", config.cross_attention_dim)), - ) - - graph, builder = _make_graph([sample, timestep, encoder_hidden_states]) + graph, builder = _make_graph() op = builder.op + sample = builder.input("sample", dtype=ir.DataType.FLOAT, shape=["batch", config.in_channels, "height", "width"]) + timestep = builder.input("timestep", dtype=ir.DataType.INT64, shape=["batch"]) + encoder_hidden_states = builder.input("encoder_hidden_states", dtype=ir.DataType.FLOAT, shape=["batch", "sequence_length", config.cross_attention_dim]) + noise_pred = module( op, sample=sample, @@ -54,7 +42,6 @@ def build( encoder_hidden_states=encoder_hidden_states, ) - noise_pred.name = "noise_pred" - graph.outputs.append(noise_pred) + builder.add_output(noise_pred, "noise_pred") return ModelPackage({"model": _make_model(graph)}, config=config) diff --git a/src/mobius/tasks/_feature_extraction.py b/src/mobius/tasks/_feature_extraction.py index 872181a11..78f7d88a4 100644 --- a/src/mobius/tasks/_feature_extraction.py +++ b/src/mobius/tasks/_feature_extraction.py @@ -37,25 +37,13 @@ def build( batch = ir.SymbolicDim("batch") seq_len = ir.SymbolicDim("sequence_len") - input_ids = ir.Value( - name="input_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), - ) - attention_mask = ir.Value( - name="attention_mask", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), - ) - token_type_ids = ir.Value( - name="token_type_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), - ) - - graph, builder = _make_graph([input_ids, attention_mask, token_type_ids]) + graph, builder = _make_graph() op = builder.op + input_ids = builder.input("input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) + attention_mask = builder.input("attention_mask", dtype=ir.DataType.INT64, shape=[batch, seq_len]) + token_type_ids = builder.input("token_type_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) + last_hidden_state = module( op, input_ids=input_ids, @@ -63,7 +51,6 @@ def build( token_type_ids=token_type_ids, ) - last_hidden_state.name = "last_hidden_state" - graph.outputs.append(last_hidden_state) + builder.add_output(last_hidden_state, "last_hidden_state") return ModelPackage({"model": _make_model(graph)}, config=config) diff --git a/src/mobius/tasks/_gemma4.py b/src/mobius/tasks/_gemma4.py index ccb73e12f..51189c06e 100644 --- a/src/mobius/tasks/_gemma4.py +++ b/src/mobius/tasks/_gemma4.py @@ -23,6 +23,7 @@ import onnx_ir as ir from onnxscript import nn +from onnxscript._internal.builder import GraphBuilder from mobius._configs import Gemma4Config from mobius._model_package import ModelPackage @@ -37,12 +38,15 @@ def _make_gemma4_kv_cache_inputs( + graph_builder: GraphBuilder, config: Gemma4Config, batch: ir.SymbolicDim, past_seq_len: ir.SymbolicDim, -) -> tuple[list[ir.Value], list[tuple[ir.Value, ir.Value]]]: +) -> list[tuple[ir.Value, ir.Value]]: """Create per-layer KV cache inputs accounting for dual head_dim and KV sharing. + Uses ``graph_builder.input()`` to create and register graph inputs directly. + Local (sliding_attention) layers use ``config.head_dim``; global (full_attention) layers use ``config.global_head_dim``. @@ -66,25 +70,23 @@ def _make_gemma4_kv_cache_inputs( f"num_hidden_layers ({config.num_hidden_layers})" ) - flat: list[ir.Value] = [] pairs: list[tuple[ir.Value, ir.Value]] = [] for i in range(num_kv_layers): layer_type = layer_types[i] if i < len(layer_types) else "sliding_attention" hd = global_head_dim if layer_type == "full_attention" else local_head_dim kv_heads = config.num_key_value_heads - past_key = ir.Value( - name=f"past_key_values.{i}.key", - shape=ir.Shape([batch, kv_heads, past_seq_len, hd]), - type=ir.TensorType(config.dtype), + past_key = graph_builder.input( + f"past_key_values.{i}.key", + dtype=config.dtype, + shape=[batch, kv_heads, past_seq_len, hd], ) - past_value = ir.Value( - name=f"past_key_values.{i}.value", - shape=ir.Shape([batch, kv_heads, past_seq_len, hd]), - type=ir.TensorType(config.dtype), + past_value = graph_builder.input( + f"past_key_values.{i}.value", + dtype=config.dtype, + shape=[batch, kv_heads, past_seq_len, hd], ) - flat.extend([past_key, past_value]) pairs.append((past_key, past_value)) - return flat, pairs + return pairs class Gemma4TextCausalLMTask(ModelTask): @@ -117,28 +119,26 @@ def build( seq_len = ir.SymbolicDim("sequence_len") past_seq_len = ir.SymbolicDim("past_sequence_len") - input_ids = ir.Value( - name="input_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), + graph, graph_builder = _make_graph() + op = graph_builder.op + + input_ids = graph_builder.input( + "input_ids", + dtype=ir.DataType.INT64, + shape=[batch, seq_len], ) - attention_mask = ir.Value( - name="attention_mask", - shape=ir.Shape([batch, "past_seq_len + seq_len"]), - type=ir.TensorType(ir.DataType.INT64), + attention_mask = graph_builder.input( + "attention_mask", + dtype=ir.DataType.INT64, + shape=[batch, "past_seq_len + seq_len"], ) - position_ids = ir.Value( - name="position_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), + position_ids = graph_builder.input( + "position_ids", + dtype=ir.DataType.INT64, + shape=[batch, seq_len], ) - graph_inputs = [input_ids, attention_mask, position_ids] - kv_inputs, past_key_values = _make_gemma4_kv_cache_inputs(config, batch, past_seq_len) - graph_inputs.extend(kv_inputs) - - graph, graph_builder = _make_graph(graph_inputs) - op = graph_builder.op + past_key_values = _make_gemma4_kv_cache_inputs(graph_builder, config, batch, past_seq_len) logits, present_key_values = module( op, @@ -147,9 +147,8 @@ def build( position_ids=position_ids, past_key_values=past_key_values, ) - logits.name = "logits" - graph.outputs.append(logits) - _register_kv_cache_outputs(graph, present_key_values) + graph_builder.add_output(logits, "logits") + _register_kv_cache_outputs(graph_builder, present_key_values) return ModelPackage({"model": _make_model(graph)}, config=config) @@ -226,34 +225,31 @@ def _build_decoder( seq_len = ir.SymbolicDim("sequence_len") past_seq_len = ir.SymbolicDim("past_sequence_len") - inputs_embeds = ir.Value( - name="inputs_embeds", - shape=ir.Shape([batch, seq_len, config.hidden_size]), - type=ir.TensorType(config.dtype), + graph, graph_builder = _make_graph(name="decoder") + op = graph_builder.op + + inputs_embeds = graph_builder.input( + "inputs_embeds", + dtype=config.dtype, + shape=[batch, seq_len, config.hidden_size], ) - attention_mask = ir.Value( - name="attention_mask", - shape=ir.Shape([batch, "past_seq_len + seq_len"]), - type=ir.TensorType(ir.DataType.INT64), + attention_mask = graph_builder.input( + "attention_mask", + dtype=ir.DataType.INT64, + shape=[batch, "past_seq_len + seq_len"], ) - position_ids = ir.Value( - name="position_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), + position_ids = graph_builder.input( + "position_ids", + dtype=ir.DataType.INT64, + shape=[batch, seq_len], ) - input_ids = ir.Value( - name="input_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), + input_ids = graph_builder.input( + "input_ids", + dtype=ir.DataType.INT64, + shape=[batch, seq_len], ) - graph_inputs = [inputs_embeds, attention_mask, position_ids, input_ids] - - kv_inputs, past_key_values = _make_gemma4_kv_cache_inputs(config, batch, past_seq_len) - graph_inputs.extend(kv_inputs) - - graph, graph_builder = _make_graph(graph_inputs, name="decoder") - op = graph_builder.op + past_key_values = _make_gemma4_kv_cache_inputs(graph_builder, config, batch, past_seq_len) logits, present_key_values = decoder( op, @@ -264,9 +260,8 @@ def _build_decoder( past_key_values=past_key_values, ) - logits.name = "logits" - graph.outputs.append(logits) - _register_kv_cache_outputs(graph, present_key_values) + graph_builder.add_output(logits, "logits") + _register_kv_cache_outputs(graph_builder, present_key_values) return _make_model(graph) @@ -293,30 +288,27 @@ def _build_vision( patch_size = config.vision.patch_size or 16 if config.vision else 16 pixel_dim = 3 * patch_size * patch_size - pixel_values = ir.Value( - name="pixel_values", - shape=ir.Shape([batch, num_patches, pixel_dim]), - type=ir.TensorType(config.dtype), + graph, graph_builder = _make_graph(name="vision_encoder") + op = graph_builder.op + + pixel_values = graph_builder.input( + "pixel_values", + dtype=config.dtype, + shape=[batch, num_patches, pixel_dim], ) - pixel_position_ids = ir.Value( - name="pixel_position_ids", - shape=ir.Shape([batch, num_patches, 2]), - type=ir.TensorType(ir.DataType.INT64), + pixel_position_ids = graph_builder.input( + "pixel_position_ids", + dtype=ir.DataType.INT64, + shape=[batch, num_patches, 2], ) - graph_inputs = [pixel_values, pixel_position_ids] - - graph, graph_builder = _make_graph(graph_inputs, name="vision_encoder") - op = graph_builder.op - image_features = vision( op, pixel_values=pixel_values, pixel_position_ids=pixel_position_ids, ) - image_features.name = "image_features" - graph.outputs.append(image_features) + graph_builder.add_output(image_features, "image_features") return _make_model(graph) @@ -348,33 +340,29 @@ def _build_audio( time = ir.SymbolicDim("time") input_size = (config.audio.input_size if config.audio else None) or 128 - input_features = ir.Value( - name="input_features", - shape=ir.Shape([batch, time, input_size]), - type=ir.TensorType(config.dtype), - ) - input_features_mask = ir.Value( - name="input_features_mask", - shape=ir.Shape([batch, time]), - type=ir.TensorType(ir.DataType.BOOL), - ) + graph, graph_builder = _make_graph(name="audio_encoder") + op = graph_builder.op - graph, graph_builder = _make_graph( - [input_features, input_features_mask], name="audio_encoder" + input_features = graph_builder.input( + "input_features", + dtype=config.dtype, + shape=[batch, time, input_size], + ) + input_features_mask = graph_builder.input( + "input_features_mask", + dtype=ir.DataType.BOOL, + shape=[batch, time], ) - op = graph_builder.op audio_features, downsampled_mask = audio( op, input_features, input_features_mask=input_features_mask, ) - audio_features.name = "audio_features" - graph.outputs.append(audio_features) + graph_builder.add_output(audio_features, "audio_features") if downsampled_mask is not None: - downsampled_mask.name = "audio_features_mask" - graph.outputs.append(downsampled_mask) + graph_builder.add_output(downsampled_mask, "audio_features_mask") return _make_model(graph) @@ -388,31 +376,29 @@ def _build_embedding( seq_len = ir.SymbolicDim("sequence_len") num_image_tokens = ir.SymbolicDim("num_image_tokens") - input_ids = ir.Value( - name="input_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), + graph, graph_builder = _make_graph(name="embedding") + op = graph_builder.op + + input_ids = graph_builder.input( + "input_ids", + dtype=ir.DataType.INT64, + shape=[batch, seq_len], ) - image_features = ir.Value( - name="image_features", - shape=ir.Shape([num_image_tokens, config.hidden_size]), - type=ir.TensorType(config.dtype), + image_features = graph_builder.input( + "image_features", + dtype=config.dtype, + shape=[num_image_tokens, config.hidden_size], ) - graph_inputs = [input_ids, image_features] audio_features_val: ir.Value | None = None if config.audio is not None: num_audio_tokens = ir.SymbolicDim("num_audio_tokens") - audio_features_val = ir.Value( - name="audio_features", - shape=ir.Shape([num_audio_tokens, config.hidden_size]), - type=ir.TensorType(config.dtype), + audio_features_val = graph_builder.input( + "audio_features", + dtype=config.dtype, + shape=[num_audio_tokens, config.hidden_size], ) - graph_inputs.append(audio_features_val) - - graph, graph_builder = _make_graph(graph_inputs, name="embedding") - op = graph_builder.op inputs_embeds = embedding( op, @@ -420,6 +406,5 @@ def _build_embedding( image_features=image_features, audio_features=audio_features_val, ) - inputs_embeds.name = "inputs_embeds" - graph.outputs.append(inputs_embeds) + graph_builder.add_output(inputs_embeds, "inputs_embeds") return _make_model(graph) diff --git a/src/mobius/tasks/_image_classification.py b/src/mobius/tasks/_image_classification.py index 44a575a97..e58610578 100644 --- a/src/mobius/tasks/_image_classification.py +++ b/src/mobius/tasks/_image_classification.py @@ -37,18 +37,13 @@ def build( image_size = getattr(config, "image_size", 224) num_channels = getattr(config, "num_channels", 3) - pixel_values = ir.Value( - name="pixel_values", - shape=ir.Shape([batch, num_channels, image_size, image_size]), - type=ir.TensorType(ir.DataType.FLOAT), - ) - - graph, builder = _make_graph([pixel_values]) + graph, builder = _make_graph() op = builder.op + pixel_values = builder.input("pixel_values", dtype=ir.DataType.FLOAT, shape=[batch, num_channels, image_size, image_size]) + last_hidden_state = module(op, pixel_values=pixel_values) - last_hidden_state.name = "last_hidden_state" - graph.outputs.append(last_hidden_state) + builder.add_output(last_hidden_state, "last_hidden_state") return ModelPackage({"model": _make_model(graph)}, config=config) diff --git a/src/mobius/tasks/_multimodal.py b/src/mobius/tasks/_multimodal.py index d96dd3e8b..f804d505c 100644 --- a/src/mobius/tasks/_multimodal.py +++ b/src/mobius/tasks/_multimodal.py @@ -52,39 +52,34 @@ def build( seq_len = ir.SymbolicDim("sequence_len") past_seq_len = ir.SymbolicDim("past_sequence_len") - input_ids = ir.Value( - name="input_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), + graph, builder = _make_graph() + op = builder.op + + input_ids = builder.input( + "input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len], ) - attention_mask = ir.Value( - name="attention_mask", - shape=ir.Shape([batch, "past_seq_len + seq_len"]), - type=ir.TensorType(ir.DataType.INT64), + attention_mask = builder.input( + "attention_mask", dtype=ir.DataType.INT64, + shape=[batch, "past_seq_len + seq_len"], ) - position_ids = ir.Value( - name="position_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), + position_ids = builder.input( + "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len], ) image_size = config.vision.image_size or 224 if config.vision else 224 - pixel_values = ir.Value( - name="pixel_values", - shape=ir.Shape([batch, 3, image_size, image_size]), - type=ir.TensorType(config.dtype), + pixel_values = builder.input( + "pixel_values", dtype=config.dtype, + shape=[batch, 3, image_size, image_size], ) audio_input_size = (config.audio.input_size if config.audio else None) or 80 - audio_features = ir.Value( - name="audio_features", - shape=ir.Shape([batch, "audio_seq_len", audio_input_size]), - type=ir.TensorType(config.dtype), + audio_features = builder.input( + "audio_features", dtype=config.dtype, + shape=[batch, "audio_seq_len", audio_input_size], ) - graph_inputs = [input_ids, attention_mask, position_ids, pixel_values, audio_features] - - kv_inputs, past_key_values = _make_kv_cache_inputs( + past_key_values = _make_kv_cache_inputs( + builder, config.num_hidden_layers, config.num_key_value_heads, config.head_dim, @@ -92,10 +87,6 @@ def build( batch, past_seq_len, ) - graph_inputs.extend(kv_inputs) - - graph, builder = _make_graph(graph_inputs) - op = builder.op logits, present_key_values = module( op, @@ -107,8 +98,7 @@ def build( past_key_values=past_key_values, ) - logits.name = "logits" - graph.outputs.append(logits) - _register_kv_cache_outputs(graph, present_key_values) + builder.add_output(logits, "logits") + _register_kv_cache_outputs(builder, present_key_values) return ModelPackage({"model": _make_model(graph)}, config=config) diff --git a/src/mobius/tasks/_object_detection.py b/src/mobius/tasks/_object_detection.py index 2d06f690d..4a3736ffa 100644 --- a/src/mobius/tasks/_object_detection.py +++ b/src/mobius/tasks/_object_detection.py @@ -38,20 +38,14 @@ def build( image_size = getattr(config, "image_size", 224) num_channels = getattr(config, "num_channels", 3) - pixel_values = ir.Value( - name="pixel_values", - shape=ir.Shape([batch, num_channels, image_size, image_size]), - type=ir.TensorType(ir.DataType.FLOAT), - ) - - graph, builder = _make_graph([pixel_values]) + graph, builder = _make_graph() op = builder.op + pixel_values = builder.input("pixel_values", dtype=ir.DataType.FLOAT, shape=[batch, num_channels, image_size, image_size]) + logits, pred_boxes = module(op, pixel_values=pixel_values) - logits.name = "logits" - pred_boxes.name = "pred_boxes" - graph.outputs.append(logits) - graph.outputs.append(pred_boxes) + builder.add_output(logits, "logits") + builder.add_output(pred_boxes, "pred_boxes") return ModelPackage({"model": _make_model(graph)}, config=config) diff --git a/src/mobius/tasks/_phi4mm_multimodal.py b/src/mobius/tasks/_phi4mm_multimodal.py index d35ab3fa2..99366bf76 100644 --- a/src/mobius/tasks/_phi4mm_multimodal.py +++ b/src/mobius/tasks/_phi4mm_multimodal.py @@ -84,22 +84,20 @@ def _build_vision( num_images = ir.SymbolicDim("num_images") image_size = (config.vision.image_size if config.vision else None) or 448 - pixel_values = ir.Value( - name="pixel_values", - shape=ir.Shape([batch, 3, image_size, image_size]), - type=ir.TensorType(config.dtype), + graph, builder = _make_graph(name="vision_encoder") + + pixel_values = builder.input( + "pixel_values", dtype=config.dtype, + shape=[batch, 3, image_size, image_size], ) - image_sizes = ir.Value( - name="image_sizes", - shape=ir.Shape([num_images, 2]), - type=ir.TensorType(ir.DataType.INT64), + image_sizes = builder.input( + "image_sizes", dtype=ir.DataType.INT64, + shape=[num_images, 2], ) - graph, builder = _make_graph([pixel_values, image_sizes], name="vision_encoder") image_features = vision(builder.op, pixel_values, image_sizes=image_sizes) - image_features.name = "image_features" - graph.outputs.append(image_features) + builder.add_output(image_features, "image_features") return _make_model(graph) def _build_speech( @@ -117,26 +115,21 @@ def _build_speech( num_audio_clips = ir.SymbolicDim("num_audio_clips") input_size = (config.audio.input_size if config.audio else None) or 80 - audio_embeds = ir.Value( - name="audio_embeds", - shape=ir.Shape([batch, audio_seq_len, input_size]), - type=ir.TensorType(config.dtype), + graph, builder = _make_graph(name="audio_encoder") + + audio_embeds = builder.input( + "audio_embeds", dtype=config.dtype, + shape=[batch, audio_seq_len, input_size], ) - audio_sizes = ir.Value( - name="audio_sizes", - shape=ir.Shape([num_audio_clips]), - type=ir.TensorType(ir.DataType.INT64), + audio_sizes = builder.input( + "audio_sizes", dtype=ir.DataType.INT64, + shape=[num_audio_clips], ) - audio_projection_mode = ir.Value( - name="audio_projection_mode", - shape=ir.Shape([]), - type=ir.TensorType(ir.DataType.INT64), + audio_projection_mode = builder.input( + "audio_projection_mode", dtype=ir.DataType.INT64, + shape=[], ) - graph, builder = _make_graph( - [audio_embeds, audio_sizes, audio_projection_mode], - name="audio_encoder", - ) speech_out = speech( builder.op, audio_embeds, @@ -144,8 +137,7 @@ def _build_speech( audio_projection_mode=audio_projection_mode, ) - speech_out.name = "audio_features" - graph.outputs.append(speech_out) + builder.add_output(speech_out, "audio_features") return _make_model(graph) def _build_embedding( @@ -159,26 +151,21 @@ def _build_embedding( num_image_tokens = ir.SymbolicDim("num_image_tokens") num_speech_tokens = ir.SymbolicDim("num_speech_tokens") - input_ids = ir.Value( - name="input_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), + graph, builder = _make_graph(name="embedding") + + input_ids = builder.input( + "input_ids", dtype=ir.DataType.INT64, + shape=[batch, seq_len], ) - image_features = ir.Value( - name="image_features", - shape=ir.Shape([num_image_tokens, config.hidden_size]), - type=ir.TensorType(config.dtype), + image_features = builder.input( + "image_features", dtype=config.dtype, + shape=[num_image_tokens, config.hidden_size], ) - audio_features = ir.Value( - name="audio_features", - shape=ir.Shape([num_speech_tokens, config.hidden_size]), - type=ir.TensorType(config.dtype), + audio_features = builder.input( + "audio_features", dtype=config.dtype, + shape=[num_speech_tokens, config.hidden_size], ) - graph, builder = _make_graph( - [input_ids, image_features, audio_features], - name="embedding", - ) inputs_embeds = embedding( builder.op, input_ids=input_ids, @@ -186,6 +173,5 @@ def _build_embedding( audio_features=audio_features, ) - inputs_embeds.name = "inputs_embeds" - graph.outputs.append(inputs_embeds) + builder.add_output(inputs_embeds, "inputs_embeds") return _make_model(graph) diff --git a/src/mobius/tasks/_qwen_image_vae.py b/src/mobius/tasks/_qwen_image_vae.py index a03b09f74..36f2a97d3 100644 --- a/src/mobius/tasks/_qwen_image_vae.py +++ b/src/mobius/tasks/_qwen_image_vae.py @@ -40,20 +40,15 @@ def _build_encoder_graph( module, config: QwenImageVAEConfig, ) -> ir.Model: - sample = ir.Value( - name="sample", - type=ir.TensorType(ir.DataType.FLOAT), - shape=ir.Shape(("batch", 3, "frames", "height", "width")), - ) - - graph, builder = _make_graph([sample], name="vae_encoder") + graph, builder = _make_graph(name="vae_encoder") op = builder.op + sample = builder.input("sample", dtype=ir.DataType.FLOAT, shape=["batch", 3, "frames", "height", "width"]) + hidden_states = module.encoder(op, sample) hidden_states = module.quant_conv(op, hidden_states) - hidden_states.name = "latent_dist" - graph.outputs.append(hidden_states) + builder.add_output(hidden_states, "latent_dist") return _make_model(graph) @@ -62,19 +57,14 @@ def _build_decoder_graph( module, config: QwenImageVAEConfig, ) -> ir.Model: - latent_sample = ir.Value( - name="latent_sample", - type=ir.TensorType(ir.DataType.FLOAT), - shape=ir.Shape(("batch", config.z_dim, "frames", "height", "width")), - ) - - graph, builder = _make_graph([latent_sample], name="vae_decoder") + graph, builder = _make_graph(name="vae_decoder") op = builder.op + latent_sample = builder.input("latent_sample", dtype=ir.DataType.FLOAT, shape=["batch", config.z_dim, "frames", "height", "width"]) + hidden_states = module.post_quant_conv(op, latent_sample) hidden_states = module.decoder(op, hidden_states) - hidden_states.name = "sample" - graph.outputs.append(hidden_states) + builder.add_output(hidden_states, "sample") return _make_model(graph) diff --git a/src/mobius/tasks/_seq2seq.py b/src/mobius/tasks/_seq2seq.py index 147f176ee..01f2152a1 100644 --- a/src/mobius/tasks/_seq2seq.py +++ b/src/mobius/tasks/_seq2seq.py @@ -18,9 +18,6 @@ _make_graph, _make_model, ) -from mobius.tasks._cache_utils import ( - _make_kv_cache_inputs, -) class Seq2SeqTask(ModelTask): @@ -58,26 +55,17 @@ def _build_encoder_graph( batch = ir.SymbolicDim("batch") seq_len = ir.SymbolicDim("sequence_len") - input_ids = ir.Value( - name="input_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), - ) - attention_mask = ir.Value( - name="attention_mask", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), - ) - - graph, builder = _make_graph([input_ids, attention_mask], name="encoder") + graph, builder = _make_graph(name="encoder") op = builder.op + input_ids = builder.input("input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) + attention_mask = builder.input("attention_mask", dtype=ir.DataType.INT64, shape=[batch, seq_len]) + encoder_hidden_states = module.encoder( op, input_ids=input_ids, attention_mask=attention_mask ) - encoder_hidden_states.name = "last_hidden_state" - graph.outputs.append(encoder_hidden_states) + builder.add_output(encoder_hidden_states, "last_hidden_state") return _make_model(graph) @@ -91,63 +79,50 @@ def _build_decoder_graph( enc_seq_len = ir.SymbolicDim("encoder_sequence_len") past_seq_len = ir.SymbolicDim("past_sequence_len") - input_ids = ir.Value( - name="input_ids", - shape=ir.Shape([batch, dec_seq_len]), - type=ir.TensorType(ir.DataType.INT64), + graph, builder = _make_graph() + op = builder.op + + input_ids = builder.input( + "input_ids", dtype=ir.DataType.INT64, shape=[batch, dec_seq_len], ) - encoder_hidden_states = ir.Value( - name="encoder_hidden_states", - shape=ir.Shape([batch, enc_seq_len, config.hidden_size]), - type=ir.TensorType(config.dtype), + encoder_hidden_states = builder.input( + "encoder_hidden_states", dtype=config.dtype, + shape=[batch, enc_seq_len, config.hidden_size], ) - attention_mask = ir.Value( - name="attention_mask", - shape=ir.Shape([batch, "past_seq_len + dec_seq_len"]), - type=ir.TensorType(ir.DataType.INT64), + attention_mask = builder.input( + "attention_mask", dtype=ir.DataType.INT64, + shape=[batch, "past_seq_len + dec_seq_len"], ) - graph_inputs = [input_ids, encoder_hidden_states, attention_mask] - num_heads = config.num_attention_heads head_dim = config.head_dim num_decoder_layers = getattr(config, "num_decoder_layers", config.num_hidden_layers) - # Self-attention KV cache - self_kv_inputs, past_self_kvs = _make_kv_cache_inputs( - num_decoder_layers, - num_heads, - head_dim, - config.dtype, - batch, - past_seq_len, - prefix="past_key_values", - ) - # Use .self. naming for seq2seq self-attention KVs - for v in self_kv_inputs: - idx = v.name.split(".")[1] - kv_type = v.name.rsplit(".", 1)[-1] - v.name = f"past_key_values.{idx}.self.{kv_type}" - graph_inputs.extend(self_kv_inputs) - - # Cross-attention KV cache - cross_kv_inputs, cross_past_kvs = _make_kv_cache_inputs( - num_decoder_layers, - num_heads, - head_dim, - config.dtype, - batch, - enc_seq_len, - prefix="past_key_values", - ) - for v in cross_kv_inputs: - idx = v.name.split(".")[1] - kv_type = v.name.rsplit(".", 1)[-1] - v.name = f"past_key_values.{idx}.cross.{kv_type}" - graph_inputs.extend(cross_kv_inputs) - - graph, builder = _make_graph(graph_inputs) - op = builder.op + # Self-attention KV cache (named past_key_values.{i}.self.key/value) + past_self_kvs: list[tuple[ir.Value, ir.Value]] = [] + for i in range(num_decoder_layers): + past_key = builder.input( + f"past_key_values.{i}.self.key", dtype=config.dtype, + shape=[batch, num_heads, past_seq_len, head_dim], + ) + past_value = builder.input( + f"past_key_values.{i}.self.value", dtype=config.dtype, + shape=[batch, num_heads, past_seq_len, head_dim], + ) + past_self_kvs.append((past_key, past_value)) + + # Cross-attention KV cache (named past_key_values.{i}.cross.key/value) + cross_past_kvs: list[tuple[ir.Value, ir.Value]] = [] + for i in range(num_decoder_layers): + past_key = builder.input( + f"past_key_values.{i}.cross.key", dtype=config.dtype, + shape=[batch, num_heads, enc_seq_len, head_dim], + ) + past_value = builder.input( + f"past_key_values.{i}.cross.value", dtype=config.dtype, + shape=[batch, num_heads, enc_seq_len, head_dim], + ) + cross_past_kvs.append((past_key, past_value)) logits, present_self_kvs, present_cross_kvs = module.decoder( op, @@ -158,17 +133,14 @@ def _build_decoder_graph( cross_past_key_values=cross_past_kvs, ) - logits.name = "logits" - graph.outputs.append(logits) + builder.add_output(logits, "logits") for i, (k, v) in enumerate(present_self_kvs): - k.name = f"present.{i}.self.key" - v.name = f"present.{i}.self.value" - graph.outputs.extend([k, v]) + builder.add_output(k, f"present.{i}.self.key") + builder.add_output(v, f"present.{i}.self.value") for i, (k, v) in enumerate(present_cross_kvs): - k.name = f"present.{i}.cross.key" - v.name = f"present.{i}.cross.value" - graph.outputs.extend([k, v]) + builder.add_output(k, f"present.{i}.cross.key") + builder.add_output(v, f"present.{i}.cross.value") return _make_model(graph) diff --git a/src/mobius/tasks/_speech_language.py b/src/mobius/tasks/_speech_language.py index 9a0b79d53..76bec8743 100644 --- a/src/mobius/tasks/_speech_language.py +++ b/src/mobius/tasks/_speech_language.py @@ -82,15 +82,14 @@ def _build_audio_encoder( mel_seq = ir.SymbolicDim("mel_sequence_len") n_mels = (config.audio.num_mel_bins if config.audio else None) or 128 - input_features = ir.Value( - name="input_features", - shape=ir.Shape([batch, n_mels, mel_seq]), - type=ir.TensorType(config.dtype), + graph, builder = _make_graph(name="audio_encoder") + + input_features = builder.input( + "input_features", dtype=config.dtype, + shape=[batch, n_mels, mel_seq], ) - graph, builder = _make_graph([input_features], name="audio_encoder") audio_features = audio_encoder(builder.op, input_features) - audio_features.name = "audio_features" - graph.outputs.append(audio_features) + builder.add_output(audio_features, "audio_features") return _make_model(graph) diff --git a/src/mobius/tasks/_speech_to_text.py b/src/mobius/tasks/_speech_to_text.py index ebed6edab..5aa5ac4e0 100644 --- a/src/mobius/tasks/_speech_to_text.py +++ b/src/mobius/tasks/_speech_to_text.py @@ -69,19 +69,17 @@ def _build_encoder( batch = ir.SymbolicDim("batch") audio_seq_len = ir.SymbolicDim("audio_seq_len") - input_features = ir.Value( - name="input_features", - shape=ir.Shape([batch, config.num_mel_bins, audio_seq_len]), - type=ir.TensorType(ir.DataType.FLOAT), - ) - - graph, builder = _make_graph([input_features], name="encoder") + graph, builder = _make_graph(name="encoder") op = builder.op + input_features = builder.input( + "input_features", dtype=ir.DataType.FLOAT, + shape=[batch, config.num_mel_bins, audio_seq_len], + ) + encoder_hidden_states = encoder(op, input_features=input_features) - encoder_hidden_states.name = "encoder_hidden_states" - graph.outputs.append(encoder_hidden_states) + builder.add_output(encoder_hidden_states, "encoder_hidden_states") return _make_model(graph) @@ -95,25 +93,22 @@ def _build_decoder( past_seq_len = ir.SymbolicDim("past_sequence_len") encoder_seq_len = ir.SymbolicDim("encoder_sequence_len") - decoder_input_ids = ir.Value( - name="decoder_input_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), + graph, builder = _make_graph() + op = builder.op + + decoder_input_ids = builder.input( + "decoder_input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len], ) - encoder_hidden_states = ir.Value( - name="encoder_hidden_states", - shape=ir.Shape([batch, encoder_seq_len, config.hidden_size]), - type=ir.TensorType(config.dtype), + encoder_hidden_states = builder.input( + "encoder_hidden_states", dtype=config.dtype, + shape=[batch, encoder_seq_len, config.hidden_size], ) - position_ids = ir.Value( - name="position_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), + position_ids = builder.input( + "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len], ) - graph_inputs = [decoder_input_ids, encoder_hidden_states, position_ids] - - kv_inputs, past_key_values = _make_kv_cache_inputs( + past_key_values = _make_kv_cache_inputs( + builder, config.num_hidden_layers, config.num_key_value_heads, config.head_dim, @@ -121,10 +116,6 @@ def _build_decoder( batch, past_seq_len, ) - graph_inputs.extend(kv_inputs) - - graph, builder = _make_graph(graph_inputs) - op = builder.op logits, present_key_values = decoder( op, @@ -134,8 +125,7 @@ def _build_decoder( past_key_values=past_key_values, ) - logits.name = "logits" - graph.outputs.append(logits) - _register_kv_cache_outputs(graph, present_key_values) + builder.add_output(logits, "logits") + _register_kv_cache_outputs(builder, present_key_values) return _make_model(graph) diff --git a/src/mobius/tasks/_ssm_causal_lm.py b/src/mobius/tasks/_ssm_causal_lm.py index 63f6f5046..4109d10e6 100644 --- a/src/mobius/tasks/_ssm_causal_lm.py +++ b/src/mobius/tasks/_ssm_causal_lm.py @@ -49,42 +49,32 @@ def _build_ssm_task( batch = ir.SymbolicDim("batch") seq_len = ir.SymbolicDim("sequence_len") - input_ids = ir.Value( - name="input_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), - ) - graph_inputs: list[ir.Value] = [input_ids] + graph, builder = _make_graph() + + input_ids = builder.input("input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) past_states: list[tuple[ir.Value, ir.Value]] = [] for i in range(config.num_hidden_layers): - conv_state = ir.Value( - name=f"past_states.{i}.conv_state", - shape=ir.Shape([batch, *conv_state_shape]), - type=ir.TensorType(config.dtype), + conv_state = builder.input( + f"past_states.{i}.conv_state", dtype=config.dtype, + shape=[batch, *conv_state_shape], ) - ssm_state = ir.Value( - name=f"past_states.{i}.ssm_state", - shape=ir.Shape([batch, *ssm_state_shape]), - type=ir.TensorType(config.dtype), + ssm_state = builder.input( + f"past_states.{i}.ssm_state", dtype=config.dtype, + shape=[batch, *ssm_state_shape], ) - graph_inputs.extend([conv_state, ssm_state]) past_states.append((conv_state, ssm_state)) - graph, builder = _make_graph(graph_inputs) logits, present_states = module( builder.op, input_ids=input_ids, past_states=past_states, ) - logits.name = "logits" - graph.outputs.append(logits) + builder.add_output(logits, "logits") for i, (conv_state, ssm_state) in enumerate(present_states): - conv_state.name = f"present.{i}.conv_state" - ssm_state.name = f"present.{i}.ssm_state" - graph.outputs.append(conv_state) - graph.outputs.append(ssm_state) + builder.add_output(conv_state, f"present.{i}.conv_state") + builder.add_output(ssm_state, f"present.{i}.ssm_state") return ModelPackage({"model": _make_model(graph)}, config=config) diff --git a/src/mobius/tasks/_tts.py b/src/mobius/tasks/_tts.py index 729269954..184c4106b 100644 --- a/src/mobius/tasks/_tts.py +++ b/src/mobius/tasks/_tts.py @@ -90,26 +90,24 @@ def _build_talker( seq_len = ir.SymbolicDim("sequence_len") past_seq_len = ir.SymbolicDim("past_sequence_len") - inputs_embeds = ir.Value( - name="inputs_embeds", - shape=ir.Shape([batch, seq_len, config.hidden_size]), - type=ir.TensorType(config.dtype), + graph, builder = _make_graph(name="talker") + + inputs_embeds = builder.input( + "inputs_embeds", dtype=config.dtype, + shape=[batch, seq_len, config.hidden_size], ) - attention_mask = ir.Value( - name="attention_mask", - shape=ir.Shape([batch, "past_seq_len + seq_len"]), - type=ir.TensorType(ir.DataType.INT64), + attention_mask = builder.input( + "attention_mask", dtype=ir.DataType.INT64, + shape=[batch, "past_seq_len + seq_len"], ) # MRoPE: 3D position_ids (3, batch, seq_len) - position_ids = ir.Value( - name="position_ids", - shape=ir.Shape([3, batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), + position_ids = builder.input( + "position_ids", dtype=ir.DataType.INT64, + shape=[3, batch, seq_len], ) - graph_inputs = [inputs_embeds, attention_mask, position_ids] - - kv_inputs, past_key_values = _make_kv_cache_inputs( + past_key_values = _make_kv_cache_inputs( + builder, config.num_hidden_layers, config.num_key_value_heads, config.head_dim, @@ -117,9 +115,7 @@ def _build_talker( batch, past_seq_len, ) - graph_inputs.extend(kv_inputs) - graph, builder = _make_graph(graph_inputs, name="talker") logits, last_hidden_state, present_key_values = talker( builder.op, inputs_embeds=inputs_embeds, @@ -128,11 +124,9 @@ def _build_talker( past_key_values=past_key_values, ) - logits.name = "logits" - last_hidden_state.name = "last_hidden_state" - graph.outputs.append(logits) - graph.outputs.append(last_hidden_state) - _register_kv_cache_outputs(graph, present_key_values) + builder.add_output(logits, "logits") + builder.add_output(last_hidden_state, "last_hidden_state") + _register_kv_cache_outputs(builder, present_key_values) return _make_model(graph) def _build_code_predictor( @@ -166,39 +160,29 @@ def _build_code_predictor( seq_len = ir.SymbolicDim("sequence_len") past_seq_len = ir.SymbolicDim("past_sequence_len") + graph, builder = _make_graph(name="code_predictor") + # Pre-embedded input in talker_hidden space (constructed by # generation loop). The model projects to cp_hidden internally. - inputs_embeds = ir.Value( - name="inputs_embeds", - shape=ir.Shape([batch, seq_len, config.hidden_size]), - type=ir.TensorType(config.dtype), + inputs_embeds = builder.input( + "inputs_embeds", dtype=config.dtype, + shape=[batch, seq_len, config.hidden_size], ) # Step index: selects which lm_head to use (0..14) - step_index = ir.Value( - name="step_index", - shape=ir.Shape([]), - type=ir.TensorType(ir.DataType.INT64), + step_index = builder.input( + "step_index", dtype=ir.DataType.INT64, shape=[], ) - attention_mask = ir.Value( - name="attention_mask", - shape=ir.Shape([batch, "past_seq_len + seq_len"]), - type=ir.TensorType(ir.DataType.INT64), + attention_mask = builder.input( + "attention_mask", dtype=ir.DataType.INT64, + shape=[batch, "past_seq_len + seq_len"], ) # 1D RoPE: 2D position_ids (batch, seq_len) - position_ids = ir.Value( - name="position_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), + position_ids = builder.input( + "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len], ) - graph_inputs = [ - inputs_embeds, - step_index, - attention_mask, - position_ids, - ] - - kv_inputs, past_key_values = _make_kv_cache_inputs( + past_key_values = _make_kv_cache_inputs( + builder, cp_num_hidden_layers, cp_num_key_value_heads, cp_head_dim, @@ -206,9 +190,7 @@ def _build_code_predictor( batch, past_seq_len, ) - graph_inputs.extend(kv_inputs) - graph, builder = _make_graph(graph_inputs, name="code_predictor") logits, present_key_values, codec_embeddings = code_predictor( builder.op, inputs_embeds=inputs_embeds, @@ -218,14 +200,12 @@ def _build_code_predictor( past_key_values=past_key_values, ) - logits.name = "logits" - graph.outputs.append(logits) + builder.add_output(logits, "logits") # Expose stacked codec embeddings for generation loop to extract. # The Identity node ensures renaming the output doesn't affect # the initializer name used for weight loading. - codec_embeddings.name = "codec_embeddings" - graph.outputs.append(codec_embeddings) - _register_kv_cache_outputs(graph, present_key_values) + builder.add_output(codec_embeddings, "codec_embeddings") + _register_kv_cache_outputs(builder, present_key_values) return _make_model(graph) def _build_embedding( @@ -238,28 +218,23 @@ def _build_embedding( text_seq = ir.SymbolicDim("text_sequence_len") codec_seq = ir.SymbolicDim("codec_sequence_len") - text_ids = ir.Value( - name="text_ids", - shape=ir.Shape([batch, text_seq]), - type=ir.TensorType(ir.DataType.INT64), + graph, builder = _make_graph(name="embedding") + + text_ids = builder.input( + "text_ids", dtype=ir.DataType.INT64, shape=[batch, text_seq], ) - codec_ids = ir.Value( - name="codec_ids", - shape=ir.Shape([batch, codec_seq]), - type=ir.TensorType(ir.DataType.INT64), + codec_ids = builder.input( + "codec_ids", dtype=ir.DataType.INT64, shape=[batch, codec_seq], ) - graph, builder = _make_graph([text_ids, codec_ids], name="embedding") text_embeds, codec_embeds = embedding( builder.op, text_ids=text_ids, codec_ids=codec_ids, ) - text_embeds.name = "text_embeds" - codec_embeds.name = "codec_embeds" - graph.outputs.append(text_embeds) - graph.outputs.append(codec_embeds) + builder.add_output(text_embeds, "text_embeds") + builder.add_output(codec_embeds, "codec_embeds") return _make_model(graph) def _build_speaker_encoder( @@ -274,15 +249,13 @@ def _build_speaker_encoder( se = tts.speaker_encoder if tts else None mel_dim = se.mel_dim if se else 128 - mel_input = ir.Value( - name="mel_input", - shape=ir.Shape([batch, mel_seq, mel_dim]), - type=ir.TensorType(config.dtype), + graph, builder = _make_graph(name="speaker_encoder") + + mel_input = builder.input( + "mel_input", dtype=config.dtype, shape=[batch, mel_seq, mel_dim], ) - graph, builder = _make_graph([mel_input], name="speaker_encoder") speaker_embedding = speaker_encoder(builder.op, mel_input) - speaker_embedding.name = "speaker_embedding" - graph.outputs.append(speaker_embedding) + builder.add_output(speaker_embedding, "speaker_embedding") return _make_model(graph) diff --git a/src/mobius/tasks/_vae.py b/src/mobius/tasks/_vae.py index 68f828884..fd40840c1 100644 --- a/src/mobius/tasks/_vae.py +++ b/src/mobius/tasks/_vae.py @@ -40,21 +40,16 @@ def _build_encoder_graph( module, config: VAEConfig, ) -> ir.Model: - sample = ir.Value( - name="sample", - type=ir.TensorType(ir.DataType.FLOAT), - shape=ir.Shape(("batch", config.in_channels, "height", "width")), - ) - - graph, builder = _make_graph([sample], name="vae_encoder") + graph, builder = _make_graph(name="vae_encoder") op = builder.op + sample = builder.input("sample", dtype=ir.DataType.FLOAT, shape=["batch", config.in_channels, "height", "width"]) + hidden_states = module.encoder(op, sample=sample) if module.quant_conv is not None: hidden_states = module.quant_conv(op, hidden_states) - hidden_states.name = "latent_dist" - graph.outputs.append(hidden_states) + builder.add_output(hidden_states, "latent_dist") return _make_model(graph) @@ -63,21 +58,16 @@ def _build_decoder_graph( module, config: VAEConfig, ) -> ir.Model: - latent_sample = ir.Value( - name="latent_sample", - type=ir.TensorType(ir.DataType.FLOAT), - shape=ir.Shape(("batch", config.latent_channels, "height", "width")), - ) - - graph, builder = _make_graph([latent_sample], name="vae_decoder") + graph, builder = _make_graph(name="vae_decoder") op = builder.op + latent_sample = builder.input("latent_sample", dtype=ir.DataType.FLOAT, shape=["batch", config.latent_channels, "height", "width"]) + hidden_states = latent_sample if module.post_quant_conv is not None: hidden_states = module.post_quant_conv(op, hidden_states) hidden_states = module.decoder(op, latent_sample=hidden_states) - hidden_states.name = "sample" - graph.outputs.append(hidden_states) + builder.add_output(hidden_states, "sample") return _make_model(graph) diff --git a/src/mobius/tasks/_video_denoising.py b/src/mobius/tasks/_video_denoising.py index 44a9a6e28..d64ee6c07 100644 --- a/src/mobius/tasks/_video_denoising.py +++ b/src/mobius/tasks/_video_denoising.py @@ -29,33 +29,13 @@ def build( module, config: CogVideoXConfig, ) -> ModelPackage: - sample = ir.Value( - name="sample", - type=ir.TensorType(ir.DataType.FLOAT), - shape=ir.Shape( - ( - "batch", - "num_frames", - config.in_channels, - "height", - "width", - ) - ), - ) - timestep = ir.Value( - name="timestep", - type=ir.TensorType(ir.DataType.INT64), - shape=ir.Shape(("batch",)), - ) - encoder_hidden_states = ir.Value( - name="encoder_hidden_states", - type=ir.TensorType(ir.DataType.FLOAT), - shape=ir.Shape(("batch", "sequence_length", config.cross_attention_dim)), - ) - - graph, builder = _make_graph([sample, timestep, encoder_hidden_states]) + graph, builder = _make_graph() op = builder.op + sample = builder.input("sample", dtype=ir.DataType.FLOAT, shape=["batch", "num_frames", config.in_channels, "height", "width"]) + timestep = builder.input("timestep", dtype=ir.DataType.INT64, shape=["batch"]) + encoder_hidden_states = builder.input("encoder_hidden_states", dtype=ir.DataType.FLOAT, shape=["batch", "sequence_length", config.cross_attention_dim]) + noise_pred = module( op, sample=sample, @@ -63,7 +43,6 @@ def build( encoder_hidden_states=encoder_hidden_states, ) - noise_pred.name = "noise_pred" - graph.outputs.append(noise_pred) + builder.add_output(noise_pred, "noise_pred") return ModelPackage({"model": _make_model(graph)}, config=config) diff --git a/src/mobius/tasks/_vision_language.py b/src/mobius/tasks/_vision_language.py index 9be0ca28c..b08722b98 100644 --- a/src/mobius/tasks/_vision_language.py +++ b/src/mobius/tasks/_vision_language.py @@ -47,49 +47,36 @@ def build( past_seq_len = ir.SymbolicDim("past_sequence_len") total_patches = ir.SymbolicDim("total_patches") - input_ids = ir.Value( - name="input_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), + graph, builder = _make_graph() + op = builder.op + + input_ids = builder.input( + "input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len], ) - attention_mask = ir.Value( - name="attention_mask", - shape=ir.Shape([batch, "past_seq_len + seq_len"]), - type=ir.TensorType(ir.DataType.INT64), + attention_mask = builder.input( + "attention_mask", dtype=ir.DataType.INT64, + shape=[batch, "past_seq_len + seq_len"], ) # MRoPE: 3D position IDs (temporal, height, width) - position_ids = ir.Value( - name="position_ids", - shape=ir.Shape([3, batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), + position_ids = builder.input( + "position_ids", dtype=ir.DataType.INT64, shape=[3, batch, seq_len], ) # Flattened image patches patch_size = config.vision.patch_size or 16 if config.vision else 16 temporal_patch_size = config.temporal_patch_size in_channels = config.vision.in_channels if config.vision else 3 pixel_dim = in_channels * temporal_patch_size * patch_size * patch_size - pixel_values = ir.Value( - name="pixel_values", - shape=ir.Shape([total_patches, pixel_dim]), - type=ir.TensorType(config.dtype), + pixel_values = builder.input( + "pixel_values", dtype=config.dtype, shape=[total_patches, pixel_dim], ) # Image grid dimensions for position embedding interpolation num_images = ir.SymbolicDim("num_images") - grid_thw = ir.Value( - name="grid_thw", - shape=ir.Shape([num_images, 3]), - type=ir.TensorType(ir.DataType.INT64), + grid_thw = builder.input( + "grid_thw", dtype=ir.DataType.INT64, shape=[num_images, 3], ) - graph_inputs = [ - input_ids, - attention_mask, - position_ids, - pixel_values, - grid_thw, - ] - - kv_inputs, past_key_values = _make_kv_cache_inputs( + past_key_values = _make_kv_cache_inputs( + builder, config.num_hidden_layers, config.num_key_value_heads, config.head_dim, @@ -97,10 +84,6 @@ def build( batch, past_seq_len, ) - graph_inputs.extend(kv_inputs) - - graph, builder = _make_graph(graph_inputs) - op = builder.op logits, present_key_values = module( op, @@ -112,8 +95,7 @@ def build( past_key_values=past_key_values, ) - logits.name = "logits" - graph.outputs.append(logits) - _register_kv_cache_outputs(graph, present_key_values) + builder.add_output(logits, "logits") + _register_kv_cache_outputs(builder, present_key_values) return ModelPackage({"model": _make_model(graph)}, config=config) diff --git a/src/mobius/tasks/_vision_language_3model.py b/src/mobius/tasks/_vision_language_3model.py index 735d7e787..35277681e 100644 --- a/src/mobius/tasks/_vision_language_3model.py +++ b/src/mobius/tasks/_vision_language_3model.py @@ -82,17 +82,15 @@ def _build_vision( batch = ir.SymbolicDim("batch") image_size = (config.vision.image_size if config.vision else None) or 224 - pixel_values = ir.Value( - name="pixel_values", - shape=ir.Shape([batch, 3, image_size, image_size]), - type=ir.TensorType(config.dtype), + graph, graph_builder = _make_graph(name="vision_encoder") + pixel_values = graph_builder.input( + "pixel_values", + dtype=config.dtype, + shape=[batch, 3, image_size, image_size], ) - - graph, graph_builder = _make_graph([pixel_values], name="vision_encoder") image_features = vision(graph_builder.op, pixel_values=pixel_values) - image_features.name = "image_features" - graph.outputs.append(image_features) + graph_builder.add_output(image_features, "image_features") return _make_model(graph) @@ -135,28 +133,25 @@ def _build_vision( in_channels = config.vision.in_channels if config.vision else 3 pixel_dim = in_channels * temporal_patch_size * patch_size * patch_size - pixel_values = ir.Value( - name="pixel_values", - shape=ir.Shape([total_patches, pixel_dim]), - type=ir.TensorType(config.dtype), + graph, graph_builder = _make_graph(name="vision_encoder") + pixel_values = graph_builder.input( + "pixel_values", + dtype=config.dtype, + shape=[total_patches, pixel_dim], ) - image_grid_thw = ir.Value( - name="image_grid_thw", - shape=ir.Shape([num_images, 3]), - type=ir.TensorType(ir.DataType.INT64), + image_grid_thw = graph_builder.input( + "image_grid_thw", + dtype=ir.DataType.INT64, + shape=[num_images, 3], ) - graph, graph_builder = _make_graph( - [pixel_values, image_grid_thw], name="vision_encoder" - ) image_features = vision( graph_builder.op, pixel_values=pixel_values, image_grid_thw=image_grid_thw, ) - image_features.name = "image_features" - graph.outputs.append(image_features) + graph_builder.add_output(image_features, "image_features") return _make_model(graph) @@ -213,15 +208,12 @@ def _build_vision( height = ir.SymbolicDim("height") width = ir.SymbolicDim("width") - pixel_values = ir.Value( - name="pixel_values", - shape=ir.Shape([batch, 3, height, width]), - type=ir.TensorType(config.dtype), + graph, graph_builder = _make_graph(name="vision_encoder") + pixel_values = graph_builder.input( + "pixel_values", + dtype=config.dtype, + shape=[batch, 3, height, width], ) - - graph_inputs = [pixel_values] - - graph, graph_builder = _make_graph(graph_inputs, name="vision_encoder") op = graph_builder.op image_features = vision( @@ -234,8 +226,7 @@ def _build_vision( # vision encoder always processes one image at a time. image_features = op.Squeeze(image_features, [0]) - image_features.name = "image_features" - graph.outputs.append(image_features) + graph_builder.add_output(image_features, "image_features") return _make_model(graph) @@ -294,59 +285,48 @@ def _build_decoder( cross_seq_len = ir.SymbolicDim("cross_sequence_len") cross_past_seq_len = ir.SymbolicDim("cross_past_seq_len") - inputs_embeds = ir.Value( - name="inputs_embeds", - shape=ir.Shape([batch, seq_len, config.hidden_size]), - type=ir.TensorType(config.dtype), + graph, graph_builder = _make_graph() + inputs_embeds = graph_builder.input( + "inputs_embeds", + dtype=config.dtype, + shape=[batch, seq_len, config.hidden_size], ) - attention_mask = ir.Value( - name="attention_mask", - shape=ir.Shape([batch, "past_seq_len + seq_len"]), - type=ir.TensorType(ir.DataType.INT64), + attention_mask = graph_builder.input( + "attention_mask", + dtype=ir.DataType.INT64, + shape=[batch, "past_seq_len + seq_len"], ) - position_ids = ir.Value( - name="position_ids", - shape=ir.Shape([batch, seq_len]), - type=ir.TensorType(ir.DataType.INT64), + position_ids = graph_builder.input( + "position_ids", + dtype=ir.DataType.INT64, + shape=[batch, seq_len], ) # Vision features: full on prefill, empty (0-length) on decode - cross_attention_states = ir.Value( - name="cross_attention_states", - shape=ir.Shape([batch, cross_seq_len, config.hidden_size]), - type=ir.TensorType(config.dtype), + cross_attention_states = graph_builder.input( + "cross_attention_states", + dtype=config.dtype, + shape=[batch, cross_seq_len, config.hidden_size], ) - graph_inputs = [ - inputs_embeds, - attention_mask, - position_ids, - cross_attention_states, - ] - # Per-layer KV cache with separate dims for self-attention # (past_seq_len) and cross-attention (cross_past_seq_len) cross_attention_layers = set(config.cross_attention_layers or []) - flat_kv: list[ir.Value] = [] past_key_values: list[tuple[ir.Value, ir.Value]] = [] for i in range(config.num_hidden_layers): psl = cross_past_seq_len if i in cross_attention_layers else past_seq_len - past_key = ir.Value( - name=f"past_key_values.{i}.key", - shape=ir.Shape([batch, config.num_key_value_heads, psl, config.head_dim]), - type=ir.TensorType(config.dtype), + past_key = graph_builder.input( + f"past_key_values.{i}.key", + dtype=config.dtype, + shape=[batch, config.num_key_value_heads, psl, config.head_dim], ) - past_value = ir.Value( - name=f"past_key_values.{i}.value", - shape=ir.Shape([batch, config.num_key_value_heads, psl, config.head_dim]), - type=ir.TensorType(config.dtype), + past_value = graph_builder.input( + f"past_key_values.{i}.value", + dtype=config.dtype, + shape=[batch, config.num_key_value_heads, psl, config.head_dim], ) - flat_kv.extend([past_key, past_value]) past_key_values.append((past_key, past_value)) - graph_inputs.extend(flat_kv) - - graph, graph_builder = _make_graph(graph_inputs) op = graph_builder.op logits, present_key_values = decoder( @@ -358,8 +338,7 @@ def _build_decoder( past_key_values=past_key_values, ) - logits.name = "logits" - graph.outputs.append(logits) - _register_kv_cache_outputs(graph, present_key_values) + graph_builder.add_output(logits, "logits") + _register_kv_cache_outputs(graph_builder, present_key_values) return _make_model(graph) From 0e2a06583d56da036e00ea45b2c6ae427cd5790b Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 1 May 2026 18:41:06 +0000 Subject: [PATCH 2/2] Fix review findings: consistent naming, public imports, lint, docstring - 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 --- src/mobius/tasks/_adapter.py | 12 ++- src/mobius/tasks/_audio_feature_extraction.py | 4 +- src/mobius/tasks/_base.py | 4 +- src/mobius/tasks/_cache_utils.py | 2 +- src/mobius/tasks/_causal_lm.py | 31 +++++--- src/mobius/tasks/_codec.py | 4 +- src/mobius/tasks/_controlnet.py | 18 ++++- src/mobius/tasks/_denoising.py | 12 ++- src/mobius/tasks/_feature_extraction.py | 8 +- src/mobius/tasks/_gemma4.py | 79 +++++++++---------- src/mobius/tasks/_image_classification.py | 6 +- src/mobius/tasks/_multimodal.py | 17 ++-- src/mobius/tasks/_object_detection.py | 6 +- src/mobius/tasks/_phi4mm_multimodal.py | 24 ++++-- src/mobius/tasks/_qwen_image_vae.py | 10 ++- src/mobius/tasks/_seq2seq.py | 26 ++++-- src/mobius/tasks/_speech_language.py | 3 +- src/mobius/tasks/_speech_to_text.py | 14 +++- src/mobius/tasks/_ssm_causal_lm.py | 6 +- src/mobius/tasks/_tts.py | 35 +++++--- src/mobius/tasks/_vae.py | 12 ++- src/mobius/tasks/_video_denoising.py | 12 ++- src/mobius/tasks/_vision_language.py | 19 +++-- src/mobius/tasks/_vision_language_3model.py | 46 +++++------ 24 files changed, 273 insertions(+), 137 deletions(-) diff --git a/src/mobius/tasks/_adapter.py b/src/mobius/tasks/_adapter.py index 17730dbfc..f586beb55 100644 --- a/src/mobius/tasks/_adapter.py +++ b/src/mobius/tasks/_adapter.py @@ -29,10 +29,18 @@ def build( # Determine input shape based on adapter type if hasattr(config, "in_channels"): # T2I-Adapter: conditioning image input - condition = builder.input("condition", dtype=ir.DataType.FLOAT, shape=["batch", config.in_channels, "height", "width"]) + condition = builder.input( + "condition", + dtype=ir.DataType.FLOAT, + shape=["batch", config.in_channels, "height", "width"], + ) else: # IP-Adapter: image embedding input - condition = builder.input("image_embeds", dtype=ir.DataType.FLOAT, shape=["batch", config.image_embed_dim]) + condition = builder.input( + "image_embeds", + dtype=ir.DataType.FLOAT, + shape=["batch", config.image_embed_dim], + ) outputs = module(op, condition) diff --git a/src/mobius/tasks/_audio_feature_extraction.py b/src/mobius/tasks/_audio_feature_extraction.py index 493183c2f..51a62b51f 100644 --- a/src/mobius/tasks/_audio_feature_extraction.py +++ b/src/mobius/tasks/_audio_feature_extraction.py @@ -31,7 +31,9 @@ def build( graph, builder = _make_graph() op = builder.op - input_values = builder.input("input_values", dtype=ir.DataType.FLOAT, shape=["batch", "time"]) + input_values = builder.input( + "input_values", dtype=ir.DataType.FLOAT, shape=["batch", "time"] + ) last_hidden_state = module(op, input_values=input_values) diff --git a/src/mobius/tasks/_base.py b/src/mobius/tasks/_base.py index 563a88a75..b09276fcb 100644 --- a/src/mobius/tasks/_base.py +++ b/src/mobius/tasks/_base.py @@ -9,8 +9,7 @@ from typing import ClassVar import onnx_ir as ir -from onnxscript import nn -from onnxscript._internal.builder import GraphBuilder +from onnxscript import GraphBuilder, nn import mobius from mobius._configs import BaseModelConfig @@ -108,6 +107,7 @@ def _make_graph( """Create an empty graph and its builder. Inputs should be added after creation via ``builder.input()``. + Outputs should be registered via ``builder.add_output()``. Returns: ``(graph, builder)`` — call ``builder.op`` to get the op handle. diff --git a/src/mobius/tasks/_cache_utils.py b/src/mobius/tasks/_cache_utils.py index 5f843be13..41cf5071c 100644 --- a/src/mobius/tasks/_cache_utils.py +++ b/src/mobius/tasks/_cache_utils.py @@ -14,7 +14,7 @@ from typing import NamedTuple import onnx_ir as ir -from onnxscript._internal.builder import GraphBuilder +from onnxscript import GraphBuilder from mobius._configs import BaseModelConfig diff --git a/src/mobius/tasks/_causal_lm.py b/src/mobius/tasks/_causal_lm.py index 9fb968002..471d350b6 100644 --- a/src/mobius/tasks/_causal_lm.py +++ b/src/mobius/tasks/_causal_lm.py @@ -6,8 +6,7 @@ from __future__ import annotations import onnx_ir as ir -from onnxscript import nn -from onnxscript._internal.builder import GraphBuilder +from onnxscript import GraphBuilder, nn from mobius._configs import ArchitectureConfig from mobius._model_package import ModelPackage @@ -121,7 +120,9 @@ def build( # --- Cache setup (static vs dynamic) --- if static: attention_mask = None - position_ids = builder.input("position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) + position_ids = builder.input( + "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len] + ) past_key_values = _make_static_cache_inputs( builder, config.num_hidden_layers, @@ -134,9 +135,13 @@ def build( else: past_seq_len = ir.SymbolicDim("past_sequence_len") attention_mask = builder.input( - "attention_mask", dtype=ir.DataType.INT64, shape=[batch, "past_seq_len + seq_len"], + "attention_mask", + dtype=ir.DataType.INT64, + shape=[batch, "past_seq_len + seq_len"], + ) + position_ids = builder.input( + "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len] ) - position_ids = builder.input("position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) # MLA attention: K/V heads equal q heads (no GQA reduction in # latent space). The ONNX Attention op is called with @@ -222,9 +227,13 @@ def build( input_ids = builder.input("input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) attention_mask = builder.input( - "attention_mask", dtype=ir.DataType.INT64, shape=[batch, "past_seq_len + seq_len"], + "attention_mask", + dtype=ir.DataType.INT64, + shape=[batch, "past_seq_len + seq_len"], + ) + position_ids = builder.input( + "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len] ) - position_ids = builder.input("position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) past_key_values = _make_hybrid_cache_inputs( builder, @@ -289,10 +298,14 @@ def _make_static_cache_inputs( # Shared inputs across all layers write_indices = builder.input( - "write_indices", dtype=ir.DataType.INT64, shape=[batch], + "write_indices", + dtype=ir.DataType.INT64, + shape=[batch], ) nonpad_kv_seqlen = builder.input( - "nonpad_kv_seqlen", dtype=ir.DataType.INT64, shape=[batch], + "nonpad_kv_seqlen", + dtype=ir.DataType.INT64, + shape=[batch], ) # Build StaticCacheState for each layer (shared indices) diff --git a/src/mobius/tasks/_codec.py b/src/mobius/tasks/_codec.py index 95b535d86..ace4f3afb 100644 --- a/src/mobius/tasks/_codec.py +++ b/src/mobius/tasks/_codec.py @@ -89,7 +89,9 @@ def _build_encoder( audio_len = ir.SymbolicDim("audio_length") graph, builder = _make_graph(name="encoder") - waveform = builder.input("waveform", dtype=ir.DataType.FLOAT, shape=[batch, 1, audio_len]) + waveform = builder.input( + "waveform", dtype=ir.DataType.FLOAT, shape=[batch, 1, audio_len] + ) codes = encoder(builder.op, waveform) diff --git a/src/mobius/tasks/_controlnet.py b/src/mobius/tasks/_controlnet.py index 5315a2e7f..4354250fc 100644 --- a/src/mobius/tasks/_controlnet.py +++ b/src/mobius/tasks/_controlnet.py @@ -31,10 +31,22 @@ def build( graph, builder = _make_graph() op = builder.op - sample = builder.input("sample", dtype=ir.DataType.FLOAT, shape=["batch", config.in_channels, "height", "width"]) + sample = builder.input( + "sample", + dtype=ir.DataType.FLOAT, + shape=["batch", config.in_channels, "height", "width"], + ) timestep = builder.input("timestep", dtype=ir.DataType.INT64, shape=["batch"]) - encoder_hidden_states = builder.input("encoder_hidden_states", dtype=ir.DataType.FLOAT, shape=["batch", "sequence_length", config.cross_attention_dim]) - controlnet_cond = builder.input("controlnet_cond", dtype=ir.DataType.FLOAT, shape=["batch", config.conditioning_channels, "cond_height", "cond_width"]) + encoder_hidden_states = builder.input( + "encoder_hidden_states", + dtype=ir.DataType.FLOAT, + shape=["batch", "sequence_length", config.cross_attention_dim], + ) + controlnet_cond = builder.input( + "controlnet_cond", + dtype=ir.DataType.FLOAT, + shape=["batch", config.conditioning_channels, "cond_height", "cond_width"], + ) down_outputs, mid_output = module( op, diff --git a/src/mobius/tasks/_denoising.py b/src/mobius/tasks/_denoising.py index 257ecef69..18c6f758a 100644 --- a/src/mobius/tasks/_denoising.py +++ b/src/mobius/tasks/_denoising.py @@ -31,9 +31,17 @@ def build( graph, builder = _make_graph() op = builder.op - sample = builder.input("sample", dtype=ir.DataType.FLOAT, shape=["batch", config.in_channels, "height", "width"]) + sample = builder.input( + "sample", + dtype=ir.DataType.FLOAT, + shape=["batch", config.in_channels, "height", "width"], + ) timestep = builder.input("timestep", dtype=ir.DataType.INT64, shape=["batch"]) - encoder_hidden_states = builder.input("encoder_hidden_states", dtype=ir.DataType.FLOAT, shape=["batch", "sequence_length", config.cross_attention_dim]) + encoder_hidden_states = builder.input( + "encoder_hidden_states", + dtype=ir.DataType.FLOAT, + shape=["batch", "sequence_length", config.cross_attention_dim], + ) noise_pred = module( op, diff --git a/src/mobius/tasks/_feature_extraction.py b/src/mobius/tasks/_feature_extraction.py index 78f7d88a4..0d32b55e7 100644 --- a/src/mobius/tasks/_feature_extraction.py +++ b/src/mobius/tasks/_feature_extraction.py @@ -41,8 +41,12 @@ def build( op = builder.op input_ids = builder.input("input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) - attention_mask = builder.input("attention_mask", dtype=ir.DataType.INT64, shape=[batch, seq_len]) - token_type_ids = builder.input("token_type_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) + attention_mask = builder.input( + "attention_mask", dtype=ir.DataType.INT64, shape=[batch, seq_len] + ) + token_type_ids = builder.input( + "token_type_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len] + ) last_hidden_state = module( op, diff --git a/src/mobius/tasks/_gemma4.py b/src/mobius/tasks/_gemma4.py index 51189c06e..532272b0c 100644 --- a/src/mobius/tasks/_gemma4.py +++ b/src/mobius/tasks/_gemma4.py @@ -22,8 +22,7 @@ from __future__ import annotations import onnx_ir as ir -from onnxscript import nn -from onnxscript._internal.builder import GraphBuilder +from onnxscript import GraphBuilder, nn from mobius._configs import Gemma4Config from mobius._model_package import ModelPackage @@ -38,14 +37,14 @@ def _make_gemma4_kv_cache_inputs( - graph_builder: GraphBuilder, + builder: GraphBuilder, config: Gemma4Config, batch: ir.SymbolicDim, past_seq_len: ir.SymbolicDim, ) -> list[tuple[ir.Value, ir.Value]]: """Create per-layer KV cache inputs accounting for dual head_dim and KV sharing. - Uses ``graph_builder.input()`` to create and register graph inputs directly. + Uses ``builder.input()`` to create and register graph inputs directly. Local (sliding_attention) layers use ``config.head_dim``; global (full_attention) layers use ``config.global_head_dim``. @@ -75,12 +74,12 @@ def _make_gemma4_kv_cache_inputs( layer_type = layer_types[i] if i < len(layer_types) else "sliding_attention" hd = global_head_dim if layer_type == "full_attention" else local_head_dim kv_heads = config.num_key_value_heads - past_key = graph_builder.input( + past_key = builder.input( f"past_key_values.{i}.key", dtype=config.dtype, shape=[batch, kv_heads, past_seq_len, hd], ) - past_value = graph_builder.input( + past_value = builder.input( f"past_key_values.{i}.value", dtype=config.dtype, shape=[batch, kv_heads, past_seq_len, hd], @@ -119,26 +118,26 @@ def build( seq_len = ir.SymbolicDim("sequence_len") past_seq_len = ir.SymbolicDim("past_sequence_len") - graph, graph_builder = _make_graph() - op = graph_builder.op + graph, builder = _make_graph() + op = builder.op - input_ids = graph_builder.input( + input_ids = builder.input( "input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len], ) - attention_mask = graph_builder.input( + attention_mask = builder.input( "attention_mask", dtype=ir.DataType.INT64, shape=[batch, "past_seq_len + seq_len"], ) - position_ids = graph_builder.input( + position_ids = builder.input( "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len], ) - past_key_values = _make_gemma4_kv_cache_inputs(graph_builder, config, batch, past_seq_len) + past_key_values = _make_gemma4_kv_cache_inputs(builder, config, batch, past_seq_len) logits, present_key_values = module( op, @@ -147,8 +146,8 @@ def build( position_ids=position_ids, past_key_values=past_key_values, ) - graph_builder.add_output(logits, "logits") - _register_kv_cache_outputs(graph_builder, present_key_values) + builder.add_output(logits, "logits") + _register_kv_cache_outputs(builder, present_key_values) return ModelPackage({"model": _make_model(graph)}, config=config) @@ -225,31 +224,31 @@ def _build_decoder( seq_len = ir.SymbolicDim("sequence_len") past_seq_len = ir.SymbolicDim("past_sequence_len") - graph, graph_builder = _make_graph(name="decoder") - op = graph_builder.op + graph, builder = _make_graph(name="decoder") + op = builder.op - inputs_embeds = graph_builder.input( + inputs_embeds = builder.input( "inputs_embeds", dtype=config.dtype, shape=[batch, seq_len, config.hidden_size], ) - attention_mask = graph_builder.input( + attention_mask = builder.input( "attention_mask", dtype=ir.DataType.INT64, shape=[batch, "past_seq_len + seq_len"], ) - position_ids = graph_builder.input( + position_ids = builder.input( "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len], ) - input_ids = graph_builder.input( + input_ids = builder.input( "input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len], ) - past_key_values = _make_gemma4_kv_cache_inputs(graph_builder, config, batch, past_seq_len) + past_key_values = _make_gemma4_kv_cache_inputs(builder, config, batch, past_seq_len) logits, present_key_values = decoder( op, @@ -260,8 +259,8 @@ def _build_decoder( past_key_values=past_key_values, ) - graph_builder.add_output(logits, "logits") - _register_kv_cache_outputs(graph_builder, present_key_values) + builder.add_output(logits, "logits") + _register_kv_cache_outputs(builder, present_key_values) return _make_model(graph) @@ -288,15 +287,15 @@ def _build_vision( patch_size = config.vision.patch_size or 16 if config.vision else 16 pixel_dim = 3 * patch_size * patch_size - graph, graph_builder = _make_graph(name="vision_encoder") - op = graph_builder.op + graph, builder = _make_graph(name="vision_encoder") + op = builder.op - pixel_values = graph_builder.input( + pixel_values = builder.input( "pixel_values", dtype=config.dtype, shape=[batch, num_patches, pixel_dim], ) - pixel_position_ids = graph_builder.input( + pixel_position_ids = builder.input( "pixel_position_ids", dtype=ir.DataType.INT64, shape=[batch, num_patches, 2], @@ -308,7 +307,7 @@ def _build_vision( pixel_position_ids=pixel_position_ids, ) - graph_builder.add_output(image_features, "image_features") + builder.add_output(image_features, "image_features") return _make_model(graph) @@ -340,15 +339,15 @@ def _build_audio( time = ir.SymbolicDim("time") input_size = (config.audio.input_size if config.audio else None) or 128 - graph, graph_builder = _make_graph(name="audio_encoder") - op = graph_builder.op + graph, builder = _make_graph(name="audio_encoder") + op = builder.op - input_features = graph_builder.input( + input_features = builder.input( "input_features", dtype=config.dtype, shape=[batch, time, input_size], ) - input_features_mask = graph_builder.input( + input_features_mask = builder.input( "input_features_mask", dtype=ir.DataType.BOOL, shape=[batch, time], @@ -359,10 +358,10 @@ def _build_audio( input_features, input_features_mask=input_features_mask, ) - graph_builder.add_output(audio_features, "audio_features") + builder.add_output(audio_features, "audio_features") if downsampled_mask is not None: - graph_builder.add_output(downsampled_mask, "audio_features_mask") + builder.add_output(downsampled_mask, "audio_features_mask") return _make_model(graph) @@ -376,15 +375,15 @@ def _build_embedding( seq_len = ir.SymbolicDim("sequence_len") num_image_tokens = ir.SymbolicDim("num_image_tokens") - graph, graph_builder = _make_graph(name="embedding") - op = graph_builder.op + graph, builder = _make_graph(name="embedding") + op = builder.op - input_ids = graph_builder.input( + input_ids = builder.input( "input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len], ) - image_features = graph_builder.input( + image_features = builder.input( "image_features", dtype=config.dtype, shape=[num_image_tokens, config.hidden_size], @@ -394,7 +393,7 @@ def _build_embedding( if config.audio is not None: num_audio_tokens = ir.SymbolicDim("num_audio_tokens") - audio_features_val = graph_builder.input( + audio_features_val = builder.input( "audio_features", dtype=config.dtype, shape=[num_audio_tokens, config.hidden_size], @@ -406,5 +405,5 @@ def _build_embedding( image_features=image_features, audio_features=audio_features_val, ) - graph_builder.add_output(inputs_embeds, "inputs_embeds") + builder.add_output(inputs_embeds, "inputs_embeds") return _make_model(graph) diff --git a/src/mobius/tasks/_image_classification.py b/src/mobius/tasks/_image_classification.py index e58610578..cb3b5dc57 100644 --- a/src/mobius/tasks/_image_classification.py +++ b/src/mobius/tasks/_image_classification.py @@ -40,7 +40,11 @@ def build( graph, builder = _make_graph() op = builder.op - pixel_values = builder.input("pixel_values", dtype=ir.DataType.FLOAT, shape=[batch, num_channels, image_size, image_size]) + pixel_values = builder.input( + "pixel_values", + dtype=ir.DataType.FLOAT, + shape=[batch, num_channels, image_size, image_size], + ) last_hidden_state = module(op, pixel_values=pixel_values) diff --git a/src/mobius/tasks/_multimodal.py b/src/mobius/tasks/_multimodal.py index f804d505c..01f7d5e35 100644 --- a/src/mobius/tasks/_multimodal.py +++ b/src/mobius/tasks/_multimodal.py @@ -56,25 +56,32 @@ def build( op = builder.op input_ids = builder.input( - "input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len], + "input_ids", + dtype=ir.DataType.INT64, + shape=[batch, seq_len], ) attention_mask = builder.input( - "attention_mask", dtype=ir.DataType.INT64, + "attention_mask", + dtype=ir.DataType.INT64, shape=[batch, "past_seq_len + seq_len"], ) position_ids = builder.input( - "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len], + "position_ids", + dtype=ir.DataType.INT64, + shape=[batch, seq_len], ) image_size = config.vision.image_size or 224 if config.vision else 224 pixel_values = builder.input( - "pixel_values", dtype=config.dtype, + "pixel_values", + dtype=config.dtype, shape=[batch, 3, image_size, image_size], ) audio_input_size = (config.audio.input_size if config.audio else None) or 80 audio_features = builder.input( - "audio_features", dtype=config.dtype, + "audio_features", + dtype=config.dtype, shape=[batch, "audio_seq_len", audio_input_size], ) diff --git a/src/mobius/tasks/_object_detection.py b/src/mobius/tasks/_object_detection.py index 4a3736ffa..6bdf5c7ed 100644 --- a/src/mobius/tasks/_object_detection.py +++ b/src/mobius/tasks/_object_detection.py @@ -41,7 +41,11 @@ def build( graph, builder = _make_graph() op = builder.op - pixel_values = builder.input("pixel_values", dtype=ir.DataType.FLOAT, shape=[batch, num_channels, image_size, image_size]) + pixel_values = builder.input( + "pixel_values", + dtype=ir.DataType.FLOAT, + shape=[batch, num_channels, image_size, image_size], + ) logits, pred_boxes = module(op, pixel_values=pixel_values) diff --git a/src/mobius/tasks/_phi4mm_multimodal.py b/src/mobius/tasks/_phi4mm_multimodal.py index 99366bf76..2e148a170 100644 --- a/src/mobius/tasks/_phi4mm_multimodal.py +++ b/src/mobius/tasks/_phi4mm_multimodal.py @@ -87,11 +87,13 @@ def _build_vision( graph, builder = _make_graph(name="vision_encoder") pixel_values = builder.input( - "pixel_values", dtype=config.dtype, + "pixel_values", + dtype=config.dtype, shape=[batch, 3, image_size, image_size], ) image_sizes = builder.input( - "image_sizes", dtype=ir.DataType.INT64, + "image_sizes", + dtype=ir.DataType.INT64, shape=[num_images, 2], ) @@ -118,15 +120,18 @@ def _build_speech( graph, builder = _make_graph(name="audio_encoder") audio_embeds = builder.input( - "audio_embeds", dtype=config.dtype, + "audio_embeds", + dtype=config.dtype, shape=[batch, audio_seq_len, input_size], ) audio_sizes = builder.input( - "audio_sizes", dtype=ir.DataType.INT64, + "audio_sizes", + dtype=ir.DataType.INT64, shape=[num_audio_clips], ) audio_projection_mode = builder.input( - "audio_projection_mode", dtype=ir.DataType.INT64, + "audio_projection_mode", + dtype=ir.DataType.INT64, shape=[], ) @@ -154,15 +159,18 @@ def _build_embedding( graph, builder = _make_graph(name="embedding") input_ids = builder.input( - "input_ids", dtype=ir.DataType.INT64, + "input_ids", + dtype=ir.DataType.INT64, shape=[batch, seq_len], ) image_features = builder.input( - "image_features", dtype=config.dtype, + "image_features", + dtype=config.dtype, shape=[num_image_tokens, config.hidden_size], ) audio_features = builder.input( - "audio_features", dtype=config.dtype, + "audio_features", + dtype=config.dtype, shape=[num_speech_tokens, config.hidden_size], ) diff --git a/src/mobius/tasks/_qwen_image_vae.py b/src/mobius/tasks/_qwen_image_vae.py index 36f2a97d3..3940753f5 100644 --- a/src/mobius/tasks/_qwen_image_vae.py +++ b/src/mobius/tasks/_qwen_image_vae.py @@ -43,7 +43,9 @@ def _build_encoder_graph( graph, builder = _make_graph(name="vae_encoder") op = builder.op - sample = builder.input("sample", dtype=ir.DataType.FLOAT, shape=["batch", 3, "frames", "height", "width"]) + sample = builder.input( + "sample", dtype=ir.DataType.FLOAT, shape=["batch", 3, "frames", "height", "width"] + ) hidden_states = module.encoder(op, sample) hidden_states = module.quant_conv(op, hidden_states) @@ -60,7 +62,11 @@ def _build_decoder_graph( graph, builder = _make_graph(name="vae_decoder") op = builder.op - latent_sample = builder.input("latent_sample", dtype=ir.DataType.FLOAT, shape=["batch", config.z_dim, "frames", "height", "width"]) + latent_sample = builder.input( + "latent_sample", + dtype=ir.DataType.FLOAT, + shape=["batch", config.z_dim, "frames", "height", "width"], + ) hidden_states = module.post_quant_conv(op, latent_sample) hidden_states = module.decoder(op, hidden_states) diff --git a/src/mobius/tasks/_seq2seq.py b/src/mobius/tasks/_seq2seq.py index 01f2152a1..0217fc9d3 100644 --- a/src/mobius/tasks/_seq2seq.py +++ b/src/mobius/tasks/_seq2seq.py @@ -59,7 +59,9 @@ def _build_encoder_graph( op = builder.op input_ids = builder.input("input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) - attention_mask = builder.input("attention_mask", dtype=ir.DataType.INT64, shape=[batch, seq_len]) + attention_mask = builder.input( + "attention_mask", dtype=ir.DataType.INT64, shape=[batch, seq_len] + ) encoder_hidden_states = module.encoder( op, input_ids=input_ids, attention_mask=attention_mask @@ -83,14 +85,18 @@ def _build_decoder_graph( op = builder.op input_ids = builder.input( - "input_ids", dtype=ir.DataType.INT64, shape=[batch, dec_seq_len], + "input_ids", + dtype=ir.DataType.INT64, + shape=[batch, dec_seq_len], ) encoder_hidden_states = builder.input( - "encoder_hidden_states", dtype=config.dtype, + "encoder_hidden_states", + dtype=config.dtype, shape=[batch, enc_seq_len, config.hidden_size], ) attention_mask = builder.input( - "attention_mask", dtype=ir.DataType.INT64, + "attention_mask", + dtype=ir.DataType.INT64, shape=[batch, "past_seq_len + dec_seq_len"], ) @@ -102,11 +108,13 @@ def _build_decoder_graph( past_self_kvs: list[tuple[ir.Value, ir.Value]] = [] for i in range(num_decoder_layers): past_key = builder.input( - f"past_key_values.{i}.self.key", dtype=config.dtype, + f"past_key_values.{i}.self.key", + dtype=config.dtype, shape=[batch, num_heads, past_seq_len, head_dim], ) past_value = builder.input( - f"past_key_values.{i}.self.value", dtype=config.dtype, + f"past_key_values.{i}.self.value", + dtype=config.dtype, shape=[batch, num_heads, past_seq_len, head_dim], ) past_self_kvs.append((past_key, past_value)) @@ -115,11 +123,13 @@ def _build_decoder_graph( cross_past_kvs: list[tuple[ir.Value, ir.Value]] = [] for i in range(num_decoder_layers): past_key = builder.input( - f"past_key_values.{i}.cross.key", dtype=config.dtype, + f"past_key_values.{i}.cross.key", + dtype=config.dtype, shape=[batch, num_heads, enc_seq_len, head_dim], ) past_value = builder.input( - f"past_key_values.{i}.cross.value", dtype=config.dtype, + f"past_key_values.{i}.cross.value", + dtype=config.dtype, shape=[batch, num_heads, enc_seq_len, head_dim], ) cross_past_kvs.append((past_key, past_value)) diff --git a/src/mobius/tasks/_speech_language.py b/src/mobius/tasks/_speech_language.py index 76bec8743..855af2b73 100644 --- a/src/mobius/tasks/_speech_language.py +++ b/src/mobius/tasks/_speech_language.py @@ -85,7 +85,8 @@ def _build_audio_encoder( graph, builder = _make_graph(name="audio_encoder") input_features = builder.input( - "input_features", dtype=config.dtype, + "input_features", + dtype=config.dtype, shape=[batch, n_mels, mel_seq], ) diff --git a/src/mobius/tasks/_speech_to_text.py b/src/mobius/tasks/_speech_to_text.py index 5aa5ac4e0..08854f035 100644 --- a/src/mobius/tasks/_speech_to_text.py +++ b/src/mobius/tasks/_speech_to_text.py @@ -73,7 +73,8 @@ def _build_encoder( op = builder.op input_features = builder.input( - "input_features", dtype=ir.DataType.FLOAT, + "input_features", + dtype=ir.DataType.FLOAT, shape=[batch, config.num_mel_bins, audio_seq_len], ) @@ -97,14 +98,19 @@ def _build_decoder( op = builder.op decoder_input_ids = builder.input( - "decoder_input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len], + "decoder_input_ids", + dtype=ir.DataType.INT64, + shape=[batch, seq_len], ) encoder_hidden_states = builder.input( - "encoder_hidden_states", dtype=config.dtype, + "encoder_hidden_states", + dtype=config.dtype, shape=[batch, encoder_seq_len, config.hidden_size], ) position_ids = builder.input( - "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len], + "position_ids", + dtype=ir.DataType.INT64, + shape=[batch, seq_len], ) past_key_values = _make_kv_cache_inputs( diff --git a/src/mobius/tasks/_ssm_causal_lm.py b/src/mobius/tasks/_ssm_causal_lm.py index 4109d10e6..efada9a7e 100644 --- a/src/mobius/tasks/_ssm_causal_lm.py +++ b/src/mobius/tasks/_ssm_causal_lm.py @@ -56,11 +56,13 @@ def _build_ssm_task( past_states: list[tuple[ir.Value, ir.Value]] = [] for i in range(config.num_hidden_layers): conv_state = builder.input( - f"past_states.{i}.conv_state", dtype=config.dtype, + f"past_states.{i}.conv_state", + dtype=config.dtype, shape=[batch, *conv_state_shape], ) ssm_state = builder.input( - f"past_states.{i}.ssm_state", dtype=config.dtype, + f"past_states.{i}.ssm_state", + dtype=config.dtype, shape=[batch, *ssm_state_shape], ) past_states.append((conv_state, ssm_state)) diff --git a/src/mobius/tasks/_tts.py b/src/mobius/tasks/_tts.py index 184c4106b..d2a20d84f 100644 --- a/src/mobius/tasks/_tts.py +++ b/src/mobius/tasks/_tts.py @@ -93,16 +93,19 @@ def _build_talker( graph, builder = _make_graph(name="talker") inputs_embeds = builder.input( - "inputs_embeds", dtype=config.dtype, + "inputs_embeds", + dtype=config.dtype, shape=[batch, seq_len, config.hidden_size], ) attention_mask = builder.input( - "attention_mask", dtype=ir.DataType.INT64, + "attention_mask", + dtype=ir.DataType.INT64, shape=[batch, "past_seq_len + seq_len"], ) # MRoPE: 3D position_ids (3, batch, seq_len) position_ids = builder.input( - "position_ids", dtype=ir.DataType.INT64, + "position_ids", + dtype=ir.DataType.INT64, shape=[3, batch, seq_len], ) @@ -165,20 +168,26 @@ def _build_code_predictor( # Pre-embedded input in talker_hidden space (constructed by # generation loop). The model projects to cp_hidden internally. inputs_embeds = builder.input( - "inputs_embeds", dtype=config.dtype, + "inputs_embeds", + dtype=config.dtype, shape=[batch, seq_len, config.hidden_size], ) # Step index: selects which lm_head to use (0..14) step_index = builder.input( - "step_index", dtype=ir.DataType.INT64, shape=[], + "step_index", + dtype=ir.DataType.INT64, + shape=[], ) attention_mask = builder.input( - "attention_mask", dtype=ir.DataType.INT64, + "attention_mask", + dtype=ir.DataType.INT64, shape=[batch, "past_seq_len + seq_len"], ) # 1D RoPE: 2D position_ids (batch, seq_len) position_ids = builder.input( - "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len], + "position_ids", + dtype=ir.DataType.INT64, + shape=[batch, seq_len], ) past_key_values = _make_kv_cache_inputs( @@ -221,10 +230,14 @@ def _build_embedding( graph, builder = _make_graph(name="embedding") text_ids = builder.input( - "text_ids", dtype=ir.DataType.INT64, shape=[batch, text_seq], + "text_ids", + dtype=ir.DataType.INT64, + shape=[batch, text_seq], ) codec_ids = builder.input( - "codec_ids", dtype=ir.DataType.INT64, shape=[batch, codec_seq], + "codec_ids", + dtype=ir.DataType.INT64, + shape=[batch, codec_seq], ) text_embeds, codec_embeds = embedding( @@ -252,7 +265,9 @@ def _build_speaker_encoder( graph, builder = _make_graph(name="speaker_encoder") mel_input = builder.input( - "mel_input", dtype=config.dtype, shape=[batch, mel_seq, mel_dim], + "mel_input", + dtype=config.dtype, + shape=[batch, mel_seq, mel_dim], ) speaker_embedding = speaker_encoder(builder.op, mel_input) diff --git a/src/mobius/tasks/_vae.py b/src/mobius/tasks/_vae.py index fd40840c1..c9dc62ee2 100644 --- a/src/mobius/tasks/_vae.py +++ b/src/mobius/tasks/_vae.py @@ -43,7 +43,11 @@ def _build_encoder_graph( graph, builder = _make_graph(name="vae_encoder") op = builder.op - sample = builder.input("sample", dtype=ir.DataType.FLOAT, shape=["batch", config.in_channels, "height", "width"]) + sample = builder.input( + "sample", + dtype=ir.DataType.FLOAT, + shape=["batch", config.in_channels, "height", "width"], + ) hidden_states = module.encoder(op, sample=sample) if module.quant_conv is not None: @@ -61,7 +65,11 @@ def _build_decoder_graph( graph, builder = _make_graph(name="vae_decoder") op = builder.op - latent_sample = builder.input("latent_sample", dtype=ir.DataType.FLOAT, shape=["batch", config.latent_channels, "height", "width"]) + latent_sample = builder.input( + "latent_sample", + dtype=ir.DataType.FLOAT, + shape=["batch", config.latent_channels, "height", "width"], + ) hidden_states = latent_sample if module.post_quant_conv is not None: diff --git a/src/mobius/tasks/_video_denoising.py b/src/mobius/tasks/_video_denoising.py index d64ee6c07..70645257f 100644 --- a/src/mobius/tasks/_video_denoising.py +++ b/src/mobius/tasks/_video_denoising.py @@ -32,9 +32,17 @@ def build( graph, builder = _make_graph() op = builder.op - sample = builder.input("sample", dtype=ir.DataType.FLOAT, shape=["batch", "num_frames", config.in_channels, "height", "width"]) + sample = builder.input( + "sample", + dtype=ir.DataType.FLOAT, + shape=["batch", "num_frames", config.in_channels, "height", "width"], + ) timestep = builder.input("timestep", dtype=ir.DataType.INT64, shape=["batch"]) - encoder_hidden_states = builder.input("encoder_hidden_states", dtype=ir.DataType.FLOAT, shape=["batch", "sequence_length", config.cross_attention_dim]) + encoder_hidden_states = builder.input( + "encoder_hidden_states", + dtype=ir.DataType.FLOAT, + shape=["batch", "sequence_length", config.cross_attention_dim], + ) noise_pred = module( op, diff --git a/src/mobius/tasks/_vision_language.py b/src/mobius/tasks/_vision_language.py index b08722b98..0fb530d27 100644 --- a/src/mobius/tasks/_vision_language.py +++ b/src/mobius/tasks/_vision_language.py @@ -51,15 +51,20 @@ def build( op = builder.op input_ids = builder.input( - "input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len], + "input_ids", + dtype=ir.DataType.INT64, + shape=[batch, seq_len], ) attention_mask = builder.input( - "attention_mask", dtype=ir.DataType.INT64, + "attention_mask", + dtype=ir.DataType.INT64, shape=[batch, "past_seq_len + seq_len"], ) # MRoPE: 3D position IDs (temporal, height, width) position_ids = builder.input( - "position_ids", dtype=ir.DataType.INT64, shape=[3, batch, seq_len], + "position_ids", + dtype=ir.DataType.INT64, + shape=[3, batch, seq_len], ) # Flattened image patches patch_size = config.vision.patch_size or 16 if config.vision else 16 @@ -67,12 +72,16 @@ def build( in_channels = config.vision.in_channels if config.vision else 3 pixel_dim = in_channels * temporal_patch_size * patch_size * patch_size pixel_values = builder.input( - "pixel_values", dtype=config.dtype, shape=[total_patches, pixel_dim], + "pixel_values", + dtype=config.dtype, + shape=[total_patches, pixel_dim], ) # Image grid dimensions for position embedding interpolation num_images = ir.SymbolicDim("num_images") grid_thw = builder.input( - "grid_thw", dtype=ir.DataType.INT64, shape=[num_images, 3], + "grid_thw", + dtype=ir.DataType.INT64, + shape=[num_images, 3], ) past_key_values = _make_kv_cache_inputs( diff --git a/src/mobius/tasks/_vision_language_3model.py b/src/mobius/tasks/_vision_language_3model.py index 35277681e..00dae849e 100644 --- a/src/mobius/tasks/_vision_language_3model.py +++ b/src/mobius/tasks/_vision_language_3model.py @@ -82,15 +82,15 @@ def _build_vision( batch = ir.SymbolicDim("batch") image_size = (config.vision.image_size if config.vision else None) or 224 - graph, graph_builder = _make_graph(name="vision_encoder") - pixel_values = graph_builder.input( + graph, builder = _make_graph(name="vision_encoder") + pixel_values = builder.input( "pixel_values", dtype=config.dtype, shape=[batch, 3, image_size, image_size], ) - image_features = vision(graph_builder.op, pixel_values=pixel_values) + image_features = vision(builder.op, pixel_values=pixel_values) - graph_builder.add_output(image_features, "image_features") + builder.add_output(image_features, "image_features") return _make_model(graph) @@ -133,25 +133,25 @@ def _build_vision( in_channels = config.vision.in_channels if config.vision else 3 pixel_dim = in_channels * temporal_patch_size * patch_size * patch_size - graph, graph_builder = _make_graph(name="vision_encoder") - pixel_values = graph_builder.input( + graph, builder = _make_graph(name="vision_encoder") + pixel_values = builder.input( "pixel_values", dtype=config.dtype, shape=[total_patches, pixel_dim], ) - image_grid_thw = graph_builder.input( + image_grid_thw = builder.input( "image_grid_thw", dtype=ir.DataType.INT64, shape=[num_images, 3], ) image_features = vision( - graph_builder.op, + builder.op, pixel_values=pixel_values, image_grid_thw=image_grid_thw, ) - graph_builder.add_output(image_features, "image_features") + builder.add_output(image_features, "image_features") return _make_model(graph) @@ -208,13 +208,13 @@ def _build_vision( height = ir.SymbolicDim("height") width = ir.SymbolicDim("width") - graph, graph_builder = _make_graph(name="vision_encoder") - pixel_values = graph_builder.input( + graph, builder = _make_graph(name="vision_encoder") + pixel_values = builder.input( "pixel_values", dtype=config.dtype, shape=[batch, 3, height, width], ) - op = graph_builder.op + op = builder.op image_features = vision( op, @@ -226,7 +226,7 @@ def _build_vision( # vision encoder always processes one image at a time. image_features = op.Squeeze(image_features, [0]) - graph_builder.add_output(image_features, "image_features") + builder.add_output(image_features, "image_features") return _make_model(graph) @@ -285,24 +285,24 @@ def _build_decoder( cross_seq_len = ir.SymbolicDim("cross_sequence_len") cross_past_seq_len = ir.SymbolicDim("cross_past_seq_len") - graph, graph_builder = _make_graph() - inputs_embeds = graph_builder.input( + graph, builder = _make_graph() + inputs_embeds = builder.input( "inputs_embeds", dtype=config.dtype, shape=[batch, seq_len, config.hidden_size], ) - attention_mask = graph_builder.input( + attention_mask = builder.input( "attention_mask", dtype=ir.DataType.INT64, shape=[batch, "past_seq_len + seq_len"], ) - position_ids = graph_builder.input( + position_ids = builder.input( "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len], ) # Vision features: full on prefill, empty (0-length) on decode - cross_attention_states = graph_builder.input( + cross_attention_states = builder.input( "cross_attention_states", dtype=config.dtype, shape=[batch, cross_seq_len, config.hidden_size], @@ -315,19 +315,19 @@ def _build_decoder( for i in range(config.num_hidden_layers): psl = cross_past_seq_len if i in cross_attention_layers else past_seq_len - past_key = graph_builder.input( + past_key = builder.input( f"past_key_values.{i}.key", dtype=config.dtype, shape=[batch, config.num_key_value_heads, psl, config.head_dim], ) - past_value = graph_builder.input( + past_value = builder.input( f"past_key_values.{i}.value", dtype=config.dtype, shape=[batch, config.num_key_value_heads, psl, config.head_dim], ) past_key_values.append((past_key, past_value)) - op = graph_builder.op + op = builder.op logits, present_key_values = decoder( op, @@ -338,7 +338,7 @@ def _build_decoder( past_key_values=past_key_values, ) - graph_builder.add_output(logits, "logits") - _register_kv_cache_outputs(graph_builder, present_key_values) + builder.add_output(logits, "logits") + _register_kv_cache_outputs(builder, present_key_values) return _make_model(graph)