[diffusion] fix: fix MiniMax-H3 dp_size>1 deadlock and cross-request audio determinism (#36398)

Co-authored-by: Kevin Mi <kevin.mi@radixark.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Cursor <cursoragent@cursor.com>
2026-08-26 14:54:08 +08:00
parent abe3aeb142
commit 3ce4f957eb
4 changed files with 68 additions and 32 deletions
@@ -34,7 +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.distributed.parallel_state import get_replica_group
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
    MINIMAX_H3_SUPPORTED_FPS,
@@ -491,7 +491,7 @@ def _write_reference_video_to_fd(command: list[str], output_fd: int) -> int:
    return os.lseek(output_fd, 0, os.SEEK_CUR)


-def _all_gather_world_objects(group: Any, value: Any) -> list[Any]:
+def _all_gather_group_objects(group: Any, value: Any) -> list[Any]:
    values = [None] * group.world_size
    torch.distributed.all_gather_object(
        values,
@@ -503,14 +503,18 @@ def _all_gather_world_objects(group: Any, value: Any) -> list[Any]:

@functools.lru_cache(maxsize=1)
 def _reference_video_host_leader() -> int:
-    group = get_world_group()
-    hostnames = _all_gather_world_objects(group, socket.gethostname())
+    group = get_replica_group()
+    hostnames = _all_gather_group_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()
+    """Decode once per host per replica; its worker ranks map the same RGB pages.
+
+    Collectives stay inside the request's pipeline replica: the world group
+    spans replicas when dp_size > 1 and only one replica runs this request.
+    """
+    group = get_replica_group()
    if (
        group.world_size <= 1
        or not sys.platform.startswith("linux")
@@ -548,7 +552,7 @@ def _decode_reference_video_shared(command: list[str]) -> tuple[Any, int]:
            owner_exception = exc
            leader_state = (None, 0, f"{type(exc).__name__}: {exc}")

-    states = _all_gather_world_objects(group, leader_state)
+    states = _all_gather_group_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),
@@ -593,7 +597,7 @@ def _decode_reference_video_shared(command: list[str]) -> tuple[Any, int]:
    except Exception as exc:
        map_error = f"{type(exc).__name__}: {exc}"

-    map_errors = _all_gather_world_objects(group, map_error)
+    map_errors = _all_gather_group_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)
@@ -3,13 +3,13 @@ from __future__ import annotations

 import functools
 from collections.abc import Mapping
-from contextlib import nullcontext
+from contextlib import contextmanager, nullcontext

 import torch

 from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
 from sglang.multimodal_gen.runtime.distributed import (
-    get_world_group,
+    get_replica_group,
    model_parallel_is_initialized,
 )
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
@@ -45,6 +45,39 @@ def _required_tensor(value, path: str) -> torch.Tensor:
    return value


+@contextmanager
+def _deterministic_audio_decode_context():
+    """Deterministic-algorithm scope for the fp32 audio-VAE decode.
+
+    Without it, cuDNN picks conv algorithms from free-workspace state, so the
+    same audio latent decodes to different bytes on a server process's first
+    request than on every later one. Deterministic algorithms with TF32 off
+    keep cuDNN speed (unlike the encode-side context, which disables cuDNN);
+    if first-request divergence ever reappears, escalate to
+    reference_encoding._AudioVAEDeterminismContext.
+    """
+    b = torch.backends
+    saved = (
+        b.cudnn.allow_tf32,
+        b.cuda.matmul.allow_tf32,
+        b.cudnn.deterministic,
+        b.cudnn.benchmark,
+    )
+    b.cudnn.allow_tf32 = False
+    b.cuda.matmul.allow_tf32 = False
+    b.cudnn.deterministic = True
+    b.cudnn.benchmark = False
+    try:
+        yield
+    finally:
+        (
+            b.cudnn.allow_tf32,
+            b.cuda.matmul.allow_tf32,
+            b.cudnn.deterministic,
+            b.cudnn.benchmark,
+        ) = saved
+
+
@functools.lru_cache(maxsize=None)
 def _cached_decode_mean_std(
    mean_values: tuple[float, ...],
@@ -306,7 +339,7 @@ class MiniMaxH3DecodingStage(DecodingStage):
                if audio_latent.is_cuda
                else nullcontext()
            )
-            with autocast_context:
+            with _deterministic_audio_decode_context(), autocast_context:
                audio_decode = self._get_vae_decode_fn(
                    audio_vae,
                    server_args,
@@ -392,12 +425,10 @@ class MiniMaxH3DecodingStage(DecodingStage):
                    canonical_frames.copy_(visual_frames)
                    visual_frames = canonical_frames

-        # DP is currently rejected by ServerArgs validation, so the world group
-        # is one request replica (TP/CFG/SP ranks), not a collection of
-        # independent requests. Decode the non-sharded audio VAE once per
-        # request and distribute its output to the ranks that decoded video.
-        world_group = get_world_group() if model_parallel_is_initialized() else None
-        is_audio_owner = world_group is None or world_group.rank_in_group == 0
+        # Audio VAE weights are replicated. Decode on replica rank 0 and broadcast
+        # only within the request's replica, excluding independent DP replicas.
+        replica_group = get_replica_group() if model_parallel_is_initialized() else None
+        is_audio_owner = replica_group is None or replica_group.rank_in_group == 0
        owner_exception = None
        owner_error = None
        audio_payload = None
@@ -407,16 +438,16 @@ class MiniMaxH3DecodingStage(DecodingStage):
            except Exception as exc:
                owner_exception = exc
                owner_error = f"{type(exc).__name__}: {exc}"
-        if world_group is not None:
-            owner_error = world_group.broadcast_object(owner_error, src=0)
+        if replica_group is not None:
+            owner_error = replica_group.broadcast_object(owner_error, src=0)
        if owner_error is not None:
            if owner_exception is not None:
                raise owner_exception
            raise RuntimeError(
                f"MiniMax H3 audio decode failed on rank 0: {owner_error}"
            )
-        if world_group is not None:
-            audio_payload = world_group.broadcast_tensor_dict(audio_payload, src=0)
+        if replica_group is not None:
+            audio_payload = replica_group.broadcast_tensor_dict(audio_payload, src=0)
        if not isinstance(audio_payload, dict):
            raise RuntimeError("MiniMax H3 audio decode produced no output payload")
        audio_waveform = _required_tensor(
@@ -8,15 +8,16 @@ from typing import Any

 def minimax_h3_replica_ctx() -> tuple[int, int]:
    from sglang.multimodal_gen.runtime.distributed.parallel_state import (
-        get_world_group,
+        get_replica_group,
        model_parallel_is_initialized,
    )

    if not model_parallel_is_initialized():
        return 1, 0
-    # ServerArgs currently rejects DP>1 and H3 rejects disaggregation, so the
-    # world group contains exactly one request replica (TP/CFG/SP ranks).
-    group = get_world_group()
+    # A request lives in one pipeline replica (TP/CFG/SP ranks); the world
+    # group spans replicas when dp_size > 1, so a world-group collective would
+    # wait on idle replicas forever.
+    group = get_replica_group()
    return int(group.world_size), int(group.rank_in_group)


@@ -26,10 +27,10 @@ def minimax_h3_replica_broadcast_extra(batch: Any, key: str) -> None:
        return

    from sglang.multimodal_gen.runtime.distributed.parallel_state import (
-        get_world_group,
+        get_replica_group,
    )

-    group = get_world_group()
+    group = get_replica_group()
    payload = {"value": batch.extra.get(key)} if rank == 0 else None
    payload = group.broadcast_tensor_dict(payload, src=0)
    value = payload.get("value") if isinstance(payload, dict) else None
@@ -45,10 +46,10 @@ def minimax_h3_replica_broadcast_error(error: str | None) -> str | None:
        return error

    from sglang.multimodal_gen.runtime.distributed.parallel_state import (
-        get_world_group,
+        get_replica_group,
    )

-    group = get_world_group()
+    group = get_replica_group()
    return group.broadcast_object(error if rank == 0 else None, src=0)


@@ -127,7 +127,7 @@ def test_video_transform_can_share_one_host_decode(monkeypatch):
        os.write(output_fd, expected.tobytes())
        return SimpleNamespace(stderr=b"")

-    monkeypatch.setattr(reference_encoding, "get_world_group", FakeGroup)
+    monkeypatch.setattr(reference_encoding, "get_replica_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()
@@ -174,7 +174,7 @@ def test_shared_video_transform_falls_back_when_proc_fd_is_blocked(monkeypatch):
            raise PermissionError("blocked by test policy")
        return real_open(path, flags)

-    monkeypatch.setattr(reference_encoding, "get_world_group", FakeGroup)
+    monkeypatch.setattr(reference_encoding, "get_replica_group", FakeGroup)
    monkeypatch.setattr(torch.distributed, "all_gather_object", all_gather_object)
    monkeypatch.setattr(subprocess, "run", run)
    monkeypatch.setattr(os, "open", guarded_open)
@@ -216,7 +216,7 @@ def test_shared_video_transform_propagates_any_host_decode_failure(monkeypatch):
            ]
        gather_index += 1

-    monkeypatch.setattr(reference_encoding, "get_world_group", FakeGroup)
+    monkeypatch.setattr(reference_encoding, "get_replica_group", FakeGroup)
    monkeypatch.setattr(torch.distributed, "all_gather_object", all_gather_object)
    monkeypatch.setattr(
        reference_encoding,