Skip to content

QMoE CUDA EP — FP4/FP8/WFP4AFP8 Quantized Mixture-of-Experts + MoE GEMM Refactor - #28467

Merged
Tianlei Wu (tianleiwu) merged 20 commits into
mainfrom
tlwu/20260511/qmoe_cuda
May 20, 2026
Merged

QMoE CUDA EP — FP4/FP8/WFP4AFP8 Quantized Mixture-of-Experts + MoE GEMM Refactor#28467
Tianlei Wu (tianleiwu) merged 20 commits into
mainfrom
tlwu/20260511/qmoe_cuda

Conversation

@tianleiwu

@tianleiwu Tianlei Wu (tianleiwu) commented May 12, 2026

Copy link
Copy Markdown
Contributor

Description

Update QMoE contrib 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-LLM MoE implementation. There is a section in moe_qmoe.md about the modifications.

Summary of Changes

New QMoE Operator

File Change
onnxruntime/core/graph/contrib_ops/contrib_defs.cc Register QMoE op schema (com.microsoft domain, opset 1)
onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc/h QMoE CUDA kernel implementation with dynamic runner selection
onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu/h Softmax top-k router, sparse mixer, zero-point pre-packing kernels
onnxruntime/contrib_ops/cuda/moe/moe_base.h Shared MoE base class updates for quantization attributes
docs/contrib_ops/cuda/moe_qmoe.md Comprehensive operator documentation (inputs, attributes, quantization formats)

MoE GEMM Refactor

File Change
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_kernels.h Unified CutlassMoeFCRunner template with FP4/FP8/WFP4AFP8 specializations
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_template_dispatch.h Three-family dispatch: Ampere GemmGrouped, TMA warp-specialized, block-scaled tensor ops
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_profiler.cc/h MoE-specific GEMM tactic profiler for auto-tuning
onnxruntime/contrib_ops/cuda/llm/moe_gemm/common.h Shared MoE GEMM types and config structs
onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/ SM80/SM90/SM120 launcher instantiations (including generated .cu files)

CUTLASS Extensions

File Change
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/arch/ Grid dependency control, TMA copy traits, multi-mem copy operations
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/collective/ Mixed-input and gated GEMM collective builders for SM90
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/ Fused MoE kernel traits/routines, MoE problem visitors, gated GEMM kernels
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/epilogue/ MoE finalize epilogue, per-row/per-col scale epilogues
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/system_barrier.h System barrier for multi-CTA synchronization

Common CUDA Utilities

  • onnxruntime/contrib_ops/cuda/llm/common/cuda_fp8_utils.cu/h — FP8 conversion, quantization, dequantization kernels
  • onnxruntime/contrib_ops/cuda/llm/common/memory_utils.cu/h — Device memory transpose, permute, type conversion utilities
  • onnxruntime/contrib_ops/cuda/llm/common/cuda_type_utils.cuh — Unified type traits for half/bfloat16/float/fp8/fp4
  • onnxruntime/contrib_ops/cuda/llm/common/quantization.h — Quantization parameter structs and helpers
  • onnxruntime/contrib_ops/cuda/llm/common/reduce_kernel_utils.cuh — Warp/block reduction primitives
  • onnxruntime/contrib_ops/cuda/llm/kernels/quantization.cuh — FP4/FP8 quantization kernels
  • onnxruntime/contrib_ops/cuda/llm/kernels/pre_quant_scale_kernel.cu/h — Pre-quantization scaling kernel

GEMM Profiler Refactor

File Change
onnxruntime/contrib_ops/cuda/llm/gemm_profiler.cc/h Refactored GEMM profiler interface for tactic selection
onnxruntime/contrib_ops/cuda/llm/cutlass_heuristic.cc/h Updated heuristics for new kernel families
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm_configs.h Extended GEMM config enums for TMA warp-specialized and gated configs

Build System

File Change
cmake/CMakeLists.txt Add ENABLE_FP4, ENABLE_FP8, ENABLE_CUDA_FP4_QMOE, ORT_QUICK_BUILD, PLACEHOLDER_KERNELS options
cmake/external/cuda_configuration.cmake FP4/FP8 capability detection based on CUDA version and SM arch
cmake/external/cutlass.cmake CUTLASS version bump
cmake/onnxruntime_providers_cuda.cmake Add MoE GEMM source files and conditional FP4/FP8 kernel compilation
cmake/onnxruntime_python.cmake Add onnxruntime_pybind_quant.cc for Python quantization bindings

Python Quantization Bindings

