Fix SkipSimplifiedLayerNormalization fallback output arity

The standard-ONNX fallback body for com.microsoft::SkipSimplifiedLayer-
Normalization returned only two outputs, [norm_out, add_out], placing the
residual sum (input_skip_bias_sum) at index 1. Per the operator spec the
four positional outputs are output(0), mean(1), inv_std_var(2) and
input_skip_bias_sum(3), and the fusion rule in rewrite_rules/_skip_norm.py
emits a 4-output node and reads outputs[3] for the residual.

When InlinePass expands that 4-output node using the 2-output fallback,
onnx_ir's replace_all_uses_with raises ValueError ("number of values and
replacements must match") because the residual at index 3 has no
replacement; even without the guard, the sum would wrongly land in the
optional mean slot at index 1.

Emit four outputs with the sum at index 3. mean and inv_std_var are
training-only outputs that the RMS (simplified) variant does not compute
and that single-output RMSNormalization cannot supply, so they are unused
Constant placeholders (pruned by RemoveUnusedNodesPass right after inline).

Verified: onnx-standard inlining now succeeds with no dangling values and
produces byte-identical logits to the native op (max abs diff 0.0) on a
tiny Qwen2 model.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>

Author: justinchuby

src/mobius/functions/skip_layer_normalization.py

@@ -94,8 +94,20 @@ def skip_simplified_layer_normalization() -> ir.Function:
         norm_out = RMSNormalization(add_out, weight, epsilon=<forwarded>)
 
     Inputs:  ``[input, skip, weight]``
-    Outputs: ``[norm_out, add_out]``
+    Outputs: ``[norm_out, mean, inv_std_var, add_out]``
     Attr:    ``epsilon`` (float)
+
+    The ``com.microsoft::SkipSimplifiedLayerNormalization`` op declares four
+    positional outputs — ``output`` (0), ``mean`` (1), ``inv_std_var`` (2) and
+    ``input_skip_bias_sum`` (3). The residual sum lives at **index 3**, so the
+    function body must expose it there for InlinePass to reconnect downstream
+    consumers correctly (see :mod:`mobius.rewrite_rules._skip_norm`, which reads
+    ``outputs[3]``). ``mean`` and ``inv_std_var`` are training-only outputs that
+    the simplified (RMS) variant does not compute and that ``RMSNormalization``
+    (single-output) cannot supply; they are emitted as unused ``Constant``
+    placeholders purely to keep the output arity aligned. Because the fusion
+    rule never wires indices 1-2 to any consumer, ``RemoveUnusedNodesPass``
+    prunes these placeholders right after the function is inlined.
     """
 
     def body(op, v_input, v_skip, v_weight):
@@ -110,9 +122,16 @@ def body(op, v_input, v_skip, v_weight):
         )
         norm_out = op.RMSNormalization(add_out, v_weight, epsilon=epsilon_attr)
 
+        # Optional mean (1) / inv_std_var (2): not produced by RMS. Emit unused
+        # placeholders so input_skip_bias_sum stays at index 3; pruned by DCE.
+        mean = op.Constant(value_floats=[0.0])
+        inv_std_var = op.Constant(value_floats=[0.0])
+
         norm_out.name = "norm_out"
+        mean.name = "mean"
+        inv_std_var.name = "inv_std_var"
         add_out.name = "add_out"
-        return norm_out, add_out
+        return norm_out, mean, inv_std_var, add_out
 
     return build_function(
         body,

src/mobius/functions/skip_layer_normalization_test.py

@@ -0,0 +1,117 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Tests for the standard-ONNX SkipLayerNormalization function bodies.
+
+These bodies are inlined by :class:`onnx_ir.passes.common.InlinePass` to
+expand the ``com.microsoft`` custom ops for EPs that do not support them.
+The key invariant is that each function's output arity and ordering match
+the ``com.microsoft`` op spec so InlinePass can reconnect downstream
+consumers by position — in particular ``input_skip_bias_sum`` must stay at
+output index 3, not collapse into the optional ``mean`` slot at index 1.
+"""
+
+from __future__ import annotations
+
+import onnx_ir as ir
+from onnx_ir.passes.common import InlinePass
+
+from mobius._builder import build_from_module
+from mobius._configs import ArchitectureConfig
+from mobius._registry import registry
+from mobius.functions import register_function_bodies
+from mobius.functions.skip_layer_normalization import (
+    skip_layer_normalization,
+    skip_simplified_layer_normalization,
+)
+
+
+def _tiny_config() -> ArchitectureConfig:
+    return ArchitectureConfig(
+        hidden_size=64,
+        intermediate_size=128,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=16,
+        num_hidden_layers=2,
+        vocab_size=256,
+        max_position_embeddings=128,
+        hidden_act="silu",
+        rms_norm_eps=1e-6,
+        rope_type="default",
+        rope_theta=10000.0,
+        pad_token_id=0,
+    )
+
+
+def _count(model: ir.Model, op_type: str) -> int:
+    return sum(1 for n in model.graph.all_nodes() if n.op_type == op_type)
+
+
+def _has_dangling_inputs(model: ir.Model) -> bool:
+    """Return True if any node consumes a value with no producer/initializer/input."""
+    known: set[ir.Value] = set(model.graph.inputs)
+    known.update(model.graph.initializers.values())
+    for node in model.graph.all_nodes():
+        known.update(node.outputs)
+    for node in model.graph.all_nodes():
+        for inp in node.inputs:
+            if inp is not None and inp.name and inp not in known:
+                return True
+    return False
+
+
+class TestSkipSimplifiedFunctionSignature:
+    def test_output_arity_matches_spec(self):
+        """SkipSimplifiedLayerNormalization must expose 4 positional outputs.
+
+        Spec order: output(0), mean(1), inv_std_var(2), input_skip_bias_sum(3).
+        """
+        fn = skip_simplified_layer_normalization()
+        assert len(fn.outputs) == 4
+        # The residual sum must live at index 3 (not index 1).
+        assert fn.outputs[3].name == "add_out"
+        assert fn.outputs[0].name == "norm_out"
+
+    def test_skip_layer_norm_output_arity(self):
+        """The non-simplified variant also exposes 4 outputs with sum at index 3."""
+        fn = skip_layer_normalization()
+        assert len(fn.outputs) == 4
+        assert fn.outputs[3].name == "add_out"
+
+
+class TestSkipSimplifiedInline:
+    def _inline_skip_norm(self, model: ir.Model) -> None:
+        register_function_bodies(model)
+
+        def criteria(func: ir.Function) -> bool:
+            return func.domain == "com.microsoft" and func.name in (
+                "SkipLayerNormalization",
+                "SkipSimplifiedLayerNormalization",
+            )
+
+        InlinePass(criteria=criteria)(model)
+
+    def test_inline_preserves_residual(self):
+        """Inlining the fallback must expand every fused op and keep the graph valid.
+
+        Regression test: with the old 2-output body, InlinePass raised
+        ``ValueError`` (output-count mismatch) because the 4-output node's
+        ``input_skip_bias_sum`` (index 3) had no replacement value. The fixed
+        4-output body reconnects the residual and leaves no dangling inputs.
+        """
+        config = _tiny_config()
+        model = build_from_module(registry.get("qwen2")(config), config)["model"]
+
+        fused = _count(model, "SkipSimplifiedLayerNormalization")
+        assert fused > 0, "expected the build pipeline to fuse Add+RMSNorm"
+
+        self._inline_skip_norm(model)
+
+        # All fused ops expanded back to Add + RMSNormalization.
+        assert _count(model, "SkipSimplifiedLayerNormalization") == 0
+        # Each expansion restores one residual Add.
+        assert _count(model, "Add") >= fused
+        # The optional mean/inv_std placeholders are unused and pruned-safe;
+        # more importantly the residual reconnected with no dangling inputs.
+        assert not _has_dangling_inputs(model)