[diffusion] optimize: optimize bit-exact h3 reference video ingress (#34563)

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
+    frame_count = payload_size // frame_bytes
-    if frame_count <= 0:
-        raise ValueError(f"reference video has no frames: {video_path}")
     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(
             {

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,

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(

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 = []