Standardize symbolic dimensions in common ORT GenAI decoder exports

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,18 +272,17 @@ 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,
         shape=[3, batch, seq_len] if mrope else [batch, seq_len],
     )
 
     # 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 [])

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]

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)