update

Author: Wanglongzhi2001

llm/server/server/engine/config.py

@@ -211,14 +211,15 @@ def get_speculate_config(self):
             SpeculateConfig: the speculate related config
         """
         speculate_config = SpeculateConfig()
-        if self.model_cfg.get("speculate_method") is not None:
-            speculate_config.speculate_method = self.model_cfg["speculate_method"]
-            speculate_config.speculate_max_draft_token_num = self.model_cfg[
+        model_cfg = self.get_model_config()
+        if model_cfg.get("speculate_method", "None") != "None":
+            speculate_config.speculate_method = str(model_cfg["speculate_method"])
+            speculate_config.speculate_max_draft_token_num = model_cfg[
                 "speculate_max_draft_token_num"]
-            speculate_config.speculate_max_ngram_size = self.model_cfg[
+            speculate_config.speculate_max_ngram_size = model_cfg[
                 "speculate_max_ngram_size"]
 
-        if speculate_config.speculate_method is not in ["none", "inference_with_reference"]:
+        if speculate_config.speculate_method not in ["None", "inference_with_reference"]:
             model_server_logger.error(f"Unsupport speculate method: {speculate_config.speculate_method}")
 
         return speculate_config
@@ -258,6 +259,6 @@ def __str__(self) -> str:
 
 @dataclass
 class SpeculateConfig:
-    speculate_method: str = None
+    speculate_method: str = "None"
     speculate_max_draft_token_num: int = 1
     speculate_max_ngram_size: int = 1

llm/server/server/engine/infer.py

@@ -48,7 +48,7 @@ def __init__(self, args):
         self.config = Config()
         self.model_cfg = self.config.get_model_config()
         self.speculate_config = self.config.get_speculate_config()
-        self.is_speculate_decoding = self.speculate_config.speculate_method is not None
+        self.is_speculate_decoding = self.speculate_config.speculate_method != "None"
         self.format_print_configuration()
 
         self.args.num_layers = self.get_value(self.model_cfg, ["num_hidden_layers", "num_layers"])
@@ -71,7 +71,7 @@ def __init__(self, args):
         self.init_inputs()
 
         if self.is_speculate_decoding:
-            logger.info(f'Using speculating decoding, method: {self.speculate_config.speculate_method}.')
+            logger.info(f'Using speculate decoding, method: {self.speculate_config.speculate_method}.')
             if self.speculate_config.speculate_method == "inference_with_reference":
                 self.proposer = InferenceWithReferenceProposer(
                     self.speculate_config.speculate_max_draft_token_num,
@@ -371,7 +371,7 @@ def step_cuda(self, seq_lens_this_time):
                     self.share_inputs['free_list'], self.share_inputs['free_list_len'],
                     self.share_inputs['input_ids'], self.share_inputs['pre_ids'],
                     self.share_inputs['step_idx'], self.share_inputs['next_tokens'],
-                    self.args.block_size, self.args.enc_dec_block_num, self.args.first_token_id
+                    self.args.block_size, self.args.enc_dec_block_num, self.args.first_token_id,
                     speculate_step_token_num)
 
     def initialize_engine_ready_check_flag(self):

llm/server/server/engine/token_processor.py

@@ -39,9 +39,9 @@ def __init__(self, cfg):
 
         self.tokens_counter = Counter()
 
-        self.is_speculate_decoding = self.cfg.get_model_config().get("speculate_method") is not None
+        self.is_speculate_decoding = self.cfg.get_speculate_config().speculate_method != "None"
         if self.is_speculate_decoding:
-            self.output_tokens = paddle.full(shape=[SPECULATE_MAX_BSZ * MAX_DRAFT_TOKENS + SPECULATE_MAX_BSZ + 2], fill_value=2, dtype="int64")
+            self.output_tokens = paddle.full(shape=[SPECULATE_MAX_BSZ * MAX_DRAFT_TOKENS + SPECULATE_MAX_BSZ + 2, 1], fill_value=2, dtype="int64")
         else:
             self.output_tokens = paddle.full(shape=[self.cfg.max_batch_size + 2, 1], fill_value=2, dtype="int64")
         self.worker = None
@@ -71,10 +71,7 @@ def run(self):
         if self.worker is not None:
             raise Exception("Worker is already running!")
 
-        if self.is_speculate_decoding:
-            self.worker = threading.Thread(target=self.process_speculate_results, args=())
-        else:
-            self.worker = threading.Thread(target=self.process_sampling_results, args=())
+        self.worker = threading.Thread(target=self.process_sampling_results, args=())
         self.worker.daemon = True
         self.worker.start()
 
@@ -86,30 +83,18 @@ def process_sampling_results(self):
             try:
                 rank_id = 0
                 is_blocking = True
-                get_output(self.output_tokens, rank_id, is_blocking)
+                if self.is_speculate_decoding:
+                    speculate_get_output(self.output_tokens, rank_id, is_blocking)
+                else:
+                    get_output(self.output_tokens, rank_id, is_blocking)
 
                 if self.output_tokens[0, 0] == -2:
                     continue
+                
                 self._process_batch_output()
             except Exception as e:
                 model_server_logger.info("while get input_data error: {0} {1}".format(e, str(traceback.format_exc())))
 
-    def process_speculate_results(self):
-        """
-        read tokens from paddle inference engine and process
-        """
-        while True:
-            try:
-                rank_id = 0
-                is_blocking = True
-                speculate_get_output(self.output_tokens, rank_id, is_blocking)
-
-                if self.output_tokens[0] == -2:
-                    continue
-                self._process_speculate_output()
-            except Exception as e:
-                model_server_logger.info("while get input_data error: {0} {1}".format(e, str(traceback.format_exc())))
-
     def postprocess(self, batch_result, exist_finished_task=False):
         """
         single post-processing function
