[diffusion] Fix component accuracy topology reuse (#33317)

python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a.py

@@ -99,9 +99,14 @@ class IpcA2AState:
         self.max_buffers = 0
         self.calls = 0
         self.rank = None
+        self.group = None
         self.failed = False
         self.inited = False
 
+    def reset(self) -> None:
+        """Drop mappings that belong to a model-parallel group being replaced."""
+        self.__init__()
+
     def _share(self, t, group):
         """Exchange `t` with the peer via torch IPC, re-opening the handle in
         the LOCAL device context (the mapping is only dereferenceable from the
@@ -136,6 +141,7 @@ class IpcA2AState:
         from torch.utils.cpp_extension import load_inline
 
         self.rank = dist.get_rank(group=group)
+        self.group = group
         dev = torch.cuda.current_device()
         if not torch.cuda.can_device_access_peer(dev, 1 - dev):
             raise _Unsupported("no peer-to-peer access between the two devices")
@@ -260,7 +266,11 @@ def ipc_a2a_ready(group) -> bool:
     from sglang.multimodal_gen.runtime.distributed import get_tp_world_size
     from sglang.multimodal_gen.runtime.platforms import current_platform
 
-    if not envs.SGLANG_DIFFUSION_IPC_A2A or IPC_A2A.failed:
+    if not envs.SGLANG_DIFFUSION_IPC_A2A:
+        return False
+    if IPC_A2A.group is not None and IPC_A2A.group is not group:
+        IPC_A2A.reset()
+    if IPC_A2A.failed:
         return False
     # TP+Ulysses groups are strided in global-rank order, while this transport
     # supports the adjacent two-device topology used by TP1+U2. Reject the

python/sglang/multimodal_gen/runtime/distributed/parallel_state.py

@@ -901,10 +901,33 @@ def destroy_model_parallel() -> None:
     """Set the groups to none and destroy them."""
     global _TP, _SP, _DP, _CFG, _PP, _VAE_DECODE, _DIT, _VAE
 
+    # The IPC transport keeps CUDA mappings associated with the current
+    # Ulysses group. Drop them before tearing down the process groups.
+    from .device_communicators.ipc_a2a import IPC_A2A
+    from .parallel_groups import PROCESS_GROUP
+
+    IPC_A2A.reset()
+
     for group in (_TP, _SP, _DP, _CFG, _PP, _VAE_DECODE):
         if group is not None:
             group.destroy()
 
+    # Ulysses and Ring groups are created separately from the SP coordinator,
+    # so GroupCoordinator.destroy() does not own or release them. Explicitly
+    # destroy them here; otherwise repeated Ulysses/Ring topology switches leak
+    # NCCL communicators and their CUDA memory.
+    destroyed_sequence_groups = []
+    for group in (PROCESS_GROUP.ULYSSES_PG, PROCESS_GROUP.RING_PG):
+        if (
+            group is not None
+            and group is not torch.distributed.group.WORLD
+            and all(group is not destroyed for destroyed in destroyed_sequence_groups)
+        ):
+            torch.distributed.destroy_process_group(group)
+            destroyed_sequence_groups.append(group)
+    PROCESS_GROUP.ULYSSES_PG = None
+    PROCESS_GROUP.RING_PG = None
+
     for group in (_DIT, _VAE):
         if group is not None:
             torch.distributed.destroy_process_group(group)

python/sglang/multimodal_gen/test/single_test_file/component_accuracy/config.py

@@ -49,8 +49,9 @@ CASE_THRESHOLDS: Dict[str, Dict[ComponentType, float]] = {
     "wan2_2_t2v_a14b_2gpu": {ComponentType.TRANSFORMER: 0.99},
     "wan2_2_t2v_a14b_teacache_2gpu": {ComponentType.TRANSFORMER: 0.99},
     "wan2_2_t2v_a14b_lora_2gpu": {ComponentType.TRANSFORMER: 0.99},
-    "zimage_image_t2i_2_gpus": {ComponentType.TRANSFORMER: 0.9935},
-    "zimage_image_t2i_2_gpus_non_square": {ComponentType.TRANSFORMER: 0.9935},
+    # H100/H200 FlashAttention runs vary between roughly 0.99338 and 0.99350.
+    "zimage_image_t2i_2_gpus": {ComponentType.TRANSFORMER: 0.993},
+    "zimage_image_t2i_2_gpus_non_square": {ComponentType.TRANSFORMER: 0.993},
 }
 
 # Active skip policy. Keep this limited to cases with current, concrete evidence

python/sglang/multimodal_gen/test/single_test_file/component_accuracy/hooks.py

@@ -362,10 +362,14 @@ def _build_transformer_hook_inputs(
         inputs["pooled_projections"] = rng.randn(
             (1, pooled_channels), device, torch.bfloat16
         )
-    if (
+    supports_text_attention_mask = (
         "encoder_attention_mask" in param_names
         or "encoder_hidden_states_mask" in param_names
-    ):
+    )
+    uses_ring_parallel = (case.server_args.ring_degree or 1) > 1
+    if supports_text_attention_mask and not uses_ring_parallel:
+        # This synthetic mask is all True and has no semantic effect. Omit it
+        # for ring-parallel cases, whose masked attention path is unsupported.
         attention_mask = torch.ones(
             1, DEFAULT_TEXT_SEQ_LEN, device=device, dtype=torch.bool
         )

python/sglang/multimodal_gen/test/single_test_file/component_accuracy/utils.py

@@ -16,8 +16,10 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
     destroy_model_parallel,
     get_classifier_free_guidance_world_size,
     get_data_parallel_world_size,
+    get_ring_parallel_world_size,
     get_sequence_parallel_world_size,
     get_tensor_model_parallel_world_size,
+    get_ulysses_parallel_world_size,
     maybe_init_distributed_environment_and_model_parallel,
     model_parallel_is_initialized,
 )
@@ -305,11 +307,15 @@ def initialize_parallel_runtime(sgl_args: ServerArgs) -> None:
     if model_parallel_is_initialized():
         current_tp = get_tensor_model_parallel_world_size()
         current_sp = get_sequence_parallel_world_size()
+        current_ulysses = get_ulysses_parallel_world_size()
+        current_ring = get_ring_parallel_world_size()
         current_dp = get_data_parallel_world_size()
         current_cfg = get_classifier_free_guidance_world_size()
         if (
             current_tp == tp_size
             and current_sp == sp_degree
+            and current_ulysses == ulysses_degree
+            and current_ring == ring_degree
             and current_dp == dp_size
             and current_cfg == cfg_degree
         ):

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

@@ -0,0 +1,45 @@
+from types import SimpleNamespace
+
+import torch.nn as nn
+
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.hooks import (
+    _build_transformer_hook_inputs,
+)
+
+
+class _TransformerWithOptionalMask(nn.Module):
+    def forward(
+        self,
+        hidden_states,
+        encoder_hidden_states,
+        timestep,
+        encoder_hidden_states_mask=None,
+    ):
+        raise NotImplementedError
+
+
+def _case(*, ring_degree: int) -> SimpleNamespace:
+    return SimpleNamespace(
+        server_args=SimpleNamespace(
+            model_path="test/model",
+            ring_degree=ring_degree,
+        )
+    )
+
+
+def test_omits_noop_attention_mask_for_ring_parallel_case():
+    inputs = _build_transformer_hook_inputs(
+        _case(ring_degree=2), _TransformerWithOptionalMask(), "cpu"
+    )
+
+    assert "encoder_hidden_states_mask" not in inputs
+    assert "encoder_attention_mask" not in inputs
+
+
+def test_keeps_attention_mask_for_non_ring_case():
+    inputs = _build_transformer_hook_inputs(
+        _case(ring_degree=1), _TransformerWithOptionalMask(), "cpu"
+    )
+
+    assert inputs["encoder_hidden_states_mask"].all()
+    assert inputs["encoder_attention_mask"].all()

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

@@ -0,0 +1,115 @@
+from contextlib import ExitStack
+from types import SimpleNamespace
+from unittest.mock import call, patch
+
+from sglang.multimodal_gen.runtime.distributed import parallel_state
+from sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a import (
+    IPC_A2A,
+)
+from sglang.multimodal_gen.runtime.distributed.parallel_groups import PROCESS_GROUP
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import (
+    initialize_parallel_runtime,
+)
+
+_UTILS = "sglang.multimodal_gen.test.single_test_file.component_accuracy.utils"
+
+
+def _server_args(*, ulysses_degree: int, ring_degree: int) -> SimpleNamespace:
+    return SimpleNamespace(
+        tp_size=1,
+        sp_degree=2,
+        ulysses_degree=ulysses_degree,
+        ring_degree=ring_degree,
+        dp_size=1,
+        cfg_parallel_degree=1,
+    )
+
+
+def _patch_current_topology(
+    stack: ExitStack, *, ulysses_degree: int, ring_degree: int
+) -> None:
+    for context_manager in (
+        patch(f"{_UTILS}.model_parallel_is_initialized", return_value=True),
+        patch(f"{_UTILS}.get_tensor_model_parallel_world_size", return_value=1),
+        patch(f"{_UTILS}.get_sequence_parallel_world_size", return_value=2),
+        patch(
+            f"{_UTILS}.get_ulysses_parallel_world_size",
+            return_value=ulysses_degree,
+        ),
+        patch(f"{_UTILS}.get_ring_parallel_world_size", return_value=ring_degree),
+        patch(f"{_UTILS}.get_data_parallel_world_size", return_value=1),
+        patch(f"{_UTILS}.get_classifier_free_guidance_world_size", return_value=1),
+    ):
+        stack.enter_context(context_manager)
+
+
+def test_reinitializes_when_sp_decomposition_changes():
+    with ExitStack() as stack:
+        _patch_current_topology(stack, ulysses_degree=2, ring_degree=1)
+        stack.enter_context(
+            patch(f"{_UTILS}.torch.distributed.is_initialized", return_value=True)
+        )
+        barrier = stack.enter_context(patch(f"{_UTILS}.torch.distributed.barrier"))
+        destroy = stack.enter_context(patch(f"{_UTILS}.destroy_model_parallel"))
+        initialize = stack.enter_context(
+            patch(f"{_UTILS}.maybe_init_distributed_environment_and_model_parallel")
+        )
+        initialize_parallel_runtime(_server_args(ulysses_degree=1, ring_degree=2))
+
+    destroy.assert_called_once_with()
+    initialize.assert_called_once_with(
+        tp_size=1,
+        sp_size=2,
+        cfg_degree=1,
+        ulysses_degree=1,
+        ring_degree=2,
+        dp_size=1,
+    )
+    assert barrier.call_count == 2
+
+
+def test_reuses_matching_sp_decomposition():
+    with ExitStack() as stack:
+        _patch_current_topology(stack, ulysses_degree=1, ring_degree=2)
+        stack.enter_context(
+            patch(f"{_UTILS}.torch.distributed.is_initialized", return_value=True)
+        )
+        destroy = stack.enter_context(patch(f"{_UTILS}.destroy_model_parallel"))
+        initialize = stack.enter_context(
+            patch(f"{_UTILS}.maybe_init_distributed_environment_and_model_parallel")
+        )
+        initialize_parallel_runtime(_server_args(ulysses_degree=1, ring_degree=2))
+
+    destroy.assert_not_called()
+    initialize.assert_not_called()
+
+
+def test_destroy_releases_sequence_parallel_subgroups_after_partial_init():
+    ulysses_group = object()
+    ring_group = object()
+
+    with ExitStack() as stack:
+        for name in (
+            "_TP",
+            "_SP",
+            "_DP",
+            "_CFG",
+            "_PP",
+            "_VAE_DECODE",
+            "_DIT",
+            "_VAE",
+        ):
+            stack.enter_context(patch.object(parallel_state, name, None))
+        stack.enter_context(patch.object(PROCESS_GROUP, "ULYSSES_PG", ulysses_group))
+        stack.enter_context(patch.object(PROCESS_GROUP, "RING_PG", ring_group))
+        reset_ipc = stack.enter_context(patch.object(IPC_A2A, "reset"))
+        destroy_group = stack.enter_context(
+            patch.object(parallel_state.torch.distributed, "destroy_process_group")
+        )
+
+        parallel_state.destroy_model_parallel()
+
+        reset_ipc.assert_called_once_with()
+        assert destroy_group.call_args_list == [call(ulysses_group), call(ring_group)]
+        assert PROCESS_GROUP.ULYSSES_PG is None
+        assert PROCESS_GROUP.RING_PG is None

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

@@ -0,0 +1,41 @@
+from unittest.mock import patch
+
+from sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a import (
+    IpcA2AState,
+    ipc_a2a_ready,
+)
+
+_IPC = "sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a"
+
+
+def test_reinitializes_ipc_transport_for_replaced_process_group():
+    state = IpcA2AState()
+    old_group = object()
+    new_group = object()
+    state.inited = True
+    state.group = old_group
+    state.calls = 7
+
+    def initialize(group):
+        state.inited = True
+        state.group = group
+
+    with (
+        patch(f"{_IPC}.IPC_A2A", state),
+        patch(f"{_IPC}.envs.SGLANG_DIFFUSION_IPC_A2A", True),
+        patch(
+            "sglang.multimodal_gen.runtime.platforms.current_platform.is_cuda",
+            return_value=True,
+        ),
+        patch(
+            "sglang.multimodal_gen.runtime.distributed.get_tp_world_size",
+            return_value=1,
+        ),
+        patch.object(state, "init", side_effect=initialize) as init,
+        patch(f"{_IPC}.torch.cuda.is_current_stream_capturing", return_value=False),
+    ):
+        assert ipc_a2a_ready(new_group)
+
+    init.assert_called_once_with(new_group)
+    assert state.group is new_group
+    assert state.calls == 0