From 8fb1d1b124df1b3c5dc69cf40610f3164557ea90 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 22 Oct 2026 20:59:32 +0000 Subject: [PATCH 1/2] Fuse INT2 FC2 routing finalization for four-row verification Make the INT2 GEMV expanded-row limit configurable and add opt-in four-row/top-10 FC2 finalization fusion. Preserve INT2 intermediate rounding and add FP16/BF16 dense-reference parity coverage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/contrib_ops/cuda/moe_qmoe.md | 16 ++++++ .../contrib_ops/cuda/llm/moe_gemm/moe_gemv.cu | 16 +++++- .../contrib_ops/cuda/moe/moe_quantization.cc | 23 +++++++-- onnxruntime/test/contrib_ops/moe_test.cc | 51 ++++++++++++++----- 4 files changed, 87 insertions(+), 19 deletions(-) diff --git a/docs/contrib_ops/cuda/moe_qmoe.md b/docs/contrib_ops/cuda/moe_qmoe.md index 2283e9fccf57b..6c99549e147fe 100644 --- a/docs/contrib_ops/cuda/moe_qmoe.md +++ b/docs/contrib_ops/cuda/moe_qmoe.md @@ -4,6 +4,22 @@ This document describes the design, schema, kernel dispatch, weight formats, and implementation status of the **MoE** (`com.microsoft::MoE`) and **QMoE** (`com.microsoft::QMoE`) operators on the CUDA execution provider. +### Four-row INT2 speculative verification + +Set `ORT_QMOE_INT2_DECODE_MAX_EXPANDED_ROWS=40` to admit four input rows with +top-10 routing to the existing packed INT2 GEMV path. Its default expanded-row +limit remains 20; other eligibility checks still apply. + +For eligible four-row/top-10 INT2 FC2 GEMV, `ORT_INT2_S4_FUSED_FINALIZE=1` +fuses the FC2 computation and routing finalization, avoiding the intermediate +FC2 output traffic and separate routing kernel launch. This is opt-in and +supports FP16 and BF16 activations. INT2 results are rounded to the activation +dtype before expert weighting, matching the unfused intermediate tensor. +Unsupported configurations retain the unfused path. + +`QMoETest_Int2CudaPackedDecodeFourRowsTop10Parity` compares fused and unfused +outputs with a dense reference for both activation types. + The CUTLASS kernels are derived from TensorRT-LLM (CUTLASS 4.4.2, commit `346018db87`) and have been significantly modified for ONNX Runtime — see [§16 Differences vs. TensorRT-LLM](#16-differences-vs-tensorrt-llm). diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.cu index 168e7980c5560..2cca5ab8c3fcc 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.cu +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.cu @@ -187,7 +187,11 @@ __device__ __forceinline__ void fused_finalize_accumulate(float* out_acc, void* for (int jj = 0; jj < WarpNum; ++jj) { val += shmem[jj * CtaN * Interleave + ii]; } - out_acc[ii] += row_scale * (val + v_bias); + if constexpr (Details::kElemsPerByteW == 4) { + out_acc[ii] += row_scale * static_cast(static_cast(val + v_bias)); + } else { + out_acc[ii] += row_scale * (val + v_bias); + } } } #endif // (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 750)) || !defined(__CUDA_ARCH__) @@ -947,8 +951,10 @@ bool is_moe_gemv_shape_supported(int sm, int64_t expanded_num_rows, int64_t n, i const int64_t max_expanded_rows = weight_bits == 2 && onnxruntime::ParseEnvironmentVariableWithDefault("ORT_QMOE_INT2_LARGE_EXPERTS", 1) == 1 - ? 20 + ? onnxruntime::ParseEnvironmentVariableWithDefault("ORT_QMOE_INT2_DECODE_MAX_EXPANDED_ROWS", 20) : kMaxProfiledExpandedRows; + ORT_ENFORCE(max_expanded_rows >= kMaxProfiledExpandedRows && max_expanded_rows <= 40, + "INT2 decode expanded-row limit must not exceed 40."); if (expanded_num_rows <= 0 || expanded_num_rows > max_expanded_rows) { return false; } @@ -1208,6 +1214,9 @@ template void launch_moe_gemv_int_symmetric( template void launch_moe_gemv_int_symmetric_fused_finalize( const half*, const cutlass::uint4b_t*, const half*, const half*, half*, const int*, const int*, const float*, int, int64_t, int64_t, int64_t, int64_t, int, int, cudaStream_t); +template void launch_moe_gemv_int_symmetric_fused_finalize( + const half*, const cutlass::uint2b_t*, const half*, const half*, half*, const int*, const int*, const float*, + int, int64_t, int64_t, int64_t, int64_t, int, int, cudaStream_t); template void launch_moe_gemv_int_symmetric_fused_finalize( const half*, const uint8_t*, const half*, const half*, half*, const int*, const int*, const float*, int, int64_t, int64_t, int64_t, int64_t, int, int, cudaStream_t); @@ -1241,6 +1250,9 @@ template void launch_moe_gemv_int_symmetric<__nv_bfloat16, uint8_t>( template void launch_moe_gemv_int_symmetric_fused_finalize<__nv_bfloat16, cutlass::uint4b_t>( const __nv_bfloat16*, const cutlass::uint4b_t*, const __nv_bfloat16*, const __nv_bfloat16*, __nv_bfloat16*, const int*, const int*, const float*, int, int64_t, int64_t, int64_t, int64_t, int, int, cudaStream_t); +template void launch_moe_gemv_int_symmetric_fused_finalize<__nv_bfloat16, cutlass::uint2b_t>( + const __nv_bfloat16*, const cutlass::uint2b_t*, const __nv_bfloat16*, const __nv_bfloat16*, __nv_bfloat16*, + const int*, const int*, const float*, int, int64_t, int64_t, int64_t, int64_t, int, int, cudaStream_t); template void launch_moe_gemv_int_symmetric_fused_finalize<__nv_bfloat16, uint8_t>( const __nv_bfloat16*, const uint8_t*, const __nv_bfloat16*, const __nv_bfloat16*, __nv_bfloat16*, const int*, const int*, const float*, int, int64_t, int64_t, int64_t, int64_t, int, int, cudaStream_t); diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc index db198230f53e5..19e4d0fd6f2f9 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc +++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc @@ -2286,15 +2286,28 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { } else { launch_fc1(static_cast(nullptr)); } - if (fc2_expert_weight_bits_ == 2) { + const bool fuse_fc2 = + fc2_expert_weight_bits_ == 2 && num_rows == 4 && k_ == 10 && + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_INT2_S4_FUSED_FINALIZE", 0) == 1 && + gemv::is_moe_gemv_fused_finalize_supported( + sm_, num_rows, k_, expanded, hidden, inter, 2, static_cast(block_size_)); + if (fuse_fc2) { + gemv::launch_moe_gemv_int_symmetric_fused_finalize( + fc1_output, static_cast(packed_int_fc2), + static_cast(p_fc2_scales), fc2_bias, reinterpret_cast(output->MutableDataRaw()), + unpermuted_row_to_permuted_row, p_exp, expert_scales, num_experts, num_rows, k_, + hidden, inter, static_cast(block_size_), sm_, stream); + } else if (fc2_expert_weight_bits_ == 2) { launch_fc2(static_cast(nullptr)); } else { launch_fc2(static_cast(nullptr)); } - ck::finalizeMoeRoutingKernelLauncher( - fc2_output, reinterpret_cast(output->MutableDataRaw()), - nullptr, expert_scales, unpermuted_row_to_permuted_row, p_r2u, expert_indices, - p_efto, num_rows, hidden, k_, num_experts, parallelism_config, false, stream); + if (!fuse_fc2) { + ck::finalizeMoeRoutingKernelLauncher( + fc2_output, reinterpret_cast(output->MutableDataRaw()), + nullptr, expert_scales, unpermuted_row_to_permuted_row, p_r2u, expert_indices, + p_efto, num_rows, hidden, k_, num_experts, parallelism_config, false, stream); + } }; if (is_fp16_) { run_packed_int_gemv(static_cast(nullptr)); diff --git a/onnxruntime/test/contrib_ops/moe_test.cc b/onnxruntime/test/contrib_ops/moe_test.cc index 72c40e7042b25..12cf1b7fa5aa9 100644 --- a/onnxruntime/test/contrib_ops/moe_test.cc +++ b/onnxruntime/test/contrib_ops/moe_test.cc @@ -2258,6 +2258,7 @@ TEST(MoETest, QMoETest_Int2CudaPrepackedRejectsRawFallback) { // path selected and once with it disabled (forcing the already-validated dense fallback from // #32743), so a stride/permutation/top-k bug in the new packed pipeline shows up as a numeric // mismatch instead of only a shape or gate mismatch. +template static std::vector RunQMoEPackedIntGemvParityCase( bool disable_packed_gemv, int64_t num_rows = 4, int64_t num_experts = 4, int64_t k = 2) { constexpr int64_t hidden_size = 512; @@ -2266,6 +2267,12 @@ static std::vector RunQMoEPackedIntGemvParityCase( constexpr int64_t bits = 2; constexpr int64_t pack_size = 8 / bits; constexpr int64_t fc1_rows = 2 * inter_size; // Fused SwiGLU doubles FC1's output rows. + auto convert = [](const std::vector& values) { + std::vector result; + result.reserve(values.size()); + for (float value : values) result.emplace_back(value); + return result; + }; uint32_t rng_state = 12345u; auto next_rand = [&rng_state]() { @@ -2328,27 +2335,27 @@ static std::vector RunQMoEPackedIntGemvParityCase( tester.AddAttribute("weights_prepacked", 0); tester.AddAttribute("block_size", block_size); - tester.AddInput("input", {num_rows, hidden_size}, ToFloat16(input)); - tester.AddInput("router_probs", {num_rows, num_experts}, ToFloat16(router_probs)); + tester.AddInput("input", {num_rows, hidden_size}, convert(input)); + tester.AddInput("router_probs", {num_rows, num_experts}, convert(router_probs)); tester.AddInput("fc1_experts_weights", {num_experts, fc1_rows, hidden_size / pack_size}, fc1_weights, /*is_initializer=*/true); - tester.AddInput("fc1_scales", {num_experts, fc1_rows, hidden_size / block_size}, - ToFloat16(fc1_scales), /*is_initializer=*/true); - tester.AddOptionalInputEdge(); + tester.AddInput("fc1_scales", {num_experts, fc1_rows, hidden_size / block_size}, + convert(fc1_scales), /*is_initializer=*/true); + tester.AddOptionalInputEdge(); tester.AddInput("fc2_experts_weights", {num_experts, hidden_size, inter_size / pack_size}, fc2_weights, /*is_initializer=*/true); - tester.AddInput("fc2_scales", {num_experts, hidden_size, inter_size / block_size}, - ToFloat16(fc2_scales), /*is_initializer=*/true); - tester.AddOptionalInputEdge(); + tester.AddInput("fc2_scales", {num_experts, hidden_size, inter_size / block_size}, + convert(fc2_scales), /*is_initializer=*/true); + tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); - tester.AddOptionalInputEdge(); - tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); // The real assertion is the packed-vs-dense parity check the caller performs on GetFetches(); // use a placeholder expected output and a huge tolerance so the framework's own comparison // never fails here. const std::vector placeholder(static_cast(num_rows * hidden_size), 0.0f); - tester.AddOutput("output", {num_rows, hidden_size}, ToFloat16(placeholder)); + tester.AddOutput("output", {num_rows, hidden_size}, convert(placeholder)); tester.SetOutputTolerance(1e6f); std::vector> execution_providers; @@ -2379,7 +2386,7 @@ static std::vector RunQMoEPackedIntGemvParityCase( const auto fetches = tester.GetFetches(); const Tensor& out = fetches[0].Get(); - const MLFloat16* data = out.Data(); + const T* data = out.Data(); std::vector result(static_cast(out.Shape().Size())); for (size_t i = 0; i < result.size(); ++i) { result[i] = data[i].ToFloat(); @@ -2414,6 +2421,26 @@ TEST(MoETest, QMoETest_Int2Cuda512ExpertsTop10Parity) { } } +TEST(MoETest, QMoETest_Int2CudaPackedDecodeFourRowsTop10Parity) { + if (!HasCudaEnvironment(800)) GTEST_SKIP(); + for (bool fused : {false, true}) { + ScopedEnvironmentVariables settings{{{"ORT_QMOE_INT2_DECODE_MAX_EXPANDED_ROWS", "40"}, + {"ORT_INT2_S4_FUSED_FINALIZE", fused ? "1" : "0"}}}; + const auto packed = RunQMoEPackedIntGemvParityCase(false, 4, 16, 10); + const auto dense = RunQMoEPackedIntGemvParityCase(true, 4, 16, 10); + ASSERT_EQ(packed.size(), dense.size()); + for (size_t i = 0; i < packed.size(); ++i) { + EXPECT_NEAR(packed[i], dense[i], 0.02f) << "mismatch at index " << i; + } + const auto packed_bf16 = RunQMoEPackedIntGemvParityCase(false, 4, 16, 10); + const auto dense_bf16 = RunQMoEPackedIntGemvParityCase(true, 4, 16, 10); + ASSERT_EQ(packed_bf16.size(), dense_bf16.size()); + for (size_t i = 0; i < packed_bf16.size(); ++i) { + EXPECT_NEAR(packed_bf16[i], dense_bf16[i], 0.02f) << "BF16 mismatch at index " << i; + } + } +} + #ifdef USE_CUDA static std::string BuildQMoEInt2DynamicRowsModel() { constexpr int64_t num_experts = 1; From b417ff420602ebf6afdf3ae781c57a61cc9c8324 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Fri, 23 Oct 2026 21:10:41 +0000 Subject: [PATCH 2/2] Strengthen INT2 fused finalize parity coverage Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/test/contrib_ops/moe_test.cc | 43 +++++++++++++++++++----- 1 file changed, 34 insertions(+), 9 deletions(-) diff --git a/onnxruntime/test/contrib_ops/moe_test.cc b/onnxruntime/test/contrib_ops/moe_test.cc index 12cf1b7fa5aa9..d4287a99dad1e 100644 --- a/onnxruntime/test/contrib_ops/moe_test.cc +++ b/onnxruntime/test/contrib_ops/moe_test.cc @@ -2260,7 +2260,8 @@ TEST(MoETest, QMoETest_Int2CudaPrepackedRejectsRawFallback) { // mismatch instead of only a shape or gate mismatch. template static std::vector RunQMoEPackedIntGemvParityCase( - bool disable_packed_gemv, int64_t num_rows = 4, int64_t num_experts = 4, int64_t k = 2) { + bool disable_packed_gemv, int64_t num_rows = 4, int64_t num_experts = 4, int64_t k = 2, + bool rounding_sensitive_fc2 = false) { constexpr int64_t hidden_size = 512; constexpr int64_t inter_size = 512; constexpr int64_t block_size = 64; @@ -2282,7 +2283,8 @@ static std::vector RunQMoEPackedIntGemvParityCase( std::vector input(static_cast(num_rows * hidden_size)); for (float& value : input) { - value = static_cast(next_rand() % 17) * 0.05f - 0.4f; + value = (static_cast(next_rand() % 17) * 0.05f - 0.4f) * + (rounding_sensitive_fc2 ? 4.0f : 1.0f); } // Favor high-numbered experts too, exercising the full expert-map range. @@ -2297,16 +2299,16 @@ static std::vector RunQMoEPackedIntGemvParityCase( } } - // Each expert's weight is a "shifted identity": row r picks out input column - // (r + expert_shift) % columns, so a bug that mixes up per-expert strides or permutation - // changes the numeric result rather than only the output shape. - auto make_expert_weights = [](int64_t rows, int64_t columns, int64_t expert_shift) { + // Each expert uses a shifted identity; the rounding-sensitive case adds a second FC2 term. + auto make_expert_weights = [](int64_t rows, int64_t columns, int64_t expert_shift, bool two_terms) { std::vector weights(static_cast(rows * columns / pack_size), 0); constexpr uint8_t base_code = static_cast(1u << (bits - 1)); for (int64_t row = 0; row < rows; ++row) { const int64_t identity_column = (row + expert_shift) % columns; + const int64_t second_column = (identity_column + 17) % columns; for (int64_t column = 0; column < columns; ++column) { - const uint8_t code = static_cast(base_code + (column == identity_column ? 1 : 0)); + const bool selected = column == identity_column || (two_terms && column == second_column); + const uint8_t code = static_cast(base_code + (selected ? 1 : 0)); const int64_t byte_index = row * (columns / pack_size) + column / pack_size; weights[static_cast(byte_index)] |= static_cast(code << ((column % pack_size) * bits)); @@ -2318,8 +2320,8 @@ static std::vector RunQMoEPackedIntGemvParityCase( std::vector fc1_weights(static_cast(num_experts * fc1_rows * hidden_size / pack_size)); std::vector fc2_weights(static_cast(num_experts * hidden_size * inter_size / pack_size)); for (int64_t e = 0; e < num_experts; ++e) { - const auto fc1_e = make_expert_weights(fc1_rows, hidden_size, e * 7); - const auto fc2_e = make_expert_weights(hidden_size, inter_size, e * 11); + const auto fc1_e = make_expert_weights(fc1_rows, hidden_size, e * 7, false); + const auto fc2_e = make_expert_weights(hidden_size, inter_size, e * 11, rounding_sensitive_fc2); std::copy(fc1_e.begin(), fc1_e.end(), fc1_weights.begin() + static_cast(fc1_e.size()) * e); std::copy(fc2_e.begin(), fc2_e.end(), fc2_weights.begin() + static_cast(fc2_e.size()) * e); } @@ -2425,6 +2427,9 @@ TEST(MoETest, QMoETest_Int2CudaPackedDecodeFourRowsTop10Parity) { if (!HasCudaEnvironment(800)) GTEST_SKIP(); for (bool fused : {false, true}) { ScopedEnvironmentVariables settings{{{"ORT_QMOE_INT2_DECODE_MAX_EXPANDED_ROWS", "40"}, + {"ORT_QMOE_INT2_LARGE_EXPERTS", "1"}, + {"ORT_ENABLE_QMOE_INT2_PREFILL", "0"}, + {"ORT_DISABLE_MOE_GEMV", "0"}, {"ORT_INT2_S4_FUSED_FINALIZE", fused ? "1" : "0"}}}; const auto packed = RunQMoEPackedIntGemvParityCase(false, 4, 16, 10); const auto dense = RunQMoEPackedIntGemvParityCase(true, 4, 16, 10); @@ -2438,6 +2443,26 @@ TEST(MoETest, QMoETest_Int2CudaPackedDecodeFourRowsTop10Parity) { for (size_t i = 0; i < packed_bf16.size(); ++i) { EXPECT_NEAR(packed_bf16[i], dense_bf16[i], 0.02f) << "BF16 mismatch at index " << i; } + + if (fused) { + // Two FC2 terms and larger activations make the per-expert FP32 dot product land between + // activation-dtype values, exercising the cast before the fused routing-weight accumulation. + const auto rounding_packed = RunQMoEPackedIntGemvParityCase(false, 4, 16, 10, true); + const auto rounding_dense = RunQMoEPackedIntGemvParityCase(true, 4, 16, 10, true); + ASSERT_EQ(rounding_packed.size(), rounding_dense.size()); + for (size_t i = 0; i < rounding_packed.size(); ++i) { + EXPECT_NEAR(rounding_packed[i], rounding_dense[i], 0.0002f) + << "rounding-sensitive FP16 mismatch at index " << i; + } + + const auto rounding_packed_bf16 = RunQMoEPackedIntGemvParityCase(false, 4, 16, 10, true); + const auto rounding_dense_bf16 = RunQMoEPackedIntGemvParityCase(true, 4, 16, 10, true); + ASSERT_EQ(rounding_packed_bf16.size(), rounding_dense_bf16.size()); + for (size_t i = 0; i < rounding_packed_bf16.size(); ++i) { + EXPECT_NEAR(rounding_packed_bf16[i], rounding_dense_bf16[i], 0.002f) + << "rounding-sensitive BF16 mismatch at index " << i; + } + } } }