diff --git a/src/mobius/tasks/_base.py b/src/mobius/tasks/_base.py index f5847f894..1473a863b 100644 --- a/src/mobius/tasks/_base.py +++ b/src/mobius/tasks/_base.py @@ -236,13 +236,13 @@ def build_decoder_from_embeds( decoder: The decoder sub-module to invoke. config: Architecture configuration. mrope: If ``True``, uses 3D MRoPE position_ids - ``[3, batch, seq_len]`` instead of the standard - ``[batch, seq_len]``. + ``[3, batch_size, sequence_length]`` instead of the standard + ``[batch_size, sequence_length]``. hybrid: If ``True``, uses hybrid KV + DeltaNet cache inputs/outputs (for Qwen3.5-VL and similar). Requires ``config.layer_types``. deepstack: If ``True`` (Qwen3-VL family with ``deepstack_visual_indexes``), adds a ``per_layer_inputs`` input - ``[batch, seq_len, D * hidden_size]`` that the decoder reshapes and + ``[batch_size, sequence_length, D * hidden_size]`` that the decoder reshapes and injects into its first ``D`` layers. Returns: @@ -259,9 +259,9 @@ def build_decoder_from_embeds( _register_linear_attention_functions, ) - batch = ir.SymbolicDim("batch") - seq_len = ir.SymbolicDim("sequence_len") - past_seq_len = ir.SymbolicDim("past_sequence_len") + batch = ir.SymbolicDim("batch_size") + seq_len = ir.SymbolicDim("sequence_length") + past_seq_len = ir.SymbolicDim("past_sequence_length") graph, builder = _make_graph() inputs_embeds = builder.input( @@ -272,10 +272,9 @@ def build_decoder_from_embeds( attention_mask = builder.input( "attention_mask", dtype=ir.DataType.INT64, - shape=[batch, "past_seq_len + seq_len"], + shape=[batch, "total_sequence_length"], ) - # MRoPE: 3D position IDs (temporal, height, width) — shape [3, batch, seq_len] - # Standard: shape [batch, seq_len] + # MRoPE: 3D position IDs (temporal, height, width). position_ids = builder.input( "position_ids", dtype=ir.DataType.INT64, @@ -283,7 +282,7 @@ def build_decoder_from_embeds( ) # DeepStack intermediate vision features, pre-scattered and flattened by - # the embedding model. Shape [batch, seq_len, D * hidden_size], matching + # the embedding model. Shape [batch_size, sequence_length, D * hidden_size], matching # ORT GenAI's generic per_layer_inputs contract. per_layer_inputs = None num_deepstack = len(getattr(config, "deepstack_visual_indexes", None) or []) diff --git a/src/mobius/tasks/_causal_lm.py b/src/mobius/tasks/_causal_lm.py index 68f0d0836..b1edefb10 100644 --- a/src/mobius/tasks/_causal_lm.py +++ b/src/mobius/tasks/_causal_lm.py @@ -40,11 +40,13 @@ class CausalLMTask(ModelTask): are concatenated internally by the Attention op. Inputs: - - input_ids: [batch, sequence_len] INT64 - - attention_mask: [batch, total_seq_len] INT64 - - position_ids: [batch, sequence_len] INT64 - - past_key_values.{i}.key: [batch, num_kv_heads, past_seq_len, head_dim] - - past_key_values.{i}.value: [batch, num_kv_heads, past_seq_len, head_dim] + - input_ids: [batch_size, sequence_length] INT64 + - attention_mask: [batch_size, total_sequence_length] INT64 + - position_ids: [batch_size, sequence_length] INT64 + - past_key_values.{i}.key: + [batch_size, num_kv_heads, past_sequence_length, head_dim] + - past_key_values.{i}.value: + [batch_size, num_kv_heads, past_sequence_length, head_dim] Outputs: - logits: FLOAT - present.{i}.key / present.{i}.value: FLOAT @@ -127,8 +129,8 @@ def build( _validate_static_cache_support(module) # --- Graph input dims --- - batch = ir.SymbolicDim("batch") - seq_len = ir.SymbolicDim("sequence_len") + batch = ir.SymbolicDim("batch_size") + seq_len = ir.SymbolicDim("sequence_length") # --- Build graph first, then create inputs via builder --- graph, builder = _make_graph() @@ -163,11 +165,11 @@ def build( cache_specs=cache_specs, ) else: - past_seq_len = ir.SymbolicDim("past_sequence_len") + past_seq_len = ir.SymbolicDim("past_sequence_length") attention_mask = builder.input( "attention_mask", dtype=ir.DataType.INT64, - shape=[batch, "past_seq_len + seq_len"], + shape=[batch, "total_sequence_length"], ) position_ids = builder.input( "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len] @@ -248,7 +250,7 @@ def build( num_kv_heads=num_kv_cache_heads, key_head_dim=kv_key_head_dim, value_head_dim=kv_value_head_dim, - total_seq_len="past_sequence_len + sequence_len", + total_seq_len="total_sequence_length", dtype=config.dtype, cache_specs=cache_specs, ) @@ -359,9 +361,9 @@ def build( module: nn.Module, config: ArchitectureConfig, ) -> ModelPackage: - batch = ir.SymbolicDim("batch") - seq_len = ir.SymbolicDim("sequence_len") - past_seq_len = ir.SymbolicDim("past_sequence_len") + batch = ir.SymbolicDim("batch_size") + seq_len = ir.SymbolicDim("sequence_length") + past_seq_len = ir.SymbolicDim("past_sequence_length") graph, builder = _make_graph() op = builder.op @@ -370,7 +372,7 @@ def build( attention_mask = builder.input( "attention_mask", dtype=ir.DataType.INT64, - shape=[batch, "past_seq_len + seq_len"], + shape=[batch, "total_sequence_length"], ) position_ids = builder.input( "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len] diff --git a/src/mobius/tasks/_task_test.py b/src/mobius/tasks/_task_test.py index 0f6a6da99..38023199e 100644 --- a/src/mobius/tasks/_task_test.py +++ b/src/mobius/tasks/_task_test.py @@ -136,9 +136,39 @@ def dims(value): assert present_dims[1] == past_dims[1] assert present_dims[3] == past_dims[3] # present covers past + current tokens. - assert present_dims[2] == "past_sequence_len + sequence_len" + assert present_dims[2] == "total_sequence_length" assert present.dtype == past.dtype + def test_dynamic_dimensions_use_ort_genai_names_after_serialization(self, tmp_path): + """Decoder ports retain the ORT GenAI symbolic contract on disk.""" + config = make_config() + module = CausalLMModel(config) + model = CausalLMTask().build(module, config)["model"] + model_path = tmp_path / "model.onnx" + ir.save(model, model_path) + model = ir.load(model_path) + inputs = {value.name: value for value in model.graph.inputs} + outputs = {value.name: value for value in model.graph.outputs} + + def dims(value): + return [dim if isinstance(dim, int) else str(dim) for dim in value.shape] + + assert dims(inputs["input_ids"]) == ["batch_size", "sequence_length"] + assert dims(inputs["position_ids"]) == ["batch_size", "sequence_length"] + assert dims(inputs["attention_mask"]) == ["batch_size", "total_sequence_length"] + assert dims(inputs["past_key_values.0.key"]) == [ + "batch_size", + config.num_key_value_heads, + "past_sequence_length", + config.head_dim, + ] + assert dims(outputs["present.0.key"]) == [ + "batch_size", + config.num_key_value_heads, + "total_sequence_length", + config.head_dim, + ] + def test_build_producer_info(self): config = make_config() module = CausalLMModel(config) @@ -1049,6 +1079,33 @@ def test_outputs_include_logits_and_kv_cache(self): assert "logits" in output_names assert "present.0.key" in output_names + def test_serialized_inputs_use_ort_genai_dimension_names(self, tmp_path): + config, decoder = self._make_decoder_module() + model = build_decoder_from_embeds(decoder, config) + model_path = tmp_path / "decoder.onnx" + ir.save(model, model_path) + model = ir.load(model_path) + inputs = {value.name: value for value in model.graph.inputs} + + def dims(value): + return [dim if isinstance(dim, int) else str(dim) for dim in value.shape] + + assert dims(inputs["inputs_embeds"]) == [ + "batch_size", + "sequence_length", + config.hidden_size, + ] + assert dims(inputs["attention_mask"]) == [ + "batch_size", + "total_sequence_length", + ] + assert dims(inputs["past_key_values.0.key"]) == [ + "batch_size", + config.num_key_value_heads, + "past_sequence_length", + config.head_dim, + ] + def test_mrope_position_ids_are_3d(self): config, decoder = self._make_decoder_module() model = build_decoder_from_embeds(decoder, config, mrope=True)