Bringing the parallel runtime up becomes a phase, not a side effect (#40345)

python/sglang/benchmark/one_batch.py

@@ -71,13 +71,13 @@ from sglang.srt.arg_groups.overrides import (
 )
 from sglang.srt.configs.hybrid_arch import mambaish_config
 from sglang.srt.configs.model_config import ModelConfig
+from sglang.srt.distributed import bootstrap
 from sglang.srt.distributed.parallel_state import (
     destroy_distributed_environment,
     destroy_model_parallel,
 )
 from sglang.srt.entrypoints.engine import _set_envs_and_config
 from sglang.srt.hardware_backend.mlx.runtime import use_mlx
-from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
 from sglang.srt.layers.moe import initialize_moe_config
 from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config
 from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config
@@ -93,6 +93,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.runtime_context import (
     SpawnRanks,
+    get_device,
     get_model,
     get_parallel,
     get_schedule,
@@ -316,18 +317,8 @@ def load_model(server_args, port_args, gpu_id, tp_rank):
     cfg = resolving_view(server_args)
     suppress_other_loggers()
     rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None
-    moe_ep_rank = tp_rank // (cfg.tp_size // cfg.ep_size)
 
     model_config = ModelConfig.from_server_args(server_args)
-    attn_tp_rank, attn_tp_size, attn_dp_rank, attn_dp_size = (
-        compute_dp_attention_world_info(
-            cfg.enable_dp_attention,
-            tp_rank,
-            cfg.tp_size,
-            cfg.dp_size,
-            cfg.attn_cp_size,
-        )
-    )
     runner_kwargs = dict(
         model_config=model_config,
         mem_fraction_static=cfg.mem_fraction_static,
@@ -336,6 +327,14 @@ def load_model(server_args, port_args, gpu_id, tp_rank):
         server_args=server_args,
     )
 
+    # Phase two: this entry has no scheduler to run it.
+    bootstrap.init_parallel_runtime(
+        server_args=server_args,
+        model_config=model_config,
+        device=get_device().device,
+        dist_port=port_args.nccl_port,
+    )
+
     _use_mlx = use_mlx()
     if _use_mlx:
         from sglang.srt.hardware_backend.mlx.model_runner_stub import (

python/sglang/srt/distributed/bootstrap.py

@@ -7,6 +7,9 @@ import torch
 import torch.distributed as dist
 
 from sglang.srt.configs.model_config import ModelConfig
+from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
+    maybe_init_shared_mooncake_transfer_engine,
+)
 from sglang.srt.distributed.gated_launch import maybe_wait_for_gated_launch
 from sglang.srt.distributed.parallel_state import (
     _tag_groups_for_flashinfer_allreduce_only,
@@ -56,92 +59,176 @@ _is_cpu_arm64 = is_host_cpu_arm64()
 _TP_ALL_TO_ALL_WARMUP_BYTES_PER_PEER = 4 << 20
 
 
-def init_torch_distributed(
+#: Set by `init_parallel`; `destroy_model_parallel` clears it, so a test that
+#: tears the groups down can build them again.
+_PARALLEL_INITIALISED = False
+
+
+def reset_parallel_initialised() -> None:
+    """Forget that the groups were built. Paired with tearing them down."""
+    global _PARALLEL_INITIALISED
+    _PARALLEL_INITIALISED = False
+
+
+def _bind_threads_if_cpu(*, device: str) -> "Optional[List[int]]":
+    """Pin OpenMP threads to this process's NUMA node, on CPU.
+
+    A precondition of the CPU group build, which reads the binding, so it is
+    done here rather than left for a caller to remember.
+    """
+    if device != "cpu":
+        return None
+    from sglang.srt.utils import numa_utils
+
+    parallel = get_parallel()
+    # With --enable-dp-attention, dp partitions the existing TP group rather
+    # than spawning additional processes, so dp_size must not be multiplied
+    # into the process count here (unlike regular DP, where dp_size * tp_size *
+    # pp_size is the true worker count).
+    dp_size = 1 if parallel.enable_dp_attention else parallel.dp_size
+    return numa_utils.init_threads_binding(
+        numa_index=get_device().gpu_id,
+        world_size=dp_size * parallel.tp_size * parallel.pp_size,
+    )
+
+
+def init_parallel_runtime(
     *,
     server_args: ServerArgs,
     model_config: ModelConfig,
     device: str,
     dist_port: int,
-    is_draft_worker: bool,
-    local_omp_cpuid: Optional[List[int]],
-):
+) -> None:
+    """Phase two of startup: bring the parallel runtime up, once.
+
+    Publish says what the topology is; this makes it exist. Nothing returns,
+    because the groups are read through the runtime context -- a caller that
+    wants one asks `get_parallel()`, in this process or any later phase.
+
+    "Runtime" rather than "groups": two things have to be in place before the
+    groups can be built, and they are done here rather than left for every
+    entry to remember. The OpenMP/NUMA binding is what the CPU group build
+    reads, and the shared Mooncake transfer engine is what the Mooncake
+    process-group backend asks for -- create that one late and a second engine
+    appears. Both are preconditions of the build, not separate work.
+
+    Runs on the target worker only. A draft worker shares its target's groups,
+    which is why this is a phase the entry runs rather than something a runner
+    does on its way up: whether the groups exist must not depend on which
+    runner happened to be constructed first.
+    """
+    global _PARALLEL_INITIALISED
+    if _PARALLEL_INITIALISED:
+        raise RuntimeError(
+            "init_parallel_runtime() ran twice in this process. The groups are built "
+            "once, before anything that reads one exists; a second build is "
+            "either a lost race between two entries or a runner trying to "
+            "bring up its own. An elastic scale-up joins an existing WORLD "
+            "through initialize_model_parallel directly and does not come "
+            "through here."
+        )
+    _PARALLEL_INITIALISED = True
+
     tic = time.perf_counter()
-    logger.info("Init torch distributed begin.")
-    parallel = get_parallel()
+    logger.info("Init parallel begin.")
 
     backend = _resolve_backend(device=device)
-
-    before_avail_memory = get_available_gpu_memory(device, get_device().gpu_id)
     if not get_parallel().enable_p2p_check:
         monkey_patch_p2p_access_check()
 
     dist_init_method = _resolve_dist_init_method(dist_port=dist_port)
     _set_all_reduce_flags()
 
-    if not is_draft_worker:
-        if device == "cpu":
-            _init_cpu_threads_env(
-                tp_size=parallel.tp_size,
-                tp_rank=parallel.tp_rank,
-                local_omp_cpuid=local_omp_cpuid,
-                dist_init_method=dist_init_method,
-            )
+    local_omp_cpuid = _bind_threads_if_cpu(device=device)
+    # Everything below allocates on the current device -- the NCCL warm-up, the
+    # mooncake all-reduce buffer -- and without this every rank on a node would
+    # pick device 0, because the default is not to reindex the visible set.
+    try:
+        torch.get_device_module(device).set_device(get_device().gpu_id)
+    except Exception:
+        logger.warning(
+            "Context: device=%s gpu_id=%s CUDA_VISIBLE_DEVICES=%s tp_rank=%s",
+            device,
+            get_device().gpu_id,
+            os.environ.get("CUDA_VISIBLE_DEVICES"),
+            get_parallel().tp_rank,
+        )
+        raise
+    maybe_init_shared_mooncake_transfer_engine(gpu_id=get_device().gpu_id)
 
-        # Only initialize the distributed environment on the target model worker.
-        # This builds the groups behind the context's live group-handle reads.
-        _init_parallel_groups(
-            backend=backend,
+    parallel = get_parallel()
+    if device == "cpu":
+        _init_cpu_threads_env(
+            tp_size=parallel.tp_size,
+            tp_rank=parallel.tp_rank,
+            local_omp_cpuid=local_omp_cpuid,
             dist_init_method=dist_init_method,
-            server_args=server_args,
-            model_config=model_config,
-            gpu_id=get_device().gpu_id,
         )
 
-        # Pre-warm NCCL/RCCL/HCCL to eliminate cold-start latency in first request
-        # Controlled by --pre-warm-nccl flag (default: enabled on AMD GPUs)
-        if get_exec().comm.pre_warm_nccl and (
-            parallel.tp_size > 1 or parallel.pp_size > 1 or parallel.moe_ep_size > 1
-        ):
-            _prewarm_nccl(
-                tp_size=parallel.tp_size,
-                pp_size=parallel.pp_size,
-                moe_ep_size=parallel.moe_ep_size,
-            )
+    _init_parallel_groups(
+        backend=backend,
+        dist_init_method=dist_init_method,
+        server_args=server_args,
+        model_config=model_config,
+        gpu_id=get_device().gpu_id,
+    )
 
-        # CUDA graph capture enables the PyNCCL communicator for TP LM-head
-        # all-to-all. Exercise that exact send/recv path before measuring
-        # pre_model_load_memory so its persistent transport allocations are
-        # included in later KV-cache sizing instead of appearing during capture.
-        if (
-            device == "cuda"
-            and get_parallel().enable_tp_lm_head_all_to_all
-            and parallel.tp_size > 1
-        ):
-            _prewarm_tp_lm_head_all_to_all()
+    # Pre-warm NCCL/RCCL/HCCL to eliminate cold-start latency in first request
+    # Controlled by --pre-warm-nccl flag (default: enabled on AMD GPUs)
+    if get_exec().comm.pre_warm_nccl and (
+        parallel.tp_size > 1 or parallel.pp_size > 1 or parallel.moe_ep_size > 1
+    ):
+        _prewarm_nccl(
+            tp_size=parallel.tp_size,
+            pp_size=parallel.pp_size,
+            moe_ep_size=parallel.moe_ep_size,
+        )
+
+    # CUDA graph capture enables the PyNCCL communicator for TP LM-head
+    # all-to-all. Exercise that exact send/recv path before measuring
+    # pre_model_load_memory so its persistent transport allocations are
+    # included in later KV-cache sizing instead of appearing during capture.
+    if (
+        device == "cuda"
+        and parallel.enable_tp_lm_head_all_to_all
+        and parallel.tp_size > 1
+    ):
+        _prewarm_tp_lm_head_all_to_all()
+
+    logger.info(f"Init parallel ends. elapsed={time.perf_counter() - tic:.2f} s")
+
+
+def measure_pre_model_load_memory(*, device: str, is_draft_worker: bool) -> float:
+    """Available memory after the groups exist and before the model loads.
+
+    Sized into the KV cache later, so it has to be taken at exactly this point
+    -- which is why it stays with the runner rather than moving into the
+    parallel phase.
+    """
+    before_avail_memory = get_available_gpu_memory(device, get_device().gpu_id)
 
     maybe_wait_for_gated_launch(
         host=get_serving().host, port=get_parallel().gated_launch_port
     )
 
-    # Draft workers reuse the target pool config and may exist on only one PP stage;
-    # including them in this WORLD reduction would deadlock on absent peers.
+    # Draft workers reuse the target pool config and may exist on only one PP
+    # stage; including them in this WORLD reduction would deadlock on absent
+    # peers.
     pre_model_load_memory = get_available_gpu_memory(
         device,
         get_device().gpu_id,
         distributed=get_world_group().world_size > 1 and not is_draft_worker,
         cpu_group=get_world_group().cpu_group,
     )
-    # Check memory for tensor parallelism
     local_gpu_memory = get_available_gpu_memory(device, get_device().gpu_id)
-    if parallel.tp_size > 1 and not is_draft_worker:
+    if get_parallel().tp_size > 1 and not is_draft_worker:
         _check_tp_memory_balance(
             pre_model_load_memory=pre_model_load_memory,
             local_gpu_memory=local_gpu_memory,
         )
-
     logger.info(
-        f"Init torch distributed ends. elapsed={time.perf_counter() - tic:.2f} s, "
-        f"mem usage={(before_avail_memory - local_gpu_memory):.2f} GB"
+        f"Memory baseline taken. mem usage="
+        f"{(before_avail_memory - local_gpu_memory):.2f} GB"
     )
     return pre_model_load_memory
 

python/sglang/srt/distributed/parallel_state.py

@@ -3240,6 +3240,9 @@ def get_moe_tensor_parallel_rank():
 
 def destroy_model_parallel():
     """Set the groups to none and destroy them."""
+    from sglang.srt.distributed.bootstrap import reset_parallel_initialised
+
+    reset_parallel_initialised()
     get_parallel().clear_stamp()
     dwdp_mgr = get_global_dwdp_manager()
     if dwdp_mgr is not None:

python/sglang/srt/managers/scheduler.py

@@ -101,6 +101,7 @@ from sglang.srt.disaggregation.utils import (
     prepare_abort,
     unified_memory_disagg_move_gate,
 )
+from sglang.srt.distributed import bootstrap
 from sglang.srt.distributed.parallel_state import (
     abort_distributed_environment,
 )
@@ -513,6 +514,18 @@ class Scheduler(
         # Init model configs
         self.init_model_config()
 
+        # Init ZBAL, switch allocator should before any torch alloc action
+        self.init_zbal_on_npu()
+
+        # The groups are the first thing that allocates, so this comes after the
+        # allocator switch above and before anything that reads a group.
+        bootstrap.init_parallel_runtime(
+            server_args=server_args,
+            model_config=self.model_config,
+            device=get_device().device,
+            dist_port=self.nccl_port,
+        )
+
         # Init metrics stats
         self.init_metrics_collector(tp_rank, pp_rank, dp_rank)
 
@@ -520,9 +533,6 @@ class Scheduler(
         self.init_ipc_channels(port_args)
         self.init_idle_sleeper()
 
-        # Init ZBAL, switch allocator should before any torch alloc action
-        self.init_zbal_on_npu()
-
         # Init PD-multiplexing context
         if self.enable_pdmux:
             self.init_pdmux()

python/sglang/srt/model_executor/model_runner.py

@@ -34,9 +34,6 @@ from sglang.srt.configs.model_config import (
 from sglang.srt.configs.update_config import adjust_config_with_unaligned_cpu_tp
 from sglang.srt.debug_utils.dumper import dumper
 from sglang.srt.distributed import bootstrap
-from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
-    maybe_init_shared_mooncake_transfer_engine,
-)
 from sglang.srt.dllm.config import DllmConfig
 from sglang.srt.elastic_ep.elastic_ep import (
     ElasticEPStateManager,
@@ -218,7 +215,6 @@ from sglang.srt.utils import (
     get_available_gpu_memory,
     is_host_cpu_arm64,
     is_npu,
-    numa_utils,
     require_gathered_buffer,
     reserve_rope_cache_for_long_sequences,
     set_cuda_arch,
@@ -402,10 +398,6 @@ class ModelRunner:
             is_draft_worker=self.is_draft_worker,
         )
 
-        # Init OpenMP threads binding for CPU
-        if self.device == "cpu":
-            self.init_threads_binding()
-
         # Set float32 matmul precision
         if get_exec().features.enable_tf32_matmul:
             torch.set_float32_matmul_precision("high")
@@ -422,11 +414,6 @@ class ModelRunner:
             )
             raise
 
-        # Initialize MooncakeTransferEngine BEFORE init_torch_distributed so
-        # that the shared TE can be passed to the Mooncake PG backend (avoids
-        # creating duplicate TransferEngines).
-        self.init_shared_mooncake_transfer_engine()
-
         # Get available memory before model loading.
         # Stored for later use by alloc_memory_pool().
         self.init_torch_distributed()
@@ -1167,13 +1154,8 @@ class ModelRunner:
         )
 
     def init_torch_distributed(self):
-        self.pre_model_load_memory = bootstrap.init_torch_distributed(
-            server_args=self.server_args,
-            model_config=self.model_config,
-            device=self.device,
-            dist_port=self.dist_port,
-            is_draft_worker=self.is_draft_worker,
-            local_omp_cpuid=self.local_omp_cpuid if self.device == "cpu" else None,
+        self.pre_model_load_memory = bootstrap.measure_pre_model_load_memory(
+            device=self.device, is_draft_worker=self.is_draft_worker
         )
         # Read once, here: a draft runner is constructed inside the scope that
         # states its topology and used outside it, so what it holds has to be
@@ -1197,9 +1179,6 @@ class ModelRunner:
         self.moe_ep_size = parallel.moe_ep_size
         self.dp_rank = parallel.dp_rank
 
-    def init_shared_mooncake_transfer_engine(self):
-        maybe_init_shared_mooncake_transfer_engine(gpu_id=self.gpu_id)
-
     def load_model(self):
         tic_total = time.perf_counter()
         before_avail_memory = get_available_gpu_memory(self.device, self.gpu_id)
@@ -1593,18 +1572,6 @@ class ModelRunner:
             phases=("prefill", "draft_prefill"),
         )
 
-    def init_threads_binding(self):
-        # With --enable-dp-attention, dp partitions the existing TP group
-        # rather than spawning additional processes, so dp_size must not be
-        # multiplied into the process count here (unlike regular DP, where
-        # dp_size * tp_size * pp_size is the true worker count).
-        parallel = get_parallel()
-        dp_size = 1 if parallel.enable_dp_attention else parallel.dp_size
-        self.local_omp_cpuid = numa_utils.init_threads_binding(
-            numa_index=self.gpu_id,
-            world_size=dp_size * parallel.tp_size * parallel.pp_size,
-        )
-
     def apply_torch_tp(self):
         model_parallel.apply_torch_tp(
             model=self.model, device=self.device, tp_size=self.tp_size