rls2.6: Fuse moegate ops for deepseekv3 (#3488)

Author: blzheng

csrc/cpu/aten/MoE.cpp

@@ -298,7 +298,8 @@ std::tuple<at::Tensor, at::Tensor> deepseek_moegate(
     const int64_t n_group,
     const int64_t topk_group,
     const int64_t n_routed_experts,
-    const int64_t top_k) {
+    const int64_t top_k,
+    c10::optional<at::Tensor> e_score_cbias) {
   RECORD_FUNCTION("ipex::deepseek_moegate", c10::ArrayRef<c10::IValue>({}));
 
   return deepseek_moegate_kernel_stub(
@@ -309,7 +310,8 @@ std::tuple<at::Tensor, at::Tensor> deepseek_moegate(
       n_group,
       topk_group,
       n_routed_experts,
-      top_k);
+      top_k,
+      e_score_cbias);
 }
 } // namespace cpu
 } // namespace torch_ipex
@@ -374,7 +376,7 @@ TORCH_LIBRARY_FRAGMENT(torch_ipex, m) {
       c10::DispatchKey::CPU,
       torch_ipex::cpu::deepseek_moe_woq);
   m.def(
-      "deepseek_moegate(Tensor hidden_states, Tensor scores, Tensor routed_scaling_factor, int n_group, int topk_group, int n_routed_experts, int top_k) -> (Tensor, Tensor)");
+      "deepseek_moegate(Tensor hidden_states, Tensor scores, Tensor routed_scaling_factor, int n_group, int topk_group, int n_routed_experts, int top_k, Tensor? e_score_cbias=None) -> (Tensor, Tensor)");
   m.impl(
       "deepseek_moegate",
       c10::DispatchKey::CPU,

csrc/cpu/aten/MoE.h

@@ -97,7 +97,8 @@ std::tuple<at::Tensor, at::Tensor> deepseek_moegate(
     const int64_t n_group,
     const int64_t topk_group,
     const int64_t n_routed_experts,
-    const int64_t top_k);
+    const int64_t top_k,
+    c10::optional<at::Tensor> e_score_cbias);
 using mixtral_moe_tpp_kernel_fn = at::Tensor (*)(
     const at::Tensor& hidden_states,
     const at::Tensor& top_x,
@@ -179,7 +180,8 @@ using deepseek_moegate_kernel_fn = std::tuple<at::Tensor, at::Tensor> (*)(
     const int64_t n_group,
     const int64_t topk_group,
     const int64_t n_routed_experts,
-    const int64_t top_k);
+    const int64_t top_k,
+    c10::optional<at::Tensor> e_score_cbias);
 IPEX_DECLARE_DISPATCH(mixtral_moe_tpp_kernel_fn, mixtral_moe_tpp_kernel_stub);
 IPEX_DECLARE_DISPATCH(deepseek_moe_tpp_kernel_fn, deepseek_moe_tpp_kernel_stub);
 IPEX_DECLARE_DISPATCH(mixtral_moe_woq_kernel_fn, mixtral_moe_woq_kernel_stub);

csrc/cpu/aten/kernels/MoEKrnl.cpp

