ep support logprob (#4089)

Author: ckl117

custom_ops/gpu_ops/get_output_msg_with_topk.cc

@@ -39,9 +39,6 @@ void GetOutputTopK(const paddle::Tensor& x,
                    int k,
                    int64_t rank_id,
                    bool wait_flag) {
-    if (rank_id > 0) {
-        return;
-    }
 
     static struct msgdata msg_rcv;
     int msg_queue_id = 1;

fastdeploy/engine/args_utils.py

@@ -400,8 +400,6 @@ def __post_init__(self):
         if self.enable_logprob:
             if self.speculative_config is not None:
                 raise NotImplementedError("Logprob does not support speculation_config.")
-            if self.enable_expert_parallel:
-                raise NotImplementedError("Logprob does not support enable_expert_parallel.")
             if not current_platform.is_cuda():
                 raise NotImplementedError("Only CUDA platform supports logprob.")
         if self.splitwise_role != "mixed":

fastdeploy/output/token_processor.py

@@ -161,24 +161,23 @@ def process_sampling_results(self):
                         continue
 
                 else:
-                    if (
+                    if self.use_logprobs:
+                        get_output_topk(
+                            self.output_tokens,
+                            self.output_scores,
+                            self.output_ranks,
+                            K,
+                            rank_id,
+                            is_blocking,
+                        )
+                    elif (
                         self.cfg.parallel_config.enable_expert_parallel
                         and self.cfg.parallel_config.data_parallel_size > 1
                     ):
                         get_output_ep(self.output_tokens, rank_id, is_blocking)
 
                     else:
-                        if self.use_logprobs:
-                            get_output_topk(
-                                self.output_tokens,
-                                self.output_scores,
-                                self.output_ranks,
-                                K,
-                                rank_id,
-                                is_blocking,
-                            )
-                        else:
-                            get_output(self.output_tokens, rank_id, is_blocking)
+                        get_output(self.output_tokens, rank_id, is_blocking)
 
                     if self.output_tokens[0, 0] == -2:
                         continue