From f0763859edbb44aaf9201ce470592b0e4d2cd7db Mon Sep 17 00:00:00 2001 From: Zhonghua Deng Date: Tue, 19 May 2026 11:47:21 +0800 Subject: [PATCH] perf(mimo-v2-epd): enable GPU image preprocess and parallel video decode (#25588) --- .../sglang/srt/managers/tokenizer_manager.py | 1 + .../srt/multimodal/processors/mimo_v2.py | 28 +++++++++- python/sglang/srt/utils/video_decoder.py | 54 ++++++++++++++++--- 3 files changed, 75 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 468947e85..0d63f608a 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/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): diff --git a/python/sglang/srt/multimodal/processors/mimo_v2.py b/python/sglang/srt/multimodal/processors/mimo_v2.py index f8c2c562f..89c2619f3 100644 --- a/python/sglang/srt/multimodal/processors/mimo_v2.py +++ b/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) diff --git a/python/sglang/srt/utils/video_decoder.py b/python/sglang/srt/utils/video_decoder.py index c82842238..0dfec642f 100644 --- a/python/sglang/srt/utils/video_decoder.py +++ b/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