onnxruntime

f42c839f - [CUDA] Add Tensor-Core BF16 INT2 QMoE GEMV kernel and 3-mask BF16 converter (#33786)

[CUDA] Add Tensor-Core BF16 INT2 QMoE GEMV kernel and 3-mask BF16 converter (#33786) ## Description Adds a Tensor-Core (`mma.sync.m16n8k16`, BF16 inputs, FP32 accumulation) grouped GEMV kernel (`MoeGemvTcInt2Kernel` / `TryLaunchInt2Bf16`) for symmetric `INT2` (`cutlass::uint2b_t`, `block_size ∈ {32, 64, 128}`) `com.microsoft::QMoE` projections, and optimizes `Int2NumericArrayConverter<bfloat16_t, N>` in `interleaved_numeric_conversion.h` using a 3-mask BF16 mantissa trick. On NVIDIA H200 (`S = 8, k = 10`, `80` routes, `H = 2560, I = 640, block_size = 64`), Nsight Compute (`ncu`) profiling shows **2.02x** faster FC1+SwiGLU (`70.56 us -> 35.01 us`) and **1.95x** faster FC2 (`47.01 us -> 24.13 us`) compared to the scalar SIMT `moe_gemv` kernel. ## Summary of Changes ### Tensor-Core BF16 INT2 QMoE GEMV & 3-Mask BF16 Converter | File | Change | |------|--------| | `onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_tc.cuh` | Implements `MoeGemvTcInt2Kernel`, `TryLaunchInt2Bf16`, `DequantPairInt2Masked`, `ComputeWordInt2`, and `DiscoverExpertRouteChunk` (grouping up to 8 routes per expert into an `mma.sync.m16n8k16` warp tile with software-pipelined activation and streaming weight/scale loads, supporting `block_size ∈ {32, 64, 128}` and sorted route discovery beyond 256 routes). | | `onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.cu` | Dispatches BF16 `cutlass::uint2b_t` launches to `fiv::moe_tc::TryLaunchInt2Bf16`. The `tensor_core_min_rows` parameter (default `33`) of the two launchers already exists on main and is reused as-is. | | `onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/interleaved_numeric_conversion.h` | Optimizes `Int2NumericArrayConverter<bfloat16_t, N>` to extract three 2-bit fields (`0x00030003`, `0x000C000C`, `0x00300030`) inside BF16's 7-bit mantissa with only two right-shifts (`>> 6` and `>> 12`) per 32-bit packed word instead of 8 shifts. | | `onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc` | Adds `int_scale_supports_packed_execution` in `QMoE::PrePack` so initializer scales are transposed once when `enable_int2_gemv_` or `enable_int2_prefill_` is active. The `ep.cuda.moe_gemv_tc_min_rows` / `ep.cuda.disable_moe_gemv_tc` controls (`ORT_MOE_GEMV_TC_MIN_ROWS` / `ORT_DISABLE_MOE_GEMV_TC`) are already parsed on main and are reused unchanged. | | `docs/contrib_ops/cuda/moe_qmoe.md` | Documents the BF16 Tensor-Core `INT2` GEMV path and session/environment controls. | | `onnxruntime/test/contrib_ops/moe_test.cc` | Adds `QMoETest_Bf16Int2TensorCoreGemvParity` (across `block_size ∈ {32, 64, 128}` and 6 shapes, including the odd-step `hidden_size = 640, inter_size = 512` case that exercises the single-K-warp SwiGLU kernel) and `QMoETest_Bf16Int2TensorCoreGemvExtremeScales`, and starts CUPTI profiling after `session.Initialize()` in `QMoETest_Int2CudaCachedScalesDecodeThenPrefill`. | ## Testing - Built `onnxruntime_provider_test` and ran `./build/direct_cuda/Release/onnxruntime_provider_test --gtest_filter="MoETest.*"` on NVIDIA H200 (`123/123` passed, including `QMoETest_Bf16Int2TensorCoreGemvParity`, `QMoETest_Bf16Int2TensorCoreGemvExtremeScales`, and `QMoETest_Int2CudaCachedScalesDecodeThenPrefill`). - Re-verified at `9ae9283` with a CUDA-EP-as-plugin build (CUDA 13.0, NVIDIA H200, `sm_90`): `MoETest.*` gives `96` passed, `17` skipped (CPU-only tests and routing-statistics tests that the CUDA plugin EP does not support), `0` failed. Nsight Systems shows the single-K-warp SwiGLU variant `MoeGemvTcInt2Kernel<true, 1, 1, G>` launched once per block size for the odd-step `{hidden_size = 640, inter_size = 512}` shape, and `compute-sanitizer --tool memcheck` reports `0` errors on `QMoETest_Bf16Int2TensorCoreGemvParity`. - Profiled `MoeGemvTcInt2Kernel` vs `moe_gemv_interleaved_swiglu_kernel` / `moe_gemv_kernel` with Nsight Compute (`ncu`) on NVIDIA H200 (`1.99x` total layer speedup at `T = 8, k = 10`). - Verified `clang-format --dry-run --Werror` and `git diff --check`. ## Checklist - [x] Tests added/updated - [x] No breaking changes (or documented in description) - [x] Documentation updated (if applicable) --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>