Fuse INT2 FC2 routing finalization for four-row verification

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).

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<float>(static_cast<Type>(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<int>("ORT_QMOE_INT2_LARGE_EXPERTS", 1) == 1
-          ? 20
+          ? onnxruntime::ParseEnvironmentVariableWithDefault<int>("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<half, uint8_t>(
 template void launch_moe_gemv_int_symmetric_fused_finalize<half, cutlass::uint4b_t>(
     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<half, cutlass::uint2b_t>(
+    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<half, uint8_t>(
     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);

onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc

@@ -2286,15 +2286,28 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const {
       } else {
         launch_fc1(static_cast<cutlass::uint4b_t*>(nullptr));
       }
-      if (fc2_expert_weight_bits_ == 2) {
+      const bool fuse_fc2 =
+          fc2_expert_weight_bits_ == 2 && num_rows == 4 && k_ == 10 &&
+          onnxruntime::ParseEnvironmentVariableWithDefault<int>("ORT_INT2_S4_FUSED_FINALIZE", 0) == 1 &&
+          gemv::is_moe_gemv_fused_finalize_supported(
+              sm_, num_rows, k_, expanded, hidden, inter, 2, static_cast<int>(block_size_));
+      if (fuse_fc2) {
+        gemv::launch_moe_gemv_int_symmetric_fused_finalize<T, cutlass::uint2b_t>(
+            fc1_output, static_cast<const cutlass::uint2b_t*>(packed_int_fc2),
+            static_cast<const T*>(p_fc2_scales), fc2_bias, reinterpret_cast<T*>(output->MutableDataRaw()),
+            unpermuted_row_to_permuted_row, p_exp, expert_scales, num_experts, num_rows, k_,
+            hidden, inter, static_cast<int>(block_size_), sm_, stream);
+      } else if (fc2_expert_weight_bits_ == 2) {
         launch_fc2(static_cast<cutlass::uint2b_t*>(nullptr));
       } else {
         launch_fc2(static_cast<cutlass::uint4b_t*>(nullptr));
       }
-      ck::finalizeMoeRoutingKernelLauncher<T, T, T>(
-          fc2_output, reinterpret_cast<T*>(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<T, T, T>(
+            fc2_output, reinterpret_cast<T*>(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<half*>(nullptr));

onnxruntime/test/contrib_ops/moe_test.cc

@@ -2258,14 +2258,22 @@ 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 <typename T = MLFloat16>
 static std::vector<float> 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;
   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<float>& values) {
+    std::vector<T> 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]() {
@@ -2275,7 +2283,8 @@ static std::vector<float> RunQMoEPackedIntGemvParityCase(
 
   std::vector<float> input(static_cast<size_t>(num_rows * hidden_size));
   for (float& value : input) {
-    value = static_cast<float>(next_rand() % 17) * 0.05f - 0.4f;
+    value = (static_cast<float>(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.
@@ -2290,16 +2299,16 @@ static std::vector<float> 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<uint8_t> weights(static_cast<size_t>(rows * columns / pack_size), 0);
     constexpr uint8_t base_code = static_cast<uint8_t>(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<uint8_t>(base_code + (column == identity_column ? 1 : 0));
+        const bool selected = column == identity_column || (two_terms && column == second_column);
+        const uint8_t code = static_cast<uint8_t>(base_code + (selected ? 1 : 0));
         const int64_t byte_index = row * (columns / pack_size) + column / pack_size;
         weights[static_cast<size_t>(byte_index)] |=
             static_cast<uint8_t>(code << ((column % pack_size) * bits));
@@ -2311,8 +2320,8 @@ static std::vector<float> RunQMoEPackedIntGemvParityCase(
   std::vector<uint8_t> fc1_weights(static_cast<size_t>(num_experts * fc1_rows * hidden_size / pack_size));
   std::vector<uint8_t> fc2_weights(static_cast<size_t>(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<int64_t>(fc1_e.size()) * e);
     std::copy(fc2_e.begin(), fc2_e.end(), fc2_weights.begin() + static_cast<int64_t>(fc2_e.size()) * e);
   }
@@ -2328,27 +2337,27 @@ static std::vector<float> RunQMoEPackedIntGemvParityCase(
   tester.AddAttribute<int64_t>("weights_prepacked", 0);
   tester.AddAttribute<int64_t>("block_size", block_size);
 
-  tester.AddInput<MLFloat16>("input", {num_rows, hidden_size}, ToFloat16(input));
-  tester.AddInput<MLFloat16>("router_probs", {num_rows, num_experts}, ToFloat16(router_probs));
+  tester.AddInput<T>("input", {num_rows, hidden_size}, convert(input));
+  tester.AddInput<T>("router_probs", {num_rows, num_experts}, convert(router_probs));
   tester.AddInput<uint8_t>("fc1_experts_weights", {num_experts, fc1_rows, hidden_size / pack_size},
                            fc1_weights, /*is_initializer=*/true);
-  tester.AddInput<MLFloat16>("fc1_scales", {num_experts, fc1_rows, hidden_size / block_size},
-                             ToFloat16(fc1_scales), /*is_initializer=*/true);
-  tester.AddOptionalInputEdge<MLFloat16>();
+  tester.AddInput<T>("fc1_scales", {num_experts, fc1_rows, hidden_size / block_size},
+                     convert(fc1_scales), /*is_initializer=*/true);
+  tester.AddOptionalInputEdge<T>();
   tester.AddInput<uint8_t>("fc2_experts_weights", {num_experts, hidden_size, inter_size / pack_size},
                            fc2_weights, /*is_initializer=*/true);
-  tester.AddInput<MLFloat16>("fc2_scales", {num_experts, hidden_size, inter_size / block_size},
-                             ToFloat16(fc2_scales), /*is_initializer=*/true);
-  tester.AddOptionalInputEdge<MLFloat16>();
+  tester.AddInput<T>("fc2_scales", {num_experts, hidden_size, inter_size / block_size},
+                     convert(fc2_scales), /*is_initializer=*/true);
+  tester.AddOptionalInputEdge<T>();
   tester.AddOptionalInputEdge<uint8_t>();
-  tester.AddOptionalInputEdge<MLFloat16>();
-  tester.AddOptionalInputEdge<MLFloat16>();
+  tester.AddOptionalInputEdge<T>();
+  tester.AddOptionalInputEdge<T>();
 
   // 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<float> placeholder(static_cast<size_t>(num_rows * hidden_size), 0.0f);
-  tester.AddOutput<MLFloat16>("output", {num_rows, hidden_size}, ToFloat16(placeholder));
+  tester.AddOutput<T>("output", {num_rows, hidden_size}, convert(placeholder));
   tester.SetOutputTolerance(1e6f);
 
   std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
@@ -2379,7 +2388,7 @@ static std::vector<float> RunQMoEPackedIntGemvParityCase(
 
   const auto fetches = tester.GetFetches();
   const Tensor& out = fetches[0].Get<Tensor>();
-  const MLFloat16* data = out.Data<MLFloat16>();
+  const T* data = out.Data<T>();
   std::vector<float> result(static_cast<size_t>(out.Shape().Size()));
   for (size_t i = 0; i < result.size(); ++i) {
     result[i] = data[i].ToFloat();
@@ -2414,6 +2423,49 @@ 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_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);
+    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<BFloat16>(false, 4, 16, 10);
+    const auto dense_bf16 = RunQMoEPackedIntGemvParityCase<BFloat16>(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;
+    }
+
+    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<BFloat16>(false, 4, 16, 10, true);
+      const auto rounding_dense_bf16 = RunQMoEPackedIntGemvParityCase<BFloat16>(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;
+      }
+    }
+  }
+}
+
 #ifdef USE_CUDA
 static std::string BuildQMoEInt2DynamicRowsModel() {
   constexpr int64_t num_experts = 1;