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

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/reference_encoding.py

@@ -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()
+    group = get_replica_group()
-    hostnames = _all_gather_world_objects(group, socket.gethostname())
+    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."""
+    """Decode once per host per replica; its worker ranks map the same RGB pages.
-    group = get_world_group()
+
+    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)

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/decoding.py

@@ -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
+        # Audio VAE weights are replicated. Decode on replica rank 0 and broadcast
-        # is one request replica (TP/CFG/SP ranks), not a collection of
+        # only within the request's replica, excluding independent DP replicas.
-        # independent requests. Decode the non-sharded audio VAE once per
+        replica_group = get_replica_group() if model_parallel_is_initialized() else None
-        # request and distribute its output to the ranks that decoded video.
+        is_audio_owner = replica_group is None or replica_group.rank_in_group == 0
-        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
         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:
+        if replica_group is not None:
-            owner_error = world_group.broadcast_object(owner_error, src=0)
+            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:
+        if replica_group is not None:
-            audio_payload = world_group.broadcast_tensor_dict(audio_payload, src=0)
+            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(

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/replica_broadcast.py

@@ -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
+    # A request lives in one pipeline replica (TP/CFG/SP ranks); the world
-    # world group contains exactly one request replica (TP/CFG/SP ranks).
+    # group spans replicas when dp_size > 1, so a world-group collective would
-    group = get_world_group()
+    # 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)
 
 

python/sglang/multimodal_gen/test/unit/test_minimax_h3_media.py

@@ -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,