[Oke Iyanuoluwa Enoch (Iyanuoluwa007)]

Describe your changes

Static-unroll MoE blocks (for example the mobius Gemma 4 fallback MoE in #2489) keep the weights of all experts in a single [E, K, N] initializer and pick one expert per MatMul:

W_k = Gather(W_all, k, axis=0)                       # scalar k
W_k = Squeeze(Gather(W_all, [k], axis=0), axes=[0])  # one-element k
y   = MatMul(x, W_k)

OnnxKQuantQuantization only quantizes MatMul nodes whose weight is an initializer, so all of these expert MatMul nodes stayed in fp16.

This PR implements option B from the issue as a new graph surgeon, UnstackGatheredMatMulWeights, which keeps k-quant focused on its job:

• For each MatMul whose weight comes from Gather(3-D initializer, constant index, axis=0), optionally through Squeeze(axes=[0]) (axes given as an input or, for opset < 13, as an attribute), the weight is replaced with a 2-D initializer holding W_all[k]. The same expert slice is created only once if several MatMul nodes use it. Negative indices are supported.
• Gather/Squeeze nodes and stacked initializers that are no longer used are removed with RemoveUnusedNodesPass.
• MatMul nodes whose expert index is only known at runtime are left unchanged.

Usage: run GraphSurgeries with {"surgeon": "UnstackGatheredMatMulWeights"} before OnnxKQuantQuantization. The surgeon is documented in docs/source/features/onnx-transformations.md.

Packing experts into a 3-D MatMulNBits style op (option A in the issue) needs ONNX Runtime kernel support and is not part of this PR.

Checklist before requesting a review

• Add unit tests for this change.
  • The rewrite is checked against ONNX Runtime outputs for scalar-index, Squeeze, and runtime-index experts.
  • An unused stacked initializer is removed, and a negative index works.
  • End to end, k-quant produces 0 MatMulNBits before the surgery and 2 after it.
• Make sure all tests can pass. (pytest test/passes/onnx/test_graph_surgeries.py: 92 passed. The 3 failures in my env need onnxruntime_genai, which is not installed, and are unrelated.)
• Update documents if necessary.
• Lint and apply fixes to your code by running lintrunner -a (ruff 0.15.11 check and format on changed files)
• Is this a user-facing change? If yes, give a description of this change to be included in the release notes.
  New UnstackGatheredMatMulWeights graph surgery lets weight-only quantizers such as k-quant quantize per-expert MatMuls in static-unroll MoE models.

(Optional) Issue link

Fixes #2489

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

🟡 Changes recommended

Generated names can collide with existing values, and overridable defaults can be incorrectly frozen during rewriting.

2 open issues

• Avoid value-name collisions across the shared graph namespace · New
• Reject overridable graph-input defaults before materializing the rewrite · New

What changed in this PR

Adds a graph surgery that exposes static MoE expert weights as initializers, enabling existing weight-only quantizers to process them and addressing #2489.

Changes:

• Materializes constant expert slices and removes unused graph components.
• Adds output-equivalence and k-quant integration tests.
• Documents configuration and recommended pass ordering.

File	Description
test/​passes/​onnx/​test_graph_surgeries.py	Tests rewriting, cleanup, negative indices, and quantization.
olive/​passes/​onnx/​graph_surgeries.py	Implements UnstackGatheredMatMulWeights.
docs/​source/​features/​onnx-transformations.md	Documents the surgeon and usage.

_{Review effort: Balanced}

[azure-pipelines]

Commenter does not have sufficient privileges for PR 2739 in repo microsoft/Olive

[Oke Iyanuoluwa Enoch (Iyanuoluwa007)]

/azp run

[azure-pipelines]

Commenter does not have sufficient privileges for PR 2739 in repo microsoft/Olive