add visual_send_bs args

Author: sangchengmeng

lightllm/models/qwen2_vl/vision_process.py

@@ -162,7 +162,7 @@ def rescale_and_normalize(
 
     def preprocess(self, image) -> Tuple[torch.Tensor, torch.Tensor]:
         image_arr = np.asarray(image, dtype=np.uint8)
-        image_data = torch.from_numpy(image_arr).permute(2, 0, 1).contiguous().to("cuda", non_blocking=True)
+        image_data = torch.from_numpy(image_arr).permute(2, 0, 1).contiguous()
         grouped_images, grouped_images_index = group_images_by_shape(
             [image_data], disable_grouping=self.disable_grouping
         )

lightllm/server/api_cli.py

@@ -355,6 +355,12 @@ def make_argument_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--visual_infer_batch_size", type=int, default=1, help="number of images to process in each inference batch"
     )
+    parser.add_argument(
+        "--visual_send_batch_size",
+        type=int,
+        default=1,
+        help="number of images embedding to send to llm process in each batch",
+    )
     parser.add_argument(
         "--visual_gpu_ids", nargs="+", type=int, default=None, help="List of GPU IDs to use, e.g., 0 1 2"
     )

lightllm/server/core/objs/start_args_type.py

@@ -77,6 +77,7 @@ class StartArgs:
     grouping_key: List[str] = field(default_factory=list)
     push_interval: int = field(default=10)
     visual_infer_batch_size: int = field(default=1)
+    visual_send_batch_size: int = field(default=10)
     visual_gpu_ids: List[int] = field(default_factory=lambda: [0])
     visual_tp: int = field(default=1)
     visual_dp: int = field(default=1)

lightllm/server/visualserver/manager.py

@@ -1,3 +1,4 @@
+import os
 import zmq
 import zmq.asyncio
 import asyncio
@@ -57,6 +58,7 @@ def __init__(
         self.trust_remote_code = args.trust_remote_code
         self.args = args
         self.visual_model_rpc_ports = visual_model_rpc_ports
+        self.send_batch_size = min(args.visual_send_batch_size, args.cache_capacity)
         self.shm_req_manager = ShmReqManager()
 
     async def wait_to_model_ready(self):
@@ -117,6 +119,18 @@ async def loop_for_fwd(self):
             else:
                 processing_group_reqs = []
                 images_need_infer = []
+                ready_to_send = []
+
+                def flush_ready(force: bool = False):
+                    if not ready_to_send:
+                        return
+                    if not force and len(ready_to_send) < self.send_batch_size:
+                        return
+
+                    for group_req_indexes in ready_to_send:
+                        self.send_to_next_module.send_pyobj(group_req_indexes, protocol=pickle.HIGHEST_PROTOCOL)
+                    ready_to_send.clear()
+
                 while len(self.waiting_reqs) > 0:
                     group_req_indexes = self.waiting_reqs.pop(0)
                     shm_req = self.shm_req_manager.get_req_obj_by_index(group_req_indexes.shm_req_indexes[0])
@@ -146,24 +160,25 @@ async def loop_for_fwd(self):
                         if len(images_need_infer) == self.infer_batch_size:
                             await self.infer_imgs(images_need_infer)
                             images_need_infer = []
-                            for _group_req_indexes in processing_group_reqs:
-                                self.send_to_next_module.send_pyobj(
-                                    _group_req_indexes, protocol=pickle.HIGHEST_PROTOCOL
-                                )
+                            ready_to_send.extend(processing_group_reqs)
                             processing_group_reqs = []
+                            flush_ready(force=False)
 
                     if len(images_need_infer) == 0:
-                        self.send_to_next_module.send_pyobj(group_req_indexes, protocol=pickle.HIGHEST_PROTOCOL)
+                        ready_to_send.append(group_req_indexes)
+                        flush_ready(force=False)
                     else:
                         processing_group_reqs.append(group_req_indexes)
 
                 if len(images_need_infer) > 0:
                     await self.infer_imgs(images_need_infer)
-                    for _group_req_indexes in processing_group_reqs:
-                        self.send_to_next_module.send_pyobj(_group_req_indexes, protocol=pickle.HIGHEST_PROTOCOL)
-                    processing_group_reqs = []
                     images_need_infer = []
 
+                    # 这些处理完 image 的 group 也 ready 了
+                    ready_to_send.extend(processing_group_reqs)
+                    processing_group_reqs = []
+                flush_ready(force=True)
+
     async def loop_for_netio_req(self):
         if not hasattr(self, "visual_recv_max_count"):
             self.visual_recv_max_count = 64