QMoE CUDA: Rename build options, refactor PrePack, add GPU kernels - #28583
Merged
Conversation
Rename cmake options to match the onnxruntime_USE_* naming convention: - onnxruntime_ENABLE_CUDA_FP4_QMOE -> onnxruntime_USE_FP4_QMOE - onnxruntime_ENABLE_CUDA_FP8_QMOE -> onnxruntime_USE_FP8_QMOE Also rename the corresponding C preprocessor defines: - ENABLE_CUDA_FP4_QMOE -> USE_FP4_QMOE - ENABLE_CUDA_FP8_QMOE -> USE_FP8_QMOE Fix CodeQL empty-except warning in test_qmoe_cuda.py by adding an explanatory comment and removing unused exception variable.
Tianlei Wu (tianleiwu)
requested review from
Akshay Sonawane (apsonawane) and
kunal-vaishnavi
May 20, 2026 17:22
Replace per-expert loop calling launch_scaled_zero_point_kernel with a single batched kernel (LaunchQMoEScaledZP4BitBatched) that uses gridDim.z for the expert dimension. This eliminates N kernel launches during PrePack and removes the now-unused fpA_intB_gemm_adaptor.h include.
Akshay Sonawane (apsonawane)
approved these changes
May 20, 2026
Tianlei Wu (tianleiwu)
enabled auto-merge (squash)
May 20, 2026 23:02
4 tasks
Tianlei Wu (tianleiwu)
added a commit
that referenced
this pull request
May 22, 2026
… fix (#28607) ## Description Follow-up to #28583. Addresses review feedback that landed after merge (input validation, redundant memset, dead branches in `PrePackComputeBias`) and fixes a pre-existing latent CUTLASS issue that surfaced as a packaging pipeline failure once MoE GEMM kernels were built with a multi-arch `CMAKE_CUDA_ARCHITECTURES` list spanning pre-Ampere and Ampere+ targets. ## Summary of Changes ### Packaging pipeline build fix | File | Change | |------|--------| | `onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/moe_cutlass_kernel.h` | Replace the unconditional `static_assert(false, ...)` in the pre-Ampere `#else` branch of `MoeFCGemm::operator()` with `CUTLASS_NOT_IMPLEMENTED()` plus a comment explaining why this is safe. | Background: `moe_gemm_kernels_*.cu` instantiate `MoeFCGemm` through `MoeGemmRunner<...>::dispatchToArch`, which contains *runtime* (not `constexpr`) `if (sm_ >= 80 && sm_ < 90)` branches. NVCC therefore instantiates the kernel for every requested device target, including pre-Sm80 device compile passes. The old `static_assert(false, ...)` fired on those passes whenever `CMAKE_CUDA_ARCHITECTURES` contained any arch below 80 (e.g. the packaging pipeline list `52-real;61-real;75-real;86-real;89-real;90-virtual`). Replacing it with `CUTLASS_NOT_IMPLEMENTED()` lets NVCC emit a runtime trap stub for pre-Sm80, while runtime dispatch in `MoeGemmRunner::dispatchToArch()` already guarantees `sm_ >= 80` before the kernel is ever launched, so the stub is unreachable in practice. ### Address PR #28583 post-merge review | File | Change | |------|--------| | `onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu` | Add `ValidateScaledZP4BitBatchedArgs` (positive `experts`/`n`/`k_blocks`, `experts ≤ 65535` for the `gridDim.z` limit) and call it from both `LaunchQMoEScaledZP4BitBatched` overloads. Matches the validation style of `LaunchQMoERepackFP4ColToRow`. | | `onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc` (`PrePackSwizzleBlockScales`) | Remove the redundant `cudaMemsetAsync` of the destination buffer. `QMoEBlockScaleInterleaveKernel`'s `(batch, row, col) -> offset` map is a bijection over the padded output extent and writes 0 for padded source positions, so every output byte is already written. Comment explains the invariant. | | `onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc` (`PrePackComputeBias`, 4-bit block-wise) | Add `ORT_ENFORCE` checks for positive shape dims and an `INT_MAX/2` bound on `packed_k_blocks` (parity with `PrePackSwizzleBlockScales` / `PrePackRepackFP4Weights`). Drop the shadowed `bool is_fp16 = is_fp16_; bool is_bf16 = !is_fp16_;` locals in favour of `is_fp16_`. Replace the dead-branch ternary `(is_fp16 \|\| is_bf16 ? 2 : 4)` with `sizeof(uint16_t)` and a clarifying comment, and remove the unreachable `else ORT_THROW(...)` (the QMoE type path is strictly FP16/BF16). | ## Testing - Built locally with CUDA 12.8 against the failing CI arch list (`-DCMAKE_CUDA_ARCHITECTURES="52-real;61-real;75-real;86-real;89-real;90-virtual"`) and confirmed `onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_kernels_bf16_bf16.cu.o` compiles cleanly (only an `sm_<75` deprecation warning, no `static_assert` failure). - Existing QMoE Python tests (`onnxruntime/test/python/transformers/test_qmoe_cuda.py`, `test_qmoe_cpu.py`) exercise the affected `PrePackSwizzleBlockScales` / `PrePackComputeBias` paths under `--config Debug` builds and continue to pass; the added `ORT_ENFORCE` checks only trigger on invalid shapes that are not produced by the supported QMoE input contract. - No behaviour change on supported devices: `dispatchToArch` already gates `MoeFCGemm` behind `sm_ >= 80`, so the new `CUTLASS_NOT_IMPLEMENTED()` stub is unreachable at runtime. ## Motivation and Context Once #28583 enabled the MoE GEMM kernels as part of the contrib CUDA build, packaging pipelines (which target a wide arch range to maximise GPU coverage) started failing on the pre-Ampere device compile passes. The kernel-side fix in this PR resolves the immediate breakage while keeping the cmake-level binary-size optimisation (per-kernel arch pinning, TensorRT-LLM style) as a follow-up — CMake's `CUDA_ARCHITECTURES` is target/directory-scoped only, so the proper way to restrict per-kernel archs is an OBJECT-library refactor, which is intentionally not in scope here. ## Checklist - [x] Tests added/updated (input validation covered by existing QMoE tests; the new `ORT_ENFORCE` checks fail loudly on out-of-contract shapes) - [x] No documentation changes needed - [x] No breaking changes - [x] Local packaging-pipeline arch list verified to compile
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
Follow-up refinements to QMoE CUDA EP (#28467): rename build options for consistency, fix CodeQL warnings, refactor
PrePackfrom 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
PrePackfunction containing 5 nested lambdas that performed weight/scale prepacking at model load time. Reviewer feedback requested:onnxruntime_ENABLE_CUDA_*cmake options toonnxruntime_USE_*for naming consistency.Key Changes
594642aENABLE_CUDA_FP4_QMOE→USE_FP4_QMOE,ENABLE_CUDA_FP8_QMOE→USE_FP8_QMOEin cmake, C++ defines, and 340+ generated .cu files594642atest_qmoe_cuda.pye8d364bPrePackTransposeAndPack,PrePackCopyToGpu,PrePackSwizzleBlockScales,PrePackRepackFP4Weights,PrePackComputeBiase8d364bQMoERepackFP4ColToRowKernel— repacks column-major FP4 packed weights to row-major on GPU (replaces per-expert CPU loop)e8d364bPrePackSwizzleBlockScalesnow calls existingLaunchQMoEBlockScaleInterleaveGPU kernel (replaces CPUSwizzleMXFPXBlockScalesToGpuloop)Impact
PrePackruns once duringInferenceSession::Initialize(), not on the inference hot path.Testing
onnxruntime_USE_FP4_QMOE=ON onnxruntime_USE_FP8_QMOE=ON(CUDA 13.0, SM90).libonnxruntime_providers_cuda.sovianm.test_qmoe_fp4_cuda.py,test_qmoe_wfp4afp8_cuda.py,test_qmoe_cuda.pycover the affected code paths.