Add mixed-width weight contract for QMoE

[David Fan (jiafatom)]

Description

Add the schema and validation contract needed for mixed-width QMoE expert weights.

• Add optional fc1_expert_weight_bits, fc2_expert_weight_bits, and fc3_expert_weight_bits attributes, inheriting from expert_weight_bits when omitted.
• Pass effective bit widths through shared validation and compute packed byte dimensions as logical_elements * bits / 8.
• Require byte-aligned weight rows while allowing zero-point rows to pad to a whole byte.
• Preserve legacy single-pack_size helper overloads and existing 2/4/8 validation diagnostics.
• Wire effective widths through CPU, CUDA, and WebGPU QMoE providers.
• Require matching FC1/FC3 widths for fused SwiGLU.
• Bypass legacy prepacking and return an explicit NOT_IMPLEMENTED status for mixed-width execution. Execution support will follow separately.
• Preserve the existing uniform-width behavior and schema version.

Attribute semantics

expert_weight_bits remains the legacy global value and the fallback for each optional projection override:

effective_fc1_bits = fc1_expert_weight_bits if present, otherwise expert_weight_bits
effective_fc2_bits = fc2_expert_weight_bits if present, otherwise expert_weight_bits
effective_fc3_bits = fc3_expert_weight_bits if present, otherwise expert_weight_bits

If an override is omitted, it is never left uninitialized: it inherits expert_weight_bits, whose schema default is 4. If all three overrides are present, each override takes precedence for its corresponding projection. All effective values are currently restricted to 2, 4, or 8 bits. Fused SwiGLU additionally requires effective FC1 and FC3 widths to match.

Examples:

• expert=4, no overrides -> effective widths (4, 4, 4), existing uniform 4-bit path.
• expert=4, fc1=2, fc2=4, fc3=2 -> effective widths (2, 4, 2), mixed-width contract path.
• expert=4, fc1=2, fc2=2, fc3=2 -> effective widths (2, 2, 2), but currently still treated as an override/mixed configuration because the effective widths differ from the legacy global value.
• For a uniform 2-bit model, use expert_weight_bits=2 and omit the overrides.

Current execution boundary

For this contract-only change, is_mixed_width means that at least one effective projection width differs from expert_weight_bits:

is_mixed_width =
    effective_fc1_bits != expert_weight_bits ||
    effective_fc2_bits != expert_weight_bits ||
    effective_fc3_bits != expert_weight_bits

This intentionally rejects even expert=4, fc1=2, fc2=2, fc3=2. Shape and zero-point validation use the effective 2-bit widths, while the existing execution and dispatch paths still use the legacy global width. Allowing execution in that state could interpret 2-bit data as 4-bit data. The explicit rejection prevents silent incorrect results.

Follow-up execution support must make prepacking, scale/zero-point interpretation, packed layouts, workspace sizing, and kernel dispatch projection-aware. Once dispatch consumes effective widths, uniformly overridden configurations can be normalized to one effective width, while genuinely mixed configurations can dispatch FC1/FC2/FC3 independently.

Testing

• Built onnxruntime_provider_test with CUDA enabled.
• Added CPU contract and invalid-shape tests.
• Added row-wise and block-wise mixed-width zero-point validation, including invalid FC2 shapes.
• Added availability-gated CUDA and WebGPU provider contract tests.
• Preserved the legacy INT2 invalid-hidden-size diagnostic.
• Covered arbitrary 3/5/6-bit packed-byte calculations in the shared helper.
• Repository lintrunner passes for the changed files.

Related to #32657.

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Copilot review overview

🟡 Changes recommended

CUDA returns the wrong status in some paths, generated schema documentation is stale, and provider-specific coverage is incomplete.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 3 · 1

Open (4)

• Test mixed-width zero-point packing and validation · New
• Allow mixed-width inputs to bypass the PrePack guard · New
• Add CUDA and WebGPU mixed-width provider tests · New
• Regenerate operator documentation for new attributes · New

What changed in this PR

Adds mixed-width QMoE weight metadata and validation while deferring execution support.

Changes:

• Adds per-FC bit-width attributes and shape validation.
• Wires effective widths through CPU, CUDA, and WebGPU.
• Adds CPU contract tests and rejects mixed-width execution.

File	Description
onnxruntime/​test/​contrib_ops/​moe_test.cc	Adds mixed-width CPU tests.
onnxruntime/​core/​graph/​contrib_ops/​contrib_defs.cc	Extends the QMoE schema.
onnxruntime/​contrib_ops/​webgpu/​moe/​qmoe.h	Parses per-FC widths.
onnxruntime/​contrib_ops/​webgpu/​moe/​qmoe.cc	Validates and rejects mixed-width execution.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_quantization.h	Stores per-FC widths.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_quantization.cc	Adds CUDA validation and prepack bypass.
onnxruntime/​contrib_ops/​cpu/​moe/​moe_quantization_cpu.h	Stores per-FC widths.
onnxruntime/​contrib_ops/​cpu/​moe/​moe_quantization_cpu.cc	Adds CPU validation and prepack bypass.
onnxruntime/​contrib_ops/​cpu/​moe/​moe_helper.h	Supports FC-specific packing dimensions.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[github-actions]

You can commit the suggested changes from lintrunner.

[Tianlei Wu (tianleiwu)]

The FC-specific validation and provider rejection paths are coherent for this contract-only change. I left two non-blocking inline comments where the public contract and its regression coverage can be tightened.