[Tianlei Wu (tianleiwu)]

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

• Tests added/updated
• No breaking changes (or documented in description)
• Documentation updated (if applicable)

[Copilot]

🟡 Changes recommended

The undefined TryParseSignedInt64 calls block compilation, and the tests lack coverage for the new odd-step single-K-warp SwiGLU variant.

2 open issues

• Replace undefined TryParseSignedInt64 calls to restore compilation · New
• Add GEMV test coverage for odd SwiGLU step counts · New

What changed in this PR

Adds a BF16 INT2 Tensor-Core GEMV path to CUDA QMoE, accelerating grouped expert projections while retaining the existing fallback paths.

Changes:

• Introduces grouped Tensor-Core GEMV kernels and optimizes BF16 INT2 conversion.
• Adds dispatch controls and expands initializer-scale prepacking.
• Adds parity and extreme-scale tests, updates profiling setup, and documents configuration options.

File	Description
onnxruntime/​test/​contrib_ops/​moe_test.cc	Adds BF16 INT2 parity and extreme-scale tests; starts profiling after initialization.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_quantization.h	Stores the Tensor-Core dispatch threshold.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_quantization.cc	Parses dispatch controls, passes thresholds, and expands scale prepacking.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemv.h	Adds threshold parameters to symmetric GEMV launch interfaces.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemv.cu	Routes eligible BF16 INT2 launches to Tensor-Core kernels.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemv_tc.cuh	Implements grouped Tensor-Core GEMV and fused SwiGLU execution.
onnxruntime/​contrib_ops/​cuda/​llm/​cutlass_extensions/​interleaved_numeric_conversion.h	Uses three masks to reduce BF16 INT2 conversion shifts.
docs/​contrib_ops/​cuda/​moe_qmoe.md	Documents the new kernel path and configuration controls.

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

The INT2 test leaves second-chunk routing and output indexing untested, while the scale-caching gate adds an unusable GPU allocation and transpose when weights are not prepacked.

2 open issues

• Avoid unused runtime scale allocation and transpose · New
• Test INT2 routing with more than eight routes per expert · New

2 resolved since last review

• Replace undefined TryParseSignedInt64 calls to restore compilation
• Add GEMV test coverage for odd SwiGLU step counts

_{Review effort: Max}