QMoE CUDA EP — FP4/FP8/WFP4AFP8 Quantized Mixture-of-Experts + MoE GEMM Refactor - #28467
Merged
Conversation
Tianlei Wu (tianleiwu)
marked this pull request as ready for review
May 15, 2026 01:08
Tianlei Wu (tianleiwu)
force-pushed
the
tlwu/20260511/qmoe_cuda
branch
from
May 17, 2026 17:01
bbc0af2 to
b3143e0
Compare
Tianlei Wu (tianleiwu)
marked this pull request as draft
May 18, 2026 07:42
Tianlei Wu (tianleiwu)
force-pushed
the
tlwu/20260511/qmoe_cuda
branch
from
May 18, 2026 17:53
6c72e8c to
16ca767
Compare
Tianlei Wu (tianleiwu)
marked this pull request as ready for review
May 19, 2026 07:05
Tianlei Wu (tianleiwu)
requested review from
Akshay Sonawane (apsonawane) and
kunal-vaishnavi
May 19, 2026 07:13
kunal-vaishnavi
approved these changes
May 20, 2026
Tianlei Wu (tianleiwu)
enabled auto-merge (squash)
May 20, 2026 05:56
Tianlei Wu (tianleiwu)
added a commit
that referenced
this pull request
May 21, 2026
…28583) ## Description Follow-up refinements to QMoE CUDA EP (#28467): rename build options for consistency, fix CodeQL warnings, refactor `PrePack` from nested lambdas into named helper methods, and replace CPU data-transformation loops with GPU kernels. ### Motivation and Context PR #28467 introduced the QMoE operator with a 373-line `PrePack` function containing 5 nested lambdas that performed weight/scale prepacking at model load time. Reviewer feedback requested: 1. Rename `onnxruntime_ENABLE_CUDA_*` cmake options to `onnxruntime_USE_*` for naming consistency. 2. Fix CodeQL empty-except warnings in test code. 3. Extract lambdas into named private methods for readability and testability. 4. Replace CPU loops (block-scale swizzle, FP4 col-to-row repack) with GPU kernels to avoid unnecessary CPU↔GPU round-trips during model loading. ## Key Changes | Commit | Scope | Description | |--------|-------|-------------| | `594642a` | Build system | Rename `ENABLE_CUDA_FP4_QMOE`→`USE_FP4_QMOE`, `ENABLE_CUDA_FP8_QMOE`→`USE_FP8_QMOE` in cmake, C++ defines, and 340+ generated .cu files | | `594642a` | Test | Fix CodeQL empty-except warning in `test_qmoe_cuda.py` | | `e8d364b` | QMoE operator | Extract 5 lambdas into private helper methods: `PrePackTransposeAndPack`, `PrePackCopyToGpu`, `PrePackSwizzleBlockScales`, `PrePackRepackFP4Weights`, `PrePackComputeBias` | | `e8d364b` | CUDA kernels | Add `QMoERepackFP4ColToRowKernel` — repacks column-major FP4 packed weights to row-major on GPU (replaces per-expert CPU loop) | | `e8d364b` | CUDA kernels | `PrePackSwizzleBlockScales` now calls existing `LaunchQMoEBlockScaleInterleave` GPU kernel (replaces CPU `SwizzleMXFPXBlockScalesToGpu` loop) | ## Impact - **No behavioral change** — all transformations produce identical output tensors. - **Model load only** — `PrePack` runs once during `InferenceSession::Initialize()`, not on the inference hot path. - **Performance**: Eliminates CPU→GPU→CPU→GPU round-trips for block-scale swizzling and FP4 weight repacking. Data stays on GPU throughout. ## Testing - Build verified with `onnxruntime_USE_FP4_QMOE=ON onnxruntime_USE_FP8_QMOE=ON` (CUDA 13.0, SM90). - All new symbols confirmed linked in `libonnxruntime_providers_cuda.so` via `nm`. - Existing `test_qmoe_fp4_cuda.py`, `test_qmoe_wfp4afp8_cuda.py`, `test_qmoe_cuda.py` cover the affected code paths.
This was referenced Jun 1, 2026
Justin Chu (justinchuby)
added a commit
to onnxruntime/mobius
that referenced
this pull request
Jun 1, 2026
Re-enables the fused `com.microsoft::MoE` op for Gemma 4 now that ORT main supports standard SwiGLU. Replaces the static-unroll fallback that was shipped in #324 as a workaround. ## Why now [microsoft/onnxruntime#28467](microsoft/onnxruntime#28467) (QMoE CUDA EP + MoE GEMM Refactor, merged on main) plumbs the existing schema attributes (`swiglu_fusion`, `activation_alpha`, `activation_beta`, `swiglu_limit`) all the way through to the kernel. Before that PR, the CUDA MoE kernel hardcoded the GPT-OSS values (`alpha=1.702, beta=1.0, limit=7.0`, interleaved) and silently produced wrong output for any model using standard SwiGLU. That was the root cause tracked in [microsoft/onnxruntime-genai#2062](microsoft/onnxruntime-genai#2062). ## Attribute set ``` activation_type = 'swiglu' activation_alpha = 1.0 # no GPT-OSS 1.702 multiplier activation_beta = 0.0 # no GPT-OSS '+1' bias on the up branch swiglu_limit = inf # no clipping swiglu_fusion = 1 # interleaved normalize_routing_weights = 1 k = top_k ``` ## Weight layout HuggingFace stores `experts.gate_up_proj` chunked as `[E, 2*inter, H]` (first `inter` rows are gate, next `inter` are up). The CPU MoE kernel still only accepts interleaved layout (`contrib_ops/cpu/moe/moe_cpu.cc:27`); the new CUDA kernel accepts either. We emit `swiglu_fusion=1` (interleaved) for maximum portability and reshape at graph-emit time via `Reshape → Transpose → Reshape` on the initializer — ORT folds the chain to a single static tensor at session load. ## Fallback `_dispatch_moe_fallback` (static per-expert unroll) is kept verbatim for EPs that don't advertise `supports_fused_moe`. The fused path is only taken when `ep_capabilities().supports_fused_moe` is true (default for the CUDA / DML / default EPs). ## Validation (H200, ORT `1.27.0.dev20260511001` which contains #28467) | Check | Result | |---|---| | 15 `gemma4` graph-construction tests | all pass | | fp16 build of `google/gemma-4-26b-a4b-it` | 30 MoE nodes emitted, all attributes correctly set | | `InferenceSession` load on CUDA EP | **12.3s** (vs 959s with the previous unrolled fallback — ~78× faster session creation) | | Prefill (B=1, S=4) | 0.65s, logits well-behaved (no NaN/Inf, top-k IDs in valid vocab range) | | lintrunner | clean | ## Compatibility - Requires ORT main (or any future release containing #28467). On older ORT releases (≤1.27 GA) the fused op will silently use the GPT-OSS hardcoded constants and produce wrong output. Users should keep the unrolled fallback build (pre-#324 reversion) until they upgrade ORT. Signed-off-by: Justin Chu <justinchu@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Justin Chu (justinchuby)
added a commit
to justinchuby/Olive
that referenced
this pull request
Jun 2, 2026
Adds a new ONNX graph pass that rewrites every com.microsoft::MoE node into a com.microsoft::QMoE node with the per-expert FC1/FC2 weight initializers quantized to symmetric int4 (default) or int8, plus corresponding fp16 scale initializers. Motivation: mobius (and similar exporters) emit the fused com.microsoft::MoE op with the per-expert weights as 3-D fp16/bf16/fp32 initializers. The existing weight-quantization passes (OnnxKQuantQuantization, OnnxBlockWiseRtnQuantization, OnnxBnb4Quantization) only target MatMul nodes, so for MoE models the per-expert weights (~80% of total parameters) stay at the model's compute dtype, leaving just ~6% size reduction after quantization. The QMoE op is the correct target for MoE weights and is supported by the CUDA + experimental CPU kernels in ORT main (PR microsoft/onnxruntime#28467). Implementation: - Walks the graph and finds every com.microsoft::MoE node whose fc1_experts_weights and fc2_experts_weights are 3-D static initializers. - For each expert, calls ORT's pybind quantize_matmul_{4,8}bits to produce per-expert int4/int8 weights + symmetric fp16 scales, then CUTLASS-prepacks them via pack_weights_for_cuda_mixed_gemm so the QMoE kernels can consume the bytes directly. - Stacks per-expert tensors along axis 0 and registers them as new initializers (uint8 weight + fp16 scale per expert). - Replaces the MoE node with a QMoE node carrying the original activation/routing attributes plus expert_weight_bits, optional block_size, and quant_type='int'. - Orphaned fp16 weight initializers are dropped. Supports per-row scales (block_size=0, default) and block-wise scales (block_size ≥ 16, must be power of two). Nodes can be selectively excluded via nodes_to_exclude. The pass requires a CUDA-enabled ONNX Runtime build because pack_weights_for_cuda_mixed_gemm is only exposed when ORT is compiled with USE_CUDA. A descriptive RuntimeError is raised at run time when the binding is unavailable. Limitations / out-of-scope: - fc3 inputs (3-fold MoE variants) are not supported and trigger a warning-skip per node. - Only symmetric int quantization (matching the kernel's preferred layout). FP4 / FP8 / WFP4AFP8 quant_types are left for a follow-up. - Calibration-aware quantization (GPTQ / AWQ) is out of scope; this pass is pure RTN. Tests: 5 unit tests covering (a) end-to-end MoE → QMoE conversion with int4 + per-row scales, (b) block-wise int4, (c) graceful skip when weights are not static initializers, (d) bits validation, and (e) block_size validation. The CUTLASS prepack helper is patched during tests so CI without onnxruntime-gpu can still exercise the graph transform. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Jayashree-mcw
pushed a commit
to Jayashree-mcw/onnxruntime
that referenced
this pull request
Jun 9, 2026
|
There is still an include for moe_kernel.h which seems to fail after this PR https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/contrib_ops/cuda/collective/sharded_moe.h#L6 |
13 tasks
This was referenced Jul 29, 2026
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
Update
QMoEcontrib operator for the CUDA EP to supports quantized Mixture-of-Experts inference with INT4, INT8, FP4 (MXFP4 e2m1), FP8 (e4m3fn), and WFP4AFP8 (mixed FP4 weight × FP8 activation) quantization formats.This also refactors the existing MoE GEMM infrastructure to support TMA warp-specialized grouped GEMM on Hopper (SM90), native MXFP4 on Blackwell (SM120), and block-scaled tensor ops on SM100+, with automatic fallback to dequantization on older architectures.
Note that this is modified from
TensorRT-LLMMoE implementation. There is a section in moe_qmoe.md about the modifications.Summary of Changes
New QMoE Operator
onnxruntime/core/graph/contrib_ops/contrib_defs.ccQMoEop schema (com.microsoft domain, opset 1)onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc/honnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu/honnxruntime/contrib_ops/cuda/moe/moe_base.hdocs/contrib_ops/cuda/moe_qmoe.mdMoE GEMM Refactor
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_kernels.hCutlassMoeFCRunnertemplate with FP4/FP8/WFP4AFP8 specializationsonnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_template_dispatch.honnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_profiler.cc/honnxruntime/contrib_ops/cuda/llm/moe_gemm/common.honnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/CUTLASS Extensions
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/arch/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/collective/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/epilogue/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/system_barrier.hCommon CUDA Utilities
onnxruntime/contrib_ops/cuda/llm/common/cuda_fp8_utils.cu/h— FP8 conversion, quantization, dequantization kernelsonnxruntime/contrib_ops/cuda/llm/common/memory_utils.cu/h— Device memory transpose, permute, type conversion utilitiesonnxruntime/contrib_ops/cuda/llm/common/cuda_type_utils.cuh— Unified type traits for half/bfloat16/float/fp8/fp4onnxruntime/contrib_ops/cuda/llm/common/quantization.h— Quantization parameter structs and helpersonnxruntime/contrib_ops/cuda/llm/common/reduce_kernel_utils.cuh— Warp/block reduction primitivesonnxruntime/contrib_ops/cuda/llm/kernels/quantization.cuh— FP4/FP8 quantization kernelsonnxruntime/contrib_ops/cuda/llm/kernels/pre_quant_scale_kernel.cu/h— Pre-quantization scaling kernelGEMM Profiler Refactor
onnxruntime/contrib_ops/cuda/llm/gemm_profiler.cc/honnxruntime/contrib_ops/cuda/llm/cutlass_heuristic.cc/honnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm_configs.hBuild System
cmake/CMakeLists.txtENABLE_FP4,ENABLE_FP8,ENABLE_CUDA_FP4_QMOE,ORT_QUICK_BUILD,PLACEHOLDER_KERNELSoptionscmake/external/cuda_configuration.cmakecmake/external/cutlass.cmakecmake/onnxruntime_providers_cuda.cmakecmake/onnxruntime_python.cmakeonnxruntime_pybind_quant.ccfor Python quantization bindingsPython Quantization Bindings
onnxruntime/python/onnxruntime_pybind_quant.cconnxruntime/python/tools/quantization/quant_utils.pysetup.pyTests
onnxruntime/test/python/transformers/test_qmoe_cuda.pyonnxruntime/test/python/transformers/test_qmoe_fp4_cuda.pyonnxruntime/test/python/transformers/test_qmoe_fp8_cuda.pyonnxruntime/test/python/transformers/test_qmoe_wfp4afp8_cuda.pyonnxruntime/test/python/transformers/test_moe_cuda.pyonnxruntime/test/contrib_ops/moe_test.ccExisting MoE Refactor
onnxruntime/contrib_ops/cuda/moe/moe.cc/h— Refactored to share base with QMoEonnxruntime/contrib_ops/cuda/moe/ft_moe/→onnxruntime/contrib_ops/cuda/llm/moe_gemm/— Relocated and rewritten MoE GEMM kernelscuda/quantization/moe_quantization.cc/hin favor of newcuda/moe/moe_quantization.cc/hTesting
python -m pytest onnxruntime/test/python/transformers/test_qmoe_cuda.py -v(requires CUDA GPU, SM75+)python -m pytest onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py -v(requires SM120+ for native, falls back on older)python -m pytest onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py -v(requires SM90+ for native)python -m pytest onnxruntime/test/python/transformers/test_qmoe_wfp4afp8_cuda.py -v(requires SM100+)python -m pytest onnxruntime/test/python/transformers/test_moe_cuda.py -vonnxruntime_test_all --gtest_filter=*MoE*Motivation and Context
Modern LLMs increasingly use Mixture-of-Experts architectures (e.g., Mixtral, DeepSeek, Phi-3.5-MoE) for efficient scaling. These models benefit significantly from weight quantization to reduce memory bandwidth and enable larger models on fewer GPUs. This PR: