From 9701cc138cc80063a911139f82d16083226faa9e Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 12 Aug 2026 20:02:49 +0800 Subject: [PATCH] [diffusion] optimize: optimize bit-exact h3 reference video ingress (#34563) --- .../minimax_h3/reference_encoding.py | 217 ++++++++++++++++-- .../minimax_h3/stages/text_encoding.py | 31 ++- .../minimax_h3/stages/visual_encoding.py | 12 +- .../test/unit/test_minimax_h3_media.py | 147 +++++++++++- 4 files changed, 381 insertions(+), 26 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/reference_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/reference_encoding.py index 311581398..155d5ff4a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/reference_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/reference_encoding.py @@ -19,6 +19,11 @@ from __future__ import annotations import functools import math +import mmap +import os +import socket +import subprocess +import sys from typing import Any import torch @@ -29,6 +34,7 @@ from sglang.multimodal_gen.configs.models.vaes.minimax_h3_audio import ( from sglang.multimodal_gen.configs.models.vaes.minimax_h3_video import ( MiniMaxH3VideoVAEArchConfig, ) +from sglang.multimodal_gen.runtime.distributed.parallel_state import get_world_group from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import ( MINIMAX_H3_SUPPORTED_FPS, ) @@ -210,8 +216,6 @@ def _load_waveform( temporary lossless file plus a second decode. """ - import subprocess - import numpy as np if max_duration_seconds is not None: @@ -361,16 +365,15 @@ def minimax_h3_decode_reference_video_frames( target_frame_count: int, fps: float = MINIMAX_H3_SUPPORTED_FPS, start_time_seconds: float = 0.0, + share_across_replicas: bool = False, ) -> Any: """Decode, transform, and truncate a reference video in one ffmpeg pass. ffmpeg applies display rotation, CFR sampling, direct Lanczos scaling, and - square-pixel normalization before writing bounded RGB24 frames to stdout. + square-pixel normalization before writing a bounded RGB24 stream. The returned array is shared by Qwen and the visual VAE, so conditioning never passes through a lossy x264 intermediate or a second video decode. """ - import subprocess - import numpy as np if target_frame_count <= 0: @@ -407,30 +410,198 @@ def minimax_h3_decode_reference_video_frames( "rawvideo", "-pix_fmt", "rgb24", - "pipe:1", ] - decoded = subprocess.run( - command, - check=True, - capture_output=True, - ) - payload = decoded.stdout - if not isinstance(payload, bytes): - raise TypeError("ffmpeg RGB24 output must be bytes") frame_bytes = target_width * target_height * 3 - if len(payload) % frame_bytes: + if share_across_replicas: + payload, payload_size = _decode_reference_video_shared(command) + else: + payload, payload_size = _decode_reference_video_local(command) + + if payload_size <= 0: + raise ValueError(f"reference video has no frames: {video_path}") + if payload_size % frame_bytes: raise ValueError( "ffmpeg returned a partial reference-video frame: " - f"{len(payload)} bytes for {target_width}x{target_height} RGB24" + f"{payload_size} bytes for {target_width}x{target_height} RGB24" ) - frame_count = len(payload) // frame_bytes - if frame_count <= 0: - raise ValueError(f"reference video has no frames: {video_path}") + frame_count = payload_size // frame_bytes return np.frombuffer(payload, dtype=np.uint8).reshape( frame_count, target_height, target_width, 3 ) +def _decode_reference_video_local(command: list[str]) -> tuple[Any, int]: + """Write one worker's RGB stream without a large stdout aggregation.""" + + # Linux workers can let ffmpeg write the exact RGB24 stream into an + # anonymous file descriptor. Mapping that output avoids communicate()'s + # chunk list and final bytes join for a several-hundred-MiB reference. + output_fd = -1 + if sys.platform.startswith("linux"): + try: + output_fd = os.memfd_create( + "sglang-h3-reference-video", + flags=os.MFD_CLOEXEC, + ) + except OSError: + output_fd = -1 + + payload: Any = b"" + payload_size = 0 + if output_fd >= 0: + try: + payload_size = _write_reference_video_to_fd(command, output_fd) + if payload_size > 0: + payload = mmap.mmap( + output_fd, + payload_size, + access=mmap.ACCESS_WRITE, + ) + finally: + os.close(output_fd) + else: + decoded = subprocess.run( + [*command, "pipe:1"], + check=True, + capture_output=True, + ) + payload = decoded.stdout + if not isinstance(payload, bytes): + raise TypeError("ffmpeg RGB24 output must be bytes") + payload_size = len(payload) + return payload, payload_size + + +def _write_reference_video_to_fd(command: list[str], output_fd: int) -> int: + subprocess.run( + [*command, f"pipe:{output_fd}"], + check=True, + pass_fds=(output_fd,), + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + ) + return os.lseek(output_fd, 0, os.SEEK_CUR) + + +def _all_gather_world_objects(group: Any, value: Any) -> list[Any]: + values = [None] * group.world_size + torch.distributed.all_gather_object( + values, + value, + group=group.cpu_group, + ) + return values + + +@functools.lru_cache(maxsize=1) +def _reference_video_host_leader() -> int: + group = get_world_group() + hostnames = _all_gather_world_objects(group, socket.gethostname()) + return hostnames.index(hostnames[group.rank_in_group]) + + +def _decode_reference_video_shared(command: list[str]) -> tuple[Any, int]: + """Decode once per host and map the same RGB pages on its worker ranks.""" + group = get_world_group() + if ( + group.world_size <= 1 + or not sys.platform.startswith("linux") + or not os.path.isdir("/proc/self/fd") + ): + return _decode_reference_video_local(command) + + leader = _reference_video_host_leader() + is_leader = group.rank_in_group == leader + + leader_fd = -1 + payload_size = 0 + owner_exception = None + leader_state = None + if is_leader: + try: + try: + leader_fd = os.memfd_create( + "sglang-h3-reference-video-shared", + flags=os.MFD_CLOEXEC, + ) + except OSError: + # Anonymous file descriptors can be disabled by a container's + # seccomp policy. Tell every host to use the unchanged local + # decode path instead of failing a valid request. + leader_state = (None, 0, None) + else: + payload_size = _write_reference_video_to_fd(command, leader_fd) + leader_state = ( + f"/proc/{os.getpid()}/fd/{leader_fd}", + payload_size, + None, + ) + except Exception as exc: + owner_exception = exc + leader_state = (None, 0, f"{type(exc).__name__}: {exc}") + + states = _all_gather_world_objects(group, leader_state) + host_states = [state for state in states if state is not None] + owner_error = next( + (state[2] for state in host_states if state[2] is not None), + None, + ) + # Every rank makes the same decision here. In particular, a decode failure + # on one host must not leave the other hosts entering the mapping collective. + if owner_error is not None: + if leader_fd >= 0: + os.close(leader_fd) + if owner_exception is not None and owner_error.endswith(str(owner_exception)): + raise owner_exception + raise RuntimeError(f"MiniMax H3 shared video decode failed: {owner_error}") + if any(state[0] is None for state in host_states): + if leader_fd >= 0: + os.close(leader_fd) + return _decode_reference_video_local(command) + if any(state[1] <= 0 for state in host_states): + if leader_fd >= 0: + os.close(leader_fd) + return b"", 0 + + state = states[leader] + if state is None: + raise RuntimeError("MiniMax H3 shared video decode returned no descriptor") + local_path, payload_size, _ = state + if payload_size <= 0: + if leader_fd >= 0: + os.close(leader_fd) + return b"", 0 + if local_path is None: + raise RuntimeError("MiniMax H3 shared video decode returned no path") + + mapping = None + map_error = None + try: + map_fd = os.open(local_path, os.O_RDWR) + try: + mapping = mmap.mmap(map_fd, payload_size, access=mmap.ACCESS_COPY) + finally: + os.close(map_fd) + except Exception as exc: + map_error = f"{type(exc).__name__}: {exc}" + + map_errors = _all_gather_world_objects(group, map_error) + failed = next((error for error in map_errors if error is not None), None) + if is_leader: + os.close(leader_fd) + if failed is not None: + if mapping is not None: + mapping.close() + # /proc fd traversal can be denied by hidepid or a container policy. + # All ranks fall back together so the optimization never makes a + # previously valid request fail or changes its RGB bytes. + return _decode_reference_video_local(command) + + if mapping is None: + raise RuntimeError("MiniMax H3 shared video mapping returned no payload") + return mapping, payload_size + + MINIMAX_H3_REFERENCE_VIDEO_ENCODE_SEED = 42 MINIMAX_H3_REFERENCE_VIDEO_PATCH_SIZE = (1, 2, 2) @@ -569,7 +740,12 @@ def _reference_video_target_frame_count( ) -def minimax_h3_prepared_reference_videos(batch: Any, plan: Any) -> dict[str, Any]: +def minimax_h3_prepared_reference_videos( + batch: Any, + plan: Any, + *, + share_across_replicas: bool = False, +) -> dict[str, Any]: """Decode the bounded reference-video RGB frames once per request. BOTH the visual-condition tokenizer and Qwen consume the same transformed @@ -629,6 +805,7 @@ def minimax_h3_prepared_reference_videos(batch: Any, plan: Any) -> dict[str, Any target_frame_count=target_frames, fps=float(plan.shape["fps"]), start_time_seconds=float(material.start_time_seconds), + share_across_replicas=share_across_replicas, ) prepared_videos.append( { diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/text_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/text_encoding.py index 007d9cc53..0180ca254 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/text_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/text_encoding.py @@ -6,8 +6,12 @@ import torch from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import ( + MINIMAX_H3_PREPARED_REFERENCE_VIDEO_EXTRA_KEY, MINIMAX_H3_TEXT_EMBEDDINGS_EXTRA_KEY, ) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.stages.replica_broadcast import ( + minimax_h3_replica_ctx, +) from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.task_profiles import ( MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES, ) @@ -19,6 +23,8 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) +_MINIMAX_H3_SINGLE_RANK_TEXT_ENCODE_EXTRA_KEY = "minimax_h3_single_rank_text_encode" + class MiniMaxH3TextEncodingStage(TextEncodingStage): deduplicated_output_fields = ("prompt_embeds", "prompt_seq_lens") @@ -54,6 +60,7 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage): minimax_h3_cleanup_temp_dirs, ) + batch.extra.pop(MINIMAX_H3_PREPARED_REFERENCE_VIDEO_EXTRA_KEY, None) minimax_h3_cleanup_temp_dirs(batch) raise return batch @@ -135,7 +142,15 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage): payload = None if dp_group.rank_in_group == owner: try: - first_result = self(first_batch, server_args) + first_batch.extra[_MINIMAX_H3_SINGLE_RANK_TEXT_ENCODE_EXTRA_KEY] = ( + True + ) + try: + first_result = self(first_batch, server_args) + finally: + first_batch.extra.pop( + _MINIMAX_H3_SINGLE_RANK_TEXT_ENCODE_EXTRA_KEY, None + ) payload = first_result.extra.get( MINIMAX_H3_TEXT_EMBEDDINGS_EXTRA_KEY ) @@ -360,7 +375,17 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage): in ("video.reference_preserve", "video_audio.reference_preserve") for material in plan.materials ): - prepared_videos = minimax_h3_prepared_reference_videos(batch, plan) + world, _ = minimax_h3_replica_ctx() + prepared_videos = minimax_h3_prepared_reference_videos( + batch, + plan, + share_across_replicas=( + world > 1 + and not bool( + batch.extra.get(_MINIMAX_H3_SINGLE_RANK_TEXT_ENCODE_EXTRA_KEY) + ) + ), + ) video_has_audio: dict[int, bool] = {} for video_index, item in enumerate((prepared_videos or {}).get("videos") or []): if item.get("condition_index") is None: @@ -512,6 +537,8 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage): pixel_values_videos=pixel_values_videos, video_grid_thw=video_grid_thw, ) + if batch.extra.get(_MINIMAX_H3_SINGLE_RANK_TEXT_ENCODE_EXTRA_KEY): + batch.extra.pop(MINIMAX_H3_PREPARED_REFERENCE_VIDEO_EXTRA_KEY, None) return { "positive": { "hidden_states": pos_hidden, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/visual_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/visual_encoding.py index d94f39979..8f907e909 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/visual_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/visual_encoding.py @@ -13,6 +13,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.condition_encoding impo ) from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import ( MINIMAX_H3_KEYFRAME_COND_ROWS_EXTRA_KEY, + MINIMAX_H3_PREPARED_REFERENCE_VIDEO_EXTRA_KEY, MINIMAX_H3_REFERENCE_IMAGE_ROWS_EXTRA_KEY, MINIMAX_H3_REFERENCE_VIDEO_ROWS_EXTRA_KEY, ) @@ -68,6 +69,11 @@ class MiniMaxH3VisualEncodingStage(ConditionEncodingStage): # No later stage will run after an encoder failure. minimax_h3_cleanup_temp_dirs(batch) raise + finally: + # Text and visual encoding are the only full-RGB consumers. Release + # the shared mapping on both success and failure before later stages + # can retain a several-hundred-MiB request-local view. + batch.extra.pop(MINIMAX_H3_PREPARED_REFERENCE_VIDEO_EXTRA_KEY, None) def _forward(self, batch: Req, server_args: ServerArgs) -> Req: from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.resolved_plan import ( @@ -283,7 +289,11 @@ class MiniMaxH3VisualEncodingStage(ConditionEncodingStage): if MINIMAX_H3_REFERENCE_VIDEO_ROWS_EXTRA_KEY in batch.extra: return - prepared = minimax_h3_prepared_reference_videos(batch, plan) + prepared = minimax_h3_prepared_reference_videos( + batch, + plan, + share_across_replicas=bool(self.video_vae.parallel_tiling), + ) videos = prepared.get("videos") if not isinstance(videos, list) or not videos: raise ValueError( diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_media.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_media.py index c40ce302f..36ecf9e21 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_media.py +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_media.py @@ -2,10 +2,13 @@ """Numerical boundaries for the one-pass Ref2VA media path.""" import json +import os import subprocess +import sys from types import SimpleNamespace import numpy as np +import pytest import torch from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3 import ( @@ -51,9 +54,14 @@ def test_video_transform_runs_once_and_qwen_samples_shared_rgb(monkeypatch): expected = np.arange(25 * 4 * 6 * 3, dtype=np.uint8).reshape(25, 4, 6, 3) commands = [] - def run(command, **_kwargs): + def run(command, **kwargs): commands.append(command) - return SimpleNamespace(stdout=expected.tobytes()) + if command[-1] == "pipe:1": + return SimpleNamespace(stdout=expected.tobytes()) + output_fd = int(command[-1].removeprefix("pipe:")) + assert kwargs["pass_fds"] == (output_fd,) + os.write(output_fd, expected.tobytes()) + return SimpleNamespace(stderr=b"") monkeypatch.setattr(subprocess, "run", run) frames = reference_encoding.minimax_h3_decode_reference_video_frames( @@ -74,8 +82,11 @@ def test_video_transform_runs_once_and_qwen_samples_shared_rgb(monkeypatch): assert command[command.index("-frames:v") + 1] == "25" assert command[command.index("-ss") + 1] == "2.25" assert command.index("-ss") < command.index("-i") - assert command[-5:] == ["-f", "rawvideo", "-pix_fmt", "rgb24", "pipe:1"] + assert command[-5:-1] == ["-f", "rawvideo", "-pix_fmt", "rgb24"] + assert command[-1].startswith("pipe:") assert "libx264" not in command + if command[-1] != "pipe:1": + assert frames.flags.writeable assert all(np.shares_memory(frame, frames) for frame in sampled["frames"]) assert [int(frame[0, 0, 0]) for frame in sampled["frames"]] == [ int(expected[index, 0, 0, 0]) for index in (0, 12, 24) @@ -83,6 +94,136 @@ def test_video_transform_runs_once_and_qwen_samples_shared_rgb(monkeypatch): assert sampled["block_timestamps"] == [0.25, 1.0] +@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="requires Linux memfd") +def test_video_transform_can_share_one_host_decode(monkeypatch): + expected = np.arange(25 * 4 * 6 * 3, dtype=np.uint8).reshape(25, 4, 6, 3) + commands = [] + + class FakeGroup: + world_size = 2 + rank_in_group = 0 + cpu_group = object() + + def barrier(self): + return None + + def all_gather_object(outputs, value, **_kwargs): + outputs[:] = [value, value] + + def run(command, **_kwargs): + commands.append(command) + output_fd = int(command[-1].removeprefix("pipe:")) + os.write(output_fd, expected.tobytes()) + return SimpleNamespace(stderr=b"") + + monkeypatch.setattr(reference_encoding, "get_world_group", FakeGroup) + monkeypatch.setattr(torch.distributed, "all_gather_object", all_gather_object) + monkeypatch.setattr(subprocess, "run", run) + reference_encoding._reference_video_host_leader.cache_clear() + + try: + frames = reference_encoding.minimax_h3_decode_reference_video_frames( + "/input/ref.mp4", + target_width=6, + target_height=4, + target_frame_count=25, + share_across_replicas=True, + ) + finally: + reference_encoding._reference_video_host_leader.cache_clear() + + assert np.array_equal(frames, expected) + assert frames.flags.writeable + assert len(commands) == 1 + assert commands[0][-1].startswith("pipe:") + + +@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="requires Linux memfd") +def test_shared_video_transform_falls_back_when_proc_fd_is_blocked(monkeypatch): + expected = np.arange(25 * 4 * 6 * 3, dtype=np.uint8).reshape(25, 4, 6, 3) + commands = [] + + class FakeGroup: + world_size = 2 + rank_in_group = 0 + cpu_group = object() + + def all_gather_object(outputs, value, **_kwargs): + outputs[:] = [value, value] + + def run(command, **_kwargs): + commands.append(command) + os.write(int(command[-1].removeprefix("pipe:")), expected.tobytes()) + return SimpleNamespace(stderr=b"") + + real_open = os.open + + def guarded_open(path, flags): + if str(path).startswith("/proc/"): + raise PermissionError("blocked by test policy") + return real_open(path, flags) + + monkeypatch.setattr(reference_encoding, "get_world_group", FakeGroup) + monkeypatch.setattr(torch.distributed, "all_gather_object", all_gather_object) + monkeypatch.setattr(subprocess, "run", run) + monkeypatch.setattr(os, "open", guarded_open) + reference_encoding._reference_video_host_leader.cache_clear() + try: + frames = reference_encoding.minimax_h3_decode_reference_video_frames( + "/input/ref.mp4", + target_width=6, + target_height=4, + target_frame_count=25, + share_across_replicas=True, + ) + finally: + reference_encoding._reference_video_host_leader.cache_clear() + + assert np.array_equal(frames, expected) + assert len(commands) == 2 + + +@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="requires Linux memfd") +def test_shared_video_transform_propagates_any_host_decode_failure(monkeypatch): + class FakeGroup: + world_size = 4 + rank_in_group = 0 + cpu_group = object() + + gather_index = 0 + + def all_gather_object(outputs, value, **_kwargs): + nonlocal gather_index + if gather_index == 0: + outputs[:] = ["host-a", "host-a", "host-b", "host-b"] + else: + outputs[:] = [ + value, + None, + (None, 0, "CalledProcessError: remote decode failed"), + None, + ] + gather_index += 1 + + monkeypatch.setattr(reference_encoding, "get_world_group", FakeGroup) + monkeypatch.setattr(torch.distributed, "all_gather_object", all_gather_object) + monkeypatch.setattr( + reference_encoding, + "_write_reference_video_to_fd", + lambda _command, _fd: 1, + ) + reference_encoding._reference_video_host_leader.cache_clear() + try: + with pytest.raises(RuntimeError, match="remote decode failed"): + reference_encoding._decode_reference_video_shared(["ffmpeg"]) + finally: + reference_encoding._reference_video_host_leader.cache_clear() + + # The failure is resolved immediately after the shared state exchange; + # no rank enters a mapping collective that another host skipped. + assert gather_index == 2 + + def test_audio_decode_is_bounded_float_pcm_without_temp_files(monkeypatch): pcm = torch.arange(8, dtype=torch.float32).numpy().tobytes() commands = []