You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Add CUDA correctness fallback for INT2 and mixed-width QMoE - #32743
Add a bounded CUDA correctness fallback for blockwise INT2 and mixed-width integer QMoE weights.
The fallback independently dequantizes FC1 and FC2 expert weights to FP16/BF16 scratch and runs the existing dense A16 CUTLASS MoE path. It supports FC1/FC2 combinations (2,4), (4,2), and (2,2), including canonical uniform INT2 via expert_weight_bits=2 without FC-specific overrides.
Fallback input contract
To keep this correctness path bounded and match the CUDA blockwise dequantizer contract, it requires:
Raw weights (weights_prepacked=0) in canonical [E,N,K/pack] layout; legacy transposed layouts are rejected.
Three-dimensional blockwise scales in [E,N,K/block_size] layout; row-wise scales are rejected.
A power-of-two block_size in [16,256].
Both FC reduction dimensions divisible by block_size.
Dequantized FC1+FC2 scratch within ep.cuda.qmoe_int_dequant_max_scratch_bytes.
Input-layout, dimension, and scratch-limit validation runs immediately after CheckInputs, before tactic profiling, workspace allocation, or stream synchronization.
Why this is a correctness fallback
This PR establishes CUDA behavior and numerical correctness without introducing the final performance kernel. Weight unpacking/dequantization and MoE execution both run on CUDA, but the implementation materializes complete FP16/BF16 expert weights in temporary scratch before invoking the existing dense A16 CUTLASS runner. It therefore uses more memory bandwidth and scratch space than a native packed INT2 path.
Follow-up #32761 pre-packs constant initializers and directly consumes packed INT2/INT4 weights for eligible decode, avoiding full dense materialization on that path. This bounded fallback remains for dynamic inputs and configurations outside the packed path; packed prefill and larger decode workloads remain follow-up work. CPU mixed-QMoE support and cross-provider parity are separate work.
Testing
Built onnxruntime_providers_cuda.
Built onnxruntime_provider_test.
Passed all 14 MoETest.QMoETest_MixedWidth* tests on NVIDIA A100 (SM80).
Added regressions for row-wise scales, legacy transposed weights, invalid block sizes, FC1/FC2 K divisibility, early scratch rejection, and canonical non-square execution.
CI follow-up: 031c236418 gates the BF16 QMoE test with the same build conditions as the BF16 MoE runner (!ORT_QUICK_BUILD && ENABLE_BF16). The four CUDA/TensorRT failures came from quick builds that intentionally omit that runner. The full A100 build still runs and passes all 12 mixed-width tests.
The WebGPU plugin failure is separate: Dawn/D3D device initialization failed in QMoETest_MixedWidthContract_WebGPU; the new commit will rerun that check.
This accepts block sizes that the invoked dequantizers cannot handle. Dequantize2Bits requires a power-of-two block size no larger than 256 and all DequantizeNBits variants require each K dimension to be divisible by it; for example, block_size=24 or hidden_size=80, block_size=32 passes here and then fails after profiling/allocation. Validate the actual kernel preconditions up front.
This issue also appears in the following locations of the same file:
The reason will be displayed to describe this comment to others. Learn more.
The bounded A16 fallback and CUDA-only tests are a useful correctness baseline, but two input shapes accepted by CheckInputs are interpreted differently by the new dequantizer: row-wise scales can cause an out-of-bounds read, and legacy weight layouts silently change the result. Please guard or support those formats before merging. The other inline comments cover an early-rejection improvement and an optional design for lower decode overhead. I have not repeated the block-size concern already raised in the existing automated review summary.
Tianlei Wu (@tianleiwu) Addressed the legacy-layout thread that GitHub no longer exposes for inline reply in 9257052114:
The dense fallback now accepts only canonical [E,N,K/pack] FC1/FC2 weight layouts and rejects legacy transposed layouts before profiling/dequantization.
Added a non-square legacy-layout rejection regression.
Added a positive canonical 64x128 CUDA fallback test to verify valid non-square layouts still execute correctly.
The full MoETest.QMoETest_MixedWidth* set passes 14/14 on A100.
Re-reviewed commit 9257052. The new validation correctly rejects incompatible block sizes, non-divisible K dimensions, row-wise scale tensors, and legacy transposed weights before profiling or allocation. The scratch-limit check is also now enforced early, and the new non-square/invalid-input regressions cover the previously identified gaps. I found no remaining blocking correctness issue. I left one non-blocking question about whether the additional mixed 8-bit combinations enabled by the fallback are intentional and supported.
Akshay Sonawane (@apsonawane) Yes, the mixed 8-bit combinations are intentional for the bounded dense fallback. The QMoE schema allows each FC-specific width to be 2, 4, or 8, and the fallback dispatches each matrix through the existing DequantizeNBits implementation, which handles all three widths before invoking the same A16 runner. Commit 38ec92f8a8 adds CUDA numerical coverage for (FC1, FC2) = (2,8), (8,2), (4,8), and (8,4); the focused test and the full 14-test MoETest.QMoETest_MixedWidth* set pass on A100.
The remaining Android minimal binary-size failure is upstream and unrelated to this CUDA-only PR. Latest main has failed the same check repeatedly: both main and this PR produce 1,596,029 bytes against the stale 1,595,392-byte threshold (637 bytes over). I opened #32772 to raise the baseline by the minimum 1 KiB increment. Once that lands, rebasing this PR will pick up the fix.
The reason will be displayed to describe this comment to others. Learn more.
Re-reviewed on 5dcede4b. My three earlier threads (2-D scale rejection, early scratch rejection before profileTactics, and the per-call dequant cost) are addressed or answered, and all three are resolved.
Things I re-verified as correct this round, so nobody has to re-derive them:
The [E*N, K] row-major dequant layout does match the grouped GEMM: TmaWarpSpecializedGroupedGemmInput::LayoutB is ColumnMajor, i.e. [K,N] col-major == [N,K] row-major. QMoETest_MixedWidthCudaCanonicalNonSquare is not transpose-invariant, so it genuinely covers this rather than passing by symmetry.
The zero-point packing convention lines up. Dequantize2BitsKernel indexes n_idx * ceil(groups_per_K/4) + kb_idx/4 (per-row byte padding), which is exactly what moe_helper::PackedByteCountWithPadding(blocks_per_row, bits) validates.
The new block_size power-of-two / k % block_size == 0 pre-checks mirror the ORT_ENFORCEs inside Dequantize2Bits, so the kernel's throwing asserts are now unreachable behind a clear error.
Relaxing expert_weight_bits to allow 2 stays bounded: fp4/nvfp4/wfp4afp8 still ORT_ENFORCE(== 4) and fp8 still ORT_ENFORCE(== 8).
PrePack early-returns for int2 and mixed-width, and the packed_fc1_weights_ == nullptr fail-loud guard correctly excludes the fallback. Separate fc3_experts_weights is already rejected outright on CUDA.
Two new findings, detailed inline:
[High]prepare_scale_zp is not short-circuited for use_int_dequant_fallback. Its zero-point branch is gated on the uniform expert_weight_bits_ rather than the per-FC width, so a legal mixed-width INT2 + zero-point model can hard-ORT_THROW from the 4-bit ZP packing check. The current asymmetric test only escapes it because blocks == 2 makes ceil(2/4) and ceil(2/2) collide.
[Medium] The moe_helper.h relaxation is a cross-EP change: the shared validator now calls [E,N,1] block-wise, but CPU QMoE's compute path still calls it row-wise (moe_quantization_cpu.cc:1511).
Plus suggestions on constructor-time validation and test coverage inline.
## 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
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.
Tianlei Wu (tianleiwu)
left a comment
Description
Add a bounded CUDA correctness fallback for blockwise INT2 and mixed-width integer QMoE weights.
The fallback independently dequantizes FC1 and FC2 expert weights to FP16/BF16 scratch and runs the existing dense A16 CUTLASS MoE path. It supports FC1/FC2 combinations
(2,4),(4,2), and(2,2), including canonical uniform INT2 viaexpert_weight_bits=2without FC-specific overrides.Fallback input contract
To keep this correctness path bounded and match the CUDA blockwise dequantizer contract, it requires:
weights_prepacked=0) in canonical[E,N,K/pack]layout; legacy transposed layouts are rejected.[E,N,K/block_size]layout; row-wise scales are rejected.block_sizein[16,256].block_size.ep.cuda.qmoe_int_dequant_max_scratch_bytes.Input-layout, dimension, and scratch-limit validation runs immediately after
CheckInputs, before tactic profiling, workspace allocation, or stream synchronization.Why this is a correctness fallback
This PR establishes CUDA behavior and numerical correctness without introducing the final performance kernel. Weight unpacking/dequantization and MoE execution both run on CUDA, but the implementation materializes complete FP16/BF16 expert weights in temporary scratch before invoking the existing dense A16 CUTLASS runner. It therefore uses more memory bandwidth and scratch space than a native packed INT2 path.
Follow-up #32761 pre-packs constant initializers and directly consumes packed INT2/INT4 weights for eligible decode, avoiding full dense materialization on that path. This bounded fallback remains for dynamic inputs and configurations outside the packed path; packed prefill and larger decode workloads remain follow-up work. CPU mixed-QMoE support and cross-provider parity are separate work.
Testing
onnxruntime_providers_cuda.onnxruntime_provider_test.MoETest.QMoETest_MixedWidth*tests on NVIDIA A100 (SM80).clang-format --dry-run --Werror.git diff --check.