@@ -126,73 +111,14 @@ def postprocess(self, batch_result, exist_finished_task=False):
             with open(result_file, "a") as f:
                 f.write("{}\n".format(result))
 
-    def _get_single_result(self, i, task_id, token_id, task):
+    def _get_single_result(self, i, task_id, token_ids, task):
         """
         processing single results
 
         Args:
             i (int): batch index
             task_id (str): task id
-            token_id (int): token id
-            task (dict): task information
-
-        Returns:
-            dict: result
-        """
-        inference_time_cost = time.time() - task["inference_start_time"]
-        task["inference_time_cost"] = inference_time_cost
-        task["tokens_all_num"] = len(self.all_tokens[i])
-        task["inference_current_step_time"] = datetime.now()
-        result = {
-            "req_id": task_id,
-            "is_end": 0,
-            "token_ids": [token_id],
-            "send_idx": self.tokens_counter[task_id],
-            "inference_time_cost": inference_time_cost,
-            "infer_seed": task["infer_seed"],
-            "return_all_tokens": task.get("return_all_tokens", False),
-        }
-
-        # get benchmark msg
-        if task.get("benchmark"):
-            keys = ["preprocess_start_time", "preprocess_end_time", "schedule_start_time",
-                    "inference_start_time", "inference_current_step_time"]
-            for key in keys:
-                if key in task:
-                    result[key] = str(task[key])
-
-        # fill some extra information
-        if token_id in task["eos_token_ids"]:
-            result["is_end"] = 1
-            result["token_ids"] = []
-            result["tokens_all_num"] = len(self.all_tokens[i]) + 1
-            result["tokens_all_ids"] = self.all_tokens[i]
-
-            info_dict = {}
-            info_dict["req_id"] = task["req_id"]
-            info_dict["input_token_num"] = len(task["input_ids"])
-            info_dict["output_token_num"] = len(self.all_tokens[i])
-            if hasattr(task, "preprocess_start_time") and hasattr(task, "preprocess_end_time"):
-                info_dict["preprocess_cost_time"] = datetime_diff(task["preprocess_start_time"],
-                                                                  task["preprocess_end_time"])
-            if hasattr(task, "preprocess_end_time") and hasattr(task, "schedule_start_time"):
-                info_dict["cache_waiting_cost_time"] = datetime_diff(task["preprocess_end_time"],
-                                                                     task["schedule_start_time"])
-            info_dict["inference_time_cost"] = task["inference_time_cost"]
-            info_dict["version"] = "4.6"
-            info_dict["timestamp"] = time.time()
-            monitor_logger.info(f"{info_dict}")
-
-        return result
-
-    def _get_speculate_result(self, i, task_id, token_ids, task):
-        """
-        processing single speculate results
-
-        Args:
-            i (int): batch index
-            task_id (str): task id
-            token_ids (int): tokens id
+            token_ids (list): token id
             task (dict): task information
 
         Returns:
@@ -220,23 +146,23 @@ def _get_speculate_result(self, i, task_id, token_ids, task):
                 if key in task:
                     result[key] = str(task[key])
 
-
-        # fill some extra information when generate eos token
+        # fill some extra information
         result["token_ids"] = []
         for token_id in token_ids:
             if token_id in task["eos_token_ids"]:
                 result["is_end"] = 1
+                result["token_ids"] = []
                 result["tokens_all_num"] = len(self.all_tokens[i]) + 1
                 result["tokens_all_ids"] = self.all_tokens[i]
 
                 info_dict = {}
                 info_dict["req_id"] = task["req_id"]
                 info_dict["input_token_num"] = len(task["input_ids"])
                 info_dict["output_token_num"] = len(self.all_tokens[i])
-                if "preprocess_start_time" in task and "preprocess_end_time" in task:
+                if hasattr(task, "preprocess_start_time") and hasattr(task, "preprocess_end_time"):
                     info_dict["preprocess_cost_time"] = datetime_diff(task["preprocess_start_time"],
                                                                     task["preprocess_end_time"])
-                if "preprocess_end_time" in task and "schedule_start_time" in task:
+                if hasattr(task, "preprocess_end_time") and hasattr(task, "schedule_start_time"):
                     info_dict["cache_waiting_cost_time"] = datetime_diff(task["preprocess_end_time"],
                                                                         task["schedule_start_time"])
                 info_dict["inference_time_cost"] = task["inference_time_cost"]
@@ -266,74 +192,36 @@ def _process_batch_output(self):
         """
         tokens = self.output_tokens.numpy()
         batch = self.output_tokens[1, 0]
