fix kernel_key for fused_gemm_epilogue_impl

Author: ooooo-create

paddle/phi/api/lib/api_custom_impl.cc

@@ -240,6 +240,21 @@ std::tuple<Tensor, Tensor> fused_gemm_epilogue_impl(
   Backend kernel_backend = Backend::UNDEFINED;
   DataLayout kernel_layout = DataLayout::UNDEFINED;
   DataType kernel_data_type = DataType::UNDEFINED;
+  if (kernel_backend == Backend::UNDEFINED ||
+      kernel_layout == DataLayout::UNDEFINED ||
+      kernel_data_type == DataType::UNDEFINED) {
+    auto kernel_key_set = ParseKernelKeyByInputArgs(x, y, bias);
+    auto kernel_key = kernel_key_set.GetHighestPriorityKernelKey();
+    if (kernel_backend == Backend::UNDEFINED) {
+      kernel_backend = kernel_key.backend();
+    }
+    if (kernel_layout == DataLayout::UNDEFINED) {
+      kernel_layout = kernel_key.layout();
+    }
+    if (kernel_data_type == DataType::UNDEFINED) {
+      kernel_data_type = kernel_key.dtype();
+    }
+  }
 #ifdef PADDLE_WITH_DISTRIBUTE
   bool run_auto_parallel = AllInputsAreDistTensor(x, y, bias);
   bool rank_is_in_current_mesh = true;
@@ -250,23 +265,6 @@ std::tuple<Tensor, Tensor> fused_gemm_epilogue_impl(
             .process_mesh();
     rank_is_in_current_mesh = phi::distributed::IsCurRankInMesh(mesh);
   }
-  if (rank_is_in_current_mesh) {
-    if (kernel_backend == Backend::UNDEFINED ||
-        kernel_layout == DataLayout::UNDEFINED ||
-        kernel_data_type == DataType::UNDEFINED) {
-      auto kernel_key_set = ParseKernelKeyByInputArgs(x, y, bias);
-      auto kernel_key = kernel_key_set.GetHighestPriorityKernelKey();
-      if (kernel_backend == Backend::UNDEFINED) {
-        kernel_backend = kernel_key.backend();
-      }
-      if (kernel_layout == DataLayout::UNDEFINED) {
-        kernel_layout = kernel_key.layout();
-      }
-      if (kernel_data_type == DataType::UNDEFINED) {
-        kernel_data_type = kernel_key.dtype();
-      }
-    }
-  }
 
   // Kernel Dispatch Body
   // Auto Parallel condition