minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Mick
2026-08-12 20:02:49 +08:00
committed by GitHub
parent 9deb6952af
commit 9701cc138c
4 changed files with 381 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/reference_encoding.py
+197 -20
View File
@@ -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(
{
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/text_encoding.py
+29 -2
View File
@@ -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
+11 -1
View File
@@ -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
+144 -3
View File
@@ -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 = []