Merge branch 'wzj_router' of https://github.com/ModelTC/lightllm into wzj_router

Author: shihaobai

lightllm/server/api_start.py

@@ -159,6 +159,7 @@ def normal_or_p_d_start(args):
         args.visual_nccl_ports = args.visual_nccl_ports[: args.visual_dp]
 
     if args.disable_chunked_prefill:
+        args.chunked_prefill_size = args.max_req_total_len
         # 普通模式下
         if args.batch_max_tokens is None:
             args.batch_max_tokens = args.max_req_total_len

lightllm/server/httpserver/manager.py

@@ -57,7 +57,6 @@ def __init__(
         self._shm_lock_pool = AtomicShmArrayLock(f"{get_unique_server_name()}_lightllm_resource_lock", 1)
         self._resource_lock = AsyncLock(self._shm_lock_pool.get_lock_context(0))
         self.node_rank = args.node_rank
-        self.transfer_lock = asyncio.Lock()  # the lock for transfer to next module in multi node mode.
         self.disable_abort = args.nnodes > 1 and args.dp == 1  # mulitnode dp=1 mode, disable abort
         self.is_multinode_tp = args.dp == 1 and args.nnodes > 1
         self.is_multinode_tp_master = args.dp == 1 and args.nnodes > 1 and args.node_rank == 0
@@ -202,25 +201,20 @@ def tokens(self, prompt, multimodal_params, samping_params: SamplingParams, kwar
 
     async def loop_for_request(self):
         assert self.args.node_rank > 0
-        tasks = []
-        self.request_order_queue = []
         while True:
             (
                 prompt,
                 sampling_params,
                 multimodal_params,
             ) = await self.multinode_req_manager.recv_pyobj()
-            self.request_order_queue.append(sampling_params.group_request_id)
             results_generator = self.generate(prompt, sampling_params, multimodal_params, None)
 
             async def generate_wrapper(results_generator):
                 async for _, _, _, _ in results_generator:
                     pass
 
-            tasks.append(asyncio.create_task(generate_wrapper(results_generator)))
-            # cleanup
-            while len(tasks) > 0 and tasks[0].done():
-                tasks.pop(0)
+            asyncio.create_task(generate_wrapper(results_generator))
+        return
 
     def alloc_req_id(self, sampling_params, is_health_req: bool = False):
         # 请求的 id 可以由外部传入，也可以由内部生成，但是由外部传入的时候，要自己保证全局唯一性
@@ -413,32 +407,13 @@ async def transfer_to_next_module_or_node(
         original_multimodal_params: MultimodalParams,
         group_req_objs: Optional[GroupReqObjs] = None,
     ):
-        # 多节点纯tp 运行模式下，master 节点需要将请求按照可控的顺序转发给slave节点，
-        # 同时转发给salve节点的时候，要保证master节点按照转发的顺序转发给next_module
-        # 所以需要锁的控制。
+        # 多节点纯tp 运行模式下，master 节点需要将请求转发给slave节点.
         if self.is_multinode_tp_master:
-            async with self.transfer_lock:
-                for sender in self.multinode_req_manager:
-                    sender.send_pyobj(
-                        (prompt, sampling_params, original_multimodal_params),
-                        protocol=pickle.HIGHEST_PROTOCOL,
-                    )
-                await self.transfer_to_next_module(group_req_objs)
-            return
-        # 多节点纯tp 的slave节点，需要按照接受到请求的顺序转发，这需要锁和排队机制来保证。
-        # self.request_order_queue 实现了一种简单的排队取出机制，这样master 和 slave
-        # 节点的请求到达各自节点的router的顺序才是一致的，才能完成同步同态调度。
-        if self.is_multinode_tp_slave:
-            while True:
-                if self.request_order_queue and self.request_order_queue[0] != group_req_objs.group_req_id:
-                    await asyncio.sleep(0.002)
-                    continue
-                else:
-                    async with self.transfer_lock:
-                        await self.transfer_to_next_module(group_req_objs)
-                        self.request_order_queue.pop(0)
-                    break
-            return
+            for sender in self.multinode_req_manager:
+                sender.send_pyobj(
+                    (prompt, sampling_params, original_multimodal_params),
+                    protocol=pickle.HIGHEST_PROTOCOL,
+                )
 
         await self.transfer_to_next_module(group_req_objs)
         return

lightllm/server/router/batch.py

@@ -54,10 +54,10 @@ def filter_out_finished_req(self, shm_req_manager: ShmReqManager):
         self.id_to_reqs = {req.request_id: req for req in self.reqs}
         return
 
-    def pop_req(self, req_id):
+    def pop_req(self, req_id) -> Req:
         self.reqs = [req for req in self.reqs if req.request_id != req_id]
-        self.id_to_reqs.pop(req_id)
-        return
+        req = self.id_to_reqs.pop(req_id)
+        return req
 
     def is_clear(self):
         return len(self.reqs) == 0

lightllm/server/router/manager.py

@@ -18,7 +18,6 @@
 from lightllm.server.core.objs.io_objs import GroupReqIndexes, AbortedReqCmd
 from lightllm.server.core.objs import ShmReqManager, StartArgs
 from .dynamic_prompt.radix_cache import RadixCacheReadOnlyClient
-from .stats import Stats
 from .shm_reqs_io_buffer import ShmReqsIOBuffer
 from lightllm.utils.log_utils import init_logger, log_time_ready
 from lightllm.server.router.token_load import TokenLoad
@@ -45,6 +44,8 @@ def __init__(self, args: StartArgs, router_port, detokenization_port, metric_por
         # 兼容多机纯tp的运行模式，这时候 1 // 2 == 0, 需要兼容
         self.dp_size_in_node = max(1, args.dp // self.nnodes)
         self.is_multinode_tp = args.nnodes > 1 and args.dp == 1
+        self.is_multinode_tp_master = self.is_multinode_tp and args.node_rank == 0
+        self.is_multinode_tp_slave = self.is_multinode_tp and args.node_rank != 0
         self.is_multinode_and_multidp = args.nnodes > 1 and args.dp > 1
         # 判断是否是保守调度，保守调度不会发生暂停 req 的情况，但是有些场景可能影响吞吐
         self.is_safe_schedule = args.router_token_ratio == 0.0
@@ -254,6 +255,8 @@ async def _step(self):
         """
         事件处理循环
         """
+        # 接受新请求，并尝试调度
+        await self._recv_new_reqs_and_schedule()
         # 判断是否有新请求加入推理
         # 激进调度满足，有新的推理batch就需要进行加入。
         # 或者延迟step的步数满足了当前条件，也需要进行新的推理batch的加入。
@@ -357,44 +360,96 @@ def _add_req(self, group_req_indexes: GroupReqIndexes):
         return
 
     def _generate_new_batch(self):
-        limit_router_queue_length = None
-        if self.is_multinode_tp:
-            # 使用 all_reduce 获取最小值
-            limit_router_queue_length = len(self.req_queue.waiting_req_list)
-            limit_router_queue_length_tensor = torch.tensor(limit_router_queue_length, dtype=torch.int32, device="cpu")
-            dist.all_reduce(limit_router_queue_length_tensor, op=dist.ReduceOp.MIN, group=self.mulitnode_group)
-            limit_router_queue_length = limit_router_queue_length_tensor.item()
-
         # 调度的时候需要考虑当前运行的batch，和调度了但是暂时还没有推理的部分请求。
         new_batch = self.req_queue.generate_new_batch(
-            Batch.merge_two_batch(self.running_batch, self.schedule_new_batch), limit_router_queue_length
+            Batch.merge_two_batch(self.running_batch, self.schedule_new_batch)
         )
         self.schedule_new_batch = Batch.merge_two_batch(self.schedule_new_batch, new_batch)
         return
 
-    async def loop_for_netio_req(self):
-        recv_max_count = 64
+    def _multinode_tp_generate_new_batch(self):
+        dist.barrier(group=self.mulitnode_group)
 
-        while True:
-            try:
-                # 一次最多从 zmq 中取 recv_max_count 个请求，防止 zmq 队列中请求数量过多导致阻塞了主循环。
-                for _ in range(recv_max_count):
-                    recv_req: GroupReqIndexes = self.recv_from_httpserver.recv_pyobj(zmq.NOBLOCK)
-                    if isinstance(recv_req, GroupReqIndexes):
-                        self._add_req(recv_req)
-                    else:
-                        assert False, f"Error Req Inf {recv_req}"
-
-                # 当队列中存在较多的请求时，将一次接受的数量上调
-                recv_max_count = min(int(recv_max_count * 1.3), 256)
-
-            except zmq.ZMQError:
-                # 当队列已经开始清空的时候，将一次接受的数量下调
-                recv_max_count = 64
+        # 调度的时候需要考虑当前运行的batch，和调度了但是暂时还没有推理的部分请求。
+        if self.is_multinode_tp_master:
+            new_batch = self.req_queue.generate_new_batch(
+                Batch.merge_two_batch(self.running_batch, self.schedule_new_batch)
+            )
+            if new_batch is not None:
+                req_ids = [req.request_id for req in new_batch.reqs]
+            else:
+                req_ids = []
+            dist.broadcast_object_list([len(req_ids)], src=0, group=self.mulitnode_group)
+            dist.broadcast_object_list(req_ids, src=0, group=self.mulitnode_group)
+            req_id_select_mark = [1 for _ in range(len(req_ids))]
+            req_id_select_mark = torch.tensor(req_id_select_mark, dtype=torch.int32, device="cpu")
+            dist.all_reduce(req_id_select_mark, op=dist.ReduceOp.MIN, group=self.mulitnode_group)
+            back_req_list = []
+            for req_id, select in zip(req_ids, req_id_select_mark.numpy()):
+                if select == 0:
+                    req = new_batch.pop_req(req_id)
+                    back_req_list.append(req)
+            self.req_queue.waiting_req_list = back_req_list + self.req_queue.waiting_req_list
+            if new_batch.is_clear():
+                new_batch = None
+        else:
+            req_nums = [None]
+            dist.broadcast_object_list(req_nums, src=0, group=self.mulitnode_group)
+            req_num = req_nums[0]
+            req_ids = [None for _ in range(req_num)]
+            dist.broadcast_object_list(req_ids, src=0, group=self.mulitnode_group)
+            all_req_id_set = set([req.request_id for req in self.req_queue.waiting_req_list])
+            req_id_select_mark = []
+            for req_id in req_ids:
+                req_id_select_mark.append(1 if req_id in all_req_id_set else 0)
+            req_id_select_mark = torch.tensor(req_id_select_mark, dtype=torch.int32, device="cpu")
+            dist.all_reduce(req_id_select_mark, op=dist.ReduceOp.MIN, group=self.mulitnode_group)
+            select_req_ids = []
+            for req_id, select in zip(req_ids, req_id_select_mark.numpy()):
+                if select == 1:
+                    select_req_ids.append(req_id)
+
+            select_reqs = []
+            for req_id in select_req_ids:
+                for req in self.req_queue.waiting_req_list:
+                    if req.request_id == req_id:
+                        select_reqs.append(req)
+
+            for req in select_reqs:
+                self.req_queue.waiting_req_list.remove(req)
+            if select_reqs:
+                new_batch = Batch(-1, reqs=select_reqs, dp_size_in_node=self.dp_size_in_node)
+            else:
+                new_batch = None
 
-            await asyncio.sleep(0.02)
+        self.schedule_new_batch = Batch.merge_two_batch(self.schedule_new_batch, new_batch)
+
+        dist.barrier(group=self.mulitnode_group)
+        return
+
+    async def _recv_new_reqs_and_schedule(self):
+        if not hasattr(self, "recv_max_count"):
+            self.recv_max_count = 64
+
+        try:
+            # 一次最多从 zmq 中取 recv_max_count 个请求，防止 zmq 队列中请求数量过多导致阻塞了主循环。
+            for _ in range(self.recv_max_count):
+                recv_req: GroupReqIndexes = self.recv_from_httpserver.recv_pyobj(zmq.NOBLOCK)
+                if isinstance(recv_req, GroupReqIndexes):
+                    self._add_req(recv_req)
+                else:
+                    assert False, f"Error Req Inf {recv_req}"
 
-            # 只有当推理侧没有发生暂停的时候，才执行新的调度
+            # 当队列中存在较多的请求时，将一次接受的数量上调
+            self.recv_max_count = min(int(self.recv_max_count * 1.3), 256)
+
+        except zmq.ZMQError:
+            # 当队列已经开始清空的时候，将一次接受的数量下调
+            self.recv_max_count = 64
+
+        if self.is_multinode_tp:
+            self._multinode_tp_generate_new_batch()
+        else:
             if self._get_paused_req_num() == 0:
                 self._generate_new_batch()
         return
@@ -436,6 +491,5 @@ def handle_exception(loop, context):
         raise
 
     pipe_writer.send("init ok")
-    loop.create_task(router.loop_for_fwd())
-    loop.run_until_complete(router.loop_for_netio_req())
+    loop.run_until_complete(router.loop_for_fwd())
     return

lightllm/server/router/model_infer/mode_backend/base_backend.py

@@ -48,7 +48,7 @@ def __init__(self) -> None:
         self.enable_decode_microbatch_overlap = get_env_start_args().enable_decode_microbatch_overlap
         self.enable_prefill_microbatch_overlap = get_env_start_args().enable_prefill_microbatch_overlap
 
-        # 控制分类的参数变量
+        # 控制 _get_classed_reqs 分类的参数变量，不同的 backend 具有可能需要不同的分类运行条件。
         self.classed_req_no_decode = False
         self.classed_req_strict_prefill = False
         pass
@@ -74,6 +74,7 @@ def init_model(self, kvargs):
         self.use_dynamic_prompt_cache = not self.args.disable_dynamic_prompt_cache
         self.eos_id: List[int] = kvargs.get("eos_id", [2])
         self.disable_cudagraph = self.args.disable_cudagraph
+        self.is_multinode_tp = self.args.nnodes > 1 and self.args.dp == 1
 
         self.logger = init_logger(__name__)
 
@@ -166,17 +167,29 @@ def init_model(self, kvargs):
                 [0 for _ in range(self.global_world_size)], dtype=torch.int32, device="cuda", requires_grad=False
             )
 
+        # 用于协同读取 ShmReqsIOBuffer 中的请求信息的通信tensor和通信组对象。
         self.node_broadcast_tensor = torch.tensor([0], dtype=torch.int32, device="cuda", requires_grad=False)
         self.node_nccl_group = create_new_group_for_current_node("nccl")
 
+        # 用于在多节点tp模式下协同读取 ShmReqsIOBuffer 中的请求信息的通信tensor和通信组对象。
+        if self.is_multinode_tp:
+            self.multinode_tp_gather_item_tensor = torch.tensor([0], dtype=torch.int32, device="cuda")
+            self.multinode_tp_all_gather_tensor = torch.tensor(
+                [0 for _ in range(self.global_world_size)], dtype=torch.int32, device="cuda", requires_grad=False
+            )
+            self.multinode_tp_nccl_group = dist.new_group(
+                [rank for rank in range(self.global_world_size)], backend="nccl"
+            )
+
         self.init_custom()
         self.shm_reqs_io_buffer = ShmReqsIOBuffer()
 
         # 开启 mtp 模式，需要完成mtp model的初始化
         if self.args.mtp_mode:
             self.init_mtp_draft_model(kvargs)
 
-        # 启动infer_loop_thread
+        # 启动infer_loop_thread, 启动两个线程进行推理，对于具备双batch推理折叠得场景
+        # 可以降低 cpu overhead，大幅提升gpu得使用率。
         self.infer_loop_thread = threading.Thread(target=self.infer_loop, daemon=True)
         self.infer_loop_thread.start()
         self.infer_loop_thread1 = threading.Thread(target=self.infer_loop, daemon=True)
@@ -239,6 +252,13 @@ def init_mtp_draft_model(self, main_kvargs: dict):
         return
 
     def _try_read_new_reqs(self):
+        if self.is_multinode_tp:
+            self._try_read_new_reqs_multinode_tp()
+        else:
+            self._try_read_new_reqs_normal()
+        return
+
+    def _try_read_new_reqs_normal(self):
         if self.is_master_in_node:
             if self.shm_reqs_io_buffer.is_ready():
                 self.node_broadcast_tensor.fill_(1)
@@ -247,16 +267,42 @@ def _try_read_new_reqs(self):
         dist.broadcast(self.node_broadcast_tensor, src=0, group=self.node_nccl_group, async_op=False)
         new_buffer_is_ready = self.node_broadcast_tensor.detach().item()
         if new_buffer_is_ready:
-            cmds: List = self.shm_reqs_io_buffer.read_obj()
-            self.shm_reqs_io_buffer.sub_state()
-            if cmds:
-                if isinstance(cmds[0], AbortedReqCmd):
-                    for obj in cmds:
-                        if obj.req_id in g_infer_context.requests_mapping:
-                            req: InferReq = g_infer_context.requests_mapping[obj.req_id]
-                            req.infer_aborted = True
-                else:
-                    self._init_reqs(reqs=cmds)
+            self._read_reqs_buffer_and_init_reqs()
+        return
+
+    def _try_read_new_reqs_multinode_tp(self):
+        """
+        多节点tp模式下,需要协调所有rank的行为同步。
+        """
+        if self.shm_reqs_io_buffer.is_ready():
+            self.multinode_tp_gather_item_tensor.fill_(1)
+        else:
+            self.multinode_tp_gather_item_tensor.fill_(0)
+        dist.all_gather_into_tensor(
+            self.multinode_tp_all_gather_tensor,
+            self.multinode_tp_gather_item_tensor,
+            group=self.multinode_tp_nccl_group,
+            async_op=False,
+        )
+        new_buffer_is_readys = self.multinode_tp_all_gather_tensor.detach().cpu().numpy()
+        new_buffer_is_ready = np.all(new_buffer_is_readys == 1)
+
+        if new_buffer_is_ready:
+            self._read_reqs_buffer_and_init_reqs()
+        return
+
+    def _read_reqs_buffer_and_init_reqs(self):
+        cmds: List = self.shm_reqs_io_buffer.read_obj()
+        self.shm_reqs_io_buffer.sub_state()
+        if cmds:
+            if isinstance(cmds[0], AbortedReqCmd):
+                for obj in cmds:
+                    obj: AbortedReqCmd = obj
+                    if obj.req_id in g_infer_context.requests_mapping:
+                        req: InferReq = g_infer_context.requests_mapping[obj.req_id]
+                        req.infer_aborted = True
+            else:
+                self._init_reqs(reqs=cmds)
         return
 
     # 一些可以复用的通用功能函数