perf(mimo-v2-epd): enable GPU image preprocess and parallel video decode (#25588)

python/sglang/srt/managers/tokenizer_manager.py

@@ -453,6 +453,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
             self.mm_receiver = create_mm_receiver(
                 self.server_args,
                 dtype=self.model_config.dtype,
+                hf_config=self.model_config.hf_config,
             )
 
     def init_metric_collector_watchdog(self):

python/sglang/srt/multimodal/processors/mimo_v2.py

@@ -359,11 +359,13 @@ class MiMoProcessor:
         pad_token_id=None,
         rope_type="rope",
         video_process_num_threads=16,
+        video_decode_num_threads=0,
         device=None,
         **kwargs,
     ):
         self.tokenizer = tokenizer
         self.video_process_num_threads = video_process_num_threads
+        self.video_decode_num_threads = video_decode_num_threads
 
         if device is None:
             self.device = None
@@ -546,6 +548,20 @@ class MiMoProcessor:
         if "sampling_rate" in audio_cfg and "audio_sampling_rate" not in kwargs:
             kwargs["audio_sampling_rate"] = audio_cfg["sampling_rate"]
 
+        image_cfg = (mm_config or {}).get("image", {})
+        if "device" in image_cfg:
+            kwargs["device"] = image_cfg["device"]
+
+        video_cfg = (mm_config or {}).get("video", {})
+        if "video_decode_num_threads" in video_cfg:
+            kwargs["video_decode_num_threads"] = video_cfg["video_decode_num_threads"]
+        else:
+            from sglang.srt.utils.common import get_int_env_var
+
+            kwargs["video_decode_num_threads"] = get_int_env_var(
+                "SGLANG_ENCODER_VIDEO_DECODE_NUM_THREADS", 0
+            )
+
         kwargs.update(overrides)
         return cls(**kwargs)
 
@@ -591,7 +607,11 @@ class MiMoProcessor:
                     f"Unsupported video input type for EPD encoder: {type(video_data)}"
                 )
 
-        vdw = VideoDecoderWrapper(video_blob, device="cpu")
+        vdw = VideoDecoderWrapper(
+            video_blob,
+            device="cpu",
+            num_decode_threads=self.video_decode_num_threads,
+        )
         try:
             video_tuple = _decode_frames_and_timestamps(
                 vdw, self.default_video_processor_kwargs
@@ -616,7 +636,11 @@ class MiMoProcessor:
             all_patches, all_grids = [], []
             for img in mm_data:
                 img_tensor, _, _ = self.get_visual_transform(
-                    img, factor=factor, min_pixels=min_pixels, max_pixels=max_pixels
+                    img,
+                    factor=factor,
+                    min_pixels=min_pixels,
+                    max_pixels=max_pixels,
+                    device=self.device,
                 )
                 patches, grid = self._flatten_visual_inputs(img_tensor, "image")
                 all_patches.append(patches)

python/sglang/srt/utils/video_decoder.py

@@ -1,6 +1,7 @@
 """Unified video decoder: torchcodec preferred, decord as fallback."""
 
 import logging
+import os
 
 import numpy as np
 
@@ -38,10 +39,16 @@ class VideoDecoderWrapper:
     All frames are returned in NHWC uint8 numpy format for consistency.
     """
 
-    def __init__(self, source, device: str = "cpu"):
+    def __init__(self, source, device: str = "cpu", num_decode_threads: int = 0):
         """source: file path (str) or video bytes.
         device: "cpu" or "cuda". GPU decoding only supported with torchcodec.
+        num_decode_threads: number of parallel decoder instances for frame
+            extraction (torchcodec only). 0 = auto (capped at 16),
+            1 = single decoder. Set > 1 to split frame indices across
+            multiple decoders in parallel threads.
         """
+        self._source = source
+        self._num_decode_threads = num_decode_threads
         self._source_bytes = source if isinstance(source, bytes) else None
         self._source_path = source if isinstance(source, str) else None
         self._tmp_path = None
@@ -49,12 +56,14 @@ class VideoDecoderWrapper:
             kwargs = {"dimension_order": "NHWC"}
             if device == "cuda" and _try_cuda_backend():
                 kwargs["device"] = "cuda"
+            self._tc_kwargs = kwargs
             try:
                 self._decoder = VideoDecoder(source, **kwargs)
             except RuntimeError:
                 if "device" in kwargs:
                     logger.warning("CUDA video decoding failed, falling back to CPU.")
                     kwargs.pop("device")
+                    self._tc_kwargs = kwargs
                     self._decoder = VideoDecoder(source, **kwargs)
                 else:
                     raise
@@ -62,7 +71,6 @@ class VideoDecoderWrapper:
             from decord import VideoReader, cpu
 
             if isinstance(source, bytes):
-                import os
                 import tempfile
 
                 fd, tmp_path = tempfile.mkstemp(suffix=".mp4")
@@ -105,6 +113,18 @@ class VideoDecoderWrapper:
         """Return frames at given indices as a torch tensor (NHWC, uint8, pinned memory)."""
         import torch
 
+        if (
+            _BACKEND == "torchcodec"
+            and self._num_decode_threads != 1
+            and len(indices) > 1
+        ):
+            num_threads = self._num_decode_threads
+            if num_threads <= 0:
+                num_threads = min(os.cpu_count() or 8, 16)
+            num_threads = min(num_threads, len(indices))
+            if num_threads > 1:
+                return self._parallel_decode(indices, num_threads)
+
         if _BACKEND == "torchcodec":
             batch = self._decoder.get_frames_at(indices)
             return batch.data.pin_memory()
@@ -112,6 +132,32 @@ class VideoDecoderWrapper:
             arr = self._decoder.get_batch(indices).asnumpy()
             return torch.from_numpy(arr).pin_memory()
 
+    def _parallel_decode(self, indices, num_threads):
+        """Decode frames using multiple VideoDecoder instances in parallel threads."""
+        from concurrent.futures import ThreadPoolExecutor, as_completed
+
+        import torch
+
+        chunks = [list(c) for c in np.array_split(indices, num_threads) if len(c) > 0]
+        source = self._source
+        kwargs = self._tc_kwargs
+
+        def _decode_chunk(chunk):
+            d = VideoDecoder(source, **kwargs)
+            return d.get_frames_at(chunk).data
+
+        with ThreadPoolExecutor(max_workers=len(chunks)) as executor:
+            future_to_idx = {
+                executor.submit(_decode_chunk, chunk): idx
+                for idx, chunk in enumerate(chunks)
+            }
+            results = [None] * len(chunks)
+            for future in as_completed(future_to_idx):
+                idx = future_to_idx[future]
+                results[idx] = future.result()
+
+        return torch.cat(results, dim=0).pin_memory()
+
     @property
     def source_bytes(self) -> bytes | None:
         """Return raw video bytes if available (needed for audio extraction)."""
@@ -119,8 +165,6 @@ class VideoDecoderWrapper:
             return self._source_bytes
         path = self._tmp_path or self._source_path
         if path is not None:
-            import os
-
             if os.path.isfile(path):
                 with open(path, "rb") as f:
                     return f.read()
@@ -129,8 +173,6 @@ class VideoDecoderWrapper:
     def close(self):
         """Explicitly clean up temporary files."""
         if self._tmp_path is not None:
-            import os
-
             if os.path.exists(self._tmp_path):
                 os.unlink(self._tmp_path)
             self._tmp_path = None