Fix nightly ut failure (#3596)

Author: blzheng

csrc/cpu/aten/DSMoE.cpp

@@ -362,6 +362,18 @@ void grouped_topk_kernel_impl(
       e_score_correction_bias.data_ptr<float>(), \
       routed_scaling_factor.data_ptr<float>());
 
+#define LAUNCH_GROUPED_TOPK_KERNEL_FP16(NE)      \
+  grouped_topk_kernel_impl<at::Half, NE>(        \
+      topk_weights.data_ptr<float>(),            \
+      topk_ids.data_ptr<int32_t>(),              \
+      gating_output.data_ptr<at::Half>(),        \
+      num_tokens,                                \
+      topk,                                      \
+      num_expert_group,                          \
+      topk_group,                                \
+      renormalize,                               \
+      e_score_correction_bias.data_ptr<float>(), \
+      routed_scaling_factor.data_ptr<float>());
 //
 std::tuple<at::Tensor, at::Tensor> grouped_topk(
     at::Tensor& hidden_states,
@@ -380,36 +392,70 @@ std::tuple<at::Tensor, at::Tensor> grouped_topk(
   TORCH_CHECK(gating_output.size(0) == num_tokens, "Number of tokens mismatch");
   auto topk_weights = at::empty({num_tokens, topk}, at::kFloat);
   auto topk_ids = at::empty_like(topk_weights, at::kInt);
-  switch (num_experts) {
-    case 1:
-      LAUNCH_GROUPED_TOPK_KERNEL(1);
-      break;
-    case 2:
-      LAUNCH_GROUPED_TOPK_KERNEL(2);
-      break;
-    case 4:
-      LAUNCH_GROUPED_TOPK_KERNEL(4);
-      break;
-    case 8:
-      LAUNCH_GROUPED_TOPK_KERNEL(8);
-      break;
-    case 16:
-      LAUNCH_GROUPED_TOPK_KERNEL(16);
-      break;
-    case 32:
-      LAUNCH_GROUPED_TOPK_KERNEL(32);
-      break;
-    case 64:
-      LAUNCH_GROUPED_TOPK_KERNEL(64);
-      break;
-    case 128:
-      LAUNCH_GROUPED_TOPK_KERNEL(128);
-      break;
-    case 256:
-      LAUNCH_GROUPED_TOPK_KERNEL(256);
-      break;
-    default:
-      TORCH_CHECK(false, "Unexpected num_experts: ", num_experts);
+  if (st == at::kBFloat16) {
+    switch (num_experts) {
+      case 1:
+        LAUNCH_GROUPED_TOPK_KERNEL(1);
+        break;
+      case 2:
+        LAUNCH_GROUPED_TOPK_KERNEL(2);
+        break;
+      case 4:
+        LAUNCH_GROUPED_TOPK_KERNEL(4);
+        break;
+      case 8:
+        LAUNCH_GROUPED_TOPK_KERNEL(8);
+        break;
+      case 16:
+        LAUNCH_GROUPED_TOPK_KERNEL(16);
+        break;
+      case 32:
+        LAUNCH_GROUPED_TOPK_KERNEL(32);
+        break;
+      case 64:
+        LAUNCH_GROUPED_TOPK_KERNEL(64);
+        break;
+      case 128:
+        LAUNCH_GROUPED_TOPK_KERNEL(128);
+        break;
+      case 256:
+        LAUNCH_GROUPED_TOPK_KERNEL(256);
+        break;
+      default:
+        TORCH_CHECK(false, "Unexpected num_experts: ", num_experts);
+    }
+  } else if (st == at::kHalf) {
+    switch (num_experts) {
+      case 1:
+        LAUNCH_GROUPED_TOPK_KERNEL_FP16(1);
+        break;
+      case 2:
+        LAUNCH_GROUPED_TOPK_KERNEL_FP16(2);
+        break;
+      case 4:
+        LAUNCH_GROUPED_TOPK_KERNEL_FP16(4);
+        break;
+      case 8:
+        LAUNCH_GROUPED_TOPK_KERNEL_FP16(8);
+        break;
+      case 16:
+        LAUNCH_GROUPED_TOPK_KERNEL_FP16(16);
+        break;
+      case 32:
+        LAUNCH_GROUPED_TOPK_KERNEL_FP16(32);
+        break;
+      case 64:
+        LAUNCH_GROUPED_TOPK_KERNEL_FP16(64);
+        break;
+      case 128:
+        LAUNCH_GROUPED_TOPK_KERNEL_FP16(128);
+        break;
+      case 256:
+        LAUNCH_GROUPED_TOPK_KERNEL_FP16(256);
+        break;
+      default:
+        TORCH_CHECK(false, "Unexpected num_experts: ", num_experts);
+    }
   }
   return std::make_tuple(topk_ids, topk_weights);
 }

intel_extension_for_pytorch/transformers/models/reference/modules/attentions.py

@@ -2701,11 +2701,11 @@ def __init__(self, module, config, sdp_module_ref, distributed=False):
                 self.hidden_size = module.q_proj.linear.weight.shape[0]
         elif hasattr(module, "o_proj"):
             if hasattr(module.o_proj, "in_features"):
-                self.hidden_size = module.q_proj.in_features
+                self.hidden_size = module.o_proj.in_features
             elif hasattr(module.o_proj, "linear") and hasattr(
                 module.o_proj.linear, "in_features"
             ):
-                self.hidden_size = module.q_proj.linear.in_features
+                self.hidden_size = module.o_proj.linear.in_features
             elif hasattr(module.o_proj, "weight"):
                 self.hidden_size = module.o_proj.weight.shape[1]
             else:

tests/cpu/hf_configs/deepseekv3/modeling_deepseek.py

@@ -409,7 +409,7 @@ def __init__(self, config):
         )
         if self.topk_method == "noaux_tc":
             self.e_score_correction_bias = nn.Parameter(
-                torch.empty((self.n_routed_experts))
+                torch.rand((self.n_routed_experts))
             )
         self.reset_parameters()
 

tests/cpu/test_ipex_optimize_transformers_nightly.py

@@ -286,6 +286,13 @@ def model_replacement_check(
             ]:
                 state_dict[weight] = torch.rand(state_dict[weight].shape)
             model.load_state_dict(state_dict)
+        elif m.name in ["deepseekv2", "deepseekv3"]:
+            model = model.to(dtype)
+            model.model.layers[
+                config.first_k_dense_replace
+            ].mlp.gate.e_score_correction_bias = torch.nn.Parameter(
+                torch.rand(config.n_routed_experts)
+            )
         elif m.name == "llava":
             model.get_vision_tower().load_model()
         elif m.name == "jamba":
@@ -390,7 +397,7 @@ def model_replacement_check(
         ):
             key_ipex = ipex_m(**input_dict)
         error_message = f"model={m.name}, deployment_mode={deployment_mode}, torchcompile={torchcompile}, return_dict={return_dict}"
-        if m.name != "mllama":
+        if m.name not in ["mllama", "deepseekv3"]:
             if return_dict:
                 assert isinstance(key_ipex, dict)
                 self.assertEqual(