fix moe_fused_gate meet nan data bug (ROCm#1086)

* fix moe_fused_gate meet nan data bug

* update asm_topksoftmax for large M

* Update csrc/kernels/moe_fused_gate.cu

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: 3 people

csrc/kernels/moe_fused_gate.cu

@@ -51,7 +51,7 @@ using float32_t  = float;
 template <typename T>
 __device__ inline bool cmp_gt(const T& a, const T& b)
 {
-    if constexpr(std::is_same<T, float16_t>::value || std::is_same<T, bfloat16_t>::value)
+    if constexpr(std::is_same<T, bfloat16_t>::value)
     {
         // at::Half (or float16_t in our native case) causes ambiguity, so we cast to float.
         return ck_tile::type_convert<float>(a) > ck_tile::type_convert<float>(b);
@@ -186,7 +186,7 @@ __device__ void moe_fused_gate_impl(void* input,
 #pragma unroll
     for(int ii = 0; ii < params.VPT; ++ii)
     {
-        row_chunk[ii] = 1.0f / (1.0f + expf(-row_chunk[ii]));
+        row_chunk[ii] = ::isnan(row_chunk[ii]) ? 0.0f : (1.0f / (1.0f + expf(-row_chunk[ii])));
     }
     // __syncthreads();
 
@@ -548,6 +548,8 @@ std::vector<at::Tensor> moe_fused_gate(at::Tensor& input,
     auto output          = topk_weights;
     auto indices         = topk_ids;
     const int out_stride = topk_ids.stride(0);
+    TORCH_CHECK(topk_weights.stride(0) == out_stride,
+                "topk_weights and topk_ids must have the same stride in dim 0");
 
     // Compute grid dimensions based on runtime value for num_expert_group.
     int64_t rows_per_warp = std::max<int64_t>(1, WARP_SIZE / num_expert_group);

csrc/py_itfs_cu/asm_topksoftmax.cu

@@ -36,7 +36,7 @@ void topk_softmax_asm(torch::Tensor& topk_weights,         // [num_tokens, topk]
     const uint num_tokens  = gating_output.numel() / num_experts;
     const uint topk        = topk_weights.size(-1);
     const uint out_stride  = topk_weights.stride(0);
-    const uint SUBM        = 4;
+    const uint SUBM        = num_tokens < 10000 ? 4 : 12;
 
     KernelArgs args;
     size_t arg_size = sizeof(args);
@@ -53,36 +53,73 @@ void topk_softmax_asm(torch::Tensor& topk_weights,         // [num_tokens, topk]
     AiterAsmKernel* impl_ptr = nullptr;
     if(num_experts == 128 && topk == 8)
     {
-        static AiterAsmKernel impl_topksoftmax_4x128x8("_ZN5aiter19topksoftmax_4x128x8E",
-                                                       "/topksoftmax/topksoftmax_4x128x8.co");
-        impl_ptr = &impl_topksoftmax_4x128x8;
+        if(SUBM == 4)
+        {
+            static AiterAsmKernel impl_topksoftmax_4x128x8("_ZN5aiter19topksoftmax_4x128x8E",
+                                                           "/topksoftmax/topksoftmax_4x128x8.co");
+            impl_ptr = &impl_topksoftmax_4x128x8;
+        }
+        else
+        {
+            static AiterAsmKernel impl_topksoftmax_12x128x8("_ZN5aiter20topksoftmax_12x128x8E",
+                                                            "/topksoftmax/topksoftmax_12x128x8.co");
+            impl_ptr = &impl_topksoftmax_12x128x8;
+        }
     }
     else if(num_experts == 256 && topk == 8)
     {
-        static AiterAsmKernel impl_topksoftmax_4x256x8("_ZN5aiter19topksoftmax_4x256x8E",
-                                                       "/topksoftmax/topksoftmax_4x256x8.co");
-        impl_ptr = &impl_topksoftmax_4x256x8;
+        if(SUBM == 4)
+        {
+            static AiterAsmKernel impl_topksoftmax_4x256x8("_ZN5aiter19topksoftmax_4x256x8E",
+                                                           "/topksoftmax/topksoftmax_4x256x8.co");
+            impl_ptr = &impl_topksoftmax_4x256x8;
+        }
+        else
+        {
+            static AiterAsmKernel impl_topksoftmax_12x256x8("_ZN5aiter20topksoftmax_12x256x8E",
+                                                            "/topksoftmax/topksoftmax_12x256x8.co");
+            impl_ptr = &impl_topksoftmax_12x256x8;
+        }
     }
     else if(num_experts == 128 && topk == 6)
     {
-        static AiterAsmKernel impl_topksoftmax_4x128x6("_ZN5aiter19topksoftmax_4x128x6E",
-                                                       "/topksoftmax/topksoftmax_4x128x6.co");
-        impl_ptr = &impl_topksoftmax_4x128x6;
+        if(SUBM == 4)
+        {
+            static AiterAsmKernel impl_topksoftmax_4x128x6("_ZN5aiter19topksoftmax_4x128x6E",
+                                                           "/topksoftmax/topksoftmax_4x128x6.co");
+            impl_ptr = &impl_topksoftmax_4x128x6;
+        }
+        else
+        {
+            static AiterAsmKernel impl_topksoftmax_12x128x6("_ZN5aiter20topksoftmax_12x128x6E",
+                                                            "/topksoftmax/topksoftmax_12x128x6.co");
+            impl_ptr = &impl_topksoftmax_12x128x6;
+        }
     }
     else if(num_experts == 256 && topk == 6)
     {
-        static AiterAsmKernel impl_topksoftmax_4x256x6("_ZN5aiter19topksoftmax_4x256x6E",
-                                                       "/topksoftmax/topksoftmax_4x256x6.co");
-        impl_ptr = &impl_topksoftmax_4x256x6;
+        if(SUBM == 4)
+        {
+            static AiterAsmKernel impl_topksoftmax_4x256x6("_ZN5aiter19topksoftmax_4x256x6E",
+                                                           "/topksoftmax/topksoftmax_4x256x6.co");
+            impl_ptr = &impl_topksoftmax_4x256x6;
+        }
+        else
+        {
+            static AiterAsmKernel impl_topksoftmax_12x256x6("_ZN5aiter20topksoftmax_12x256x6E",
+                                                            "/topksoftmax/topksoftmax_12x256x6.co");
+            impl_ptr = &impl_topksoftmax_12x256x6;
+        }
     }
     else
     {
-        TORCH_CHECK(false,
-                    __func__,
-                    " only support num_experts/topk in [128/6, 128/8, 256/6, 256/8], but get num_experts: ",
-                    num_experts,
-                    " , topk: ",
-                    topk);
+        TORCH_CHECK(
+            false,
+            __func__,
+            " only support num_experts/topk in [128/6, 128/8, 256/6, 256/8], but get num_experts: ",
+            num_experts,
+            " , topk: ",
+            topk);
     }
 
     const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));