[diffusion] fix: decouple encoder parallelism from the dit parallel layout (#34713)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

python/sglang/multimodal_gen/configs/models/encoders/base.py

@@ -78,9 +78,9 @@ class EncoderConfig(ModelConfig):
 
     # Parallel folding: during the encoding stage the whole DiT replica is idle,
     # so TP-shard the encoder across those otherwise-unused GPUs instead of
-    # running it on a single rank. None = replicated, else the group to fold
-    # over ("sp"|"world"); resolved by finalize_encoder_folding.
-    parallel_folding_mode: Literal["sp", "world"] | None = None
+    # leaving it on the DiT TP group. None keeps that TP group; the other modes
+    # override it with a wider group selected by finalize_encoder_folding.
+    parallel_folding_mode: Literal["sp", "world", "replica"] | None = None
 
 
 @dataclass

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

@@ -12,6 +12,8 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
     get_dp_group,
     get_dp_rank,
     get_dp_world_size,
+    get_encoder_data_parallel_group,
+    get_replica_group,
     get_sp_group,
     get_sp_parallel_rank,
     get_sp_world_size,
@@ -44,6 +46,8 @@ __all__ = [
     "get_world_size",
     # Data parallel group
     "get_dp_group",
+    "get_replica_group",
+    "get_encoder_data_parallel_group",
     "get_dp_rank",
     "get_dp_world_size",
     # Sequence parallel group

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

@@ -66,10 +66,18 @@ _SP: SequenceParallelGroupCoordinator | None = None
 _PP: PipelineGroupCoordinator | None = None
 _CFG: GroupCoordinator | None = None
 _DP: GroupCoordinator | None = None
+# all ranks serving one pipeline replica (every dim except dp); with
+# dp_size == 1 it coincides with the world group
+_REPLICA: GroupCoordinator | None = None
+# Corresponding TP lanes across the replicated encoder copies in one pipeline
+# replica. None means the encoder has only one TP copy, so batch DP cannot run.
+_ENCODER_DP: GroupCoordinator | None = None
 _VAE_DECODE: GroupCoordinator | None = None
 _DIT: ProcessGroup | None = None
 _VAE: ProcessGroup | None = None
 _VAE_DECODE_PARALLEL_AXES = "tp-sp-pp-cfg"
+_REPLICA_PARALLEL_AXES = "tp-sp-pp-cfg"
+_ENCODER_DP_PARALLEL_AXES = "sp-pp-cfg"
 
 TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])
 
