Add OnnxMoEQuantization pass (com.microsoft::MoE → QMoE)

[justinchuby]

Problem

Mobius and other ONNX exporters emit MoE blocks as fused com.microsoft::MoE nodes whose per-expert FC1/FC2 weights are 3-D fp16/bf16/fp32 initializers. The existing weight-quantization passes (OnnxKQuantQuantization, OnnxBlockWiseRtnQuantization, OnnxBnb4Quantization) only target standalone MatMul nodes, so:

• The per-expert weights — which are typically ~80% of total parameters in a MoE model — stay at the model's compute dtype after running any of those passes.
• For Gemma 4 26B-A4B, this means int4 quantization yields only ~6% file-size reduction (~48 GB → ~45 GB).

The correct target for these weights is com.microsoft::QMoE, which is now supported by both the CUDA and (experimental) CPU kernels in ORT main after microsoft/onnxruntime#28467. There is no existing pass in either Olive or ORT that performs the MoE → QMoE rewrite.

Solution

New OnnxMoEQuantization pass that walks the graph and rewrites every com.microsoft::MoE node:

1. Validates that fc1_experts_weights and fc2_experts_weights are 3-D static initializers.
2. For each expert, calls ORT's pybind quantize_matmul_{4,8}bits to produce per-expert quantized weights and symmetric fp16 scales, then CUTLASS-prepacks them via pack_weights_for_cuda_mixed_gemm so the QMoE kernels can read them directly.
3. Stacks the per-expert tensors along axis 0 and registers them as new initializers (uint8 weights + fp16 scales).
4. Replaces the MoE node with a QMoE node, carrying over all original routing/activation attributes plus expert_weight_bits, optional block_size, and quant_type='int'.
5. Drops the orphan fp16 weight initializers.

Configuration

Param	Default	Notes
bits	4	4 (int4) or 8 (int8)
block_size	0	0 = per-row scales; otherwise must be a power of two ≥ 16
nodes_to_exclude	None	MoE node names to leave unquantized
force_arch	80	CUTLASS prepacking target SM (80 = Ampere, 90 = Hopper)

Example

{
  "input_model": {"type": "OnnxModel", "model_path": "decoder/model.onnx"},
  "passes": {
    "moe_quant": {
      "type": "OnnxMoEQuantization",
      "bits": 4,
      "block_size": 0
    }
  }
}

Runtime requirement

The pass uses pack_weights_for_cuda_mixed_gemm from ORT, which is only exported when ORT is built with USE_CUDA. A descriptive RuntimeError is raised at run time when the binding is unavailable, telling the user to install onnxruntime-gpu >= 1.28 (or a recent nightly).

Limitations / out-of-scope

• fc3 inputs (3-fold MoE variants) trigger a warning-skip per node.
• Symmetric integer quantization only (matching the QMoE kernel's preferred layout). fp4 / fp8 / wfp4afp8 quant_types are left for a follow-up.
• No calibration-aware quantization (GPTQ / AWQ). This pass is pure RTN.

Tests

5 unit tests in test/passes/onnx/test_moe_quantization.py:

Test	Coverage
test_moe_to_qmoe_conversion	End-to-end MoE→QMoE with int4 + per-row scales; checks node replacement, input slot order, attribute preservation, and uint8/fp16 shape/dtype
test_moe_to_qmoe_blockwise	block_size=16 produces 3-D scales [E, N, K/block_size] and emits the block_size attribute
test_moe_to_qmoe_skip_when_not_initializer	Dynamic-weight MoE nodes are skipped with a warning, not converted
test_invalid_bits_rejected	bits ∉ {4, 8} fails fast
test_invalid_block_size_rejected	non-power-of-two block_size fails fast

The CUTLASS prepack helper is patched with a pass-through stub during tests so CPU-only CI can exercise the graph transform. All 5 pass locally + the existing 416 test/passes/onnx tests still pass + lintrunner is clean.

[Copilot]

Pull request overview

This PR introduces a new ONNX graph-rewrite pass to quantize fused MoE blocks exported as com.microsoft::MoE by rewriting them into com.microsoft::QMoE, enabling int4/int8 weight-only quantization of per-expert FC weights using ONNX Runtime’s quantization + CUTLASS prepack helpers.

Changes:

• Added OnnxMoEQuantization pass to convert MoE → QMoE, quantize per-expert FC1/FC2 weights, and register packed uint8 weights + fp16 scales.
• Added unit tests covering successful conversion, block-wise scales, skip behavior when weights are dynamic, and config validation.
• Registered the new pass in olive_config.json with CUDA EP support and int4/int8 precision metadata.

Reviewed changes

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

File	Description
olive/passes/onnx/moe_quantization.py	Implements the MoE→QMoE rewrite and per-expert quantization/packing pipeline.
test/passes/onnx/test_moe_quantization.py	Adds targeted unit tests for the new pass (conversion, blockwise scales, skips, config validation).
olive/olive_config.json	Registers the new pass and its supported providers/accelerators/precisions.

[justinchuby]

Follow-ups filed:

• com.microsoft::QMoE should prepack int4/int8 weights in PrePack(), like MatMulNBits does onnxruntime#28748 — proposed fix on the ORT side: com.microsoft::QMoE should prepack int4/int8 weights in PrePack(), like com.microsoft::MatMulNBits already does. That would let this pass quantize on a CPU-only host with no special steps. The current asymmetry is most likely an oversight inherited from the TensorRT-LLM-flavoured fpA_intB code, since MatMulNBits does it correctly.
• OnnxMoEQuantization: port pack_weights_for_cuda_mixed_gemm to numpy / cupy so the pass runs on CPU-only hosts #2492 — tracks a numpy + optional-cupy port of pack_weights_for_cuda_mixed_gemm for this pass, so it can run on CPU-only hosts without waiting for the ORT-side fix to ship.

[justinchuby]

Sent the proposed ORT-side fix as microsoft/onnxruntime#28749. Once that lands and ships in a release, this pass can drop the pack_weights_for_cuda_mixed_gemm dependency entirely — the Olive pass would only need quantize_matmul_{4,8}bits (which is in every ORT build, CPU or CUDA).

[titaiwangms]

Related: #2583 proposes resuming jambayk/moe-quant to add MoE-aware weight quantization at the PyTorch/HF level (storage-only QuantTensor representation for fused-3D and ModuleList-style expert weights), as a complement to this ONNX-graph-level MoE work. Cross-linking for context in case an implementing agent picks either up.