decode for mtp

Author: shihaobai

lightllm/server/router/manager.py

@@ -62,6 +62,8 @@ def __init__(self, args, router_port, detokenization_port, metric_port):
         # 初始化 radix_cache_client 用于读取 prompt cache 的管理信息
         self.radix_cache_client = None
 
+        self.spec_step = args.spec_step
+
         # 共享变量，用于存储router端调度分析得到的机器负载信息
         self.shared_token_load = TokenLoad(f"{get_unique_server_name()}_shared_token_load", self.dp_size_in_node)
         for dp_index in range(self.dp_size_in_node):
@@ -386,8 +388,9 @@ async def _prefill_batch(self, batch: Batch):
         self.overlap_event.set()
         await self.model_rpc_client.prefill(reqs)
         batch.filter_out_finished_req(self.shm_req_manager)
-        # 发个None包触发一下detokenization
-        self.send_to_detokenization.send_pyobj(None, protocol=pickle.HIGHEST_PROTOCOL)
+        # 发spec_step + 1 个 None包触发一下detokenization
+        for _ in range(self.spec_step + 1):
+            self.send_to_detokenization.send_pyobj(None, protocol=pickle.HIGHEST_PROTOCOL)
 
         logger.debug(f"Prefill Batch: {batch.simple_log()} \n")
         self.metric_client.histogram_observe(
@@ -403,8 +406,9 @@ async def _decode_batch(self, batch: Batch):
         # 在 self.is_multinode_and_multidp 为 True 时，传入的 batch 对象可能为 None。
         if batch is not None:
             batch.filter_out_finished_req(self.shm_req_manager)
-        # 发个None包触发一下detokenization
-        self.send_to_detokenization.send_pyobj(None, protocol=pickle.HIGHEST_PROTOCOL)
+        # 发spec_step + 1 个 None包触发一下detokenization
+        for _ in range(self.spec_step + 1):
+            self.send_to_detokenization.send_pyobj(None, protocol=pickle.HIGHEST_PROTOCOL)
         self.metric_client.histogram_observe(
             "lightllm_batch_inference_duration_bucket", time.time() - start_time, "decode"
         )

lightllm/server/router/model_infer/mode_backend/continues_batch/impl_mtp.py

@@ -33,8 +33,6 @@ def update_draft_token_mem_indexes(draft_token_memindex_map, run_reqs, mem_index
 class ContinuesBatchWithMTPBackend(ModeBackend):
     def __init__(self) -> None:
         super().__init__()
-        self.accepted_cnt = 0
-        self.all_cnt = 0
 
     # 支持双模型
     def init_model(self, kvargs):
@@ -83,8 +81,6 @@ def init_model(self, kvargs):
         self.mtp_draft_token_memindex_map = torch.full(
             (max_req_num,), fill_value=IS_NONE, dtype=torch.int32, device="cpu"
         )
-        self.draft_accept_count = torch.zeros((max_req_num,), dtype=torch.int32, device="cpu")
-        self.main_step = 0
 
     def prefill(self, reqs: List[Tuple]):
         self._init_reqs(reqs, init_req_obj=False)
@@ -103,8 +99,6 @@ def decode(self):
                 prefill_reqs, is_chuncked_mode=False, is_multimodal=self.is_multimodal
             )
             model_output = self.model.forward(model_input)
-            self.main_step += 1
-            device0_print(f"main_step: {self.main_step}")
 
             self._overlap_req_init_and_filter(
                 uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True
@@ -134,9 +128,6 @@ def decode(self):
             model_output = self.model.forward(model_input)
             assert model_output.logits.shape[0] % 2 == 0
 
-            self.main_step += 1
-            device0_print(f"main_step: {self.main_step}")
-
             self._overlap_req_init_and_filter(
                 uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True
             )
@@ -165,7 +156,6 @@ def decode(self):
                 is_chuncked_mode=False,
                 do_filter_finished_reqs=False,
             )
-
             # spec decode: MTP
             draft_model_input = copy.deepcopy(model_input)
             draft_model_input.input_ids = torch.tensor(next_token_ids, dtype=torch.int64, device="cuda")
@@ -191,8 +181,6 @@ def verify(self, next_token_ids0, run_reqs):
             if self.draft_token_id_map[req.req_idx] == next_token_ids0[i]:
                 accepted_reqs.append(req)
                 accepted_index.append(i)
-                self.draft_accept_count[req.req_idx] += 1
-                device0_print(f"draft_accept_count: {self.draft_accept_count[req.req_idx]}")
                 self.main_draft_token_memindex_map[req.req_idx] = IS_NONE
             else:
                 need_free_mem_indexes.append(self.main_draft_token_memindex_map[req.req_idx])