Description

Add a new QMoE contrib operator for the CUDA EP that 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.

Summary of Changes

New QMoE Operator

MoE GEMM Refactor

CUTLASS Extensions

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

Build System

Python Quantization Bindings

Tests

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.