Name the two widths of the WORLD group (#40070)

python/sglang/srt/disaggregation/ascend/transfer_engine.py

@@ -65,7 +65,8 @@ class AscendTransferEngine(MooncakeTransferEngine):
             # through all_gather to avoid conflicts with rdma initialization.
             tmp_tensor = torch.zeros(1, device="npu")
             output_tensor_list = [
-                torch.empty_like(tmp_tensor) for _ in range(get_parallel().world_size)
+                torch.empty_like(tmp_tensor)
+                for _ in range(get_parallel().launch_world_size)
             ]
             torch.distributed.all_gather(
                 output_tensor_list,

python/sglang/srt/elastic_ep/elastic_ep.py

@@ -92,7 +92,7 @@ class ElasticEPStateManager:
 
         if get_exec().moe.elastic_ep_backend is not None:
             world_size = torch.distributed.get_world_size()
-            active_rank_capacity = get_parallel().max_ep_size or world_size
+            active_rank_capacity = get_parallel().max_world_size
             assert active_rank_capacity >= world_size, (
                 f"--max-ep-size ({active_rank_capacity}) must be >= "
                 f"world_size ({world_size})."

python/sglang/srt/elastic_ep/expert_backup_client.py

@@ -50,7 +50,7 @@ class ExpertBackupClient:
         self.buffer_size = 0
         self.use_backup = False
         local_ip = get_local_ip_auto()
-        all_ips = [None] * get_parallel().world_size
+        all_ips = [None] * get_parallel().launch_world_size
         torch.distributed.all_gather_object(
             all_ips, local_ip, group=get_parallel().world_group.cpu_group
         )
@@ -59,14 +59,14 @@ class ExpertBackupClient:
         for i in range(self.engine_num):
             self.recv_list[i] = context.socket(zmq.SUB)
             self.recv_list[i].connect(
-                f"tcp://{all_ips[i * get_parallel().world_size // get_parallel().nnodes]}:{PORT_BASE + i * 2 + 1}"
+                f"tcp://{all_ips[i * get_parallel().launch_world_size // get_parallel().nnodes]}:{PORT_BASE + i * 2 + 1}"
             )
             self.recv_list[i].setsockopt(zmq.SUBSCRIBE, b"")
 
             # Synchronization channel to notify the manager when this client is ready.
             self.ready_sockets[i] = context.socket(zmq.PUSH)
             self.ready_sockets[i].connect(
-                f"tcp://{all_ips[i * get_parallel().world_size // get_parallel().nnodes]}:{PORT_BASE + i * 2}"
+                f"tcp://{all_ips[i * get_parallel().launch_world_size // get_parallel().nnodes]}:{PORT_BASE + i * 2}"
             )
             sock_send(self.ready_sockets[i], UpdateExpertBackupReq())
 

python/sglang/srt/managers/scheduler.py

@@ -5665,7 +5665,7 @@ class Scheduler(
 
         old_ep_size = ElasticEPStateManager.get_effective_ep_size()
         new_ep_size = recv_req.new_ep_size
-        max_ep_size = get_parallel().max_ep_size or old_ep_size
+        max_ep_size = get_parallel().max_world_size
 
         logger.debug(
             "[Elastic EP][scale] request received: new_ep_size=%d "

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

@@ -60,13 +60,17 @@ def _resolve_elastic_world_dp_size(
 
     live_dp_size = get_parallel().attn_dp_size
     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
+    # served width has to stay under.
     world_size = torch.distributed.get_world_size(group)
 
     if live_dp_size != effective_ep_size:
         raise RuntimeError(
             "[Elastic EP] WORLD MLP sync dp_size is out of sync: "
             f"rank={torch.distributed.get_rank(group)} "
-            f"live_dp_size={live_dp_size} effective_ep_size={effective_ep_size} "
+            f"live_dp_size={live_dp_size} "
+            f"effective_ep_size={effective_ep_size} "
             f"world_size={world_size} server_args_dp_size={dp_size} "
             f"local_num_tokens={local_num_tokens} "
             f"local_forward_mode={local_forward_mode}"

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -2195,7 +2195,7 @@ class KVCacheConfigurator:
         available_gpu_memory = get_available_gpu_memory(
             self.device,
             self.gpu_id,
-            distributed=get_parallel().world_group.world_size > 1,
+            distributed=get_parallel().launch_world_size > 1,
             cpu_group=get_parallel().world_group.cpu_group,
         )
 

python/sglang/srt/mem_cache/pool_host/base.py

@@ -31,20 +31,21 @@ _WRITE_BACK_STAGING_PAGE_CHUNK = 64
 def ranks_per_host() -> int:
     """Number of ranks of this job running on the same machine as this one.
 
-    Derived as world_size // nnodes: the launcher slices ranks uniformly
-    across nodes (resolution asserts divisibility), so no hostname collective
-    is needed — a collective here would have to be issued the same number of
-    times on every rank, and ranks build different numbers of host pools.
+    Derived as the launch width // nnodes: the launcher slices ranks
+    uniformly across nodes (resolution asserts divisibility), so no hostname
+    collective is needed — a collective here would have to be issued the same
+    number of times on every rank, and ranks build different numbers of host
+    pools.
     """
     if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
         return 1
     try:
-        world_group = get_parallel().world_group
+        launch_world_size = get_parallel().launch_world_size
     except AssertionError:
         return 1
-    if world_group.world_size == 1:
+    if launch_world_size == 1:
         return 1
-    return max(world_group.world_size // get_parallel().nnodes, 1)
+    return max(launch_world_size // get_parallel().nnodes, 1)
 
 
 def host_memory_budget_bytes() -> int:

python/sglang/srt/model_executor/model_runner_components/attention_backend_setup.py

@@ -146,7 +146,7 @@ def build_attention_backends(*, model_runner: ModelRunner) -> AttentionBackends:
             model_runner.device,
             model_runner.gpu_id,
             get_parallel().world_group.rank_in_group,
-            get_parallel().world_group.world_size,
+            get_parallel().launch_world_size,
             get_parallel().world_group.cpu_group,
         )
 

python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py

@@ -227,7 +227,7 @@ def refresh_deep_gemm_layout_memory_budget(
     available_memory_gb = get_available_gpu_memory(
         model_runner.device,
         model_runner.gpu_id,
-        distributed=world_group.world_size > 1,
+        distributed=get_parallel().launch_world_size > 1,
         cpu_group=world_group.cpu_group,
     )
     budget_bytes = set_masked_standard_layout_memory_budget(

python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py

@@ -59,7 +59,7 @@ def compute_post_capture_kv_resize(
     free_gb = get_available_gpu_memory(
         model_runner.device,
         model_runner.gpu_id,
-        distributed=get_parallel().world_group.world_size > 1,
+        distributed=get_parallel().launch_world_size > 1,
         cpu_group=get_parallel().world_group.cpu_group,
     )
     headroom_gb = model_runner.pre_model_load_memory * (

python/sglang/srt/model_loader/weight_utils.py

@@ -1008,7 +1008,7 @@ def _prefetch_all_checkpoints(
     if torch.distributed.is_initialized():
         world_group = get_parallel().world_group
         local_rank = world_group.local_rank
-        local_world_size = world_group.local_size or world_group.world_size
+        local_world_size = world_group.local_size or get_parallel().launch_world_size
     else:
         local_rank = 0
         local_world_size = 1

python/sglang/srt/runtime_context.py

@@ -117,9 +117,57 @@ def _parallel_config_leaves() -> frozenset:
 # this process's attention-DP rank.
 _MISSING_READ = object()
 
+
+class Live(msgspec.Struct, frozen=True):
+    """How a rank / group / world width is answered, and what it means.
+
+    `source` is the canonical getter's name in `parallel_state`, a callable
+    taking the context, or `None` for a name only a stamp can answer.
+
+    Most entries in the table below are a bare getter name: a rank or a group
+    handle is its own explanation. This shape is for a name whose meaning is
+    not in its getter, and it carries the prose with the declaration rather
+    than in a second table keyed by the same names.
+    """
+
+    source: Any = None
+    doc: str = ""
+
+
 _LIVE_READS: dict = {
-    "world_size": "get_world_size",
-    "world_rank": "get_world_rank",
+    # Two widths of the WORLD group: what it was built at, and what it has
+    # room for. Both are properties of the group itself. How much of that room
+    # is currently serving is elastic-EP state, owned by `ElasticEPStateManager`
+    # and asked of it directly -- a width that lives somewhere else does not
+    # become a WORLD fact by being readable from here.
+    "launch_world_size": Live(
+        source="get_world_size",
+        doc=(
+            "Width the WORLD group was built at: `len(ranks)`, frozen when the "
+            "coordinator was constructed. What every startup reader wants -- "
+            "memory accounting, KV cache sizing, graph capture, weight loading "
+            "-- and what a scale-up leaves behind rather than updates."
+        ),
+    ),
+    "max_world_size": Live(
+        source=lambda self: self.max_ep_size or self.launch_world_size,
+        doc=(
+            "Ranks the WORLD group has room for: `--max-ep-size` when it is "
+            "set, otherwise the launch width. This is the ceiling the process "
+            "group was pre-allocated to -- mooncake sizes its active-rank mask "
+            "to it -- which is why `init_distributed_environment` takes it "
+            "under this name. Whether the group can grow at all is a separate "
+            "question, answered by the leaf being set rather than by this width."
+        ),
+    ),
+    "launch_world_rank": Live(
+        source="get_world_rank",
+        doc=(
+            "This process's rank in the WORLD group as built. Frozen with the "
+            "coordinator, exactly like `launch_world_size`, and named for the "
+            "same reason: a scale-up does not renumber it."
+        ),
+    ),
     "tp_rank": "get_tensor_model_parallel_rank",
     "pp_rank": "get_pipeline_model_parallel_rank",
     "moe_ep_rank": "get_moe_expert_parallel_rank",
@@ -195,9 +243,15 @@ def derive_parallel_widths(
     the arithmetic lives here rather than being read back off the group
     coordinators.
 
-    `world_size` is not among them: it is not a quotient, and `get_world_size()`
-    answers with the live WORLD group, which stays right through an elastic
-    scale-up that a value fixed at group build would not survive.
+    The world widths are not among them: neither is a quotient, and they are
+    not one number. `launch_world_size` is what the WORLD group was built at
+    and is frozen there -- `GroupCoordinator.world_size` is `len(ranks)`, so it
+    does not move when mooncake admits ranks into an expandable WORLD;
+    `max_world_size` is what that group has room for. How much of that room is
+    serving right now is elastic-EP state and is asked of its owner. Deriving
+    either from the leaves would be wrong in a further way: on a scale joiner it
+    would answer with the joining cohort's own `tp * pp`, while that process's
+    WORLD spans `ep_join_rank_offset + tp * pp`.
     """
     return {
         "attn_dp_size": attn_dp_size,
@@ -335,10 +389,11 @@ class ParallelContext:
             return getattr(config, name)
         live = _LIVE_READS.get(name, _MISSING_READ)
         if live is not _MISSING_READ:
-            if isinstance(live, str):
-                return getattr(_ps(), live)()
-            if live is not None:
-                return live(self)
+            source = live.source if isinstance(live, Live) else live
+            if isinstance(source, str):
+                return getattr(_ps(), source)()
+            if source is not None:
+                return source(self)
             raise RuntimeError(
                 f"parallel rank {name!r} is not available: it is computed from "
                 "this process's `tp_rank` when the attention topology is "
@@ -422,8 +477,15 @@ def _install_parallel_properties() -> None:
     chain rather than one per kind of name.
     """
     docs = {name: decl.doc for name, decl in _derived_widths().items()}
+    docs.update(
+        {
+            name: live.doc
+            for name, live in _LIVE_READS.items()
+            if isinstance(live, Live) and live.doc
+        }
+    )
 
-    for name in list(docs) + list(_LIVE_READS):
+    for name in list(_derived_widths()) + list(_LIVE_READS):
 
         def getter(self, _name=name):
             return self._read(_name)