File Change
onnxruntime/python/onnxruntime_pybind_quant.cc C++ pybind module for MoE weight preprocessing (quantize, pack, preprocess)
onnxruntime/python/tools/quantization/quant_utils.py FP4/FP8 quantization utilities
setup.py Include new pybind module in package build

Tests

File Change
onnxruntime/test/python/transformers/test_qmoe_cuda.py INT4/INT8 QMoE tests (Phi3 topology, SwiGLU, blockwise, asymmetric)
onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py MXFP4 QMoE tests
onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py FP8 QMoE tests
onnxruntime/test/python/transformers/test_qmoe_wfp4afp8_cuda.py WFP4AFP8 mixed-precision QMoE tests
onnxruntime/test/python/transformers/test_moe_cuda.py Updated existing MoE tests for refactored infrastructure
onnxruntime/test/contrib_ops/moe_test.cc C++ MoE unit tests updated

Existing MoE Refactor

  • onnxruntime/contrib_ops/cuda/moe/moe.cc/h — Refactored to share base with QMoE
  • onnxruntime/contrib_ops/cuda/moe/ft_moe/onnxruntime/contrib_ops/cuda/llm/moe_gemm/ — Relocated and rewritten MoE GEMM kernels
  • Removed old cuda/quantization/moe_quantization.cc/h in favor of new cuda/moe/moe_quantization.cc/h

Testing

  • INT4/INT8 QMoE: python -m pytest onnxruntime/test/python/transformers/test_qmoe_cuda.py -v (requires CUDA GPU, SM75+)
  • FP4 QMoE: python -m pytest onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py -v (requires SM120+ for native, falls back on older)
  • FP8 QMoE: python -m pytest onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py -v (requires SM90+ for native)
  • WFP4AFP8 QMoE: python -m pytest onnxruntime/test/python/transformers/test_qmoe_wfp4afp8_cuda.py -v (requires SM100+)
  • Existing MoE: python -m pytest onnxruntime/test/python/transformers/test_moe_cuda.py -v
  • C++ MoE tests: Build with CUDA EP enabled, run onnxruntime_test_all --gtest_filter=*MoE*
  • All tests compare QMoE output against PyTorch reference implementations with configurable tolerance

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:

  1. Adds native low-precision MoE support — FP4 and FP8 quantized weights avoid the dequantization overhead of INT4/INT8 on supported hardware (Hopper, Blackwell).
  2. Introduces WFP4AFP8 — A novel mixed-precision mode where weights are MXFP4 and activations are dynamically quantized to FP8, enabling 2× weight compression with minimal accuracy loss on Blackwell GPUs.
  3. Refactors MoE GEMM infrastructure — The previous FasterTransformer-derived MoE GEMM code is replaced with a modern CUTLASS 4.x-based dispatch system supporting three kernel families across SM75–SM120+.
  4. Adds auto-tuning — The GEMM profiler enables runtime tactic selection for optimal performance across different expert sizes and batch configurations.

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 wasn't able to review this pull request because it exceeds the maximum number of lines (20,000). Try reducing the number of changed lines and requesting a review from Copilot again.

@tianleiwu
Tianlei Wu (tianleiwu) marked this pull request as ready for review May 15, 2026 01:08
@tianleiwu
Tianlei Wu (tianleiwu) marked this pull request as draft May 18, 2026 07:42
Comment thread onnxruntime/test/python/transformers/test_qmoe_cuda.py Fixed
Comment thread onnxruntime/test/python/transformers/test_qmoe_cuda.py Dismissed
Comment thread onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py Fixed
Comment thread onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py Fixed
Comment thread onnxruntime/test/python/transformers/test_qmoe_wfp4afp8_cuda.py Fixed
Comment thread cmake/CMakeLists.txt
Comment thread docs/contrib_ops/cuda/moe_qmoe.md
Comment thread docs/contrib_ops/cuda/moe_qmoe.md
Comment thread onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc
Comment thread onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu
Comment thread onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc
Comment thread onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc
Comment thread onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc
@tianleiwu
Tianlei Wu (tianleiwu) enabled auto-merge (squash) May 20, 2026 05:56
@tianleiwu
Tianlei Wu (tianleiwu) merged commit 548ab6e into main May 20, 2026
91 of 95 checks passed
@tianleiwu
Tianlei Wu (tianleiwu) deleted the tlwu/20260511/qmoe_cuda branch 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.
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
@SuperSandro2000

Copy link
Copy Markdown

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

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.

6 participants