[DO NOT MERGE] Copy of #2739

docs/source/features/onnx-transformations.md

@@ -2583,6 +2583,33 @@ No parameters required.
 }
 ```
 
+### `UnstackGatheredMatMulWeights`
+
+#### Description
+
+Static-unroll Mixture-of-Experts (MoE) blocks often keep the weights of all experts in a single `[E, K, N]` initializer and select one expert per `MatMul` with a constant index, either `Gather(W_all, k, axis=0)` with a scalar index or `Gather(W_all, [k], axis=0)` followed by `Squeeze(axes=[0])`.
+
+Weight-only quantizers such as `OnnxKQuantQuantization` only quantize `MatMul` nodes whose weight is an initializer, so these per-expert `MatMul` nodes are left unquantized. This surgery replaces each gathered weight with a 2-D initializer holding `W_all[k]` so a following quantization pass picks it up. `Gather`/`Squeeze` nodes and stacked initializers that are no longer used are removed. `MatMul` nodes whose expert index is computed at runtime are left unchanged.
+
+Run it before `OnnxKQuantQuantization` (or another weight-only `MatMul` quantizer) on models with static-unroll MoE experts.
+
+#### Configurations
+
+No parameters required.
+
+#### Example
+
+```json
+{
+    "type": "GraphSurgeries",
+    "surgeries": [
+        {
+            "surgeon": "UnstackGatheredMatMulWeights"
+        }
+    ]
+}
+```
+
 
 ## ORT Performance Tuning
 

olive/passes/onnx/graph_surgeries.py

@@ -3434,6 +3434,134 @@ def _remove_from_graph(cls, graph: ir.Graph) -> int:
         return removed
 
 
+class UnstackGatheredMatMulWeights(Surgeon):
+    """Materialize MatMul weights gathered from a stacked 3-D initializer as 2-D initializers.
+
+    Static-unroll Mixture-of-Experts blocks often keep the weights of all experts in a
+    single ``[E, K, N]`` initializer and select one expert per MatMul with a constant index::
+
+        W_k = Gather(W_all, k, axis=0)                       # scalar k -> [K, N]
+        W_k = Squeeze(Gather(W_all, [k], axis=0), axes=[0])  # 1-element k -> [K, N]
+        y = MatMul(x, W_k)
+
+    Weight-only quantizers such as ``OnnxKQuantQuantization`` only quantize MatMuls whose
+    weight is an initializer, so these per-expert MatMuls stay unquantized. This surgery
+    replaces each gathered weight with a 2-D initializer holding ``W_all[k]``, so that a
+    following quantization pass picks them up. Gather/Squeeze nodes and stacked initializers
+    left without consumers are removed.
+
+    Only constant indices on axis 0 are rewritten. MatMuls whose expert index is computed at
+    runtime are left unchanged.
+
+    When to use:
+        Run **before** ``OnnxKQuantQuantization`` (or another weight-only MatMul quantizer)
+        on models with static-unroll MoE experts.
+    """
+
+    def call_ir(self, model: ir.Model) -> ir.Model:
+        graph = model.graph
+        unstacked_weights: dict[tuple[str, int], ir.Value] = {}
+        # all values share one namespace (see _unique_value_name); collect it once since a model can have
+        # thousands of experts
+        taken_names = set(graph.initializers)
+        taken_names.update(value.name for value in graph.inputs if value.name)
+        taken_names.update(value.name for value in graph.outputs if value.name)
+        for graph_node in graph:
+            taken_names.update(output.name for output in graph_node.outputs if output.name)
+        rewritten = 0
+        for node in list(graph):
+            if node.op_type != "MatMul" or node.domain not in ("", "ai.onnx") or len(node.inputs) != 2:
+                continue
+            weight = node.inputs[1]
+            if weight is None or weight.is_initializer():
+                continue
+
+            gathered = self._get_gathered_weight(weight)
+            if gathered is None:
+                continue
+            stacked, index = gathered
+
+            key = (stacked.name, index)
+            if key not in unstacked_weights:
+                name = self._get_unique_name(taken_names, f"{stacked.name}_{index}")
+                taken_names.add(name)
+                array = np.ascontiguousarray(stacked.const_value.numpy()[index])
+                value = ir.Value(name=name, const_value=ir.tensor(array, name=name))
+                graph.register_initializer(value)
+                unstacked_weights[key] = value
+            node.replace_input_with(1, unstacked_weights[key])
+            rewritten += 1
+
+        if rewritten:
+            RemoveUnusedNodesPass()(model)
+            logger.debug("Unstacked %d gathered MatMul weights into %d initializers", rewritten, len(unstacked_weights))
+        return model
+
+    @classmethod
+    def _get_gathered_weight(cls, weight: ir.Value) -> tuple[ir.Value, int] | None:
+        """Return the stacked initializer and constant expert index that produce ``weight``, if any."""
+        producer = weight.producer()
+        squeezed = False
+        if producer is not None and producer.op_type == "Squeeze" and producer.domain in ("", "ai.onnx"):
+            if not cls._squeezes_only_leading_axis(producer):
+                return None
+            squeezed = True
+            producer = producer.inputs[0].producer() if producer.inputs[0] is not None else None
+
+        if producer is None or producer.op_type != "Gather" or producer.domain not in ("", "ai.onnx"):
+            return None
+        if producer.attributes.get_int("axis", 0) != 0:
+            return None
+
+        stacked, indices = producer.inputs
+        # an initializer that is also a graph input is an overridable default, not a constant
+        if stacked is None or not stacked.is_initializer() or stacked.is_graph_input() or stacked.const_value is None:
+            return None
+        if len(stacked.const_value.shape) != 3:
+            return None
+
+        if indices is None or indices.is_graph_input():
+            return None
+        index_tensor = ir.convenience.get_const_tensor(indices)
+        if index_tensor is None:
+            return None
+        index_array = index_tensor.numpy()
+        # a scalar index drops the expert axis, a one-element index keeps it and needs the Squeeze
+        if (squeezed and index_array.shape != (1,)) or (not squeezed and index_array.ndim != 0):
+            return None
+
+        num_experts = stacked.const_value.shape[0]
+        index = int(index_array.reshape(-1)[0])
+        if not -num_experts <= index < num_experts:
+            return None
+        return stacked, index % num_experts
+
+    @staticmethod
+    def _squeezes_only_leading_axis(squeeze: ir.Node) -> bool:
+        if len(squeeze.inputs) > 1 and squeeze.inputs[1] is not None:
+            # overriding default axes could change the rank of the weight, so they must be constant
+            if squeeze.inputs[1].is_graph_input():
+                return False
+            axes_tensor = ir.convenience.get_const_tensor(squeeze.inputs[1])
+            if axes_tensor is None:
+                return False
+            axes = axes_tensor.numpy().reshape(-1).tolist()
+        else:
+            # opset < 13 keeps axes as an attribute
+            axes = list(squeeze.attributes.get_ints("axes", []))
+        # rank of the Gather output is 3, so -3 also refers to the leading axis
+        return axes in ([0], [-3])
+
+    @staticmethod
+    def _get_unique_name(taken_names: set[str], base_name: str) -> str:
+        if base_name not in taken_names:
+            return base_name
+        index = 1
+        while f"{base_name}_{index}" in taken_names:
+            index += 1
+        return f"{base_name}_{index}"
+
+
 class RenameInputDims(Surgeon):
     """Rename / promote a dimension in an input tensor's shape to a named symbolic dim.
 

test/passes/onnx/test_graph_surgeries.py

@@ -3464,3 +3464,226 @@ def test_remove_memcpy_chained(tmp_path):
     x = np.random.randn(4, 8).astype(np.float32)
     result = sess.run(["output"], {"input": x})[0]
     np.testing.assert_allclose(result, np.maximum(x, 0), atol=1e-6)
+
+
+def _make_stacked_expert_matmul_model(model_path: Path, weights: np.ndarray):
+    """Build x -> 3 MatMuls with expert weights gathered from one stacked [E, K, N] initializer.
+
+    expert0: Gather with a scalar Constant index (no Squeeze)
+    expert1: Gather with a one-element initializer index followed by Squeeze(axes=[0])
+    expert_dyn: Gather with a runtime index from a graph input, which must stay unchanged
+    """
+    num_experts, k, n = weights.shape
+    nodes = [
+        helper.make_node(
+            "Constant", [], ["idx0"], name="Idx0", value=numpy_helper.from_array(np.array(0, dtype=np.int64))
+        ),
+        helper.make_node("Gather", ["W_all", "idx0"], ["w0"], name="Gather0", axis=0),
+        helper.make_node("MatMul", ["x", "w0"], ["y0"], name="MatMul0"),
+        helper.make_node("Gather", ["W_all", "idx1"], ["w1_3d"], name="Gather1", axis=0),
+        helper.make_node("Squeeze", ["w1_3d", "squeeze_axes"], ["w1"], name="Squeeze1"),
+        helper.make_node("MatMul", ["x", "w1"], ["y1"], name="MatMul1"),
+        helper.make_node("Gather", ["W_all", "expert_idx"], ["w_dyn"], name="GatherDyn", axis=0),
+        helper.make_node("MatMul", ["x", "w_dyn"], ["y_dyn"], name="MatMulDyn"),
+    ]
+    graph = helper.make_graph(
+        nodes=nodes,
+        name="StackedExperts",
+        inputs=[
+            helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, k]),
+            helper.make_tensor_value_info("expert_idx", TensorProto.INT64, []),
+        ],
+        outputs=[helper.make_tensor_value_info(name, TensorProto.FLOAT, [2, n]) for name in ("y0", "y1", "y_dyn")],
+        initializer=[
+            numpy_helper.from_array(weights, "W_all"),
+            numpy_helper.from_array(np.array([num_experts - 1], dtype=np.int64), "idx1"),
+            numpy_helper.from_array(np.array([0], dtype=np.int64), "squeeze_axes"),
+        ],
+    )
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
+    model.ir_version = 9
+    onnx.save(model, model_path)
+
+
+def test_unstack_gathered_matmul_weights_materializes_constant_expert_weights(tmp_path):
+    rng = np.random.default_rng(0)
+    weights = rng.standard_normal((3, 4, 5)).astype(np.float32)
+    model_path = tmp_path / "model.onnx"
+    _make_stacked_expert_matmul_model(model_path, weights)
+
+    p = create_pass_from_dict(
+        GraphSurgeries, {"surgeries": [{"surgeon": "UnstackGatheredMatMulWeights"}]}, disable_search=True
+    )
+    output_model = p.run(ONNXModelHandler(model_path=str(model_path)), str(tmp_path / "onnx"))
+    output_model_def = output_model.load_model()
+
+    # constant-index MatMuls now take 2-D initializers, the runtime-index one is untouched
+    initializers = {init.name: numpy_helper.to_array(init) for init in output_model_def.graph.initializer}
+    matmuls = {node.name: node for node in output_model_def.graph.node if node.op_type == "MatMul"}
+    np.testing.assert_array_equal(initializers[matmuls["MatMul0"].input[1]], weights[0])
+    np.testing.assert_array_equal(initializers[matmuls["MatMul1"].input[1]], weights[2])
+    assert matmuls["MatMulDyn"].input[1] == "w_dyn"
+    # the Gather/Squeeze feeding constant experts are removed, the stacked weight stays for the dynamic Gather
+    node_names = {node.name for node in output_model_def.graph.node}
+    assert {"Gather0", "Gather1", "Squeeze1"}.isdisjoint(node_names)
+    assert "GatherDyn" in node_names
+    assert "W_all" in initializers
+
+    # outputs are unchanged
+    x = rng.standard_normal((2, 4)).astype(np.float32)
+    feeds = {"x": x, "expert_idx": np.array(1, dtype=np.int64)}
+    expected = InferenceSession(str(model_path), providers=["CPUExecutionProvider"]).run(None, feeds)
+    actual = InferenceSession(output_model.model_path, providers=["CPUExecutionProvider"]).run(None, feeds)
+    for expected_output, actual_output in zip(expected, actual):
+        np.testing.assert_allclose(actual_output, expected_output, rtol=1e-6, atol=1e-6)
+
+
+def test_unstack_gathered_matmul_weights_removes_unused_stacked_initializer(tmp_path):
+    weights = np.arange(2 * 3 * 4, dtype=np.float32).reshape(2, 3, 4)
+    nodes = [
+        helper.make_node("Gather", ["W_all", "idx"], ["w"], name="Gather", axis=0),
+        helper.make_node("MatMul", ["x", "w"], ["y"], name="MatMul"),
+    ]
+    graph = helper.make_graph(
+        nodes=nodes,
+        name="SingleExpert",
+        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3])],
+        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 4])],
+        initializer=[
+            numpy_helper.from_array(weights, "W_all"),
+            numpy_helper.from_array(np.array(-1, dtype=np.int64), "idx"),
+        ],
+    )
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
+    model.ir_version = 9
+    model_path = tmp_path / "model.onnx"
+    onnx.save(model, model_path)
+
+    p = create_pass_from_dict(
+        GraphSurgeries, {"surgeries": [{"surgeon": "UnstackGatheredMatMulWeights"}]}, disable_search=True
+    )
+    output_model_def = p.run(ONNXModelHandler(model_path=str(model_path)), str(tmp_path / "onnx")).load_model()
+
+    initializers = {init.name: numpy_helper.to_array(init) for init in output_model_def.graph.initializer}
+    assert "W_all" not in initializers
+    assert [node.op_type for node in output_model_def.graph.node] == ["MatMul"]
+    # negative index -1 selects the last expert
+    np.testing.assert_array_equal(initializers[output_model_def.graph.node[0].input[1]], weights[-1])
+
+
+def test_unstack_gathered_matmul_weights_lets_kquant_quantize_expert_matmuls(tmp_path):
+    from olive.hardware.accelerator import AcceleratorSpec
+    from olive.passes.onnx.kquant_quantization import OnnxKQuantQuantization
+
+    rng = np.random.default_rng(0)
+    weights = rng.standard_normal((3, 64, 32)).astype(np.float32)
+    model_path = tmp_path / "model.onnx"
+    _make_stacked_expert_matmul_model(model_path, weights)
+    accelerator_spec = AcceleratorSpec(accelerator_type="CPU", execution_provider="CPUExecutionProvider")
+
+    def count_matmul_nbits(model_handler):
+        return sum(node.op_type == str(OpType.MatMulNBits) for node in model_handler.load_model().graph.node)
+
+    kquant_config = {"bits": 4, "block_size": 32}
+    kquant = create_pass_from_dict(
+        OnnxKQuantQuantization, kquant_config, disable_search=True, accelerator_spec=accelerator_spec
+    )
+    # without the surgery no expert MatMul has an initializer weight
+    assert count_matmul_nbits(kquant.run(ONNXModelHandler(model_path=str(model_path)), tmp_path / "q0.onnx")) == 0
+
+    surgeries = create_pass_from_dict(
+        GraphSurgeries, {"surgeries": [{"surgeon": "UnstackGatheredMatMulWeights"}]}, disable_search=True
+    )
+    unstacked = surgeries.run(ONNXModelHandler(model_path=str(model_path)), str(tmp_path / "unstacked"))
+    kquant = create_pass_from_dict(
+        OnnxKQuantQuantization, kquant_config, disable_search=True, accelerator_spec=accelerator_spec
+    )
+    # the two constant-index experts are quantized, the runtime-index one is not
+    assert count_matmul_nbits(kquant.run(unstacked, tmp_path / "q1.onnx")) == 2
+
+
+def _run_unstack_surgery(model: onnx.ModelProto, tmp_path: Path) -> ONNXModelHandler:
+    model_path = tmp_path / "model.onnx"
+    onnx.save(model, model_path)
+    p = create_pass_from_dict(
+        GraphSurgeries, {"surgeries": [{"surgeon": "UnstackGatheredMatMulWeights"}]}, disable_search=True
+    )
+    return p.run(ONNXModelHandler(model_path=str(model_path)), str(tmp_path / "onnx"))
+
+
+def test_unstack_gathered_matmul_weights_avoids_value_name_collisions(tmp_path):
+    weights = np.arange(2 * 3 * 4, dtype=np.float32).reshape(2, 3, 4)
+    nodes = [
+        helper.make_node("Gather", ["W_all", "idx"], ["w"], name="Gather", axis=0),
+        helper.make_node("MatMul", ["x", "w"], ["y"], name="MatMul"),
+        # a node output already uses the name the expert-0 initializer would get
+        helper.make_node("Identity", ["x"], ["W_all_0"], name="Identity"),
+    ]
+    graph = helper.make_graph(
+        nodes=nodes,
+        name="NameCollision",
+        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3])],
+        outputs=[
+            helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 4]),
+            helper.make_tensor_value_info("W_all_0", TensorProto.FLOAT, [1, 3]),
+        ],
+        initializer=[
+            numpy_helper.from_array(weights, "W_all"),
+            numpy_helper.from_array(np.array(0, dtype=np.int64), "idx"),
+        ],
+    )
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
+    model.ir_version = 9
+
+    output_model = _run_unstack_surgery(model, tmp_path)
+    output_model_def = output_model.load_model()
+    onnx.checker.check_model(output_model_def)
+
+    matmul = next(node for node in output_model_def.graph.node if node.op_type == "MatMul")
+    assert matmul.input[1] != "W_all_0"
+    initializers = {init.name: numpy_helper.to_array(init) for init in output_model_def.graph.initializer}
+    np.testing.assert_array_equal(initializers[matmul.input[1]], weights[0])
+    x = np.ones((1, 3), dtype=np.float32)
+    y, identity = InferenceSession(output_model.model_path, providers=["CPUExecutionProvider"]).run(None, {"x": x})
+    np.testing.assert_allclose(y, x @ weights[0])
+    np.testing.assert_array_equal(identity, x)
+
+
+@pytest.mark.parametrize("overridable", ["W_all", "idx1", "squeeze_axes"])
+def test_unstack_gathered_matmul_weights_skips_overridable_defaults(overridable, tmp_path):
+    rng = np.random.default_rng(0)
+    weights = rng.standard_normal((3, 4, 5)).astype(np.float32)
+    nodes = [
+        helper.make_node("Gather", ["W_all", "idx1"], ["w1_3d"], name="Gather1", axis=0),
+        helper.make_node("Squeeze", ["w1_3d", "squeeze_axes"], ["w1"], name="Squeeze1"),
+        helper.make_node("MatMul", ["x", "w1"], ["y1"], name="MatMul1"),
+    ]
+    initializers = {
+        "W_all": weights,
+        "idx1": np.array([2], dtype=np.int64),
+        "squeeze_axes": np.array([0], dtype=np.int64),
+    }
+    # an initializer that is also a graph input is a default the caller can override at inference time
+    graph = helper.make_graph(
+        nodes=nodes,
+        name="OverridableDefault",
+        inputs=[
+            helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 4]),
+            helper.make_tensor_value_info(
+                overridable,
+                helper.np_dtype_to_tensor_dtype(initializers[overridable].dtype),
+                list(initializers[overridable].shape),
+            ),
+        ],
+        outputs=[helper.make_tensor_value_info("y1", TensorProto.FLOAT, None)],
+        initializer=[numpy_helper.from_array(value, name) for name, value in initializers.items()],
+    )
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
+    model.ir_version = 9
+
+    output_model_def = _run_unstack_surgery(model, tmp_path).load_model()
+
+    # nothing is materialized, so overriding the default still affects the MatMul
+    matmul = next(node for node in output_model_def.graph.node if node.op_type == "MatMul")
+    assert matmul.input[1] == "w1"
+    assert {node.name for node in output_model_def.graph.node} == {"Gather1", "Squeeze1", "MatMul1"}