Merge pull request #1 from ckl117/logprobs

infer engine support base logprobs

Author: zhenwenDang

custom_ops/gpu_ops/get_output_msg_with_topk.cc

@@ -24,7 +24,7 @@
 #endif
 
 #define MAX_BSZ 512
-#define K 10
+#define K 20
 
 struct msgdata {
     long mtype;

custom_ops/gpu_ops/save_output_msg_with_topk.cc

@@ -23,8 +23,8 @@
 #define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name)
 #endif
 
-#define MAX_BSZ 128
-#define K 10
+#define MAX_BSZ 512
+#define K 20
 // #define SAVE_WITH_OUTPUT_DEBUG
 
 struct msgdata {
@@ -35,22 +35,15 @@ struct msgdata {
 
 void SaveOutMmsgTopK(const paddle::Tensor& x,
                      const paddle::Tensor& scores,
-                     const paddle::Tensor& topk_ids,
-                     const paddle::Tensor& topk_scores,  // [bsz, k]
                      const paddle::Tensor& not_need_stop,
-                     int k,
                      int64_t rank_id) {
     if (rank_id > 0) {
         return;
     }
     auto x_cpu = x.copy_to(paddle::CPUPlace(), false);
     auto scores_cpu = scores.copy_to(paddle::CPUPlace(), false);
-    auto topk_ids_cpu = topk_ids.copy_to(paddle::CPUPlace(), false);
-    auto topk_scores_cpu = topk_scores.copy_to(paddle::CPUPlace(), false);
     int64_t* x_data = x_cpu.data<int64_t>();
     float* scores_data = scores_cpu.data<float>();
-    int64_t* topk_ids_data = topk_ids_cpu.data<int64_t>();
-    float* topk_scores_data = topk_scores_cpu.data<float>();
     static struct msgdata msg_sed;
     int msg_queue_id = 1;
     if (const char* inference_msg_queue_id_env_p =
@@ -106,20 +99,14 @@ void SaveOutMmsgTopK(const paddle::Tensor& x,
     msg_sed.mtext[0] = not_need_stop_data ? inference_msg_id_from_env
                                           : -inference_msg_id_from_env;
     int bsz = x.shape()[0];
+    int token_num = x.shape()[1];
+    int k = token_num - 1;
     msg_sed.mtext[1] = bsz;
     for (int i = 0; i < bsz; i++) {
-        for (int j = 0; j < k + 1; j++) {
+        for (int j = 0; j < token_num; j++) {
             const int64_t offset = i * (K + 1) + j;
-            if (j == 0) {
-                msg_sed.mtext[offset + 2] = (int)x_data[i];
-                msg_sed.mtext_f[offset] = scores_data[i];
-            } else if (j <= k + 1) {
-                msg_sed.mtext[offset + 2] = (int)topk_ids_data[i * k + j - 1];
-                msg_sed.mtext_f[offset] = topk_scores_data[i * k + j - 1];
-            } else {
-                msg_sed.mtext[offset + 2] = -1;
-                msg_sed.mtext_f[offset] = 0.0;
-            }
+            msg_sed.mtext[offset + 2] = (int)x_data[i * token_num + j];
+            msg_sed.mtext_f[offset] = scores_data[i * token_num + j];
         }
     }
 #ifdef SAVE_WITH_OUTPUT_DEBUG
@@ -139,8 +126,8 @@ void SaveOutMmsgTopK(const paddle::Tensor& x,
 }
 
 PD_BUILD_STATIC_OP(save_output_topk)
-    .Inputs({"x", "scores", "topk_ids", "topk_scores", "not_need_stop"})
-    .Attrs({"k: int", "rank_id: int64_t"})
+    .Inputs({"x", "scores", "not_need_stop"})
+    .Attrs({"rank_id: int64_t"})
     .Outputs({"x_out"})
     .SetInplaceMap({{"x", "x_out"}})
     .SetKernelFn(PD_KERNEL(SaveOutMmsgTopK));

custom_ops/setup_ops_base.py

@@ -22,6 +22,7 @@
             "gpu_ops/save_with_output_msg.cc",
             "gpu_ops/get_output.cc",
             "gpu_ops/get_output_msg_with_topk.cc",
+            "gpu_ops/save_output_msg_with_topk.cc",
             "gpu_ops/transfer_output.cc",
             "cpu_ops/rebuild_padding.cc",
         ],

fastdeploy/model_executor/layers/sample/meta_data.py

@@ -42,3 +42,4 @@ class SamplingMetadata:
 
     top_p: paddle.Tensor
     top_k: Optional[paddle.Tensor] = None
+    max_num_logprobs: Optional[int] = None

fastdeploy/model_executor/layers/sample/sampler.py

@@ -29,6 +29,7 @@
     apply_penalty_multi_scores, apply_speculative_penalty_multi_scores,
     top_p_sampling)
 from fastdeploy.platforms import current_platform
+from fastdeploy.worker.output import LogprobsTensors, SamplerOutput
 
 
 class SamplerProcessor:
@@ -189,14 +190,65 @@ def pre_process(self, skip_idx_list: List[int] = []):
         """ pre process before running """
         self.processor.pre_process(skip_idx_list)
 
+    def compute_logprobs(self, logits: paddle.Tensor) -> paddle.Tensor:
+        """
+        """
+        return F.log_softmax(logits, axis=-1)
+
+    def gather_logprobs(
+        self,
+        logprobs: paddle.Tensor,
+        num_logprobs: int,
+        token_ids: paddle.Tensor,
+    ) -> LogprobsTensors:
+        """
+        Gather logprobs for topk and sampled/prompt token.
+
+        Args:
+          logprobs: (num tokens) x (vocab) tensor
+          num_logprobs: minimum number of logprobs to
+                        retain per token
+          token_ids: prompt tokens (if prompt logprobs)
+                     or sampled tokens (if sampled
+                     logprobs); 1D token ID tensor
+                     with (num tokens) elements
+                     Must be int64.
+
+        Returns:
+          Top-k int indices tensor, (num tokens) x (num_logprobs + 1)
+          Top-k float logprobs tensor, (num tokens) x (num_logprobs + 1)
+          Sampled token rank tensor, (num tokens)
+        """
+        assert token_ids.dtype == paddle.int64
+        # Find the topK values.
+        token_logprobs = paddle.take_along_axis(logprobs, token_ids, axis=-1)
+        if num_logprobs >= 1:
+            topk_logprobs, topk_indices = paddle.topk(logprobs,
+                                                    num_logprobs,
+                                                    axis=-1)
+            indices = paddle.concat([token_ids, topk_indices], axis=1)
+            top_logprobs = paddle.concat([token_logprobs, topk_logprobs], axis=1)
+        else:
+            indices = token_ids
+            top_logprobs = token_logprobs
+
+        # Compute the ranks of the actual token.
+        token_ranks = (logprobs >= token_logprobs).sum(-1)
+
+        return LogprobsTensors(indices, top_logprobs, token_ranks)
+
     def forward_cuda(
         self,
         logits: paddle.Tensor,
         sampling_metadata: SamplingMetadata,
         skip_idx_list: List[int] = [],
-    ) -> paddle.Tensor:
+    ) -> SamplerOutput:
         """
         """
+        num_logprobs = sampling_metadata.max_num_logprobs
+        if num_logprobs is not None:
+            raw_logprobs = self.compute_logprobs(logits)
+
         logits = self.processor.apply_token_mask(logits, skip_idx_list)
 
         logits = apply_penalty_multi_scores(
@@ -216,8 +268,19 @@ def forward_cuda(
 
         _, next_tokens = top_p_sampling(probs, sampling_metadata.top_p)
 
+        logprobs_tensors = None if num_logprobs is None else \
+            self.gather_logprobs(raw_logprobs, num_logprobs, token_ids=next_tokens)
+
         self.processor.update_output_tokens(next_tokens, skip_idx_list)
-        return next_tokens
+
+        sampler_output = SamplerOutput(
+            # The sampled tokens are expanded to 2D tensor with shape
+            # [num_requests, 1], where each row represents one generated
+            # token per request.
+            sampled_token_ids=next_tokens,
+            logprobs_tensors=logprobs_tensors,
+        )
+        return sampler_output
 
 
 class SpeculativeSampler(nn.Layer):

fastdeploy/model_executor/pre_and_post_process.py

@@ -32,8 +32,9 @@
         speculate_save_output, speculate_set_value_by_flags_and_idx,
         speculate_step_paddle, speculate_step_system_cache,
         speculate_update_v3, step_paddle, step_system_cache, update_inputs,
-        step_reschedule)
-from fastdeploy.worker.output import ModelOutputData
+        step_reschedule, save_output_topk)
+from fastdeploy.worker.output import (ModelOutputData, ModelRunnerOutput,
+                                      SamplerOutput)
 
 DISABLE_RECOVER = (envs.FD_DISABLED_RECOVER == "1")
 
@@ -109,10 +110,10 @@ def pre_process(
             cu_seqlens_k, output_cum_offsets, output_padding_offset)
 
 
-def post_process_normal(sampled_token_ids: paddle.Tensor,
+def post_process_normal(sampler_output: SamplerOutput,
                         model_output: ModelOutputData,
                         save_each_rank: bool = False,
-                        skip_save_output: bool = False) -> None:
+                        skip_save_output: bool = False) -> ModelRunnerOutput:
     """ Post-processing steps after completing a single token generation. """
     # 1. Set stop value
     paddle.assign(
@@ -130,7 +131,8 @@ def post_process_normal(sampled_token_ids: paddle.Tensor,
         model_output.stop_flags,
     )
     # TODO(gongshaotian): Add use_stop_seqs
-    set_stop_value_multi_ends(sampled_token_ids, model_output.stop_flags,
+    set_stop_value_multi_ends(sampler_output.sampled_token_ids,
+                              model_output.stop_flags,
                               model_output.seq_lens_this_time,
                               model_output.eos_token_id,
                               model_output.next_tokens, False)  # multi ends
@@ -145,18 +147,26 @@ def post_process_normal(sampled_token_ids: paddle.Tensor,
             model_output.seq_lens_decoder,
             model_output.input_ids,
             model_output.stop_nums,
-            sampled_token_ids,
+            sampler_output.sampled_token_ids,
             model_output.is_block_step,
         )
     # 3. Transmit the model's output and stop generation signal via message queue.
     #    In the future, we will abandon this approach.
     if not skip_save_output:
-        save_output(
-            sampled_token_ids,
-            model_output.not_need_stop,
-            model_output.mp_rank,
-            save_each_rank,  # save_each_rank
-        )
+        if sampler_output.logprobs_tensors is None:
+            save_output(
+                sampler_output.sampled_token_ids,
+                model_output.not_need_stop,
+                model_output.mp_rank,
+                save_each_rank,  # save_each_rank
+            )
+        else:
+            save_output_topk(
+                sampler_output.logprobs_tensors.logprob_token_ids,
+                sampler_output.logprobs_tensors.logprobs,
+                model_output.not_need_stop,
+                model_output.mp_rank,
+            )
 
 
 def post_process_specualate(model_output, skip_save_output: bool = False):
@@ -201,7 +211,7 @@ def post_process_specualate(model_output, skip_save_output: bool = False):
     )
 
 
-def post_process(sampled_token_ids: paddle.Tensor,
+def post_process(sampler_output: SamplerOutput,
                  model_output: ModelOutputData,
                  save_each_rank: bool = False,
                  speculative_decoding: bool = False,
@@ -210,7 +220,7 @@ def post_process(sampled_token_ids: paddle.Tensor,
     if speculative_decoding:
         post_process_specualate(model_output, skip_save_output)
     else:
-        post_process_normal(sampled_token_ids, model_output, save_each_rank,
+        post_process_normal(sampler_output, model_output, save_each_rank,
                             skip_save_output)
 
 

fastdeploy/worker/gpu_model_runner.py

@@ -582,6 +582,7 @@ def _prepare_inputs(self) -> None:
             min_dec_lens=self.share_inputs["min_dec_len"],
             bad_words_token_ids=self.share_inputs["bad_tokens"],
             eos_token_ids=self.share_inputs["eos_token_id"],
+            max_num_logprobs=None,
         )
 
     def load_model(self) -> None:
@@ -786,15 +787,15 @@ def _dummy_run(self,
                     self.share_inputs["step_idx"],
                     self.share_inputs["stop_flags"],
                 )
-                sampled_token_ids = self.sampler(logits,
+                sampler_output = self.sampler(logits,
                                                  self.sampling_metadata)
                 if self.parallel_config.tensor_parallel_degree > 1:
-                    paddle.distributed.broadcast(sampled_token_ids, 0)
+                    paddle.distributed.broadcast(sampler_output.sampled_token_ids, 0)
             else:
                 self.sampler(logits, self.sampling_metadata,
                              self.parallel_config.max_model_len,
                              self.share_inputs)
-                sampled_token_ids = None
+                sampler_output = None
                 if self.parallel_config.tensor_parallel_degree > 1:
                     paddle.distributed.broadcast(
                         self.share_inputs["accept_tokens"], 0)
@@ -834,7 +835,7 @@ def _dummy_run(self,
                 accept_num=self.share_inputs["accept_num"]
                 if self.speculative_decoding else None)
 
-            post_process(sampled_token_ids=sampled_token_ids,
+            post_process(sampler_output=sampler_output,
                          model_output=model_output_data,
                          speculative_decoding=self.speculative_decoding,
                          skip_save_output=True)
@@ -1021,18 +1022,18 @@ class at the server level, which is too granular for ModelRunner.
                 self.share_inputs["step_idx"],
                 self.share_inputs["stop_flags"],
             )
-            sampled_token_ids = self.sampler(
+            sampler_output = self.sampler(
                 logits,
                 self.sampling_metadata,
                 skip_idx_list,
             )
             if self.parallel_config.tensor_parallel_degree > 1:
-                paddle.distributed.broadcast(sampled_token_ids, 0)
+                paddle.distributed.broadcast(sampler_output.sampled_token_ids, 0)
 
         else:
             self.sampler(logits, self.sampling_metadata,
                          self.parallel_config.max_model_len, self.share_inputs)
-            sampled_token_ids = None
+            sampler_output = None
             if self.parallel_config.tensor_parallel_degree > 1:
                 paddle.distributed.broadcast(
                     self.share_inputs["accept_tokens"], 0)
@@ -1075,7 +1076,7 @@ class at the server level, which is too granular for ModelRunner.
             skip_save_output = True
         else:
             skip_save_output = False
-        post_process(sampled_token_ids=sampled_token_ids,
+        post_process(sampler_output=sampler_output,
                      model_output=model_output_data,
                      save_each_rank=self.parallel_config.use_ep,
                      speculative_decoding=self.speculative_decoding,