fix chunked prefill

Author: shihaobai

lightllm/server/api_cli.py

@@ -102,7 +102,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
         help="the port for multinode http manager, default is 20000",
     )
     parser.add_argument(
-        "--multinode_router_gloo_port",
+        "--multinode_router_nccl_port",
         type=int,
         default=20001,
         help="the gloo port for multinode router, default is 20001",

lightllm/server/router/manager.py

@@ -77,12 +77,10 @@ def __init__(self, args, router_port, detokenization_port, model_rpc_ports, metr
         self.send_to_detokenization.connect(f"{args.zmq_mode}127.0.0.1:{detokenization_port}")
         self.model_rpc_ports = model_rpc_ports
 
-        self.multinode_req_manager = None
-        self.multinode_req_queue_lock = asyncio.Lock()
         if args.nnodes > 1:
             self.mulitnode_group = dist.init_process_group(
-                backend="gloo",
-                init_method=f"tcp://{args.nccl_host}:{args.multinode_router_gloo_port}",
+                backend="nccl",
+                init_method=f"tcp://{args.nccl_host}:{args.multinode_router_nccl_port}",
                 world_size=args.nnodes,
                 rank=args.node_rank,
             )
@@ -217,8 +215,7 @@ async def add_req(self, group_req_indexes: GroupReqIndexes):
             req_group.append(req)
 
             logger.info(f"router recive req id {req.request_id} cost time {time.time() - req.start_time} s")
-        async with self.multinode_req_queue_lock:
-            self.req_queue.extend(req_group)
+        self.req_queue.extend(req_group)
         self.send_to_detokenization.send_pyobj(group_req_indexes, protocol=pickle.HIGHEST_PROTOCOL)
         return
 

lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py

@@ -54,7 +54,7 @@ def post_handel(self, run_reqs: List[InferReq], next_token_ids, next_token_logpr
 
             req_obj.cur_kv_len = len(req_obj.get_chuncked_input_token_ids())
             if req_obj.cur_kv_len < req_obj.get_cur_total_len():
-                if self.tp_rank < self.dp_size:
+                if self.local_tp_rank < self.dp_size:
                     req_obj.shm_req.shm_cur_kv_len = req_obj.cur_kv_len
                 continue
 
@@ -67,7 +67,7 @@ def post_handel(self, run_reqs: List[InferReq], next_token_ids, next_token_logpr
             if req_obj.finish_status.is_finished() or req_obj.shm_req.router_aborted:
                 finished_req_ids.append(req_obj.shm_req.request_id)
 
-            if self.tp_rank < self.dp_size:
+            if self.local_tp_rank < self.dp_size:
                 # shm_cur_kv_len shm_cur_output_len 是 router 调度进程需要读的信息
                 # finish_token_index finish_status candetoken_out_len 是
                 # detokenization 进程需要的信息，注意这些变量的写入顺序避免异步协同问题。

lightllm/server/router/req_queue/base_queue.py

@@ -30,27 +30,7 @@ def __init__(self, args, router, dp_index, dp_size) -> None:
     @property
     def waiting_req_id_list(self):
         return [req.request_id for req in self.waiting_req_list]
-    
-    def all_contain(self, id_list: List[int]):
-        return all([_id in self.waiting_req_id_list for _id in id_list])
-
-    def arrange(self, id_list: List[int]):
-        id2index_list = {}
-        for index, req in enumerate(self.waiting_req_list):
-           id2index_list[req.request_id] = index
-        return [self.waiting_req_list[id2index_list[id]] for id in id_list]
-    
-    def pop_list(self, req_list):
-        id2index_list = {}
-        for index, req in enumerate(self.waiting_req_list):
-           id2index_list[req.request_id] = index
-        remove_target = []
-        for req in req_list:
-            remove_target.append(self.waiting_req_list[id2index_list[req.request_id]])
-        for req in remove_target:
-            self.waiting_req_list.remove(req)
-        return
-    
+
     def append(self, req: Req):
         req.sample_params.suggested_dp_index = self.dp_index
         self.waiting_req_list.append(req)