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 98f1f72e8..6bc1bbddb 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 @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/decoding.py index 29ab53965..6c20d41c6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/decoding.py +++ b/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 - # 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( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/replica_broadcast.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/replica_broadcast.py index d3499642e..5367b5d8e 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/replica_broadcast.py +++ b/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 - # 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) 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 f1fdabe91..48862032c 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 @@ -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,