@@ -292,7 +292,6 @@ at::Tensor mixtral_moe_woq_kernl_impl(
 
 template <typename T>
 std::tuple<at::Tensor, at::Tensor> deepseek_moegate_kernel(
-    const at::Tensor& hidden_states,
     const at::Tensor& scores,
     const at::Tensor& routed_scaling_factor,
     const int64_t n_group,
@@ -302,7 +301,7 @@ std::tuple<at::Tensor, at::Tensor> deepseek_moegate_kernel(
   auto group_size = n_routed_experts / n_group;
   auto n = scores.size(0);
   auto h = scores.size(1);
-  auto group_scores = at::empty({n, n_group}, hidden_states.options());
+  auto group_scores = at::empty({n, n_group}, scores.options());
   auto group_scores_ptr = group_scores.data_ptr<T>();
   auto scores_ptr = scores.data_ptr<T>();
 #pragma omp parallel for collapse(2)
@@ -319,7 +318,7 @@ std::tuple<at::Tensor, at::Tensor> deepseek_moegate_kernel(
   }
 
   auto group_idx = std::get<1>(group_scores.topk(topk_group, -1, true, false));
-  auto tmp_scores = at::zeros_like(scores, hidden_states.options());
+  auto tmp_scores = at::zeros_like(scores, scores.options());
   auto group_idx_ptr = group_idx.data_ptr<int64_t>();
   auto tmp_scores_ptr = tmp_scores.data_ptr<T>();
   T scale = routed_scaling_factor.item<T>();
@@ -339,17 +338,117 @@ std::tuple<at::Tensor, at::Tensor> deepseek_moegate_kernel(
   return std::make_tuple(topk, topk_weight);
 }
 
+template <typename T>
+std::tuple<at::Tensor, at::Tensor> deepseekv3_moegate_kernel(
+    const at::Tensor& scores,
+    const at::Tensor& routed_scaling_factor,
+    const int64_t n_group,
+    const int64_t topk_group,
+    const int64_t n_routed_experts,
+    const int64_t top_k,
+    const at::Tensor& e_score_cbias) {
+  auto group_size = n_routed_experts / n_group;
+  auto n = scores.size(0);
+  auto h = scores.size(1);
+  auto scores_for_choice = at::empty({n, n_group, group_size}, at::kFloat);
+  auto scores_ptr = scores.data_ptr<T>();
+  auto scores_for_choice_ptr = scores_for_choice.data_ptr<float>();
+  auto scores_for_choice_stride0 = scores_for_choice.stride(0);
+  auto e_score_cbias_ptr = e_score_cbias.data_ptr<float>();
+#pragma omp parallel for collapse(2)
+  for (auto i = 0; i < n; i++) {
+    for (auto j = 0; j < n_group; j++) {
+      auto k_start = j * group_size;
+      auto k_end = k_start + group_size;
+      for (auto k = k_start; k < k_end; k++) {
+        scores_for_choice_ptr[i * scores_for_choice_stride0 + k] =
+            scores_ptr[i * h + k] + e_score_cbias_ptr[k];
+      }
+    }
+  }
+  auto group_scores =
+      std::get<0>(scores_for_choice.topk(2, -1, true, false)).sum(-1);
+  auto group_idx = std::get<1>(group_scores.topk(topk_group, -1, true, false));
+  auto tmp_scores = at::zeros_like(scores, at::kFloat);
+  auto group_idx_ptr = group_idx.data_ptr<int64_t>();
+  auto tmp_scores_ptr = tmp_scores.data_ptr<float>();
+#pragma omp parallel for collapse(2)
+  for (auto i = 0; i < n; i++) {
+    for (auto j = 0; j < topk_group; j++) {
+      auto selected_idx = group_idx_ptr[i * topk_group + j];
+      auto k_start = selected_idx * group_size;
+      auto k_end = k_start + group_size;
+      for (auto k = k_start; k < k_end; k++) {
+        tmp_scores_ptr[i * h + k] =
+            scores_for_choice_ptr[i * scores_for_choice_stride0 + k];
+      }
+    }
+  }
+  auto topk = std::get<1>(tmp_scores.topk(top_k, -1, true, false));
+  auto topk_weight = scores.gather(1, topk);
+  return std::make_tuple(topk, topk_weight);
+}
+
 std::tuple<at::Tensor, at::Tensor> deepseek_moegate_kernel_impl(
     const at::Tensor& hidden_states,
     const at::Tensor& scores,
     const at::Tensor& routed_scaling_factor,
     const int64_t n_group,
     const int64_t topk_group,
     const int64_t n_routed_experts,
-    const int64_t top_k) {
+    const int64_t top_k,
+    c10::optional<at::Tensor> e_score_cbias) {
+  if (e_score_cbias.has_value()) { // deepseekv3
+    if (hidden_states.scalar_type() == at::ScalarType::Float) {
+      return deepseekv3_moegate_kernel<float>(
+          scores,
+          routed_scaling_factor,
+          n_group,
+          topk_group,
+          n_routed_experts,
+          top_k,
+          e_score_cbias.value());
+    } else if (hidden_states.scalar_type() == at::ScalarType::BFloat16) {
+      return deepseekv3_moegate_kernel<at::BFloat16>(
+          scores,
+          routed_scaling_factor,
+          n_group,
+          topk_group,
+          n_routed_experts,
+          top_k,
+          e_score_cbias.value());
+    } else if (hidden_states.scalar_type() == at::ScalarType::Half) {
+      return deepseekv3_moegate_kernel<at::Half>(
+          scores,
+          routed_scaling_factor,
+          n_group,
+          topk_group,
+          n_routed_experts,
+          top_k,
+          e_score_cbias.value());
+    }
+    auto n = hidden_states.size(0);
+    auto group_size = n_routed_experts / n_group;
+    auto scores_for_choice =
+        scores.view({n, -1}) + e_score_cbias.value().unsqueeze(0);
+    auto group_scores = std::get<0>(
+        scores_for_choice.view({n, n_group, -1}).topk(2, -1, true, false));
+    group_scores = group_scores.sum(-1);
+    auto group_idx =
+        std::get<1>(group_scores.topk(topk_group, -1, true, false));
+    auto group_mask = at::zeros_like(group_scores);
+    group_mask.scatter_(1, group_idx, 1);
+    auto score_mask = group_mask.unsqueeze(-1)
+                          .expand({n, n_group, group_size})
+                          .reshape({n, -1});
+    auto tmp_scores =
+        scores_for_choice.masked_fill(~score_mask.to(at::kBool), 0.0);
+    auto topk = std::get<1>(tmp_scores.topk(top_k, -1, true, false));
+    auto topk_weight = scores.gather(1, topk);
+    return std::make_tuple(topk, topk_weight.to(hidden_states.scalar_type()));
+  }
   if (hidden_states.scalar_type() == at::ScalarType::Float) {
     return deepseek_moegate_kernel<float>(
-        hidden_states,
         scores,
         routed_scaling_factor,
         n_group,
@@ -358,7 +457,14 @@ std::tuple<at::Tensor, at::Tensor> deepseek_moegate_kernel_impl(
         top_k);
   } else if (hidden_states.scalar_type() == at::ScalarType::BFloat16) {
     return deepseek_moegate_kernel<at::BFloat16>(
-        hidden_states,
+        scores,
+        routed_scaling_factor,
+        n_group,
+        topk_group,
+        n_routed_experts,
+        top_k);
+  } else if (hidden_states.scalar_type() == at::ScalarType::Half) {
+    return deepseek_moegate_kernel<at::Half>(
         scores,
         routed_scaling_factor,
         n_group,

intel_extension_for_pytorch/transformers/models/reference/models.py

@@ -5875,27 +5875,16 @@ def Deepseek_MoEGate_forward(self, hidden_states):
             self.top_k,
         )
     elif self.topk_method == "noaux_tc":
-        # TODO: fuse the following ops.
-        n = hidden_states.size(0)
-        scores_for_choice = scores.view(n, -1) + self.e_score_correction_bias.unsqueeze(
-            0
-        )
-        group_scores = (
-            scores_for_choice.view(n, self.n_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
-        )  # [n, n_group]
-        group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[
-            1
-        ]  # [n, top_k_group]
-        group_mask = torch.zeros_like(group_scores)  # [n, n_group]
-        group_mask.scatter_(1, group_idx, 1)  # [n, n_group]
-        score_mask = (
-            group_mask.unsqueeze(-1)
-            .expand(n, self.n_group, self.n_routed_experts // self.n_group)
-            .reshape(n, -1)
-        )  # [n, e]
-        tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(), 0.0)  # [n, e]
-        _, topk_idx = torch.topk(tmp_scores, k=self.top_k, dim=-1, sorted=False)
-        topk_weight = scores.gather(1, topk_idx)
+        topk_idx, topk_weight = torch.ops.torch_ipex.deepseek_moegate(
+            hidden_states,
+            scores,
+            torch.tensor(self.routed_scaling_factor),
+            self.n_group,
+            self.topk_group,
+            self.n_routed_experts,
+            self.top_k,
+            torch.tensor(self.e_score_correction_bias),
+        )
 
     # norm gate to sum 1
     if self.top_k > 1 and self.norm_topk_prob:

tests/cpu/test_cpu_ops.py

@@ -2245,7 +2245,10 @@ def moe_gate(scores):
             topk_weight = topk_weight * routed_scaling_factor
             return topk_idx, topk_weight
 
-        for dtype in [torch.float32, torch.bfloat16]:
+        dtypes = [torch.float32, torch.bfloat16]
+        if core.onednn_has_fp16_support():
+            dtypes.append(torch.float16)
+        for dtype in dtypes:
             hidden_states = torch.rand(10, 2560, dtype=dtype)
             weight = torch.rand(16, 2560, dtype=dtype)
             logits = torch.nn.functional.linear(
@@ -2267,6 +2270,66 @@ def moe_gate(scores):
                 self.assertEqual(topk_idx_ref, topk_idx_ipex)
                 self.assertEqual(topk_weight_ref, topk_weight_ipex)
 
+    def test_deepseekv3_moegate(self):
+        n_group = 8
+        topk_group = 3
+        n_routed_experts = 16
+        top_k = 6
+        routed_scaling_factor = 16.0
+        e_score_correction_bias = torch.rand(n_routed_experts)
+
+        def moe_gate(scores):
+            n, h = scores.shape
+            scores_for_choice = scores.view(n, -1) + e_score_correction_bias.unsqueeze(
+                0
+            )
+            group_scores = (
+                scores_for_choice.view(n, n_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
+            )  # [n, n_group]
+            group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[
+                1
+            ]  # [n, top_k_group]
+            group_mask = torch.zeros_like(group_scores)  # [n, n_group]
+            group_mask.scatter_(1, group_idx, 1)  # [n, n_group]
+            score_mask = (
+                group_mask.unsqueeze(-1)
+                .expand(n, n_group, n_routed_experts // n_group)
+                .reshape(n, -1)
+            )  # [n, e]
+            tmp_scores = scores_for_choice.masked_fill(
+                ~score_mask.bool(), 0.0
+            )  # [n, e]
+            _, topk_idx = torch.topk(tmp_scores, k=top_k, dim=-1, sorted=False)
+            topk_weight = scores.gather(1, topk_idx)
+
+            return topk_idx, topk_weight
+
+        dtypes = [torch.float32, torch.bfloat16]
+        if core.onednn_has_fp16_support():
+            dtypes.append(torch.float16)
+        for dtype in dtypes:
+            hidden_states = torch.rand(10, 2560, dtype=dtype)
+            weight = torch.rand(16, 2560, dtype=dtype)
+            logits = torch.nn.functional.linear(
+                hidden_states.type(torch.float32), weight.type(torch.float32), None
+            )
+            scores = logits.sigmoid()
+            enable_autocast = dtype == torch.bfloat16
+            with torch.no_grad(), torch.cpu.amp.autocast(enabled=enable_autocast):
+                topk_idx_ref, topk_weight_ref = moe_gate(scores)
+                topk_idx_ipex, topk_weight_ipex = torch.ops.torch_ipex.deepseek_moegate(
+                    hidden_states,
+                    scores.to(dtype),
+                    torch.tensor(routed_scaling_factor),
+                    n_group,
+                    topk_group,
+                    n_routed_experts,
+                    top_k,
+                    torch.tensor(e_score_correction_bias),
+                )
+                self.assertEqual(topk_idx_ref, topk_idx_ipex)
+                self.assertEqual(topk_weight_ref, topk_weight_ipex)
+
 
 if __name__ == "__main__":
     test = unittest.main()