Repository navigation
Add mixed-width weight contract for QMoE - #32697
Merged
Merged
Conversation
Contributor
There was a problem hiding this comment.
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
Open (4)
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.
Tianlei Wu (tianleiwu)
left a comment
Contributor
There was a problem hiding this comment.
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.
Tianlei Wu (tianleiwu)
approved these changes
Sep 21, 2026
David Fan (jiafatom)
deleted the
jiafa/int2-qmoe-mixed-width-contract
branch
September 21, 2026 20:21
This was referenced Sep 21, 2026
3 tasks done
Ti-Tai Wang (titaiwangms)
added a commit
to microsoft/Olive
that referenced
this pull request
Sep 23, 2026
## Summary - qualify deterministic RTN checkpoints with default INT4 weights, INT2 fused expert gate/up banks, and INT4 down banks - validate actual `QuantTensor` bit widths, packed payload geometry, scales, zero points, and odd packed zero-point ceilings - cover exact and regex overrides for Qwen3 and Qwen3.5 source prefixes - verify save/reload persistence and exact eager pre/post-reload parity - document reproducible Qwen3 and Qwen3.5 recipes ## Implementation result Existing Olive fused-MoE quantization, override resolution, packing, and persistence already support this baseline. No production quantization code changes were required; this PR adds qualification tests and documentation. The fixture uses a Mobius/ORT-compatible layout: - hidden size: 32 - expert intermediate size: 48 - group size: 16 - FC1 gate/up: INT2 - FC2 down: INT4 ## Validation - `python -m pytest test/passes/pytorch/test_rtn_mixed_qmoe.py -q --tb=short` — 6 passed - Ruff format and lint passed - `git diff --check` passed ## Runtime boundary microsoft/onnxruntime#32697 has merged the mixed-width QMoE schema and validation contract. Mixed-width execution is not implemented yet, so this PR qualifies checkpoint production and persistence rather than ORT runtime inference. Addresses #2638. --------- Signed-off-by: titaiwangms <titaiwang@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Ti-Tai Wang (titaiwangms)
added a commit
to onnxruntime/mobius
that referenced
this pull request
Sep 28, 2026
## Status update (September 28, 2026) - microsoft/onnxruntime#32743 merged as `ee5f6e7` on September 24, and microsoft/onnxruntime#32761 merged as `a11b4e5` on September 28. The latter optimizes packed decode for `(2,2)`, `(2,4)`, and `(4,2)`; `(4,8)`/`(8,4)` still use the bounded dense correctness fallback. These merges do **not** automatically update Mobius CI's installed ORT wheel. - The last successful Mobius GPU fast-integration run (September 26) installed `onnxruntime-gpu==1.30.0.dev20260812001` from the CUDA 12 feed. The opt-in QMoE CUDA tests in this PR are not part of the normal CI selections, so the CI checks do not qualify the merged ORT runtime for mixed-width QMoE. The CPU ORT GenAI E2E lane is separately pinned to ORT 1.29.0. - The separate experimental manual RTN recipe in microsoft/olive-recipes#637 exported the full pinned Qwen3-30B-A3B checkpoint into an approximately 18 GB ONNX package with 24 `(4,8)` and 24 `(4,4)` QMoE layers. A bounded CUDA `model-qa.py` run returned `391` for `17 * 23` with `ORT_QMOE_INT_DEQUANT_MAX_SCRATCH_BYTES=2147483648`, and a 57-question MMLU smoke run completed. This used a locally built **earlier** ORT commit `50b8fcb`, not the newly merged `a11b4e5`; it does not establish full-logit parity, production throughput, benchmark-scale quality, or validation of the final merged runtime. ## Summary - follow the projection-specific QMoE bit-width contract added by microsoft/onnxruntime#32697 - ingest Olive fused expert checkpoints with FC1/FC2 INT2/INT4, INT4/INT8, INT8/INT4, or uniform INT8 weights under a model-wide INT4 fallback - allocate and validate FC1/FC2 packed weights, scales, and zero points independently, including adjacent layers with different expert widths - preserve legacy uniform INT4 QMoE emission and prevent mixed geometry from being incorrectly rewrite-fused - support authoritative Qwen3, Qwen3.5, and Qwen3.5-VL checkpoint prefixes - preserve pinned GGUF route fingerprints when the new QMoE source aliases are empty ## Contract Supported FC1/FC2 expert widths are `(4,4)`, `(2,4)`, `(4,8)`, `(8,4)`, and `(8,8)` with a model-wide INT4 fallback. Differing-width nodes emit projection-specific bit attributes, including FC3=FC1 for fused SwiGLU. Uniform INT8 nodes emit `expert_weight_bits=8` with no projection-specific overrides. A model-wide INT8 fallback is not enabled. Checkpoint preprocessing fails closed for missing, partial, stale, malformed, or unsupported expert sidecars. Plain and regex overrides follow Olive producer semantics. ## Validation - `python -m pytest src/mobius/models/moe_test.py src/mobius/components/_moe_test.py src/mobius/models/qwen35_test.py src/mobius/integrations/gguf/_builder_core_test.py src/mobius/integrations/gguf/_moe_cohort_test.py -q --tb=short` — 399 passed - `python -m pytest src/mobius/rewrite_rules/_qmoe_fusion_test.py src/mobius/_weight_utils_test.py -q --tb=short` — 172 passed - Opt-in `MOBIUS_QMOE_CUDA_TEST=1 python -m pytest tests/integration/qmoe_cuda_test.py -q`: **2 passed** on A100 with isolated ORT CUDA build at merged commit `ee5f6e7`. This exports a small Qwen3-MoE model using Olive-style fused sidecars, runs adjacent (4,8)/(8,4) and uniform (8,8) layers, checks full FP16 logits against an independent dequantized PyTorch reference (`rtol=atol=0.01`), and confirms each QMoE executes on CUDA via profiling. The full graph assigns shape/control nodes to CPU; it does not claim an all-CUDA graph. - Opt-in `MOBIUS_QMOE_REAL_SLICE_CUDA_TEST=1 python -m pytest tests/integration/qmoe_cuda_test.py -k real_qwen3 -q`: **2 passed** on A100 using layer-0 experts 0–3 and sliced embedding/router weights from pinned Qwen3-30B-A3B. Olive RTN tensor producer created mixed (4,8)/(8,4) sidecars; both reloaded exports ran QMoE on CUDA and matched independently unpacked full logits at `rtol=atol=0.01` (maximum absolute errors 0.005148/0.004774). Reduced four-expert top-2 routing and synthetic surrounding decoder are intentional; no full-checkpoint parity is claimed. - touched-file Ruff lint and formatting, `python -m mypy src/mobius/_weight_utils.py`, and `git diff --check` passed ## Runtime boundary microsoft/onnxruntime#32743 provides a bounded CUDA correctness fallback for `(4,8)` and `(8,4)` despite its earlier summary listing only 2/4-bit cases. It dequantizes expert weights into FP16/BF16 scratch and requires canonical raw weights, 3-D scales, block size 16–256, divisible reduction dimensions, and scratch within its configured cap. Uniform (8,8) uses the native packed INT8 CUDA path: the tested tiny model needed block size divisible by 32 and both FC reduction dimensions divisible by 64. The synthetic test uses small weights with zero attention to isolate routed experts. The separate real-slice test validates selected checkpoint experts via Olive RTN tensor production, not complete Olive pass orchestration or the original 128-expert top-8 model. Full-checkpoint parity, prefill/decode generation parity, throughput, and peak VRAM have **not** been validated. For Qwen3-30B-A3B, FP16 expert scratch alone is approximately 1.125 GiB per MoE layer, above the fallback default 1 GiB: configure `ep.cuda.qmoe_int_dequant_max_scratch_bytes` (or the runtime's corresponding environment override) and measure before claiming deployability at production size. CPU mixed-width execution is not qualified. microsoft/onnxruntime#32761 adds a packed decode subset for 2/4-bit combinations, not the `(4,8)` INT8 path. Addresses #735. --------- Signed-off-by: titaiwang <titaiwang@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.


Description
Add the schema and validation contract needed for mixed-width QMoE expert weights.
fc1_expert_weight_bits,fc2_expert_weight_bits, andfc3_expert_weight_bitsattributes, inheriting fromexpert_weight_bitswhen omitted.logical_elements * bits / 8.pack_sizehelper overloads and existing 2/4/8 validation diagnostics.NOT_IMPLEMENTEDstatus for mixed-width execution. Execution support will follow separately.Attribute semantics
expert_weight_bitsremains the legacy global value and the fallback for each optional projection override: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.expert_weight_bits=2and omit the overrides.Current execution boundary
For this contract-only change,
is_mixed_widthmeans that at least one effective projection width differs fromexpert_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
onnxruntime_provider_testwith CUDA enabled.Related to #32657.