fix moe stage2 when prebuild aot (#1018)

Author: lalala-sh

csrc/ck_gemm_moe_2stages_codegen/gemm_moe_ck2stages.cu

@@ -38,7 +38,7 @@ MoeKernel moe_dispatch(std::string &kernelName, int block_m, int inter_dim, at::
     }
     else
     {
-        return moe_stage2_heuristic_dispatch(block_m, inter_dim, x_dtype, w_dtype, y_dtype, act_op, quant_type, mul_routed_weight);
+        return moe_stage2_heuristic_dispatch(block_m, inter_dim, x_dtype, w_dtype, y_dtype, 0, quant_type, mul_routed_weight);
     }
 }
 

csrc/ck_gemm_moe_2stages_codegen/gen_instances.py

@@ -59,6 +59,9 @@
 """
 
 heuristic_dispatch_end = """
+    TORCH_CHECK(
+        false,
+        "Unsupported kernel config for moe heuristic dispatch");
 }}
 
 """
@@ -227,7 +230,6 @@
     if (dtype_checker<{A0DataType}>{{}}(x_dtype)
         && dtype_checker<{B0DataType}>{{}}(w_dtype)
         && dtype_checker<{EDataType}>{{}}(y_dtype)
-        && {ActOP} == act_op
         && {MulRoutedWeight} == mul_routed_weight_stage
         && {Quant} == quant)
     {{
@@ -261,7 +263,6 @@
     if (dtype_checker<{A0DataType}>{{}}(x_dtype)
         && dtype_checker<{B0DataType}>{{}}(w_dtype)
         && dtype_checker<{EDataType}>{{}}(y_dtype)
-        && {ActOP} == act_op
         && {MulRoutedWeight} == mul_routed_weight_stage
         && {Quant} == quant)
     {{
@@ -295,7 +296,6 @@
     if (dtype_checker<{A0DataType}>{{}}(x_dtype)
         && dtype_checker<{B0DataType}>{{}}(w_dtype)
         && dtype_checker<{EDataType}>{{}}(y_dtype)
-        && {ActOP} == act_op
         && {MulRoutedWeight} == mul_routed_weight_stage
         && {Quant} == quant)
     {{
@@ -331,7 +331,6 @@
     if (dtype_checker<{A0DataType}>{{}}(x_dtype)
         && dtype_checker<{B0DataType}>{{}}(w_dtype)
         && dtype_checker<{EDataType}>{{}}(y_dtype)
-        && {ActOP} == act_op
         && {MulRoutedWeight} == mul_routed_weight_stage
         && {Quant} == quant)
     {{
@@ -389,7 +388,6 @@
     if (dtype_checker<{A0DataType}>{{}}(x_dtype)
         && dtype_checker<{B0DataType}>{{}}(w_dtype)
         && dtype_checker<{EDataType}>{{}}(y_dtype)
-        && {ActOP} == act_op
         && {MulRoutedWeight} == mul_routed_weight_stage
         && {Quant} == quant)
     {{