use aync load for encoder_server (#15456)

Signed-off-by: liuanqi <liuanqi6@xiaomi.com>
Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com>
Co-authored-by: liuanqi <liuanqi6@xiaomi.com>

python/sglang/srt/disaggregation/encode_server.py

@@ -1,4 +1,5 @@
 import asyncio
+import concurrent.futures
 import ctypes
 import logging
 import multiprocessing as mp
@@ -17,7 +18,6 @@ import zmq.asyncio
 from fastapi import FastAPI
 from fastapi.responses import ORJSONResponse, Response
 from transformers import AutoImageProcessor
-from transformers.image_utils import load_images
 
 from sglang.srt.configs.device_config import DeviceConfig
 from sglang.srt.configs.load_config import LoadConfig
@@ -38,7 +38,14 @@ from sglang.srt.server_args import (
     ServerArgs,
     set_global_server_args_for_scheduler,
 )
-from sglang.srt.utils import get_local_ip_auto, get_zmq_socket, random_uuid
+from sglang.srt.utils import (
+    get_local_ip_auto,
+    get_zmq_socket,
+    load_audio,
+    load_image,
+    load_video,
+    random_uuid,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -46,6 +53,10 @@ rid_lock = asyncio.Lock()
 rid_to_receive_endpoint: Dict[str, List[str]] = dict()
 rid_to_receive_count: Dict[str, int] = dict()
 
+use_image_processor_gpu = (
+    int(os.getenv("SGLANG_ENCODER_IMAGE_PROCESSOR_USE_GPU", "0")) == 1
+)
+
 
 class TensorWrapper:
     """Wrapper to keep tensor alive while exposing buffer for zero-copy."""
@@ -138,6 +149,8 @@ class MMEncoder:
 
         torch.get_device_module(self.device).set_device(self.gpu_id)
 
+        self.use_image_processor_gpu = use_image_processor_gpu
+
         init_distributed_environment(
             world_size=server_args.tp_size,
             rank=rank,
@@ -159,6 +172,10 @@ class MMEncoder:
         self.mm_cache = MultiModalStaticCache(embedding_cache_size * 1024 * 1024)
         self.mm_cache_lock = asyncio.Lock()
 
+        self.io_executor = concurrent.futures.ThreadPoolExecutor(
+            max_workers=int(os.environ.get("SGLANG_ENCODER_MM_LOAD_WORKERS", 4))
+        )
+
         if schedule_path is not None:
             self.schedule_socket = get_zmq_socket(
                 self.context, zmq.PULL, schedule_path, True
@@ -182,10 +199,102 @@ class MMEncoder:
 
         logger.info(f"rank {rank} init finish ")
 
-    async def _encode(self, mm_items) -> torch.Tensor:
-        images = load_images(mm_items)
+    def _load_single_item(
+        self,
+        data,
+        modality: Modality,
+        frame_count_limit=None,
+        audio_sample_rate: Optional[int] = None,
+        discard_alpha_channel=True,
+    ):
+        """
+        Load a single multimodal data.
+        If data is precomputed, returns directly.
+        Static method that can be pickled for multiprocessing"""
+        if isinstance(data, dict):
+            return data
+        try:
+            if modality == Modality.IMAGE:
+                img, _ = load_image(data)
+                if discard_alpha_channel and img.mode != "RGB":
+                    img = img.convert("RGB")
+                return img
+            elif modality == Modality.VIDEO:
+                return load_video(data, frame_count_limit)
+            elif modality == Modality.AUDIO:
+                return load_audio(data, audio_sample_rate)
 
-        images_input = self.image_processor(images=images)
+        except Exception as e:
+            raise RuntimeError(f"Error while loading data {data}: {e}")
+
+    def submit_data_loading_tasks(self, items, modalities):
+        futures = []
+        task_info = []
+
+        for data, modality in zip(items, modalities):
+            if modality is not None:
+                futures.append(
+                    self.io_executor.submit(
+                        self._load_single_item,
+                        data,
+                        modality,
+                    )
+                )
+                task_info.append((modality, data))
+        return futures, task_info
+
+    async def _flatten_and_load_images(self, mm_items):
+        """
+        Flatten mm_items structure, load images concurrently, and restore original structure.
+
+        Returns:
+            Same structure as load_images would return
+        """
+        # Handle single image (not a list)
+        if not isinstance(mm_items, (list, tuple)):
+            futures, _ = self.submit_data_loading_tasks([mm_items], [Modality.IMAGE])
+            return await asyncio.wrap_future(futures[0])
+
+        # Handle nested list (list of lists)
+        if len(mm_items) > 0 and isinstance(mm_items[0], (list, tuple)):
+            # Flatten nested structure
+            flat_data = []
+            flat_indices = []  # Track which group each item belongs to
+            for group_idx, image_group in enumerate(mm_items):
+                for item in image_group:
+                    flat_data.append(item)
+                    flat_indices.append(group_idx)
+
+            # Submit all tasks concurrently
+            futures, _ = self.submit_data_loading_tasks(
+                flat_data, [Modality.IMAGE] * len(flat_data)
+            )
+
+            # Wait for all tasks to complete asynchronously
+            async_futures = [asyncio.wrap_future(f) for f in futures]
+            results = await asyncio.gather(*async_futures)
+
+            # Restore nested structure
+            nested_results = [[] for _ in range(len(mm_items))]
+            for idx, result in zip(flat_indices, results):
+                nested_results[idx].append(result)
+
+            return nested_results
+
+        # Handle simple list
+        else:
+            futures, _ = self.submit_data_loading_tasks(
+                mm_items, [Modality.IMAGE] * len(mm_items)
+            )
+            # Wait for all tasks to complete asynchronously
+            async_futures = [asyncio.wrap_future(f) for f in futures]
+            return await asyncio.gather(*async_futures)
+
+    async def _encode(self, mm_items) -> torch.Tensor:
+        images = await self._flatten_and_load_images(mm_items)
+
+        kwargs = {"device": self.device} if self.use_image_processor_gpu else {}
+        images_input = self.image_processor(images=images, **kwargs)
         feature = images_input["pixel_values"]
         mm_item = MultimodalDataItem.from_dict(
             {