@@ -187,6 +195,8 @@ def init_parallel_group_coordinator(
         "sequence",
         "classifier_free_guidance",
         "vae_decode",
+        "replica",
+        "encoder_data",
     ], f"parallel_mode {parallel_mode} is not supported"
     if parallel_mode == "pipeline":
         return PipelineGroupCoordinator(
@@ -204,19 +214,19 @@ def init_parallel_group_coordinator(
             **kwargs,
         )
     else:
+        group_name = {
+            "tensor": "tp_group",
+            "vae_decode": "vae_decode_group",
+            "replica": "replica_group",
+            "encoder_data": "encoder_dp_group",
+        }.get(parallel_mode, "cfg_group")
         return GroupCoordinator(
             group_ranks=group_ranks,
             local_rank=local_rank,
             torch_distributed_backend=backend,
             use_device_communicator=parallel_mode != "tensor",
             use_srt_custom_allreduce=parallel_mode == "tensor",
-            group_name=(
-                "tp_group"
-                if parallel_mode == "tensor"
-                else (
-                    "vae_decode_group" if parallel_mode == "vae_decode" else "cfg_group"
-                )
-            ),
+            group_name=group_name,
         )
 
 
@@ -228,6 +238,12 @@ def _get_vae_decode_group_ranks(
     return rank_generator.get_ranks(_VAE_DECODE_PARALLEL_AXES)
 
 
+def _get_encoder_data_parallel_group_ranks(
+    rank_generator: RankGenerator,
+) -> list[list[int]]:
+    return rank_generator.get_ranks(_ENCODER_DP_PARALLEL_AXES)
+
+
 def get_tp_group() -> GroupCoordinator:
     assert _TP is not None, "tensor model parallel group is not initialized"
     return _TP
@@ -324,6 +340,35 @@ def get_dp_group() -> GroupCoordinator:
     return _DP
 
 
+def get_replica_group() -> GroupCoordinator:
+    """The ranks that share this rank's request batch — its pipeline replica.
+
+    Replica-wide encoder folding uses this group rather than the world group,
+    which spans replicas when dp_size > 1. The encoder batch-DP group is the
+    orthogonal non-TP subset of this replica. A process that only initialized
+    the world group has exactly one replica by construction, so the world group
+    doubles as it."""
+    if _REPLICA is not None:
+        return _REPLICA
+    return get_world_group()
+
+
+def get_encoder_data_parallel_group() -> GroupCoordinator | None:
+    """Encoder copies that may split a batch inside one pipeline replica.
+
+    Native encoders use the DiT TP group when they are not folded. The batch-DP
+    group is therefore the orthogonal non-TP part of the replica: ranks with the
+    same TP coordinate receive different batch slices, while every rank in one
+    TP group receives the same slice. A pure-TP replica has only one encoder
+    copy and returns None.
+    """
+    if _ENCODER_DP is not None:
+        return _ENCODER_DP
+    if not model_parallel_is_initialized():
+        return get_world_group()
+    return None
+
+
 # xDiT
 def initialize_model_parallel(
     data_parallel_size: int = 1,
@@ -486,6 +531,64 @@ def initialize_model_parallel(
     )
     _sync_srt_tp_group()
 
+    global _REPLICA
+    assert _REPLICA is None, "replica group is already initialized"
+    if data_parallel_size == 1:
+        _REPLICA = get_world_group()
+    else:
+        replica_axis_groups = [
+            group
+            for degree, group in (
+                (tensor_parallel_degree, _TP),
+                (sequence_parallel_degree, _SP),
+                (pipeline_parallel_degree, _PP),
+                (classifier_free_guidance_degree, _CFG),
+            )
+            if degree > 1
+        ]
+        if len(replica_axis_groups) <= 1:
+            # One nontrivial axis already is the replica group. If every axis
+            # is size one, any singleton axis group represents the replica.
+            _REPLICA = replica_axis_groups[0] if replica_axis_groups else _TP
+        else:
+            _REPLICA = init_parallel_group_coordinator(
+                group_ranks=rank_generator.get_ranks(_REPLICA_PARALLEL_AXES),
+                local_rank=get_world_group().local_rank,
+                backend=backend,
+                parallel_mode="replica",
+            )
+
+    global _ENCODER_DP
+    assert _ENCODER_DP is None, "encoder data parallel group is already initialized"
+    encoder_dp_degree = (
+        sequence_parallel_degree
+        * pipeline_parallel_degree
+        * classifier_free_guidance_degree
+    )
+    if tensor_parallel_degree == 1:
+        # With no encoder TP sharding, every replica rank owns a full encoder.
+        _ENCODER_DP = _REPLICA
+    elif encoder_dp_degree > 1:
+        nontrivial_groups = [
+            group
+            for degree, group in (
+                (sequence_parallel_degree, _SP),
+                (pipeline_parallel_degree, _PP),
+                (classifier_free_guidance_degree, _CFG),
+            )
+            if degree > 1
+        ]
+        if len(nontrivial_groups) == 1:
+            # Reuse the existing axis group in the common TP x SP case.
+            _ENCODER_DP = nontrivial_groups[0]
+        else:
+            _ENCODER_DP = init_parallel_group_coordinator(
+                group_ranks=_get_encoder_data_parallel_group_ranks(rank_generator),
+                local_rank=get_world_group().local_rank,
+                backend=backend,
+                parallel_mode="encoder_data",
+            )
+
     global _VAE_DECODE
     assert _VAE_DECODE is None, "VAE decode parallel group is already initialized"
     _VAE_DECODE = init_parallel_group_coordinator(
@@ -920,7 +1023,7 @@ def init_vae_group(
 
 def destroy_model_parallel() -> None:
     """Set the groups to none and destroy them."""
-    global _TP, _SP, _DP, _CFG, _PP, _VAE_DECODE, _DIT, _VAE
+    global _TP, _SP, _DP, _CFG, _PP, _VAE_DECODE, _DIT, _VAE, _REPLICA, _ENCODER_DP
 
     _clear_srt_tp_group()
     # The IPC transport keeps CUDA mappings associated with the current
@@ -930,9 +1033,25 @@ def destroy_model_parallel() -> None:
 
     IPC_A2A.reset()
 
-    for group in (_TP, _SP, _DP, _CFG, _PP, _VAE_DECODE):
-        if group is not None:
+    destroyed_groups = []
+    for group in (
+        _TP,
+        _SP,
+        _DP,
+        _CFG,
+        _PP,
+        _REPLICA,
+        _ENCODER_DP,
+        _VAE_DECODE,
+    ):
+        # Replica-derived groups may alias each other or the world group.
+        if (
+            group is not None
+            and group is not _WORLD
+            and all(group is not destroyed for destroyed in destroyed_groups)
+        ):
             group.destroy()
+            destroyed_groups.append(group)
 
     # Ulysses and Ring groups are created separately from the SP coordinator,
     # so GroupCoordinator.destroy() does not own or release them. Explicitly
@@ -954,4 +1073,6 @@ def destroy_model_parallel() -> None:
         if group is not None:
             torch.distributed.destroy_process_group(group)
 
-    _TP, _SP, _DP, _CFG, _PP, _VAE_DECODE, _DIT, _VAE = (None,) * 8
+    _TP, _SP, _DP, _CFG, _PP, _VAE_DECODE, _DIT, _VAE, _REPLICA, _ENCODER_DP = (
+        None,
+    ) * 10

python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py

@@ -15,6 +15,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
     QwenImageEditPipelineConfig,
 )
 from sglang.multimodal_gen.runtime.distributed import (
+    get_encoder_data_parallel_group,
     get_local_torch_device,
 )
 from sglang.multimodal_gen.runtime.distributed.parallel_state import (
@@ -420,17 +421,19 @@ class TextEncoderLoader(ComponentLoader):
             model_config,
             component_model_path,
         )
+        encoder_dp_group = get_encoder_data_parallel_group()
+        prefer_dp = (
+            server_args.batching_max_size > 1
+            and encoder_dp_group is not None
+            and encoder_dp_group.world_size > 1
+            and issubclass(model_cls, TextEncoder)
+            and model_cls.supports_dp_encode
+        )
         # real dims are populated now; resolve fold vs replicate
         finalize_encoder_folding(
             encoder_config,
             server_args.encoder_parallel,
-            prefer_dp=(
-                server_args.batching_max_size > 1
-                and (server_args.tp_size or 1) == 1
-                and (server_args.dp_size or 1) == 1
-                and issubclass(model_cls, TextEncoder)
-                and model_cls.supports_dp_encode
-            ),
+            prefer_dp=prefer_dp,
         )
         encoder_dtype = server_args.pipeline_config.text_encoder_precisions[
             encoder_index

python/sglang/multimodal_gen/runtime/models/encoders/base.py

@@ -14,6 +14,7 @@ from sglang.multimodal_gen.configs.models.encoders import (
     TextEncoderConfig,
 )
 from sglang.multimodal_gen.runtime.distributed import (
+    get_replica_group,
     get_sp_group,
     get_tp_group,
     get_world_group,
@@ -36,6 +37,10 @@ def get_folding_tp_group(config: EncoderConfig):
     if mode == "world":
         # the whole single-replica DiT (all GPUs), regardless of tp/sp/cfg.
         return get_world_group()
+    if mode == "replica":
+        # the ranks serving this rank's pipeline replica (== world when
+        # dp_size is 1); the shape-independent choice for explicit folding
+        return get_replica_group()
     if mode is None:
         return get_tp_group()
     raise ValueError(f"Unsupported encoder folding mode: {mode!r}")

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

@@ -23,7 +23,7 @@ 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"
+_MINIMAX_H3_SINGLE_COPY_TEXT_ENCODE_EXTRA_KEY = "minimax_h3_single_copy_text_encode"
 
 
 class MiniMaxH3TextEncodingStage(TextEncodingStage):
@@ -108,13 +108,14 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage):
         batches: list[Req],
         server_args: ServerArgs,
     ) -> list[Req]:
-        """Distribute independent H3 presentations over replicated encoders.
+        """Distribute independent H3 presentations over encoder copies.
 
         H3 presentations have variable multimodal layouts, so they cannot be
         stacked into the generic text batch without changing padding/kernels.
-        Assigning one complete request to a rank preserves the exact
-        single-request encoder path, then broadcasts that request's native
-        payload to the other ranks.
+        Assigning one complete request to a copy preserves the exact
+        single-request encoder path. A copy may itself span a TP group; the
+        orthogonal encoder-DP group then broadcasts that request's native
+        payload to the other copies.
         """
         grouped = self._group_requests_by_fingerprint(
             batches,
@@ -142,14 +143,14 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage):
             payload = None
             if dp_group.rank_in_group == owner:
                 try:
-                    first_batch.extra[_MINIMAX_H3_SINGLE_RANK_TEXT_ENCODE_EXTRA_KEY] = (
+                    first_batch.extra[_MINIMAX_H3_SINGLE_COPY_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
+                            _MINIMAX_H3_SINGLE_COPY_TEXT_ENCODE_EXTRA_KEY, None
                         )
                     payload = first_result.extra.get(
                         MINIMAX_H3_TEXT_EMBEDDINGS_EXTRA_KEY
@@ -191,7 +192,7 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage):
         self._dp_choice_logged = True
         logger.info(
             "encoder_parallel: distributing %d independent MiniMax H3 "
-            "presentations over %d replicated encoder ranks",
+            "presentations over %d encoder copies",
             batch_size,
             world_size,
         )
@@ -381,7 +382,7 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage):
                 share_across_replicas=(
                     world > 1
                     and not bool(
-                        batch.extra.get(_MINIMAX_H3_SINGLE_RANK_TEXT_ENCODE_EXTRA_KEY)
+                        batch.extra.get(_MINIMAX_H3_SINGLE_COPY_TEXT_ENCODE_EXTRA_KEY)
                     )
                 ),
             )
@@ -536,7 +537,7 @@ 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):
+        if batch.extra.get(_MINIMAX_H3_SINGLE_COPY_TEXT_ENCODE_EXTRA_KEY):
             batch.extra.pop(MINIMAX_H3_PREPARED_REFERENCE_VIDEO_EXTRA_KEY, None)
         return {
             "positive": {

python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py

@@ -17,8 +17,8 @@ import torch
 from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput
 from sglang.multimodal_gen.configs.pipeline_configs.base import TextConditioningOutput
 from sglang.multimodal_gen.runtime.distributed import (
+    get_encoder_data_parallel_group,
     get_local_torch_device,
-    get_world_group,
 )
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
 from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
@@ -46,12 +46,13 @@ logger = init_logger(__name__)
 
 
 def _data_parallel_text_encode(forward_fn, forward_kwargs: dict, group):
-    """each rank encodes its 1/world_size batch slice, then all-gathers
+    """Each encoder copy encodes a batch slice, then all-gathers the outputs.
 
-    every rank runs the full unsharded encoder on its slice, so each row is
-    computed by the same kernels as the replicated forward; the batch is padded
-    to a multiple of world_size and padding rows are dropped after the gather.
-    Requires a TextEncoder (BaseEncoderOutput) -- see _text_encode_dp_group.
+    An encoder copy may itself be TP-sharded. Corresponding TP ranks use the
+    same batch-DP rank, so every rank in that TP group receives the same slice;
+    the orthogonal batch-DP groups then gather the replicated encoder outputs.
+    The batch is padded to a multiple of the DP degree and padding rows are
+    dropped after the gather. Requires a TextEncoder (BaseEncoderOutput).
     """
     world = group.world_size
     rank = group.rank_in_group
@@ -515,29 +516,32 @@ class TextEncodingStage(ConditionEncodingStage):
     ):
         """group to data-parallel a batched text-encode over, or None
 
-        requires a replicated encoder (tp==1, dp==1, not folded): each rank
-        would otherwise redundantly encode the whole batch. Also requires a
-        TextEncoder, whose forward returns BaseEncoderOutput -- the gather needs
-        to know which fields carry the batch, and a raw transformers encoder
-        returns its own output type (e.g. Qwen2_5_VLCausalLMOutputWithPast).
+        DP splits the request batch across encoder copies and all-gathers the
+        outputs. A non-folded native encoder uses the DiT TP group, so this uses
+        the orthogonal non-TP ranks inside the same pipeline replica. It never
+        mixes requests across pipeline replicas and composes with DiT TP.
         """
-        policy = server_args.encoder_parallel
-        if (
-            policy not in ("auto", "dp")
-            # isinstance first: the loader can return a raw transformers
-            # encoder, which carries no such attribute
-            or not isinstance(text_encoder, TextEncoder)
-            or not text_encoder.supports_dp_encode
-            or (server_args.tp_size or 1) != 1
-            or (server_args.dp_size or 1) != 1
-            or encoder_config.parallel_folding_mode is not None
-        ):
+        if server_args.encoder_parallel not in ("auto", "dp"):
             return None
-        group = get_world_group()
-        if group.world_size <= 1:
+        # A folded encoder has one TP copy spanning its folding group, so there
+        # are no independent copies over which to split the batch.
+        if encoder_config.parallel_folding_mode is not None:
+            return None
+        # the gather rebuilds a BaseEncoderOutput, which only a TextEncoder
+        # forward produces -- a raw transformers encoder returns its own output
+        # type (e.g. Qwen2_5_VLCausalLMOutputWithPast). isinstance first: the
+        # loader can return such an encoder, which carries no dp attribute.
+        if not isinstance(text_encoder, TextEncoder):
+            return None
+        if not text_encoder.supports_dp_encode:
+            return None
+        group = get_encoder_data_parallel_group()
+        if group is None or group.world_size <= 1:
             return None
         # explicit dp trusts the operator on an unmeasured topology; auto does not
-        measured = policy == "dp" or group_has_measured_topology(group)
+        measured = server_args.encoder_parallel == "dp" or group_has_measured_topology(
+            group
+        )
         if not encoder_dp_worthwhile(encoder_config, batch_size, measured):
             return None
         self._log_dp_choice(batch_size, group.world_size)
@@ -548,7 +552,7 @@ class TextEncodingStage(ConditionEncodingStage):
             return
         self._dp_choice_logged = True
         logger.info(
-            "encoder_parallel: data-parallel text encode over %d ranks "
+            "encoder_parallel: data-parallel text encode over %d encoder copies "
             "(batch %d). Measured 1.9x on the encode stage at batch 2/4/8 "
             "(2xH100, T5-XXL width) with max_abs_diff=0 against the replicated "
             "forward.",

python/sglang/multimodal_gen/runtime/server_args/server_args.py

@@ -710,8 +710,20 @@ class ServerArgs(DisaggServerArgsMixin):
         replica_size = (self.num_gpus or tp_size) // dp_size
 
         fold_world = dp_size == 1 and not self.disagg_mode and replica_size > tp_size
+        # An explicit fold policy is honored for any multi-rank replica —
+        # pure-TP (replica == tp) and dp>1 shapes included. Encoder layout is
+        # independent of the DiT's parallelism; the replica group is simply
+        # "the ranks that share this batch". `auto` keeps its conservative
+        # proposals below and never folds a pure-TP replica.
+        fold_replica = (
+            self.encoder_parallel == "fold"
+            and not self.disagg_mode
+            and replica_size > 1
+        )
 
-        if fold_world:
+        if fold_replica:
+            mode = "replica"
+        elif fold_world:
             mode = "world"
         elif tp_size == 1 and sp_degree > 1:
             # Preserve prior behavior for dp>1 / disaggregated SP runs.
@@ -1891,8 +1903,9 @@ class ServerArgs(DisaggServerArgsMixin):
                 "`auto` folds encoders wide enough to benefit (best "
                 "single-request latency) and data-parallels eligible native "
                 "text encoders at batch>1; `fold` always tensor-parallels the "
-                "encoder weights; `dp` never folds and splits the batch across "
-                "ranks (best batched throughput; requires TP=1 and DP=1); "
+                "encoder weights across the replica; `dp` never folds and "
+                "splits the batch across encoder copies inside each replica, "
+                "composing with encoder TP (best batched throughput); "
                 "`replicate` disables both. The default is `auto`."
             ),
         )
@@ -3171,10 +3184,6 @@ class ServerArgs(DisaggServerArgsMixin):
             raise ValueError("batching_max_size must be >= 1")
         if self.batching_delay_ms < 0:
             raise ValueError("batching_delay_ms must be >= 0")
-        if self.encoder_parallel == "dp" and (
-            (self.tp_size or 1) != 1 or (self.dp_size or 1) != 1
-        ):
-            raise ValueError("encoder_parallel=dp requires tp_size=1 and dp_size=1")
 
     def _set_default_attention_backend(self) -> None:
         """Configure ROCm defaults when users do not specify an attention backend."""

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

@@ -68,10 +68,25 @@ def _proposed_mode(tp, sp, cfg, dp=1, disagg=False, num_gpus=None, policy="auto"
 
 
 def test_pure_tp_not_folded():
-    # replica == tp: encoder already uses every replica GPU; nothing to fold.
+    # replica == tp under auto: each rank keeps a full encoder replica; auto
+    # never proposes folding here (the trade-off stays opt-in, see below).
     assert _proposed_mode(tp=2, sp=1, cfg=1) is None
 
 
+def test_explicit_fold_proposes_replica_for_any_shape():
+    # Encoder layout is independent of the DiT's parallelism: an operator's
+    # explicit fold is honored on every multi-rank replica instead of being
+    # silently ignored (pure TP) or narrowed to the SP group (dp > 1).
+    assert _proposed_mode(tp=2, sp=1, cfg=1, policy="fold") == "replica"
+    assert _proposed_mode(tp=1, sp=2, cfg=1, policy="fold") == "replica"
+    assert (
+        _proposed_mode(tp=2, sp=2, cfg=1, dp=2, num_gpus=8, policy="fold") == "replica"
+    )
+    # single-GPU replica: nothing to shard over
+    assert _proposed_mode(tp=1, sp=1, cfg=1, policy="fold") is None
+    assert _proposed_mode(tp=1, sp=1, cfg=1, dp=2, num_gpus=2, policy="fold") is None
+
+
 def test_single_gpu_not_folded():
     assert _proposed_mode(tp=1, sp=1, cfg=1) is None
 
@@ -121,10 +136,14 @@ def test_all_encoders_get_the_same_proposed_mode():
     assert img.parallel_folding_mode == "world"
 
 
-def test_adjust_proposes_regardless_of_policy():
-    # adjust reads the parallelism only; finalize owns the policy decision.
-    for policy in ("auto", "fold", "dp", "replicate"):
+def test_adjust_proposal_policy_dependence():
+    # adjust reads the parallelism only for auto/dp/replicate; finalize owns
+    # those policy decisions. An explicit fold is the one exception: it widens
+    # the proposal to the whole replica, because finalize cannot conjure a
+    # group that was never proposed (pure-TP replicas had none at all).
+    for policy in ("auto", "dp", "replicate"):
         assert _proposed_mode(tp=1, sp=2, cfg=1, policy=policy) == "world", policy
+    assert _proposed_mode(tp=1, sp=2, cfg=1, policy="fold") == "replica"
 
 
 def test_no_policy_touches_batching_max_size():

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

@@ -0,0 +1,119 @@
+"""Unit test for the batched text-encode data-parallel gate.
+
+_text_encode_dp_group decides whether a batched encode may split across
+encoder copies. It composes with DiT TP through the orthogonal encoder-DP
+group and stays off for folded encoders, single-copy replicas, and non-DP
+policies. Pure logic, no GPU / distributed init.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+from sglang.multimodal_gen.configs.models.encoders import TextEncoderConfig
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    _get_encoder_data_parallel_group_ranks,
+)
+from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder
+from sglang.multimodal_gen.runtime.pipelines_core.stages import text_encoding as _te_mod
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.distributed import RankGenerator
+
+
+def _enc(hidden=4096, heads=64, inter=10240, folding_mode=None):
+    enc = TextEncoderConfig()
+    enc.hidden_size = hidden
+    enc.num_attention_heads = heads
+    enc.intermediate_size = inter
+    enc.parallel_folding_mode = folding_mode
+    return enc
+
+
+def _gate(
+    monkeypatch,
+    policy="dp",
+    tp=1,
+    dp=1,
+    folding_mode=None,
+    batch_size=4,
+    world_size=4,
+    group_available=True,
+    measured=True,
+    supports_dp=True,
+):
+    group = SimpleNamespace(world_size=world_size) if group_available else None
+    monkeypatch.setattr(_te_mod, "get_encoder_data_parallel_group", lambda: group)
+    monkeypatch.setattr(_te_mod, "group_has_measured_topology", lambda group: measured)
+    encoder = Mock(spec=TextEncoder)
+    encoder.supports_dp_encode = supports_dp
+    server_args = SimpleNamespace(encoder_parallel=policy, tp_size=tp, dp_size=dp)
+    self = SimpleNamespace(_log_dp_choice=lambda batch, world: None)
+    return _te_mod.TextEncodingStage._text_encode_dp_group(
+        self, server_args, _enc(folding_mode=folding_mode), batch_size, encoder
+    )
+
+
+def test_dp_engages_for_replicated_encoder(monkeypatch):
+    assert _gate(monkeypatch) is not None
+
+
+def test_dit_tp_composes_with_encoder_dp_group(monkeypatch):
+    # TP ranks jointly run one encoder copy; the orthogonal group distributes
+    # batch slices across those TP-sharded copies.
+    assert _gate(monkeypatch, tp=4) is not None
+    assert _gate(monkeypatch, policy="auto", tp=4) is not None
+
+
+def test_pure_tp_has_no_independent_encoder_copy(monkeypatch):
+    # A pure-TP replica already uses every GPU for one encoder copy.
+    assert _gate(monkeypatch, tp=4, group_available=False) is None
+
+
+def test_folded_encoder_blocks_dp(monkeypatch):
+    # folding shards the weights over the folding group; a rank cannot
+    # encode alone, whatever the policy says
+    assert _gate(monkeypatch, folding_mode="world") is None
+    assert _gate(monkeypatch, tp=4, folding_mode="world") is None
+
+
+def test_multi_replica_uses_per_replica_encoder_group(monkeypatch):
+    # dp>1 no longer blocks: the encoder-DP group stays within one pipeline
+    # replica. A replica with one encoder copy uses its normal TP forward.
+    assert _gate(monkeypatch, dp=2) is not None
+    assert _gate(monkeypatch, dp=2, world_size=1) is None
+
+
+def test_policy_and_encoder_gates(monkeypatch):
+    assert _gate(monkeypatch, policy="replicate") is None
+    assert _gate(monkeypatch, policy="fold") is None
+    assert _gate(monkeypatch, supports_dp=False) is None
+    assert _gate(monkeypatch, world_size=1) is None
+    assert _gate(monkeypatch, batch_size=1) is None  # unbatched: not worthwhile
+    # explicit dp trusts the operator on an unmeasured topology; auto does not
+    assert _gate(monkeypatch, policy="dp", measured=False) is not None
+    assert _gate(monkeypatch, policy="auto", measured=False) is None
+
+
+def _validate_batching(encoder_parallel, tp, dp):
+    args = SimpleNamespace(
+        batching_mode="dynamic",
+        batching_max_size=1,
+        batching_delay_ms=0,
+        encoder_parallel=encoder_parallel,
+        tp_size=tp,
+        dp_size=dp,
+    )
+    ServerArgs._validate_batching(args)
+
+
+def test_server_args_accepts_dp_for_any_parallel_shape():
+    # Encoder DP composes with TP and stays within each pipeline replica.
+    _validate_batching("dp", tp=4, dp=1)
+    _validate_batching("dp", tp=1, dp=2)
+    _validate_batching("dp", tp=4, dp=2)
+
+
+def test_encoder_dp_groups_are_tp_orthogonal_and_replica_local():
+    ranks = _get_encoder_data_parallel_group_ranks(
+        RankGenerator(tp=2, sp=2, pp=1, cfg=1, dp=2, order="tp-sp-pp-cfg-dp")
+    )
+    assert ranks == [[0, 2], [1, 3], [4, 6], [5, 7]]