Check the topology identities where the layout is written, and build at the published widths (#40340)

python/sglang/srt/arg_groups/arg_utils.py

@@ -110,7 +110,7 @@ class Derived(msgspec.Struct, frozen=True):
     Every declaration carries ``fn`` today, the parallel quotients included:
     they are a function of the configured leaves, so they are computed at
     publish like the rest. What is special about them is not how they are
-    computed but that a stamp can move one afterwards -- an elastic scale-up
+    computed but that a stamp can move one afterwards -- ``initialize_dp_attention``
     restamps ``attn_dp_size`` -- which ``ParallelContext`` answers above the
     published leaf.
     """

python/sglang/srt/disaggregation/encoder/server.py

@@ -589,10 +589,32 @@ class MMEncoder:
             distributed_init_method=dist_init_method,
             local_rank=rank,
         )
-        initialize_model_parallel(
-            tensor_model_parallel_size=get_parallel().tp_size,
-            attention_context_model_parallel_size=get_parallel().attn_cp_size,
+        # The encoder serves the vision tower on a world of its own: `tp_size`
+        # ranks wide, with no pipeline, no expert or MoE-DP dimension and no
+        # decode context parallelism, whatever the generation side published.
+        # That has always been the layout it builds; stating it is what stops
+        # the context from answering with the other side's topology while these
+        # groups answer with this one.
+        parallel = get_parallel()
+        attn_cp_size = parallel.attn_cp_size
+        attn_tp_size = parallel.tp_size // attn_cp_size
+        attn_cp_rank, attn_tp_rank = divmod(rank, attn_tp_size)
+        parallel.override_permanently(
+            tp_rank=rank,
+            pp_size=1,
+            pp_rank=0,
+            attn_dp_size=1,
+            attn_dp_rank=0,
+            attn_tp_size=attn_tp_size,
+            attn_tp_rank=attn_tp_rank,
+            attn_cp_rank=attn_cp_rank,
+            attn_dcp_size=1,
+            moe_ep_size=1,
+            moe_ep_rank=0,
+            moe_dp_size=1,
+            moe_tp_size=parallel.tp_size,
         )
+        initialize_model_parallel()
         initialize_dp_attention(server_args, self.model_config)
 
         self.model = load_model(

python/sglang/srt/distributed/bootstrap.py

@@ -106,15 +106,6 @@ def init_torch_distributed(
             server_args=server_args,
             model_config=model_config,
             gpu_id=ps.gpu_id,
-            tp_rank=ps.tp_rank,
-            tp_size=ps.tp_size,
-            pp_rank=ps.pp_rank,
-            pp_size=ps.pp_size,
-            attn_dp_size=ps.attn_dp_size,
-            attn_cp_size=ps.attn_cp_size,
-            moe_ep_size=ps.moe_ep_size,
-            moe_dp_size=ps.moe_dp_size,
-            dcp_size=ps.attn_dcp_size,
         )
 
         # Pre-warm NCCL/RCCL/HCCL to eliminate cold-start latency in first request
@@ -255,19 +246,13 @@ def _init_parallel_groups(
     server_args: ServerArgs,
     model_config: ModelConfig,
     gpu_id: int,
-    tp_rank: int,
-    tp_size: int,
-    pp_rank: int,
-    pp_size: int,
-    attn_dp_size: int,
-    attn_cp_size: int,
-    moe_ep_size: int,
-    moe_dp_size: int,
-    dcp_size: int,
 ) -> None:
+    parallel = get_parallel()
+    tp_size, pp_size = parallel.tp_size, parallel.pp_size
+    tp_rank, pp_rank = parallel.tp_rank, parallel.pp_rank
     is_ep_joiner = get_exec().moe.is_ep_joiner
     is_scale_joiner = get_exec().moe.is_ep_scale_joiner
-    rank_offset = get_parallel().ep_join_rank_offset if is_scale_joiner else 0
+    rank_offset = parallel.ep_join_rank_offset if is_scale_joiner else 0
     world_size = (
         rank_offset + tp_size * pp_size if is_scale_joiner else tp_size * pp_size
     )
@@ -285,14 +270,6 @@ def _init_parallel_groups(
         max_world_size=get_parallel().max_ep_size,
     )
     initialize_model_parallel(
-        tensor_model_parallel_size=tp_size,
-        attention_data_parallel_size=attn_dp_size,
-        pipeline_model_parallel_size=pp_size,
-        expert_model_parallel_size=moe_ep_size,
-        attention_context_model_parallel_size=attn_cp_size,
-        moe_data_model_parallel_size=moe_dp_size,
-        decode_context_parallel_size=dcp_size,
-        shared_experts_tensor_parallel_size=get_parallel().shared_experts_tp_size,
         duplicate_tp_group=get_disagg().enable_pdmux,
         enable_symm_mem=get_exec().comm.enable_symm_mem,
         # Only WORLD is extended during scale-up. The joiner's model-parallel

python/sglang/srt/distributed/parallel_state.py

@@ -50,6 +50,7 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo
 )
 from sglang.srt.platforms.device_mixin import _DEVICE_TO_DISTRIBUTED_BACKEND
 from sglang.srt.runtime_context import (
+    _validate_parallel,
     derive_parallel_widths,
     get_global_dwdp_manager,
     get_parallel,
@@ -2515,44 +2516,41 @@ def init_distributed_environment(
 
 
 def initialize_model_parallel(
-    tensor_model_parallel_size: int = 1,
-    expert_model_parallel_size: int = 1,
-    pipeline_model_parallel_size: int = 1,
-    attention_data_parallel_size: int = 1,
-    attention_context_model_parallel_size: int = 1,
-    moe_data_model_parallel_size: int = 1,
-    decode_context_parallel_size: int = 1,
     backend: Optional[str] = None,
     duplicate_tp_group: bool = False,
     enable_symm_mem: bool = False,
     recovered_rank: bool = False,
     rank_offset: int = 0,
     max_world_size: Optional[int] = None,
-    shared_experts_tensor_parallel_size: Optional[int] = None,
 ) -> None:
     """
-    Initialize model parallel groups.
+    Initialize model parallel groups at the published widths.
 
-    Arguments:
-        tensor_model_parallel_size: number of GPUs used for tensor model
-            parallelism.
-        expert_model_parallel_size: number of GPUs used for expert model
-            parallelism.
-        pipeline_model_parallel_size: number of GPUs used for pipeline model
-            parallelism.
-        attention_data_parallel_size: number of GPUs used for attention data
-            parallelism.
-        attention_context_model_parallel_size: number of GPUs used for attention context
-            parallelism.
-        moe_data_model_parallel_size: number of GPUs used for moe data
-            parallelism.
-        decode_context_parallel_size: number of GPUs used for decode context
-            parallelism, which splits the KV cache across GPUs within each
-            tensor-parallel group during decoding. Must be a divisor of
-            tensor_model_parallel_size and is currently only supported on the
-            AMD HIP platform.
-        shared_experts_tensor_parallel_size: optional shared-expert TP width.
-            Must divide attention TP; subgroups never cross attention replicas.
+    Every width comes from the runtime context rather than from an argument:
+    the configuration already says how wide each dimension is, and a caller
+    that translates it again is a second place for the two to disagree. A
+    process that needs a narrower layout than the one it published -- the
+    media encoder is the case in the tree -- states that layout on the context
+    first, so what it builds and what it answers stay the same thing.
+
+    The remaining arguments are not topology. `backend` is decided by the
+    device, `duplicate_tp_group` and `enable_symm_mem` by other namespaces, and
+    `recovered_rank` / `rank_offset` / `max_world_size` describe this
+    particular join rather than the layout being joined.
+
+    The widths this reads:
+        tp_size: GPUs used for tensor model parallelism.
+        moe_ep_size: GPUs used for expert model parallelism.
+        pp_size: GPUs used for pipeline model parallelism.
+        attn_dp_size: GPUs used for attention data parallelism.
+        attn_cp_size: GPUs used for attention context parallelism.
+        moe_dp_size: GPUs used for MoE data parallelism.
+        attn_dcp_size: GPUs used for decode context parallelism, which splits
+            the KV cache across GPUs within each tensor-parallel group during
+            decoding. Must be a divisor of `tp_size` and is currently only
+            supported on the AMD HIP platform.
+        shared_experts_tp_size: optional shared-expert TP width. Must divide
+            attention TP; subgroups never cross attention replicas.
 
     Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we
     use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize
@@ -2589,6 +2587,16 @@ def initialize_model_parallel(
     assert torch.distributed.is_initialized()
     backend = backend or torch.distributed.get_backend(get_world_group().device_group)
 
+    parallel = get_parallel()
+    tensor_model_parallel_size = parallel.tp_size
+    expert_model_parallel_size = parallel.moe_ep_size
+    pipeline_model_parallel_size = parallel.pp_size
+    attention_data_parallel_size = parallel.attn_dp_size
+    attention_context_model_parallel_size = parallel.attn_cp_size
+    moe_data_model_parallel_size = parallel.moe_dp_size
+    decode_context_parallel_size = parallel.attn_dcp_size
+    shared_experts_tensor_parallel_size = parallel.shared_experts_tp_size
+
     # Joiners construct their local TP/PP layout in global rank space.
     world_size: int = (
         tensor_model_parallel_size * pipeline_model_parallel_size
@@ -2933,6 +2941,12 @@ def initialize_model_parallel(
             max_world_size=max_world_size,
         )
 
+    # The groups just built and the configuration they were built from are two
+    # accounts of one layout. Check them against each other here, where the
+    # disagreement is still attributable, rather than letting a collective run
+    # on the wrong peers.
+    _validate_parallel(get_parallel(), "group build")
+
 
 def create_custom_parallel_group(
     group_ranks: List[int], backend: str = "gloo"
@@ -3094,10 +3108,18 @@ def patch_tensor_parallel_group(tp_group: GroupCoordinator, *, owns_attention: b
         narrowed.update(
             attn_tp_size=tp_group.world_size,
             attn_tp_rank=tp_group.rank_in_group,
+            attn_tp_group=tp_group,
             attn_dp_size=1,
             attn_dp_rank=0,
             attn_cp_size=1,
             attn_cp_rank=0,
+            attn_cp_group=None,
+            moe_ep_size=1,
+            moe_ep_rank=0,
+            moe_ep_group=None,
+            moe_dp_size=1,
+            moe_tp_size=tp_group.world_size,
+            moe_tp_rank=tp_group.rank_in_group,
         )
     try:
         with get_parallel().override(**narrowed):

python/sglang/srt/layers/dp_attention.py

@@ -15,11 +15,9 @@ from sglang.srt.arg_groups.model_override_base import (
 )
 from sglang.srt.distributed import (
     GroupCoordinator,
-    get_attn_tensor_model_parallel_world_size,
 )
 from sglang.srt.distributed import get_moe_dp_group as _get_moe_dp_group
 from sglang.srt.distributed import (
-    get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
 )
 from sglang.srt.distributed.device_communicators.pynccl_allocator import (
@@ -49,6 +47,33 @@ if TYPE_CHECKING:
     from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 
 
+def dp_gather_width() -> int:
+    """How many replicas the DP sync gathers over.
+
+    The attention-DP replicas, except after an elastic-EP scale-up, when the
+    gather spans the expanded WORLD -- whose width is the `dp_size` the
+    scale-up published. Read from the context either way: a scoped width has
+    to reach this, which is the whole reason the name has one home.
+    """
+    parallel = get_parallel()
+    return parallel.dp_size if world_dp_gather_enabled() else parallel.attn_dp_size
+
+
+def dp_gather_slot() -> int:
+    """This process's index in the list the DP sync just gathered.
+
+    The gather spans the attention-DP replicas, except after an elastic-EP
+    scale-up, when it spans the expanded WORLD and the joining cohort is
+    numbered from its offset. Which list was gathered is what the flag below
+    says, so the index is read from there rather than kept as a second name on
+    the topology.
+    """
+    parallel = get_parallel()
+    if world_dp_gather_enabled():
+        return parallel.tp_rank + parallel.ep_join_rank_offset
+    return parallel.attn_dp_rank
+
+
 def world_dp_gather_enabled() -> bool:
     """Whether DP gathers should use expanded WORLD after joiner admission."""
     dp = get_flags().dp
@@ -60,9 +85,13 @@ def enable_joiner_all_gather():
 
 
 def update_dp_attention_post_scale(new_dp_size: int, new_dp_rank: int):
-    get_parallel().override_permanently(
-        attn_dp_size=new_dp_size, attn_dp_rank=new_dp_rank
-    )
+    """Point the DP gather at the expanded WORLD.
+
+    The widths themselves are not written here: the caller scales `dp_size` on
+    the published bag, and the gather reads its width and this process's slot
+    from there. The arguments are the values the caller is about to publish,
+    kept so the log says which scale-up this was.
+    """
     get_flags().dp.use_world_group_for_gather = True
     logger.debug(
         "[Elastic EP] dp_attention switched to WORLD: dp_size=%d dp_rank=%d",
@@ -92,7 +121,7 @@ class DpPaddingMode(IntEnum):
     def get_dp_padding_mode(
         cls, is_extend_in_batch, global_num_tokens: List[int]
     ) -> DpPaddingMode:
-        dp_size = get_parallel().attn_dp_size
+        dp_size = dp_gather_width()
 
         # (trangdough) pplx-kernels a2a is a symmetric collective: every EP rank
         # must dispatch the same number of tokens or the device-side handshake
@@ -374,14 +403,13 @@ def initialize_dp_attention(
     dp.enabled = enable_dp_attention
 
     tp_rank = get_parallel().tp_rank
-    tp_size = get_tensor_model_parallel_world_size()
+    tp_size = get_parallel().tp_size
 
     _, _, attn_dp_rank, attn_dp_size = compute_dp_attention_world_info(
         enable_dp_attention, tp_rank, tp_size, dp_size, attn_cp_size
     )
 
     if get_exec().moe.elastic_ep_backend is not None and get_parallel().max_ep_size:
-        attn_dp_rank = tp_rank + get_parallel().ep_join_rank_offset
         # Reads the resolution, not a bag: this runs under
         # `initialize_dp_attention`, which the weight-cache daemon calls from
         # `_init_distributed` -- and other callers reach it from processes
@@ -414,7 +442,10 @@ def is_allocation_symmetric() -> bool:
 
 def get_dp_local_info(forward_batch: ForwardBatch) -> Tuple[torch.Tensor, torch.Tensor]:
     # `get_dp_local_info` is only called in global DP gather and scatter. We use global DP rank here.
-    dp_rank = get_parallel().attn_dp_rank
+    # The slot in the list that was gathered. A scale-up widens that list
+    # to WORLD, and this process's index in it is not its index among the
+    # launch replicas.
+    dp_rank = dp_gather_slot()
 
     if forward_batch.dp_local_start_pos is None:
         cumtokens = torch.cumsum(forward_batch.global_num_tokens_gpu, dim=0)
@@ -438,7 +469,10 @@ def get_dp_local_slice_cpu(
     # CPU (start, length) slice for DP-local data in a rank-padded buffer.
     # Returns Python ints (no D2H sync) and handles the cuda-graph-padded layout.
     global_num_tokens = forward_batch.global_num_tokens_cpu
-    dp_rank = get_parallel().attn_dp_rank
+    # The slot in the list that was gathered. A scale-up widens that list
+    # to WORLD, and this process's index in it is not its index among the
+    # launch replicas.
+    dp_rank = dp_gather_slot()
     local_num_tokens = global_num_tokens[dp_rank]
     if can_run_graph:
         local_start_pos = dp_rank * cuda_graph_batch
@@ -514,7 +548,7 @@ def _dp_gather_via_all_reduce(
         NUM_GPUS_PER_NODE = 8
         if (
             not local_tokens.dtype.is_floating_point
-            and get_tensor_model_parallel_world_size() <= NUM_GPUS_PER_NODE
+            and get_parallel().tp_size <= NUM_GPUS_PER_NODE
         ):
             from sglang.srt.distributed.parallel_state import inplace_all_reduce
 
@@ -534,7 +568,7 @@ def _dp_gather_via_all_gather(
 ):
     use_world = world_dp_gather_enabled()
 
-    if get_attn_tensor_model_parallel_world_size() == 1:
+    if get_parallel().attn_tp_size == 1:
         if use_world:
             torch.distributed.all_gather_into_tensor(
                 global_tokens,
@@ -548,9 +582,9 @@ def _dp_gather_via_all_gather(
     if not is_partial:
         if get_parallel().attn_tp_rank != 0:
             local_tokens.fill_(0)
-    scattered_local_tokens = local_tokens.tensor_split(
-        get_attn_tensor_model_parallel_world_size()
-    )[get_parallel().attn_tp_rank]
+    scattered_local_tokens = local_tokens.tensor_split(get_parallel().attn_tp_size)[
+        get_parallel().attn_tp_rank
+    ]
     get_parallel().attn_tp_group.reduce_scatter_tensor(
         scattered_local_tokens, local_tokens
     )
@@ -721,8 +755,8 @@ def is_dp_gatherv_active() -> bool:
     return (
         _USE_DP_GATHERV
         and not world_dp_gather_enabled()
-        and get_attn_tensor_model_parallel_world_size() == 1
-        and get_tensor_model_parallel_world_size() == get_parallel().attn_dp_size
+        and get_parallel().attn_tp_size == 1
+        and get_parallel().tp_size == get_parallel().attn_dp_size
         and not _DpGatheredBufferWrapper.is_dp_max_padding()
     )
 
@@ -881,12 +915,12 @@ def dp_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor):
         if sizes is not None:
             get_parallel().tp_group.reduce_scatterv(input, output=output, sizes=sizes)
             return
-    if get_tensor_model_parallel_world_size() == get_parallel().attn_dp_size:
+    if get_parallel().tp_size == get_parallel().attn_dp_size:
         get_parallel().tp_group.reduce_scatter_tensor(output, input)
     else:
-        scattered_local_tokens = input.tensor_split(
-            get_tensor_model_parallel_world_size()
-        )[get_parallel().tp_rank]
+        scattered_local_tokens = input.tensor_split(get_parallel().tp_size)[
+            get_parallel().tp_rank
+        ]
         get_parallel().tp_group.reduce_scatter_tensor(scattered_local_tokens, input)
         get_parallel().attn_tp_group.all_gather_into_tensor(
             output, scattered_local_tokens

python/sglang/srt/layers/sampler.py

@@ -112,7 +112,12 @@ class Sampler(nn.Module):
         self.cp_sync_group = None
         if is_dp_attention_enabled():
             self.tp_sync_group = get_parallel().attn_tp_group.device_group
-            self.cp_sync_group = get_parallel().attn_cp_group.device_group
+            # Only when there is more than one context shard to reconcile. The
+            # sync below already short-circuits on that, and a model running on
+            # one shard -- a speculative draft, under the scope that says so --
+            # has no context-parallel communicator to name.
+            if get_parallel().attn_cp_size > 1:
+                self.cp_sync_group = get_parallel().attn_cp_group.device_group
 
         self.rl_on_policy_target = get_exec().deterministic.rl_on_policy_target
         # In RL on-policy mode, deterministic inference is automatically enabled.

python/sglang/srt/managers/scheduler_components/dp_attn.py

@@ -10,7 +10,7 @@ from sglang.srt.configs.model_config import ModelConfig
 from sglang.srt.distributed.parallel_state_wrapper import ParallelState
 from sglang.srt.environ import envs
 from sglang.srt.layers.cp.utils import get_cp_strategy
-from sglang.srt.layers.dp_attention import world_dp_gather_enabled
+from sglang.srt.layers.dp_attention import dp_gather_width, world_dp_gather_enabled
 from sglang.srt.layers.moe.utils import get_moe_a2a_backend
 from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.scheduler_components.recv_skipper import (
@@ -58,7 +58,7 @@ def _resolve_elastic_world_dp_size(
 
     from sglang.srt.elastic_ep.elastic_ep import ElasticEPStateManager
 
-    live_dp_size = get_parallel().attn_dp_size
+    live_dp_size = dp_gather_width()
     effective_ep_size = ElasticEPStateManager.get_effective_ep_size()
     # The group's own membership, not the width it was built at: this is the
     # one number an out-of-process join moves, and it is the upper bound the

python/sglang/srt/model_executor/forward_batch_info.py

@@ -45,6 +45,7 @@ from sglang.srt.kv_canary.req_to_expected_token_ids_manager import (
 )
 from sglang.srt.layers.dp_attention import (
     DpPaddingMode,
+    dp_gather_slot,
     set_dp_buffer_len,
     set_is_extend_in_batch,
     world_dp_gather_enabled,
@@ -1165,9 +1166,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
         if self.global_num_tokens_cpu is not None:
             # DP / MLP-sync path: per-DP padded width.
             if require_mlp_tp_gather():
-                num_tokens_per_dp = self.global_num_tokens_cpu[
-                    get_parallel().attn_dp_rank
-                ]
+                num_tokens_per_dp = self.global_num_tokens_cpu[dp_gather_slot()]
             else:
                 num_tokens_per_dp = self.global_num_tokens_cpu[0]
         else:
@@ -1520,7 +1519,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             buffer_len = sum(global_num_tokens)
 
         if len(global_num_tokens) > 1:
-            num_tokens = global_num_tokens[get_parallel().attn_dp_rank]
+            num_tokens = global_num_tokens[dp_gather_slot()]
         else:
             num_tokens = global_num_tokens[0]
 

python/sglang/srt/models/inkling_common/moe.py

@@ -673,9 +673,15 @@ class InklingSharedFusedMoE(FusedMoE):
     ) -> None:
         # FusedMoE.__init__ reads get_parallel() once and caches it on self, so
         # scoping the override to just this call is sufficient for the module's lifetime.
+        # The shared experts are replicated rather than sharded, so there is no
+        # expert-parallel communication here and no group to name: a width of
+        # one with the wider group still installed would describe a layout that
+        # does not exist.
         with get_parallel().override(
             moe_ep_size=1,
             moe_ep_rank=0,
+            moe_ep_group=None,
+            moe_dp_size=1,
             moe_tp_size=get_parallel().tp_size,
             moe_tp_rank=get_parallel().tp_rank,
         ):

python/sglang/srt/models/kimi_k3.py

@@ -22,7 +22,6 @@ from sglang.srt.configs.kimi_k3 import KimiK3Config
 from sglang.srt.configs.kimi_linear import KimiLinearConfig
 from sglang.srt.distributed import (
     divide,
-    get_shared_experts_tp_group,
     tensor_model_parallel_all_reduce,
 )
 from sglang.srt.distributed.device_communicators.pynccl_allocator import (
@@ -582,7 +581,7 @@ class KimiK3MoE(nn.Module):
             shared_experts_tp_kwargs = dict(tp_rank=0, tp_size=1)
         elif self._shared_experts_tp_comm:
             group = (
-                get_shared_experts_tp_group()
+                parallel.shared_experts_tp_group
                 if requested_shared_tp is not None
                 else parallel.attn_tp_group
             )

python/sglang/srt/runtime_context.py

@@ -218,6 +218,7 @@ _LIVE_READS: dict = {
     "moe_tp_group": "get_moe_tp_group",
     "attn_tp_group": "get_attn_tp_group",
     "attn_cp_group": "get_attn_cp_group",
+    "shared_experts_tp_group": "get_shared_experts_tp_group",
     "dcp_group": "get_dcp_group",
 }
 
@@ -448,14 +449,138 @@ class SpawnRanks(msgspec.Struct, frozen=True):
     the rank cannot say which replica this is. `None` means "no controller",
     which is an answer rather than an absence, and it is recorded as one.
 
-    Nothing else belongs here. A device index, for instance, is a placement
-    decision rather than a position -- the launcher may reindex it, and Ray
-    assigns it from its own allocator -- so it stays an argument to whoever
-    was handed it.
+    `gpu_id` is the device the parent picked for this process. It is not a
+    position in any group -- reindexing narrows the visible devices before the
+    spawn, and Ray allocates from its own pool -- but it is the same kind of
+    fact: something only the entry that spawned the process can state. `None`
+    for a process that runs on no device.
     """
 
     world_rank: int
     dp_rank: Optional[int] = None
+    gpu_id: Optional[int] = None
+
+
+_RANK_AND_WIDTH = (
+    ("tp_rank", "tp_size"),
+    ("pp_rank", "pp_size"),
+    ("attn_tp_rank", "attn_tp_size"),
+    ("attn_dp_rank", "attn_dp_size"),
+    ("attn_cp_rank", "attn_cp_size"),
+    ("moe_ep_rank", "moe_ep_size"),
+)
+
+# `moe_dp` is absent because `initialize_model_parallel` aliases the MoE-DP
+# group to the attention-CP group when the latter is wider: there the group and
+# the name are two facts, which is the same reason `moe_dp_rank` is left off the
+# record at publish.
+_WIDTH_AND_GROUP = (
+    ("tp_size", "tp_group"),
+    ("pp_size", "pp_group"),
+    ("attn_tp_size", "attn_tp_group"),
+    ("attn_cp_size", "attn_cp_group"),
+    ("moe_ep_size", "moe_ep_group"),
+)
+
+_UNREADABLE = object()
+
+
+def _validate_parallel(parallel, source: str) -> None:
+    """Fail on a topology that cannot describe a real process layout.
+
+    Every identity holds unconditionally: a width and a rank are both
+    plausible small integers whichever way they are wrong, so an inconsistent
+    set is not caught by anything downstream -- it surfaces as a hang or a
+    wrong answer in a collective, far from the write. Stating one leaf without
+    the quotients that follow from it leaves the namespace describing no real
+    layout, and the caller that did so is the one that has to say what it meant.
+
+    Names that cannot be read are skipped rather than treated as zero: a
+    process that has published nothing can still stamp a rank, and a group that
+    has not been built answers nothing at all.
+    """
+
+    def read(name):
+        """A width or a rank, or `_UNREADABLE` for anything these identities
+        cannot be stated about -- an absent name, `None`, or a stand-in a test
+        put in a group's place. Booleans are integers in Python and are not
+        widths, so they are out too."""
+        try:
+            value = getattr(parallel, name)
+        except Exception:
+            return _UNREADABLE
+        if isinstance(value, bool) or not isinstance(value, int):
+            return _UNREADABLE
+        return value
+
+    problems = []
+
+    for rank_name, size_name in _RANK_AND_WIDTH:
+        rank, size = read(rank_name), read(size_name)
+        if _UNREADABLE in (rank, size):
+            continue
+        if not 0 <= rank < size:
+            problems.append(
+                f"0 <= {rank_name} < {size_name}\n  {rank} is not a rank of {size}"
+            )
+
+    terms = ("tp_size", "attn_tp_size", "attn_dp_size", "attn_cp_size")
+    tp_size, a_tp, a_dp, a_cp = (read(n) for n in terms)
+    if _UNREADABLE not in (tp_size, a_tp, a_dp, a_cp):
+        if tp_size != a_tp * a_dp * a_cp:
+            problems.append(
+                "tp_size == attn_tp_size * attn_dp_size * attn_cp_size\n"
+                f"  {tp_size} != {a_tp} * {a_dp} * {a_cp} (= {a_tp * a_dp * a_cp})"
+            )
+
+    moe_terms = ("tp_size", "moe_ep_size", "moe_dp_size", "moe_tp_size")
+    tp_size, m_ep, m_dp, m_tp = (read(n) for n in moe_terms)
+    if _UNREADABLE not in (tp_size, m_ep, m_dp, m_tp):
+        if tp_size != m_ep * m_dp * m_tp:
+            problems.append(
+                "tp_size == moe_ep_size * moe_dp_size * moe_tp_size\n"
+                f"  {tp_size} != {m_ep} * {m_dp} * {m_tp} (= {m_ep * m_dp * m_tp})"
+            )
+
+    layout_terms = (
+        "tp_rank",
+        "attn_dp_rank",
+        "attn_cp_rank",
+        "attn_tp_rank",
+        "attn_cp_size",
+        "attn_tp_size",
+    )
+    tp_rank, r_dp, r_cp, r_tp, w_cp, w_tp = (read(n) for n in layout_terms)
+    if _UNREADABLE not in (tp_rank, r_dp, r_cp, r_tp, w_cp, w_tp):
+        laid_out = (r_dp * w_cp + r_cp) * w_tp + r_tp
+        if tp_rank != laid_out:
+            problems.append(
+                "tp_rank == (attn_dp_rank * attn_cp_size + attn_cp_rank)"
+                " * attn_tp_size + attn_tp_rank\n"
+                f"  {tp_rank} != ({r_dp} * {w_cp} + {r_cp})"
+                f" * {w_tp} + {r_tp} (= {laid_out})"
+            )
+
+    for size_name, group_name in _WIDTH_AND_GROUP:
+        size = read(size_name)
+        if size is _UNREADABLE:
+            continue
+        try:
+            group = getattr(parallel, group_name)
+        except Exception:
+            continue
+        built = getattr(group, "world_size", _UNREADABLE)
+        if isinstance(built, int) and not isinstance(built, bool) and built != size:
+            problems.append(
+                f"{group_name}.world_size == {size_name}\n"
+                f"  built {built}, configured {size}"
+            )
+
+    if problems:
+        raise ValueError(
+            f"parallel topology is inconsistent (set by {source}):\n"
+            + "\n".join(problems)
+        )
 
 
 class ParallelContext:
@@ -561,7 +686,13 @@ class ParallelContext:
         unknown = set(values) - _parallel_fields()
         if unknown:
             raise ValueError(f"unknown parallel field(s): {sorted(unknown)}")
+        saved = dict(self._stamp)
         self._stamp.update(values)
+        try:
+            _validate_parallel(self, "override_permanently")
+        except Exception:
+            self._stamp = saved
+            raise
 
     def clear_stamp(self) -> None:
         """Drop every stamped name, ranks included."""
@@ -576,6 +707,11 @@ class ParallelContext:
             raise ValueError(f"unknown parallel field(s): {sorted(unknown)}")
         saved = dict(self._overrides)
         self._overrides.update(kwargs)
+        try:
+            _validate_parallel(self, "override")
+        except Exception:
+            self._overrides = saved
+            raise
         try:
             yield self
         finally:
@@ -1776,6 +1912,8 @@ def publish(
             ),
         )
     _CONTEXT._publish_role = role
+    if ranks is not None and ranks.gpu_id is not None:
+        _CONTEXT.override("spawn", gpu_id=ranks.gpu_id)
     if ranks is not None:
         # The placement, worked out here rather than carried: the widths are on
         # the bag a moment ago, and `world_rank` fixes the rest. A read of any
@@ -1804,8 +1942,13 @@ def publish(
         # "no controller" rather than an absence.
         placement["dp_rank"] = ranks.dp_rank
         placement["launch_world_rank"] = ranks.world_rank
+        placement.update(_attention_ranks(parallel, placement["tp_rank"]))
+        # One stamp, not two: the identities are checked on every write, and a
+        # half-placed process satisfies none of them.
         parallel.override_permanently(**placement)
-        _stamp_attention_ranks(parallel, placement["tp_rank"])
+        # Publish established the whole layout, so every identity applies here,
+        # not just the ones the stamp happened to name.
+        _validate_parallel(parallel, "publish")
     if _ROLE_NS_MODE == "record":
         # The '-' marker distinguishes a zero-read role from a process where
         # recording never ran (signal teardown skips atexit).
@@ -1820,7 +1963,7 @@ def publish(
     return _CONTEXT
 
 
-def _stamp_attention_ranks(parallel, tp_rank: int) -> None:
+def _attention_ranks(parallel, tp_rank: int) -> dict:
     """Place this process in the attention topology, from the configuration.
 
     The widths are already on the bag -- `publish` computed them a moment ago --
@@ -1829,8 +1972,9 @@ def _stamp_attention_ranks(parallel, tp_rank: int) -> None:
     that never initialises distributed, which is what `ParallelState` provided
     by being a plain frozen record.
 
-    It is a stamp rather than a bag leaf because it is a per-process fact, and
-    nothing about the configuration distinguishes one rank from another.
+    These are stamped rather than written as bag leaves because they are
+    per-process facts, and nothing about the configuration distinguishes one
+    rank from another.
     """
     attn_tp_rank, attn_dp_rank = derive_attention_ranks(
         tp_rank=tp_rank,
@@ -1838,7 +1982,7 @@ def _stamp_attention_ranks(parallel, tp_rank: int) -> None:
         attn_cp_size=parallel.attn_cp_size,
         enable_dp_attention=parallel.enable_dp_attention,
     )
-    parallel.override_permanently(attn_tp_rank=attn_tp_rank, attn_dp_rank=attn_dp_rank)
+    return {"attn_tp_rank": attn_tp_rank, "attn_dp_rank": attn_dp_rank}
 
 
 def assert_published(server_args, *, role: str) -> RuntimeContext:

python/sglang/srt/weight_cache/daemon.py

@@ -232,16 +232,7 @@ class WeightCacheDaemon:
                 moe_a2a_backend=self.moe_a2a_backend,
             )
 
-        initialize_model_parallel(
-            tensor_model_parallel_size=self.tp_size,
-            pipeline_model_parallel_size=self.pp_size,
-            expert_model_parallel_size=self.ep_size,
-            attention_data_parallel_size=(
-                self.dp_size if self.enable_dp_attention else 1
-            ),
-            attention_context_model_parallel_size=self.attn_cp_size,
-            moe_data_model_parallel_size=self.moe_dp_size,
-        )
+        initialize_model_parallel()
 
         # Initialize DP attention state (required by some models like Qwen3 MoE)
         from sglang.srt.layers.dp_attention import initialize_dp_attention

python/sglang/test/layer_ut_utils.py

@@ -7,6 +7,8 @@ import os
 
 import torch
 
+from sglang.test.test_utils import publish_build_topology
+
 
 def init_single_process_dist(master_port: int = 29632, backend: str = "gloo"):
     """world=1 dist + model-parallel groups; srt layers require them even
@@ -29,12 +31,8 @@ def init_single_process_dist(master_port: int = 29632, backend: str = "gloo"):
     if not model_parallel_is_initialized():
         # kwargs only: a positional backend would land in the
         # attention_data_parallel_size slot and explode on int // str.
-        initialize_model_parallel(
-            tensor_model_parallel_size=1,
-            expert_model_parallel_size=1,
-            pipeline_model_parallel_size=1,
-            backend=backend,
-        )
+        publish_build_topology(tp_size=1, ep_size=1, pp_size=1)
+        initialize_model_parallel(backend=backend)
 
 
 def make_tp1_column_parallel_linear(

python/sglang/test/test_utils.py

@@ -2084,6 +2084,28 @@ def published_topology(role: str = "test", *, ranks=None, **server_args_fields):
         reset_context()
 
 
+def publish_build_topology(*, world_rank: int = 0, **server_args_fields):
+    """State the widths `initialize_model_parallel` is about to build at.
+
+    The build reads every width from the runtime context, so a test that wants
+    a particular topology publishes it here rather than passing it in -- the
+    same door production uses, which also keeps the derived widths honest.
+
+    Unlike `published_topology` this is not a scope: the groups it is about to
+    build outlive any block, so the configuration describing them has to as
+    well. Callers that tear the groups down are already resetting the process.
+    """
+    from sglang.srt.runtime_context import SpawnRanks, publish, reset_context
+    from sglang.srt.server_args import ServerArgs
+
+    reset_context()
+    publish(
+        ServerArgs(model_path="dummy", **server_args_fields),
+        role="test",
+        ranks=SpawnRanks(world_rank=world_rank),
+    )
+
+
 _GPU_IDLE_TIMEOUT_SECS = 30.0
 _GPU_IDLE_POLL_INTERVAL_SECS = 2.0
 _GPU_IDLE_USED_MEMORY_THRESHOLD = 2 << 30  # 2 GiB