Add UnstackGatherSqueezeWeight surgeon for static-unroll MoE quantization

[justinchuby]

Fixes #2489.

Problem

When mobius exports models with statically-unrolled MoE blocks (e.g. Gemma 4 26B-A4B), each per-expert MatMul looks like:

fc1 = Squeeze(Gather(W_3d, [const_idx], axis=0), [0])
y   = MatMul(x, Transpose(fc1))

The MatMul's B input is not a graph-level initializer, so weight-quantization passes (OnnxKQuantQuantization, OnnxBlockWiseRtnQuantization) skip it via the node.inputs[1].is_initializer() check. The model's expert weights stay at compute dtype (fp16), giving only ~6% size reduction instead of the expected 4-bit ratio.

Mobius emits this fallback because the fused com.microsoft::MoE op's SwiGLU mode is hardcoded for GPT-OSS semantics and is incompatible with the standard SwiGLU used by models like Gemma 4 26B — see microsoft/onnxruntime#28738.

Solution

Add a new GraphSurgeries surgeon UnstackGatherSqueezeWeight that:

1. Detects Gather(W_3d_init, [const_int], axis=0) → Squeeze([0]) chains where the data is a 3-D initializer and the index is a constant scalar.
2. Materialises W_3d[const_int] as a new 2-D initializer (<W_3d_name>__slice_<idx>).
3. Folds any downstream Transpose into a transposed initializer (__T suffix).
4. Rewires consumers and removes orphaned nodes / 3-D initializers.

After running this surgery, the existing OnnxKQuantQuantization and OnnxBlockWiseRtnQuantization passes pick up every per-expert MatMul with no changes to their own logic.

Usage

{
  "input_model": {"type": "OnnxModel", "model_path": "decoder/model.onnx"},
  "passes": {
    "unstack": {
      "type": "GraphSurgeries",
      "surgeries": [{"surgeon": "UnstackGatherSqueezeWeight"}]
    },
    "kquant": {"type": "OnnxKQuantQuantization", "bits": 4, "block_size": 32}
  }
}

Test

New unit test test_unstack_gather_squeeze_weight builds a tiny model with two Gather → Squeeze → Transpose → MatMul chains over a single 3-D initializer and verifies:

• All Gather/Squeeze/Transpose ops are removed.
• New per-slice transposed initializers are created with the correct values.
• The original 3-D initializer is dropped.
• Numerical parity with the unrewritten model holds via InferenceSession on CPU EP.

All 83 test_graph_surgeries.py tests pass.

[Copilot]

Pull request overview

Adds a new UnstackGatherSqueezeWeight surgeon to GraphSurgeries that materialises per-slice 2-D initializers from Gather(W_3d, [const], axis=0) → Squeeze([0]) chains (with optional downstream Transpose folding). This unblocks weight quantization for static-unroll MoE models such as Gemma 4 26B-A4B where per-expert MatMul B inputs are not graph-level initializers, so passes like OnnxKQuantQuantization previously skipped them.

Changes:

• New UnstackGatherSqueezeWeight surgeon in olive/passes/onnx/graph_surgeries.py that detects the Gather→Squeeze chain over a 3-D initializer, slices it at the constant index, optionally folds a downstream Transpose, rewires consumers, and cleans up orphaned nodes/initializers.
• Helpers _fold_transpose_of_initializer and _scalar_int_list for static-int resolution from both initializers and inline Constant nodes.
• Unit test test_unstack_gather_squeeze_weight building a tiny two-chain model and asserting node/initializer rewrites plus numerical parity via InferenceSession.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
olive/passes/onnx/graph_surgeries.py	Implements the new UnstackGatherSqueezeWeight surgeon and helpers.
test/passes/onnx/test_graph_surgeries.py	Adds a unit test validating slice/transpose materialisation and numerical parity.