feat(qwen35-moe): emit fused com.microsoft::QMoE for int4 A3B decode - #451
Open
justinchuby wants to merge 2 commits into
Open
feat(qwen35-moe): emit fused com.microsoft::QMoE for int4 A3B decode#451justinchuby wants to merge 2 commits into
justinchuby wants to merge 2 commits into
Conversation
Qwen3.5-MoE (e.g. Qwen3.6-35B-A3B) exported its MoE as dense_fallback: every routed expert un-fused into per-expert MatMulNBits MLPs + Equal/ReduceSum masking, so native decode computed all 256 experts/token instead of the top-k=8 active. onnx-genai's CUDA QMoE kernel already does sparse top-k decode; the only gap was the export. - _supported_qmoe_quantization: accept quant_method="olive" (blk32 int4 integer-affine) alongside gptq/awq, and require a power-of-two block_size>=16 (the CUDA QMoE constraint) so unrunnable configs fall back to dense. - preprocess_olive_weights: preserve all leading dims when reshaping .qweight so fused expert-major tensors [E,N,packed_K] -> [E,N,n_blocks,blob_size] feed the QMoE repacker; 2-D linears are byte-identical (no behavior change). - Qwen35MoECausalLMModel.preprocess_weights: guard the per-expert un-fuse behind use_qmoe; when QMoE is supported, keep fused expert-major weights and route through pack_qmoe_expert_weights (target_moe_path=".mlp") to emit one QMoE node per MoE layer. Mirrors DeepSeek-V3's use_qmoe path. - deepseek: use the shared _supported_qmoe_quantization predicate so preprocess and MoELayer never disagree (DRY, avoids block_size drift). Byte-equivalent to dense_fallback: same (q-zp)*scale affine, same softmax-top-k routing, same SwiGLU (swiglu_fusion=2). Pure layout repack, no requantization. Tests: Olive QMoE emission (fused node, zero per-expert MatMulNBits storm), non-pow2 block_size -> dense fallback, and byte-exact fused-expert repack vs the dense per-expert reference (dequant atol=0 + SwiGLU forward parity). Model-level Qwen35 preprocess routing (QMoE params bind; dense path still un-fuses). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| @@ -0,0 +1,123 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR updates MoE export for Qwen3.5/3.6-MoE (A3B) to emit a fused com.microsoft::QMoE node per MoE layer when the quantization config matches the CUDA QMoE ABI, enabling sparse top‑k expert execution for int4 decode instead of the dense per‑expert fallback graph.
Changes:
- Expand the QMoE-eligibility predicate to include
quant_method="olive"and enforce CUDA kernel constraints (power-of-twogroup_size >= 16) so unsupported configs fall back to the dense representation. - Update Olive weight preprocessing to preserve leading dimensions when reshaping
.qweight, enabling expert-major fused tensors to be repacked into QMoE ABI layout. - Wire Qwen3.5-MoE and DeepSeek to use the shared QMoE predicate, and add targeted tests validating QMoE emission and expert-major packing.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
src/mobius/models/qwen35.py |
Gate MoE expert un-fusing behind use_qmoe and pack expert-major weights into QMoE parameters. |
src/mobius/models/qwen35_test.py |
New build/weight tests focused on Qwen3.5-MoE QMoE emission and packed-weight binding. |
src/mobius/models/deepseek.py |
Use shared _supported_qmoe_quantization predicate for consistent QMoE enablement. |
src/mobius/components/_moe.py |
Extend QMoE support to Olive and add block-size constraints to avoid emitting unrunnable QMoE nodes. |
src/mobius/components/_moe_test.py |
Add tests for Olive QMoE emission, fallback behavior, and byte-exact packing parity. |
src/mobius/_weight_utils.py |
Preserve leading dims when reshaping Olive .qweight to support fused expert-major tensors. |
Comment on lines
+403
to
+408
| Accepts the integer-affine int4 block schemes whose ``(q - zero_point) * | ||
| scale`` dequantization is byte-identical to ``MatMulNBits`` and to the | ||
| ``com.microsoft::QMoE`` kernel: GPTQ, AWQ, and Olive RTN | ||
| (``quant_method="olive"``, blk32 int4). The CUDA QMoE kernel requires a | ||
| power-of-two ``block_size >= 16``, so configs outside that range fall back | ||
| to the portable dense representation instead of emitting an unrunnable node. |
Comment on lines
+6
to
+10
| Focuses on the ``com.microsoft::QMoE`` emission path: when the quantization | ||
| config matches the native QMoE ABI (blk32 int4 Olive/GPTQ/AWQ), | ||
| :meth:`Qwen35MoECausalLMModel.preprocess_weights` keeps the fused expert-major | ||
| tensors and repacks them into ``fc1``/``fc2`` QMoE parameters (mirroring | ||
| DeepSeek-V3), instead of un-fusing into a per-expert dense fallback. All tiny |
Comment on lines
+380
to
+384
| # Unpack fused float expert weights into per-expert tensors for the | ||
| # dense fallback. When ``use_qmoe`` is set the fused quantized | ||
| # tensors arrive as ``.qweight``/``.scales``/``.qzeros`` (which do | ||
| # not match these suffixes) and are kept expert-major for | ||
| # ``pack_qmoe_expert_weights`` below. |
Add fuse_dense_moe_to_qmoe, an ONNX-to-ONNX rewrite that recognises the per-expert dense-fallback MoE subgraph (router MatMulNBits -> TopK -> Softmax, per-expert Equal/Cast/Mul/ReduceSum masking, E x3 expert MatMulNBits MLPs and the left-folded weighted sum) and rewrites it to a single com.microsoft::QMoE node per MoE layer. The existing int4 MatMulNBits expert weights and zero-points are reused byte-for-byte (expert-major concat/flatten layout transform, no requant); float16 scales are upcast to the float32 the QMoE kernel requires (lossless). The router logits are Cast to float32 and passed as router_probs with normalize_routing_weights=1, matching TopKGate. Any shared expert is preserved and the final Add(routed, shared) is rewired onto the QMoE output. Handles the merged-graph variant that inserts a Cast between the gate MatMulNBits and TopK, and between TopK values and Softmax. Tests build a tiny Qwen35-MoE-shaped dense graph and assert the E x3 MatMulNBits storm collapses to one QMoE node, the packed weights/zero-points are bit-identical to the concatenated per-expert tensors, scales are a lossless float32 upcast, and a reference forward matches the dense forward. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.
Suppressed comments (1)
src/mobius/rewrite_rules/_qmoe_fusion.py:220
- Dense-fallback graphs may wire the constant expert id on either side of Equal (Equal is commutative).
_collect_expertscurrently assumes the constant is alwaysequal.inputs[1], so it will silently miss experts (and skip fusing) when the constant isequal.inputs[0]and indices areequal.inputs[1].
indices = self.topk.outputs[1]
for equal in _consumers_of_type(indices, "Equal"):
expert_id = _scalar_int(equal.inputs[1])
if expert_id is None:
continue
Comment on lines
+314
to
+315
| bits, block_size, has_zp = _expert_geometry(down_nodes[0]) | ||
|
|
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.
Goal
Realize Qwen3.5/3.6-MoE (35B-A3B) A3B decode speed by exporting a fused
com.microsoft::QMoEnode per MoE layer instead of thedense_fallbackrepresentation (per-expert
MatMulNBitsMLPs +Equal/ReduceSummasking).The dense_fallback graph computes all 256 experts every token (full 35B HBM
weight traffic). onnx-genai's CUDA QMoE kernel already does sparse top-k=8
decode (reads only k experts' weights, ~32x less MoE weight traffic), so the
only change on the critical path was the mobius export.
Changes
_supported_qmoe_quantization(components/_moe.py): acceptquant_method="olive"(blk32 int4 integer-affine) alongsidegptq/awq, andrequire a power-of-two
block_size >= 16(the CUDA QMoE kernel constraint) sonon-conforming configs fall back to the portable dense representation rather
than emitting an unrunnable node.
preprocess_olive_weights(_weight_utils.py): preserve all leading dimswhen reshaping
.qweight, so fused expert-major tensors[E, N, packed_K] -> [E, N, n_blocks, blob_size]feed the QMoE repacker. 2-Dlinears reshape identically (
[N, n_blocks, blob_size]) — no behavior change.Qwen35MoECausalLMModel.preprocess_weights(models/qwen35.py): guard theper-expert un-fuse behind
use_qmoe. When QMoE is supported, keep the fusedexpert-major weights and route through
pack_qmoe_expert_weights(..., target_moe_path=".mlp")to emit one QMoE node per MoE layer instead ofun-fusing. Mirrors DeepSeek-V3's
use_qmoepath.deepseek.py: use the shared_supported_qmoe_quantizationpredicate foruse_qmoeso the repacked weights and the emitted graph never disagree (DRY;avoids drift from the new
block_sizeguard).Correctness
Byte-equivalent to
dense_fallback: identical(q - zero_point) * scaleint4affine, identical softmax-top-k routing, identical SwiGLU (
swiglu_fusion=2,gate = first half of
fc1, up = second half). This is a pure layout repack —no requantization, no accuracy change.
Tests (all green via the conda
onnxenv)test_int4_olive_moe_emits_expert_major_qmoe: Olive MoE emits a single fusedQMoEnode, zero per-expertMatMulNBits.test_olive_group_size_not_power_of_two_falls_back_to_dense: non-pow2block_size-> dense fallback (kernel would reject).test_olive_fused_expert_packing_is_byte_exact_vs_dense: fused Olive expertweights (asymmetric int4, per-block zero-points) repack to the QMoE ABI with
dequant
atol=0, and a loop-over-experts SwiGLU forward matches the denseper-expert reference within fp32 tolerance.
models/qwen35_test.py: Olivepreprocess_weightspacksfc1/fc2QMoEparams that bind to real model parameters and leak no per-expert tensors; the
unquantized path still un-fuses to the dense fallback.
ruffclean;_moe_test.py,qwen35_test.py,deepseek_test.py,_qwen35_mtp_test.py,_weight_utils_test.pyall pass (132 tests).Structural proof (tiny 4-layer Qwen3.5-MoE, E=16/k=4, blk32 int4 Olive)
com.microsoft::QMoEMatMulEqual(mask)ReduceSum(mask)TopKOne QMoE node per MoE layer; the per-expert MatMul storm and Equal/ReduceSum
masking are eliminated. Extrapolates to the real 35B: 40 MoE layers -> 40 QMoE
nodes, ~20,480 per-expert expert MatMulNBits removed.
35B regen status (follow-up)
The existing dense artifact was produced via mobius float export ->
Olive ONNX-graph
onnxkquantquantization(q4_k_m). This PR implements themobius olive-import QMoE path (
quant_method="olive", pre-packed.qweight). A full 35B QMoE artifact therefore needs an Olive-RTN-quantizedQwen3.6-35B HF checkpoint (fused expert
.qweight/.scales/.qzeroswithquantization_config.quant_method="olive"), which is not on disk yet — the onlylocal copy is the float bf16 source. The tiny proof + byte-exact tests validate
the path end to end; the 35B regen is a compute step (RTN-quantize the 35B, then
re-export) that a measurement pass can pick up.
Please review and merge — I do not self-merge mobius.
Update — dense-fallback → QMoE graph rewrite (reuse existing int4 weights)
To get a measured 35B QMoE number without re-quantizing a 70GB HF checkpoint,
this PR also adds an ONNX-to-ONNX rewrite that fuses the existing
dense_fallback artifact's per-expert
MatMulNBitsstorm intoQMoEnodes,reusing the int4 weights byte-for-byte.
New:
mobius.rewrite_rules.fuse_dense_moe_to_qmoerewrite_rules/_qmoe_fusion.py— a topology-driven pass (anchored on each MoETopK) that recognises the dense_fallback subgraph and rewrites it to onecom.microsoft::QMoEnode per layer:MatMulNBitslogits areCastto float32 and passed asrouter_probswithnormalize_routing_weights=1(matchesTopKGate:Softmax(TopK(logits))). Walks through an optionalCastbetween the gate andTopK(themergedgraph variant).fc1 = concat(flatten(gate_w), flatten(up_w))andfc2 = flatten(down_w), stacked expert-major. Packeduint8weights andzero-points are copied bit-identically (same low-nibble-first, K-contiguous
packing as the QMoE kernel — verified against
qmoe.rs);float16scales areupcast to the
float32the kernel requires (value-exact, lossless).activation_type=swiglu,swiglu_fusion=2,k=8,expert_weight_bits=4,block_size=32,quant_type=int.shared_expert+shared_expert_gate) is leftdense; the final
Add(routed_sum, shared)is rewired onto the QMoE output.Dead dense-routing nodes and now-unused initializers are removed.
Tests (
rewrite_rules/_qmoe_fusion_test.py, 6 cases)On a tiny Qwen35-MoE-shaped dense graph (incl. a shared expert):
E x3 MatMulNBitsstorm collapses to 1 QMoE node (TopK/Softmax/Equal/ReduceSumremoved); packed weights + zero-points are bit-identical to theconcatenated per-expert tensors; scales are a lossless float32 upcast; and a
reference forward on the rewritten graph matches the dense forward within fp32
tolerance.
Applied to the real 35B artifact (A/B, no requant)
Rewrote
qwen36-35b-a3b-artifacts/{decoder,merged}/model.onnxintoqwen36-35b-a3b-qmoe-artifacts/(dense artifact left untouched):The
MatMulNBitsdrop is exactly40 x 256 x 3 = 30,720(the full expertstorm); the remaining 391 are the router gate + shared-expert projections. Graph
integrity re-checked (0 dangling inputs, 0 topo violations). Byte-exactness
re-verified on real experts (layer 0, experts 0/1/127/255): fc1/fc2 weights and
zero-points bit-identical, scales fp16→fp32 lossless.