-        tokens = tokens[2:batch + 2]
+        if not self.is_speculate_decoding:
+            tokens = tokens[2:batch + 2]
+        else:
+            accept_num = tokens[2:batch + 2]
 
         batch_result = list()
         exist_finished_task = False
         for i in range(batch):
             if self.resource_manager.stop_flags[i]:
                 continue
 
-            token_id = int(tokens[i, 0])
-            if token_id < 0:
+            if not self.is_speculate_decoding:
+                token_ids = [int(tokens[i, 0])]
+            else:
+                token_ids = tokens[2 + SPECULATE_MAX_BSZ + i * MAX_DRAFT_TOKENS: 2 + SPECULATE_MAX_BSZ + i * MAX_DRAFT_TOKENS + accept_num[i, 0], 0].tolist()
+            
+            if any(token_id < 0 for token_id in token_ids):
                 continue
 
             task = self.resource_manager.tasks_list[i]
 
             task_id = task["req_id"]
-            result = self._get_single_result(i, task_id, token_id, task)
-
-            self.tokens_counter[task_id] += 1
-            if token_id not in task["eos_token_ids"]:
-                self.all_tokens[i].append(token_id)
+            result = self._get_single_result(i, task_id, token_ids, task)
 
-            self.number_of_output_tokens += 1
-            if token_id in task["eos_token_ids"]:
-                self._recycle_resources(task_id, i, task)
-                model_server_logger.info("req_id: {0} finished".format(task_id))
-                model_server_logger.info(f"{self.resource_manager.info()}")
-                exist_finished_task = True
-            batch_result.append(result)
-
-        self.postprocess(batch_result, exist_finished_task)
-
-    def _process_speculate_output(self):
-        """
-        batch post-processing function
-        """
-        tokens = self.output_tokens.numpy()
-        batch = self.output_tokens[1]
-        output_token_msg_id = int(self.output_tokens[0])
-        accept_num = tokens[2 : batch + 2]
-        batch_result = list()
-        # 用于判断当前此批结果中是否存在已完成的任务
-        exist_finished_task = False
-        prefill_mode = False
-        tasks_prefill = []
-        
-        for i in range(batch):
-            # 对应task如若已结束，跳过
-            if self.resource_manager.stop_flags[i]:
-                continue
-
-            token_ids = tokens[2 + SPECULATE_MAX_BSZ + i * MAX_DRAFT_TOKENS: 2 + SPECULATE_MAX_BSZ + i * MAX_DRAFT_TOKENS + accept_num[i]].tolist()
-            # 跳过非法token
-            if len(token_ids) == 0 or token_ids[-1] == 0:
-                continue
-
-            task = self.resource_manager.tasks_list[i]
-
-            # 将会移至data server解决
-            task_id = task["req_id"]
-            result = self._get_speculate_result(i, task_id, token_ids, task)
-            
             for token_id in token_ids:
                 self.tokens_counter[task_id] += 1
                 if token_id not in task["eos_token_ids"]:
                     self.all_tokens[i].append(token_id)
 
                 self.number_of_output_tokens += 1
-                # 生成结束符时，重置相应变量
                 if token_id in task["eos_token_ids"]:
                     self._recycle_resources(task_id, i, task)
                     model_server_logger.info("req_id: {0} finished".format(task_id))
@@ -342,7 +230,6 @@ def _process_speculate_output(self):
                     break
             batch_result.append(result)
 
-        # 后处理函数调用
         self.postprocess(batch_result, exist_finished_task)
 
 
@@ -365,29 +252,17 @@ def process_sampling_results(self):
         while self._is_running:
             try:
                 rank_id = 0
-                get_output(self.output_tokens, rank_id, self._is_blocking)
+                if self.is_speculate_decoding:
+                    speculate_get_output(self.output_tokens, rank_id, self._is_blocking)
+                else:
+                    get_output(self.output_tokens, rank_id, self._is_blocking)
 
                 if self.output_tokens[0, 0] == -2:
                     continue
                 self._process_batch_output()
             except Exception as e:
                 model_server_logger.info("while get input_data error: {0} {1}".format(e, str(traceback.format_exc())))
 
-    def process_speculate_results(self):
-        """
-        read tokens from paddle inference engine and process
-        """
-        while self._is_running:
-            try:
-                rank_id = 0
-                speculate_get_output(self.output_tokens, rank_id, self._is_blocking)
-
-                if self.output_tokens[0] == -2:
-                    continue
-                self._process_speculate_output()
-            except Exception as e:
-                model_server_logger.info("while get input_data error: {0} {1}".format(e, str(traceback.format_exc())))
-
     def stop(self):
         """
         stop warm up thread