onnxruntime
ee5f6e7c - Add CUDA correctness fallback for INT2 and mixed-width QMoE (#32743)

Commit
11 days ago
Add CUDA correctness fallback for INT2 and mixed-width QMoE (#32743) ### Description Add a bounded CUDA correctness fallback for blockwise INT2 and mixed-width integer QMoE weights. The fallback independently dequantizes FC1 and FC2 expert weights to FP16/BF16 scratch and runs the existing dense A16 CUTLASS MoE path. It supports FC1/FC2 combinations `(2,4)`, `(4,2)`, and `(2,2)`, including canonical uniform INT2 via `expert_weight_bits=2` without FC-specific overrides. ### Fallback input contract To keep this correctness path bounded and match the CUDA blockwise dequantizer contract, it requires: - Raw weights (`weights_prepacked=0`) in canonical `[E,N,K/pack]` layout; legacy transposed layouts are rejected. - Three-dimensional blockwise scales in `[E,N,K/block_size]` layout; row-wise scales are rejected. - A power-of-two `block_size` in `[16,256]`. - Both FC reduction dimensions divisible by `block_size`. - Dequantized FC1+FC2 scratch within `ep.cuda.qmoe_int_dequant_max_scratch_bytes`. Input-layout, dimension, and scratch-limit validation runs immediately after `CheckInputs`, before tactic profiling, workspace allocation, or stream synchronization. ### Why this is a correctness fallback This PR establishes CUDA behavior and numerical correctness without introducing the final performance kernel. Weight unpacking/dequantization and MoE execution both run on CUDA, but the implementation materializes complete FP16/BF16 expert weights in temporary scratch before invoking the existing dense A16 CUTLASS runner. It therefore uses more memory bandwidth and scratch space than a native packed INT2 path. Follow-up #32761 pre-packs constant initializers and directly consumes packed INT2/INT4 weights for eligible decode, avoiding full dense materialization on that path. This bounded fallback remains for dynamic inputs and configurations outside the packed path; packed prefill and larger decode workloads remain follow-up work. CPU mixed-QMoE support and cross-provider parity are separate work. ### Testing - Built `onnxruntime_providers_cuda`. - Built `onnxruntime_provider_test`. - Passed all 14 `MoETest.QMoETest_MixedWidth*` tests on NVIDIA A100 (SM80). - Added regressions for row-wise scales, legacy transposed weights, invalid block sizes, FC1/FC2 K divisibility, early scratch rejection, and canonical non-square execution. - Passed `clang-format --dry-run --Werror`. - Passed `git diff --check`.
Author
Parents
Loading