Add fused MoE support to SelectiveMixedPrecision heuristics

[Ti-Tai Wang (titaiwangms)]

Describe your changes

Add explicit, fail-closed fused-MoE support to the fixed SelectiveMixedPrecision heuristics.

• Add moe: true support to high_precision_mlp_down and high_precision_mlp_down_qkv for explicitly recognized Qwen3/Qwen3.5 fused K-last expert layouts.
• Resolve the routed expert output projection through LayerWrapper and canonical iter_quant_targets identities, emitting whole per-layer experts.down_proj overrides rather than per-expert precision.
• Keep high_precision_lm_head as a legal MoE no-op, and reject score-based algorithms with moe: true until their parameter-level scoring is implemented separately.
• Add mixed_precision_info.requires_moe and require an explicit moe: true on the first MoE-capable RTN, GPTQ, or KQuant consumer. Validate that required expert overrides survive skip/component selection and are actually selected or already materialized before mutating parameters.
• Preserve AutoClip and follow-up category-pass composition, prevent persisted metadata from enabling/disabling the consumer’s MoE setting, and reject malformed or stale MoE metadata.
• Keep routers and shared-expert gates unquantized; shared-expert MLP weights remain at the default precision. ModuleList experts, unknown/transposed layouts, score-based MoE selection, and ONNX/Mobius export are intentionally out of scope.

This is the first implementation stage for #2638. A follow-up will generalize SNR/IQE scoring to fused parameter targets; KLD-gradient requires separate routing/gradient/memory design.

Checklist before requesting a review

• Add unit tests for this change.
• Make sure all tests can pass.
• Update documents if necessary.
• Lint and apply fixes to your code by running lintrunner -a.
• Is this a user-facing change? If yes, give a description of this change to be included in the release notes.

Release note: SelectiveMixedPrecision fixed heuristics can now plan higher precision for supported fused Qwen MoE expert output projections with explicit downstream MoE opt-in and fail-closed validation.

(Optional) Issue link

Part of #2638.

[Copilot]

Pull request overview

This PR extends Olive’s SelectiveMixedPrecision fixed heuristics to support explicitly recognized fused-MoE (Qwen3/Qwen3.5) expert output projections, and enforces a “double opt-in” so MoE-capable quantizers must explicitly enable MoE when consuming a plan that requires it.

Changes:

• Add moe: true support to fixed SelectiveMixedPrecision MLP heuristics by resolving fused expert down_proj targets via LayerWrapper + canonical iter_quant_targets identities, and emit mixed_precision_info.requires_moe when expert overrides are produced.
• Add fail-closed validation and consumer-side enforcement in PyTorch quantization utilities so required MoE overrides can’t be silently skipped (unless already materialized by a compatible prior Olive checkpoint).
• Add unit tests and documentation covering supported MoE layouts, rejection paths, and the double-opt-in workflow.

Reviewed changes

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

Show a summary per file

File	Description
olive/passes/pytorch/selective_mixed_precision.py	Adds moe config, MoE-aware fixed-heuristic targeting, and emits requires_moe metadata.
olive/passes/pytorch/quant_utils.py	Validates/propagates requires_moe, enforces consumer opt-in, and delays parameter mutation until after validation.
olive/common/hf/wrapper.py	Adds fail-closed LayerWrapper.get_expert_output for supported fused-expert output parameter resolution.
test/passes/pytorch/test_selective_mixed_precision.py	Adds MoE planning and end-to-end RTN consumption tests (including rejection and hybrid cases).
test/passes/pytorch/test_quant_utils.py	Adds tests for requires_moe enforcement, malformed metadata rejection, and required-target survival checks.
test/common/test_hf_wrapper.py	Adds tests validating LayerWrapper.get_expert_output behavior and fail-closed errors.
docs/source/features/quantization.md	Documents MoE selective mixed precision support and the required double opt-in behavior.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟡 Changes recommended

Stale targets and mismatched materialized precision can be accepted, while failure paths may mutate tied embeddings before validation.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (2)

olive/passes/pytorch/quant_utils.py:669

• Name-only membership does not prove that an already-materialized target fulfills the mixed-precision override. For example, the added test installs experts.down_proj at 4 bits while metadata requires 8 bits, and this path accepts it because the name is present. Compare the existing QuantTensor/effective existing qcfg settings with the required override before treating it as fulfilled; otherwise a follow-up pass can silently preserve the wrong precision.

    materialized_names = already_quantized_names if has_compatible_existing_config else set()
    if config.moe is True:
        materialized_names = materialized_names | new_target_names
    missing = sorted(required_names - materialized_names)

olive/passes/pytorch/quant_utils.py:894

• This does not actually defer parameter mutation until after validation because wrapper.maybe_untie_word_embeddings() is still called earlier at lines 751-752. If a required MoE target then fails validation, tied embeddings and configuration have already been mutated. Remove the earlier call and keep this post-validation call so the fail-closed path remains mutation-free.

    # Everything above is discovery/validation. Mutate parameters only after all required
    # final targets have been proven present.
    if fresh_qcfg.lm_head or fresh_qcfg.embeds:
        wrapper.maybe_untie_word_embeddings()

• Files reviewed: 7/7 changed files
• Comments generated: 1
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

The documentation implies support for an export path that the PR explicitly leaves out of scope.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

• Files reviewed: 7/7 changed files
• Comments generated: 1
• Review effort level: Balanced

[Copilot]

🔵 Needs a closer look

Stateful MoE quantization, checkpoint reuse, and multi-pass composition warrant final human validation.

Review details

• Files reviewed: 7/7 changed files
• Comments generated: 0 new
• Review effort level: Balanced