Skip to content

Add mixed-width weight contract for QMoE - #32697

Merged
David Fan (jiafatom) merged 6 commits into
mainfrom
jiafa/int2-qmoe-mixed-width-contract
Sep 21, 2026
Merged

David Fan (jiafatom) merged 6 commits into
mainfrom
jiafa/int2-qmoe-mixed-width-contract

Conversation

@jiafatom

@jiafatom David Fan (jiafatom) commented Sep 19, 2026 •

Copy link
Copy Markdown
Contributor

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.

Copilot AI balanced review requested due to automatic review settings September 19, 2026 16:09

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can commit the suggested changes from lintrunner.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 Medium severity · 1 Low severity

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.

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can commit the suggested changes from lintrunner.

@tianleiwu Tianlei Wu (tianleiwu) left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@jiafatom
David Fan (jiafatom) merged commit 090f6ba into main Sep 21, 2026
92 of 93 checks passed
@jiafatom
David Fan (jiafatom) deleted the jiafa/int2-qmoe-mixed-width-contract branch September 21, 2026 20:21
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

⌂ 换网址