kenerl: add kernel for moe permutation with mask map (#433)

Author: guocuimi

src/kernels/moe/CMakeLists.txt

@@ -8,6 +8,7 @@ cc_library(
     topk_softmax_kernel.cu
     grouped_topk_sigmoid_kernel.cu
     permutation_index_kernel.cu
+    permutation_mask_kernel.cu
   DEPS
     cutlass
     glog::glog

src/kernels/moe/permutation_index_kernel.cu

@@ -99,20 +99,20 @@ __global__ void permute_kernel(
     const int topk,
     const int dim) {
   // one block corresponds to one token
-  const int token_idx = blockIdx.x;
+  const int t_idx = blockIdx.x;
   const int tid = threadIdx.x;
 
   // frag for load/store
   float4 frag_ls;
 
   static constexpr int kFragSize = 16 / sizeof(T);
   // tokens: [n_tokens, dim]
-  const T* token_base = tokens + token_idx * dim;
+  const T* token_base = tokens + t_idx * dim;
   for (int i = tid * kFragSize; i < dim; i += blockDim.x * kFragSize) {
     // load fragment into frag_ls (float4)
     frag_ls = __ldlu(reinterpret_cast<const float4*>(token_base + i));
 
-    int idx = token_idx;
+    int idx = t_idx;
     for (int k_idx = 0; k_idx < topk; ++k_idx) {
       // row_id_map: [topk, n_tokens] => idx in permuted tokens
       const int p_idx = row_id_map[idx];
@@ -145,7 +145,7 @@ __global__ void unpermute_kernel(
 
   // load prob into shared memory for the token
   // let first topk thread to load probs
-  for (int i = tid; i < topk; i += blockDim.x * blockDim.y) {
+  for (int i = tid; i < topk; i += blockDim.x) {
     s_probs[i] = probs[(t_idx * topk) + i];
   }
   __syncthreads();
@@ -281,10 +281,8 @@ std::tuple<torch::Tensor, torch::Tensor> permute_with_index_map(
 
   const auto type = tokens.scalar_type();
 
-  auto permuted_tokens = torch::empty({n_permuted_tokens, dim},
-                                      torch::dtype(type).device(torch::kCUDA));
-  auto row_id_map = torch::empty(
-      {n_tokens * topk}, torch::dtype(torch::kInt32).device(torch::kCUDA));
+  auto permuted_tokens = torch::empty({n_permuted_tokens, dim}, options);
+  auto row_id_map = torch::empty({n_tokens * topk}, int32_options);
 
   auto* stream = at::cuda::getCurrentCUDAStream().stream();
 
@@ -321,17 +319,17 @@ std::tuple<torch::Tensor, torch::Tensor> permute_with_index_map(
 torch::Tensor unpermute_with_index_map(
     torch::Tensor permuted_tokens,  // [n_permuted_tokens, dim]
     torch::Tensor row_id_map,       // [topk, n_tokens] => dst row
-    torch::Tensor probs,            // [n_tokens, topk]
-    int64_t n_tokens,
-    int64_t topk) {
+    torch::Tensor probs             // [n_tokens, topk]
+) {
   const auto dim = permuted_tokens.size(1);
+  const auto n_tokens = probs.size(0);
+  const auto topk = probs.size(1);
   const auto type = permuted_tokens.scalar_type();
 
-  // [n_tokens, dim]
-  auto tokens = torch::empty(
-      {n_tokens, dim},
-      torch::dtype(type).device(torch::kCUDA).requires_grad(false));
+  const auto options = permuted_tokens.options();
 
+  // [n_tokens, dim]
+  auto tokens = torch::empty({n_tokens, dim}, options);
   auto* stream = at::cuda::getCurrentCUDAStream().stream();
 
 #define LAUNCH_UNPERMUTE_KERNEL(DType)                                   \

src/kernels/moe/permutation_kernel_test.cu

@@ -12,14 +12,26 @@ namespace llm {
 namespace kernel::moe {
 // forward declare the kernel function
 std::tuple<torch::Tensor, torch::Tensor> permute_with_index_map(
-    torch::Tensor tokens,
-    torch::Tensor indices);
-
-torch::Tensor unpermute_with_index_map(torch::Tensor permuted_tokens,
-                                       torch::Tensor row_id_map,
-                                       torch::Tensor probs,
-                                       int64_t n_tokens,
-                                       int64_t topk);
+    torch::Tensor tokens,  // [n_tokens, dim]
+    torch::Tensor indices  // [n_tokens, topk]
+);
+
+torch::Tensor unpermute_with_index_map(
+    torch::Tensor permuted_tokens,  // [n_permuted_tokens, dim]
+    torch::Tensor row_id_map,       // [topk, n_tokens]
+    torch::Tensor probs             // [n_tokens, topk]
+);
+
+std::tuple<torch::Tensor, torch::Tensor> permute_with_mask_map(
+    torch::Tensor tokens,       // [n_tokens, dim]
+    torch::Tensor routing_map,  // [n_tokens, n_experts]
+    int64_t topk);
+
+torch::Tensor unpermute_with_mask_map(
+    torch::Tensor permuted_tokens,  // [n_permuted_tokens, dim]
+    torch::Tensor row_id_map,       // [n_experts, n_tokens]
+    torch::Tensor probs             // [n_tokens, n_experts]
+);
 
 }  // namespace kernel::moe
 
@@ -66,6 +78,45 @@ torch::Tensor unpermute_index_ref(
   return tokens.sum(/*dim=*/1);
 }
 
+std::tuple<torch::Tensor, torch::Tensor> permute_mask_ref(
+    const torch::Tensor& tokens,      // [n_tokens, dim]
+    const torch::Tensor& routing_map  // [n_tokens, n_experts]
+) {
+  const auto n_tokens = routing_map.size(0);
+  const auto n_experts = routing_map.size(1);
+  const auto options = tokens.options();
+
+  // [n_experts, n_tokens]
+  auto token_indices = torch::arange(n_tokens, options.dtype(torch::kLong))
+                           .unsqueeze(/*dim=*/0)
+                           .expand({n_experts, n_tokens});
+
+  // [n_permuted_tokens] original token indices, sorted by expert idx
+  auto sorted_indices = token_indices.masked_select(/*mask=*/routing_map.t());
+  auto permuted_tokens = tokens.index_select(
+      /*dim=*/0, /*index=*/sorted_indices);
+  return {permuted_tokens, sorted_indices};
+}
+
+torch::Tensor unpermute_mask_ref(
+    const torch::Tensor& permuted_tokens,  // [n_permuted_tokens, dim]
+    const torch::Tensor& permuted_probs,   // [n_permuted_tokens]
+    const torch::Tensor& sorted_incices,   // [n_permuted_tokens]
+    int64_t n_tokens) {
+  const auto dim = permuted_tokens.size(1);
+  const auto options = permuted_tokens.options();
+  // [n_tokens, dim]
+  auto tokens = torch::zeros({n_tokens, dim}, options);
+  // [n_permuted_tokens] => [n_permuted_tokens, dim]
+  auto index = sorted_incices.unsqueeze(/*dim=*/1).expand({-1, dim});
+  // reduce sum over experts
+  tokens.scatter_add_(
+      /*dim=*/0,
+      /*index=*/index,
+      /*src=*/permuted_tokens * permuted_probs.unsqueeze(/*dim=*/1));
+  return tokens;
+}
+
 }  // namespace
 
 class PermuteTest
@@ -101,7 +152,7 @@ TEST_P(PermuteTest, Index) {
   EXPECT_TRUE(torch::allclose(permuted_tokens, ref_permuted_tokens));
 
   auto unpermute_out = kernel::moe::unpermute_with_index_map(
-      permuted_tokens, sorted_indices, probs, n_tokens, topk);
+      permuted_tokens, sorted_indices, probs);
 
   auto ref_unpermute_out = unpermute_index_ref(
       ref_permuted_tokens, ref_sorted_indices, probs, n_tokens, topk);
@@ -111,6 +162,49 @@ TEST_P(PermuteTest, Index) {
       torch::allclose(tokens, unpermute_out, /*rtol=*/1e-2, /*atol=*/1e-2));
 }
 
+TEST_P(PermuteTest, Mask) {
+  const auto [dtype, n_tokens, dim, n_experts, topk] = GetParam();
+
+  const auto options = torch::dtype(dtype).device(torch::kCUDA);
+
+  const auto tokens = torch::randn({n_tokens, dim}, options);
+  const auto gating_logit = torch::randn({n_tokens, n_experts}, options);
+
+  auto [weights, indices] = gating_logit.topk(topk, /*dim=*/-1);
+  //   auto probs = weights.softmax(/*dim=*/-1);
+
+  // construct dense routing map and probs
+  auto probs = torch::zeros_like(gating_logit)
+                   .scatter(
+                       /*dim=*/1, /*index=*/indices, /*value=*/1.0 / topk);
+  auto routing_map = torch::zeros_like(gating_logit, torch::kInt)
+                         .scatter(
+                             /*dim=*/1, /*index=*/indices, /*value=*/1)
+                         .to(torch::kBool);
+
+  auto [permuted_tokens, row_id_map] =
+      kernel::moe::permute_with_mask_map(tokens, routing_map, topk);
+
+  auto [ref_permuted_tokens, ref_row_id_map] =
+      permute_mask_ref(tokens, routing_map);
+
+  EXPECT_TRUE(torch::allclose(permuted_tokens, ref_permuted_tokens));
+
+  auto unpermute_out =
+      kernel::moe::unpermute_with_mask_map(permuted_tokens, row_id_map, probs);
+
+  auto ref_permuted_probs = probs.t().masked_select(/*mask=*/routing_map.t());
+  auto ref_unpermute_out = unpermute_mask_ref(
+      ref_permuted_tokens, ref_permuted_probs, ref_row_id_map, n_tokens);
+  EXPECT_TRUE(torch::allclose(
+      unpermute_out, ref_unpermute_out, /*rtol=*/1e-2, /*atol=*/1e-2));
+
+  EXPECT_TRUE(torch::allclose(tokens,
+                              unpermute_out,
+                              /*rtol=*/1e-2,
+                              /*atol=*/1e-2));
+}
+
 INSTANTIATE_TEST_SUITE_P(
     Moe,
     PermuteTest,
@@ -123,4 +217,4 @@ INSTANTIATE_TEST_SUITE_P(
                        ::testing::Values(1, 2, 4)            // topk
                        ));
 
-}  // namespace llm
+}  // namespace llm