From 3a5f52e14419b7f1d25d7a3833ebaecf10d484b4 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Fri, 18 Sep 2026 23:52:37 -0700 Subject: [PATCH] Record a process's placement at publish, not at group build (#40071) --- python/sglang/benchmark/one_batch.py | 20 +- .../sglang/srt/constrained/grammar_manager.py | 6 +- python/sglang/srt/disaggregation/decode.py | 2 +- python/sglang/srt/disaggregation/prefill.py | 6 +- python/sglang/srt/distributed/bootstrap.py | 1 + .../sglang/srt/distributed/parallel_state.py | 9 +- .../layers/deep_gemm_wrapper/compile_utils.py | 4 +- python/sglang/srt/layers/dp_attention.py | 16 +- python/sglang/srt/managers/scheduler.py | 165 +++++++----- .../managers/scheduler_components/dp_attn.py | 4 +- .../kv_events_publisher.py | 5 +- .../scheduler_components/load_inquirer.py | 6 +- .../scheduler_components/metrics_reporter.py | 2 +- .../scheduler_components/output_streamer.py | 6 +- .../scheduler_components/profiler_manager.py | 33 +-- .../scheduler_components/request_receiver.py | 12 +- .../sglang/srt/managers/scheduler_pp_mixin.py | 50 ++-- python/sglang/srt/managers/tp_worker.py | 8 +- .../srt/mem_cache/kv_cache_configurator.py | 4 +- .../sglang/srt/model_executor/model_runner.py | 36 +-- .../srt/model_executor/runner/base_runner.py | 4 +- .../runner/flashinfer_autotune.py | 5 +- .../srt/observability/metrics_collector.py | 3 +- python/sglang/srt/ray/scheduler_actor.py | 23 +- python/sglang/srt/runtime_context.py | 205 ++++++++++++++- python/sglang/srt/rust_server/server.py | 2 +- .../srt/speculative/dflash_worker_v2.py | 13 +- .../dspark_components/dspark_worker_v2.py | 7 +- .../sglang/srt/speculative/eagle_worker_v2.py | 2 + .../speculative/frozen_kv_mtp_worker_v2.py | 5 +- .../multi_layer_eagle_worker_v2.py | 5 +- .../srt/speculative/standalone_worker_v2.py | 8 +- python/sglang/srt/utils/profile_utils.py | 25 +- python/sglang/srt/weight_cache/daemon.py | 14 +- python/sglang/test/test_utils.py | 42 +++ .../unit/constrained/test_grammar_manager.py | 31 ++- .../unit/managers/test_auxiliary_output.py | 19 +- .../test_output_streamer_customized_info.py | 3 + .../unit/managers/test_pp_cp_rank_offsets.py | 27 +- .../test_scheduler_chunked_abort_race.py | 12 +- .../managers/test_scheduler_hicache_events.py | 8 +- .../unit/managers/test_scheduler_timeouts.py | 12 +- .../test_startup_weight_load.py | 7 +- .../rust/shared/test_rust_server_extension.py | 5 +- .../test_forward_pass_metrics.py | 8 +- test/registered/unit/test_runtime_context.py | 246 ++++++++++++++++++ .../unit/utils/test_weight_versions.py | 12 +- 47 files changed, 917 insertions(+), 231 deletions(-) diff --git a/python/sglang/benchmark/one_batch.py b/python/sglang/benchmark/one_batch.py index a4d7d0de3..b305bbb30 100644 --- a/python/sglang/benchmark/one_batch.py +++ b/python/sglang/benchmark/one_batch.py @@ -93,10 +93,12 @@ from sglang.srt.model_executor.cuda_graph_config import ( 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_model, get_parallel, get_schedule, publish, + spawn_world_rank, ) from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.server_args import PortArgs, ServerArgs @@ -707,7 +709,15 @@ def correctness_test( gpu_id, tp_rank, ): - publish(server_args, role="scheduler") + # With the placement this process was spawned with, so a rank read here + # does not need a process group -- the same bundle the runner is handed. + publish( + server_args, + role="scheduler", + ranks=SpawnRanks( + world_rank=spawn_world_rank(server_args, tp_rank=tp_rank, pp_rank=0) + ), + ) # Configure the logger configure_logger(server_args, prefix=f" TP{tp_rank}") @@ -912,7 +922,13 @@ def latency_test( cfg = resolving_view(server_args) # `main` runs this inline for tp_size == 1 and spawns it per rank otherwise; # a spawned child arrives with nothing published. - publish(server_args, role="scheduler") + publish( + server_args, + role="scheduler", + ranks=SpawnRanks( + world_rank=spawn_world_rank(server_args, tp_rank=tp_rank, pp_rank=0) + ), + ) initialize_moe_config() initialize_fp8_gemm_config() initialize_fp4_gemm_config() diff --git a/python/sglang/srt/constrained/grammar_manager.py b/python/sglang/srt/constrained/grammar_manager.py index 7886e7060..5d5d33d29 100644 --- a/python/sglang/srt/constrained/grammar_manager.py +++ b/python/sglang/srt/constrained/grammar_manager.py @@ -14,7 +14,7 @@ from sglang.srt.constrained.base_grammar_backend import ( from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject from sglang.srt.distributed.communication_tags import P2PTag from sglang.srt.environ import envs -from sglang.srt.runtime_context import get_serving +from sglang.srt.runtime_context import get_parallel, get_serving from sglang.srt.sampling.sampling_params import ( get_request_reasoning_end_token_ids, ) @@ -53,8 +53,8 @@ class GrammarManager: self.grammar_sync_size = scheduler.dp_tp_group.world_size self.grammar_sync_entry = scheduler.dp_tp_group.first_rank self.is_grammar_sync_entry = scheduler.dp_tp_group.is_first_rank - self.pp_rank = scheduler.ps.pp_rank - self.pp_size = scheduler.ps.pp_size + self.pp_rank = get_parallel().pp_rank + self.pp_size = get_parallel().pp_size self.pp_group = scheduler.pp_group self.grammar_pp_sync_work_list = [] diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index ebf388519..e1f50d201 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -393,7 +393,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): self.bootstrap_port = bootstrap_port self.max_total_num_tokens = max_total_num_tokens self.pp_rank = pp_rank - self.pp_size = scheduler.ps.pp_size + self.pp_size = get_parallel().pp_size self.num_reserved_decode_tokens = num_reserved_decode_tokens self.transfer_backend = transfer_backend # Queue for requests pending pre-allocation diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index e1db8cd18..6fc2269f8 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -997,7 +997,7 @@ class SchedulerDisaggregationPrefillMixin: KVPoll.Failed, ): logger.warning_once( - f"PP rank {self.ps.pp_rank}: unexpected poll state {poll} for rid {req.rid} " + f"PP rank {get_parallel().pp_rank}: unexpected poll state {poll} for rid {req.rid} " f"from consensus; treating as undone", ) undone_reqs.append(req) @@ -1077,7 +1077,7 @@ class SchedulerDisaggregationPrefillMixin: ) -> Optional[Exception]: """Conclude an inflight request whose KV transfer failed.""" error_message = ( - f"Prefill transfer failed for request rank={self.ps.tp_rank} " + f"Prefill transfer failed for request rank={get_parallel().tp_rank} " f"{req.rid=} {req.bootstrap_room=}" ) exc: Optional[Exception] = None @@ -1142,7 +1142,7 @@ class SchedulerDisaggregationPrefillMixin: def handle_bootstrap_failure(self: Scheduler, req: Req) -> None: self.clear_pending_chunk_send(req) error_message = ( - f"Prefill bootstrap failed for request rank={self.ps.tp_rank} " + f"Prefill bootstrap failed for request rank={get_parallel().tp_rank} " f"{req.rid=} {req.bootstrap_room=}" ) is_propagated = False diff --git a/python/sglang/srt/distributed/bootstrap.py b/python/sglang/srt/distributed/bootstrap.py index da0997fa0..d45c48d39 100644 --- a/python/sglang/srt/distributed/bootstrap.py +++ b/python/sglang/srt/distributed/bootstrap.py @@ -99,6 +99,7 @@ def init_torch_distributed( ) # 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, dist_init_method=dist_init_method, diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index e59f903ce..1d03cdc64 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -3029,7 +3029,14 @@ def patch_pipeline_parallel_group(pp_group: GroupCoordinator): global _PP _PP = pp_group try: - yield + # `pp_size` is a configured leaf: unlike the rank and the handle it + # does not follow the group being swapped, so the scope has to name it. + with get_parallel().override( + pp_size=pp_group.world_size, + pp_rank=pp_group.rank_in_group, + pp_group=pp_group, + ): + yield finally: _PP_STATE_PATCHED = False _PP = old_pp_group diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py index eb3e5988d..4db8a4a92 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py @@ -509,7 +509,7 @@ def pp_parallel_deep_gemm_warmup(runner) -> None: logger.info( "PP-parallel DeepGEMM warmup start " "(pp_rank=%d, tp_rank=%d, batch_sizes=%s, disagg=%s).", - model_runner.ps.pp_rank, + get_parallel().pp_rank, model_runner.ps.tp_rank, batch_sizes, disagg_mode, @@ -538,5 +538,5 @@ def pp_parallel_deep_gemm_warmup(runner) -> None: logger.info( "PP-parallel DeepGEMM warmup done in %.2fs (pp_rank=%d).", time.perf_counter() - t0, - model_runner.ps.pp_rank, + get_parallel().pp_rank, ) diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index d2bc67116..4919c34cd 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/python/sglang/srt/layers/dp_attention.py @@ -27,6 +27,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( ) from sglang.srt.environ import envs from sglang.srt.runtime_context import ( + derive_attention_ranks, derive_attention_widths, get_device, get_exec, @@ -349,15 +350,12 @@ def compute_dp_attention_world_info( dp_size=dp_size, enable_dp_attention=enable_dp_attention, ) - attn_tp_rank = tp_rank % attn_tp_size - - if not enable_dp_attention: - attn_dp_rank = 0 - else: - # Rank layout is (dp, cp, tp) where tp is the fastest-changing dim: - # tp_rank = (attn_dp_rank * attn_cp_size + attn_cp_rank) * attn_tp_size + attn_tp_rank - attn_dp_rank = tp_rank // (attn_tp_size * attn_cp_size) - + attn_tp_rank, attn_dp_rank = derive_attention_ranks( + tp_rank=tp_rank, + attn_tp_size=attn_tp_size, + attn_cp_size=attn_cp_size, + enable_dp_attention=enable_dp_attention, + ) return attn_tp_rank, attn_tp_size, attn_dp_rank, attn_dp_size diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 2daac51cd..a25ca6e0d 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -29,6 +29,7 @@ from http import HTTPStatus from typing import TYPE_CHECKING, Any, Deque, Dict, List, Optional, Set, Tuple, Union from sglang.srt.runtime_context import ( + SpawnRanks, attention_backends, get_context, get_device, @@ -44,6 +45,7 @@ from sglang.srt.runtime_context import ( get_serving, get_spec, publish, + spawn_world_rank, ) from sglang.srt.utils.common import suppress_noisy_warnings # isort: skip @@ -595,9 +597,9 @@ class Scheduler( enable_metrics=get_observability().enable_metrics, enable_kv_cache_events=bool( get_observability().kv_events_config - and self.ps.pp_rank == 0 - and self.ps.attn_tp_rank == 0 - and self.ps.attn_cp_rank == 0 + and get_parallel().pp_rank == 0 + and get_parallel().attn_tp_rank == 0 + and get_parallel().attn_cp_rank == 0 ), ps=self.ps, tp_group=self.tp_group, @@ -759,8 +761,8 @@ class Scheduler( return rank = ( - self.ps.dp_rank - if self.ps.dp_rank is not None + get_parallel().dp_rank + if get_parallel().dp_rank is not None else self.tp_group.rank_in_group ) logger.info("HCCL DP prewarm start: rank=%s", rank) @@ -775,10 +777,10 @@ class Scheduler( if _is_npu: from sglang.srt.hardware_backend.npu.utils import init_zbal - if self.ps.pp_size > 1: + if get_parallel().pp_size > 1: logger.error("only zbal mix mode support pp_size > 1!") init_zbal( - self.ps.tp_size, self.ps.gpu_id, self.ps.tp_rank + get_parallel().tp_size, self.ps.gpu_id, get_parallel().tp_rank ) # only switch allocator if is mix mode def init_model_config(self): @@ -811,9 +813,9 @@ class Scheduler( def init_ipc_channels(self, port_args: PortArgs): is_rank_zero = ( - self.ps.pp_rank == 0 - and self.ps.attn_tp_rank == 0 - and self.ps.attn_cp_rank == 0 + get_parallel().pp_rank == 0 + and get_parallel().attn_tp_rank == 0 + and get_parallel().attn_cp_rank == 0 ) self.ipc_channels = SchedulerIpcChannels.create( port_args=port_args, @@ -824,7 +826,7 @@ class Scheduler( skip_tokenizer_init=self.skip_tokenizer_init, metrics_enabled=get_observability().enable_metrics and ( - self.ps.attn_tp_rank == 0 + get_parallel().attn_tp_rank == 0 or get_observability().enable_metrics_for_all_schedulers ), enable_scripted_runtime=envs.SGLANG_TEST_SCRIPTED_RUNTIME.get(), @@ -837,7 +839,7 @@ class Scheduler( return self.recv_from_tokenizer = self.ipc_channels.recv_from_tokenizer - dp_rank = self.ps.dp_rank if self.ps.dp_rank is not None else 0 + dp_rank = get_parallel().dp_rank if get_parallel().dp_rank is not None else 0 try: self.load_snapshot_writer = create_load_snapshot_writer( port_args, @@ -850,9 +852,9 @@ class Scheduler( def init_idle_sleeper(self) -> None: if ( - self.ps.pp_rank == 0 - and self.ps.attn_tp_rank == 0 - and self.ps.attn_cp_rank == 0 + get_parallel().pp_rank == 0 + and get_parallel().attn_tp_rank == 0 + and get_parallel().attn_cp_rank == 0 and get_device().sleep_on_idle ): self.idle_sleeper = IdleSleeper( @@ -1032,8 +1034,8 @@ class Scheduler( if ( envs.SGLANG_ENABLE_PP_SPEC.get() - and self.ps.pp_size > 1 - and self.ps.pp_rank != self.ps.pp_size - 1 + and get_parallel().pp_size > 1 + and get_parallel().pp_rank != get_parallel().pp_size - 1 ): # PP+spec: the draft model (MTP layer) needs final hidden states and # the lm_head, both of which live on the last PP stage only. @@ -1130,7 +1132,7 @@ class Scheduler( # Match run_batch / _pp_launch_batch so warmup allocations stay reusable. forward_stream = ( model_runner.forward_stream - if self.enable_overlap or self.ps.pp_size > 1 or use_mlx() + if self.enable_overlap or get_parallel().pp_size > 1 or use_mlx() else self.schedule_stream ) with device_module.stream(forward_stream): @@ -1208,7 +1210,7 @@ class Scheduler( get_context().override( "scheduler.pp_max_micro_batch_size_default", pp_max_micro_batch_size=max( - self.max_running_requests // self.ps.pp_size, 1 + self.max_running_requests // get_parallel().pp_size, 1 ), ) @@ -1239,7 +1241,7 @@ class Scheduler( self.startup_available_gpu_memory_gb = get_available_gpu_memory( self.device, self.ps.gpu_id, empty_cache=False ) - if self.ps.tp_rank == 0: + if get_parallel().tp_rank == 0: logger.info( f"max_total_num_tokens={self.max_total_num_tokens}, " f"chunked_prefill_size={get_schedule().chunked_prefill_size}, " @@ -1344,7 +1346,7 @@ class Scheduler( def maybe_init_dynamic_chunk_sizer(self) -> None: """Profile a PP prefill latency model that sizes chunks per stage.""" self.dynamic_chunk_sizer: Optional[DynamicChunkSizer] = None - if not (get_schedule().enable_dynamic_chunking and self.ps.pp_size > 1): + if not (get_schedule().enable_dynamic_chunking and get_parallel().pp_size > 1): return sizer = DynamicChunkSizer( model_runner=self.tp_worker.model_runner, @@ -1359,7 +1361,7 @@ class Scheduler( device=self.device, pp_group=self.pp_group, world_group=self.world_group, - pp_rank=self.ps.pp_rank, + pp_rank=get_parallel().pp_rank, ) if sizer.profile_and_fit(): self.dynamic_chunk_sizer = sizer @@ -1422,7 +1424,7 @@ class Scheduler( else: self.prefill_delayer = PrefillDelayer( dp_size=self.ps.dp_size, - attn_tp_size=self.ps.attn_tp_size, + attn_tp_size=get_parallel().attn_tp_size, cpu_group=self.tp_cpu_group, device_group=self.tp_group.device_group, metrics_collector=( @@ -1433,7 +1435,7 @@ class Scheduler( max_delay_passes=get_schedule().prefill_delayer_max_delay_passes, token_usage_low_watermark=get_schedule().prefill_delayer_token_usage_low_watermark, device=self.tp_group.device, - debug_log_enabled=self.ps.attn_tp_rank == 0, + debug_log_enabled=get_parallel().attn_tp_rank == 0, ) # NOTE: preemption is enabled by default for priority scheduling. @@ -1464,7 +1466,7 @@ class Scheduler( # Init recv skipper and input blocker self.recv_skipper = SchedulerRecvSkipper.maybe_create() self.input_blocker = ( - SchedulerInputBlocker(noop=self.ps.attn_tp_rank != 0) + SchedulerInputBlocker(noop=get_parallel().attn_tp_rank != 0) if get_bool_env_var("SGLANG_ENABLE_COLOCATED_BATCH_GEN") else None ) @@ -1553,7 +1555,7 @@ class Scheduler( self.disagg_decode_transfer_queue = DecodeTransferQueue( gloo_group=self.attn_tp_cpu_group, req_to_metadata_buffer_idx_allocator=self.req_to_metadata_buffer_idx_allocator, - tp_rank=self.ps.tp_rank, + tp_rank=get_parallel().tp_rank, metadata_buffers=self.disagg_metadata_buffers, scheduler=self, tree_cache=self.tree_cache, @@ -1570,13 +1572,13 @@ class Scheduler( transfer_queue=self.disagg_decode_transfer_queue, tree_cache=self.tree_cache, gloo_group=self.attn_tp_cpu_group, - tp_rank=self.ps.tp_rank, - tp_size=self.ps.tp_size, + tp_rank=get_parallel().tp_rank, + tp_size=get_parallel().tp_size, dp_size=get_parallel().dp_size, gpu_id=self.ps.gpu_id, bootstrap_port=get_disagg().disaggregation_bootstrap_port, max_total_num_tokens=self.max_total_num_tokens, - pp_rank=self.ps.pp_rank, + pp_rank=get_parallel().pp_rank, num_reserved_decode_tokens=get_disagg().num_reserved_decode_tokens, transfer_backend=self.transfer_backend, ) @@ -1602,16 +1604,16 @@ class Scheduler( draft_token_to_kv_pool=draft_token_to_kv_pool, req_to_metadata_buffer_idx_allocator=self.req_to_metadata_buffer_idx_allocator, metadata_buffers=self.disagg_metadata_buffers, - tp_rank=self.ps.tp_rank, - tp_size=self.ps.tp_size, + tp_rank=get_parallel().tp_rank, + tp_size=get_parallel().tp_size, gpu_id=self.ps.gpu_id, bootstrap_port=get_disagg().disaggregation_bootstrap_port, gloo_group=self.attn_tp_cpu_group, max_total_num_tokens=self.max_total_num_tokens, scheduler=self, scheduler_stage_metrics=self.scheduler_stage_metrics, - pp_rank=self.ps.pp_rank, - pp_size=self.ps.pp_size, + pp_rank=get_parallel().pp_rank, + pp_size=get_parallel().pp_size, transfer_backend=self.transfer_backend, ) # The prefill requests that are in the middle of kv sending @@ -1638,8 +1640,8 @@ class Scheduler( self.server_args, dtype=self.model_config.dtype, hf_config=self.model_config.hf_config, - pp_rank=self.ps.pp_rank, - tp_rank=self.ps.tp_rank, + pp_rank=get_parallel().pp_rank, + tp_rank=get_parallel().tp_rank, tp_group=self.tp_group, scheduler=self, ) @@ -1895,7 +1897,9 @@ class Scheduler( if self.device == "cpu": self.schedule_stream.synchronize = lambda: None # No-op for CPU - elif (is_cuda() or _is_hip) and (self.enable_overlap or self.ps.pp_size > 1): + elif (is_cuda() or _is_hip) and ( + self.enable_overlap or get_parallel().pp_size > 1 + ): # CUDA/HIP streams come from a fixed round-robin pool. Redraw if this # stream aliases forward_stream, which would eliminate scheduler # overlap. Only CUDA/HIP streams expose a ``cuda_stream`` handle; @@ -2090,9 +2094,9 @@ class Scheduler( """ local_reqs = [] if ( - self.ps.pp_rank == 0 - and self.ps.attn_tp_rank == 0 - and self.ps.attn_cp_rank == 0 + get_parallel().pp_rank == 0 + and get_parallel().attn_tp_rank == 0 + and get_parallel().attn_cp_rank == 0 ): local_reqs = self._poll_timeout_aborts() recv_reqs = self.request_receiver.recv_requests(local_reqs=local_reqs) @@ -2290,9 +2294,9 @@ class Scheduler( and with it the server-process duties a Python ``TokenizerManager`` would otherwise own (e.g. serving the PD KV bootstrap registry).""" return envs.SGLANG_RUST_SERVER.get() and ( - self.ps.pp_rank == 0 - and self.ps.attn_tp_rank == 0 - and self.ps.attn_cp_rank == 0 + get_parallel().pp_rank == 0 + and get_parallel().attn_tp_rank == 0 + and get_parallel().attn_cp_rank == 0 ) def maybe_init_rust_server(self) -> None: @@ -2417,10 +2421,10 @@ class Scheduler( self.kv_events_publisher = SchedulerKvEventsPublisher( kv_events_config=get_observability().kv_events_config, ps=self.ps, - attn_tp_rank=self.ps.attn_tp_rank, - attn_cp_rank=self.ps.attn_cp_rank, + attn_tp_rank=get_parallel().attn_tp_rank, + attn_cp_rank=get_parallel().attn_cp_rank, attn_dp_rank=self.ps.attn_dp_rank, - dp_rank=self.ps.dp_rank, + dp_rank=get_parallel().dp_rank, tree_cache=self.tree_cache, send_metrics_from_scheduler=self.ipc_channels.send_metrics_from_scheduler, max_running_requests=self.max_running_requests, @@ -3398,7 +3402,7 @@ class Scheduler( if (timeout_s := envs.SGLANG_REQ_RUNNING_TIMEOUT.get()) > 0: deadline = time.perf_counter() - timeout_s - if self.ps.pp_size == 1: + if get_parallel().pp_size == 1: inflight_batches = [self.running_batch, self.last_batch] else: inflight_batches = [*self.running_mbs, *self.mbs] @@ -4460,7 +4464,7 @@ class Scheduler( batch.input_ids = None self._copy_auxiliary_output_to_cpu(batch, batch_result) elif not batch.spec_algorithm.is_none(): - is_verify_round = self.ps.pp_size > 1 and not ( + is_verify_round = get_parallel().pp_size > 1 and not ( batch.forward_mode.is_extend() or batch.is_extend_in_batch ) # The relayed tree is what the requests carry between rounds; @@ -4548,7 +4552,10 @@ class Scheduler( # Only the last PP rank owns real results requiring D2H; other ranks # consume device tensors rebuilt from the output ring. batch_result.copy_done = self.device_module.Event() - if batch_result.has_sampled_token_ids and self.ps.pp_size == 1: + if ( + batch_result.has_sampled_token_ids + and get_parallel().pp_size == 1 + ): batch_result.copy_to_cpu( return_logprob=batch.return_logprob, return_hidden_states=batch.return_hidden_states, @@ -4685,7 +4692,7 @@ class Scheduler( ): return # PP transports the device output to the first rank before copying it. - if self.ps.pp_size > 1: + if get_parallel().pp_size > 1: return if result.copy_done is not None: raise RuntimeError( @@ -5024,7 +5031,7 @@ class Scheduler( return idle def _pp_microbatches_drained(self) -> bool: - if self.ps.pp_size == 1: + if get_parallel().pp_size == 1: return True return all(x.is_empty() for x in self.running_mbs) and all( mb is None or mb.is_empty() for mb in self.mbs @@ -5265,10 +5272,10 @@ class Scheduler( if_success = False break elif k == "pp_max_micro_batch_size" and ( - v > self.max_running_requests // self.ps.pp_size or v < 1 + v > self.max_running_requests // get_parallel().pp_size or v < 1 ): logging.warning( - f"Updating {k} to {v} is rejected because it is out of the valid range [1, {self.max_running_requests // self.ps.pp_size}]." + f"Updating {k} to {v} is rejected because it is out of the valid range [1, {self.max_running_requests // get_parallel().pp_size}]." ) if_success = False break @@ -5382,7 +5389,7 @@ class Scheduler( ) def collect_inflight_reqs(self) -> Set[Req]: - if self.ps.pp_size == 1: + if get_parallel().pp_size == 1: inflight_batches = [self.running_batch, self.last_batch] else: inflight_batches = [*self.running_mbs, *self.mbs] @@ -5472,7 +5479,7 @@ class Scheduler( if hasattr(req.disagg_kv_sender, "abort"): req.disagg_kv_sender.abort() - if self.ps.pp_size > 1: + if get_parallel().pp_size > 1: prepare_abort(req, "Aborted by AbortReq.") # Abort in-flight requests @@ -5488,7 +5495,7 @@ class Scheduler( if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid): logger.debug(f"Abort prealloc queue request. {decode_req.req.rid=}") decode_req.kv_receiver.abort() - if self.ps.pp_size > 1: + if get_parallel().pp_size > 1: prepare_abort(decode_req.req, "Aborted by AbortReq.") # Abort requests waiting for kvcache to release tree cache @@ -5821,7 +5828,11 @@ class Scheduler( output = self.session_controller.open(recv_req) if output.success and self.enable_session_radix_cache: self.tree_cache.open_radix_session(recv_req.session_id) - if self.ps.pp_rank == 0 and self.ps.tp_rank == 0 and self.ps.attn_cp_rank == 0: + if ( + get_parallel().pp_rank == 0 + and get_parallel().tp_rank == 0 + and get_parallel().attn_cp_rank == 0 + ): return output return None @@ -5922,6 +5933,20 @@ def _dispatch_event_loop_once(scheduler: Scheduler): scheduler.event_loop_normal_disagg_decode() +def resolve_spawn_dp_rank(dp_rank: Optional[int]) -> Optional[int]: + """The `dp_rank` this process was spawned with, in either of its two forms. + + A router does not pass it as an argument, it sets `SGLANG_DP_RANK`. Both + forms are the launcher naming this process's place, so both have to be in + hand before `publish` records the placement -- resolving one of them after + would leave the context answering `None` for a process that has a rank. + """ + if dp_rank is None and "SGLANG_DP_RANK" in os.environ: + # [For Router] if env var "SGLANG_DP_RANK" exist, set dp_rank to the value of the env var + return int(os.environ["SGLANG_DP_RANK"]) + return dp_rank + + def configure_scheduler_process( server_args: ServerArgs, gpu_id: int, @@ -5934,18 +5959,15 @@ def configure_scheduler_process( display_tp_rank: Optional[int] = None, display_dp_rank: Optional[int] = None, display_moe_ep_rank: Optional[int] = None, -) -> Optional[int]: +) -> None: """Configure scheduler worker logging and process title. - display_* ranks are cosmetic; runtime ranks stay local. + display_* ranks are cosmetic; runtime ranks stay local. `dp_rank` arrives + already resolved -- see `resolve_spawn_dp_rank`. """ kill_itself_when_parent_died() # Generate the logger prefix - if dp_rank is None and "SGLANG_DP_RANK" in os.environ: - # [For Router] if env var "SGLANG_DP_RANK" exist, set dp_rank to the value of the env var - dp_rank = int(os.environ["SGLANG_DP_RANK"]) - shown_dp = display_dp_rank if display_dp_rank is not None else dp_rank shown_tp = display_tp_rank if display_tp_rank is not None else tp_rank shown_moe_ep = ( @@ -5987,8 +6009,6 @@ def configure_scheduler_process( if numa_node is not None: numa_bind_to_node(numa_node) - return dp_rank - def run_scheduler_process( server_args: ServerArgs, @@ -6007,9 +6027,20 @@ def run_scheduler_process( ): # Load plugins so hooks can override Scheduler and its dependencies. load_plugins() - # Publish before anything in this process reads configuration. - publish(server_args, role="scheduler") - dp_rank = configure_scheduler_process( + dp_rank = resolve_spawn_dp_rank(dp_rank) + # Publish before anything in this process reads configuration, with the + # placement the launcher decided: from here on a rank read is answered + # without a process group, which is what every reader needs before + # `init_torch_distributed` has run. + publish( + server_args, + role="scheduler", + ranks=SpawnRanks( + world_rank=spawn_world_rank(server_args, tp_rank=tp_rank, pp_rank=pp_rank), + dp_rank=dp_rank, + ), + ) + configure_scheduler_process( server_args, gpu_id, tp_rank, diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index 82770601b..a822d380e 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/python/sglang/srt/managers/scheduler_components/dp_attn.py @@ -548,8 +548,8 @@ class SchedulerDPAttnAdapter: local_batch, model_runner=self.model_runner, dp_size=get_parallel().dp_size, - attn_tp_size=self.ps.attn_tp_size, - attn_cp_size=self.ps.attn_cp_size, + attn_tp_size=get_parallel().attn_tp_size, + attn_cp_size=get_parallel().attn_cp_size, tp_group=self.tp_group, get_idle_batch=self.get_idle_batch, disable_cuda_graph=cuda_graph_fully_disabled(), diff --git a/python/sglang/srt/managers/scheduler_components/kv_events_publisher.py b/python/sglang/srt/managers/scheduler_components/kv_events_publisher.py index b483a9ac9..94d24d391 100644 --- a/python/sglang/srt/managers/scheduler_components/kv_events_publisher.py +++ b/python/sglang/srt/managers/scheduler_components/kv_events_publisher.py @@ -19,6 +19,7 @@ from sglang.srt.disaggregation.kv_events import ( select_kv_publisher_dp_rank, ) from sglang.srt.managers.io_struct import hook_custom_types, sock_send +from sglang.srt.runtime_context import get_parallel if TYPE_CHECKING: from sglang.srt.distributed.parallel_state_wrapper import ParallelState @@ -68,7 +69,7 @@ class SchedulerKvEventsPublisher: self.kv_event_publisher = EventPublisherFactory.create( kv_events_config, select_kv_publisher_dp_rank( - self.ps.attn_dp_size, self.ps.attn_dp_rank, self.ps.dp_rank + self.ps.attn_dp_size, self.ps.attn_dp_rank, get_parallel().dp_rank ), ) @@ -87,7 +88,7 @@ class SchedulerKvEventsPublisher: kv_metrics.gpu_cache_usage_perc = self.get_stats().token_usage kv_metrics.gpu_prefix_cache_hit_rate = self.get_stats().cache_hit_rate kv_metrics.data_parallel_rank = ( - self.ps.dp_rank if self.ps.dp_rank is not None else 0 + get_parallel().dp_rank if get_parallel().dp_rank is not None else 0 ) if not self.send_metrics_from_scheduler.closed: diff --git a/python/sglang/srt/managers/scheduler_components/load_inquirer.py b/python/sglang/srt/managers/scheduler_components/load_inquirer.py index bb1f5d47a..104a52c73 100644 --- a/python/sglang/srt/managers/scheduler_components/load_inquirer.py +++ b/python/sglang/srt/managers/scheduler_components/load_inquirer.py @@ -14,7 +14,7 @@ from sglang.srt.managers.load_snapshot import ( QueueMetrics, SpeculativeMetrics, ) -from sglang.srt.runtime_context import get_lora +from sglang.srt.runtime_context import get_lora, get_parallel if TYPE_CHECKING: from sglang.srt.distributed.parallel_state_wrapper import ParallelState @@ -215,7 +215,9 @@ class SchedulerLoadInquirer: decode_moments = list(totals) if totals[0] > 0 else None return LoadSnapshot( - dp_rank=int(self.ps.dp_rank) if self.ps.dp_rank is not None else 0, + dp_rank=int(get_parallel().dp_rank) + if get_parallel().dp_rank is not None + else 0, timestamp=time.time(), num_running_reqs=num_running_reqs, num_waiting_reqs=num_waiting_reqs, diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index 317e4e9aa..f8be4ab2a 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py @@ -309,7 +309,7 @@ class SchedulerMetricsReporter: if ( get_observability().enable_forward_pass_metrics and self.scheduler.ps.attn_tp_rank == 0 - and self.scheduler.ps.pp_rank == self.scheduler.ps.pp_size - 1 + and get_parallel().pp_rank == get_parallel().pp_size - 1 ): from sglang.srt.observability.forward_pass_metrics import ( _FpmPublisherThread, diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py index 15659cf2d..27250f484 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py @@ -32,7 +32,7 @@ from sglang.srt.managers.schedule_batch import ( Req, ) from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache -from sglang.srt.runtime_context import get_observability, get_serving +from sglang.srt.runtime_context import get_observability, get_parallel, get_serving from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.utils.weight_versions import compute_weight_version_spans @@ -205,7 +205,7 @@ class SchedulerOutputStreamer: # Send to detokenizer payload = acc.to_payload( - dp_rank=self.ps.dp_rank, + dp_rank=get_parallel().dp_rank, is_idle_batch=is_idle_batch, ) if payload is not None: @@ -232,7 +232,7 @@ class SchedulerOutputStreamer: def _maybe_log_time_stats(self, *, req: Req) -> None: if ( req.finished() - and self.ps.attn_tp_rank == 0 + and get_parallel().attn_tp_rank == 0 and get_observability().enable_request_time_stats_logging ): req.log_time_stats() diff --git a/python/sglang/srt/managers/scheduler_components/profiler_manager.py b/python/sglang/srt/managers/scheduler_components/profiler_manager.py index 53387b6aa..ecc181a8c 100644 --- a/python/sglang/srt/managers/scheduler_components/profiler_manager.py +++ b/python/sglang/srt/managers/scheduler_components/profiler_manager.py @@ -21,7 +21,7 @@ from sglang.srt.managers.io_struct import ProfileReq, ProfileReqOutput, ProfileR from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.model_executor.step_span_utils import set_detailed_annotations_enabled from sglang.srt.platforms import current_platform -from sglang.srt.runtime_context import get_device +from sglang.srt.runtime_context import get_device, get_parallel from sglang.srt.utils import is_mps, is_npu from sglang.srt.utils.profile_merger import ProfileMerger from sglang.srt.utils.profile_utils import ProfileManager @@ -207,10 +207,13 @@ class SchedulerProfilerManager: self.rpd_profile_path = os.path.join( self.torch_profiler_output_dir, - "rpd-" + str(time.time()) + f"-TP-{self.ps.tp_rank}" + ".trace.json.gz", + "rpd-" + + str(time.time()) + + f"-TP-{get_parallel().tp_rank}" + + ".trace.json.gz", ) - if self.ps.tp_rank == 0: + if get_parallel().tp_rank == 0: import sqlite3 from rocpd.schema import RocpdSchema @@ -282,13 +285,13 @@ class SchedulerProfilerManager: if not self.merge_profiles: return "" - if self.ps.tp_rank != 0: + if get_parallel().tp_rank != 0: return "" - if self.ps.dp_size > 1 and self.ps.dp_rank != 0: + if self.ps.dp_size > 1 and get_parallel().dp_rank != 0: return "" - if self.ps.pp_size > 1 and self.ps.pp_rank != 0: + if get_parallel().pp_size > 1 and get_parallel().pp_rank != 0: return "" - if self.ps.moe_ep_size > 1 and self.ps.moe_ep_rank != 0: + if get_parallel().moe_ep_size > 1 and get_parallel().moe_ep_rank != 0: return "" try: @@ -336,15 +339,15 @@ class SchedulerProfilerManager: self.torch_profiler.stop() if not _is_npu: # Build filename with only non-zero ranks to maintain backward compatibility - filename_parts = [self.profile_id, f"TP-{self.ps.tp_rank}"] + filename_parts = [self.profile_id, f"TP-{get_parallel().tp_rank}"] # Only add other ranks if parallelism is enabled (size > 1) if self.ps.dp_size > 1: - filename_parts.append(f"DP-{self.ps.dp_rank}") - if self.ps.pp_size > 1: - filename_parts.append(f"PP-{self.ps.pp_rank}") - if self.ps.moe_ep_size > 1: - filename_parts.append(f"EP-{self.ps.moe_ep_rank}") + filename_parts.append(f"DP-{get_parallel().dp_rank}") + if get_parallel().pp_size > 1: + filename_parts.append(f"PP-{get_parallel().pp_rank}") + if get_parallel().moe_ep_size > 1: + filename_parts.append(f"EP-{get_parallel().moe_ep_rank}") filename = ( stage_prefix @@ -364,7 +367,7 @@ class SchedulerProfilerManager: self.rpd_profiler.flush() torch.distributed.barrier(self.dp_tp_cpu_group) - if self.ps.tp_rank == 0: + if get_parallel().tp_rank == 0: from sglang.srt.utils.rpd_utils import rpd_to_chrome_trace rpd_to_chrome_trace("trace.rpd", self.rpd_profile_path) @@ -376,7 +379,7 @@ class SchedulerProfilerManager: self.torch_profiler_output_dir, stage_prefix + str(time.time()) - + f"-TP-{self.ps.tp_rank}-memory" + + f"-TP-{get_parallel().tp_rank}-memory" + stage_suffix + ".pickle", ) diff --git a/python/sglang/srt/managers/scheduler_components/request_receiver.py b/python/sglang/srt/managers/scheduler_components/request_receiver.py index 5970fb768..79c88a17f 100644 --- a/python/sglang/srt/managers/scheduler_components/request_receiver.py +++ b/python/sglang/srt/managers/scheduler_components/request_receiver.py @@ -109,7 +109,7 @@ class SchedulerRequestReceiver: recv_reqs = self._broadcast_reqs_across_ranks(recv_reqs, local_reqs) - if self.ps.pp_rank == 0: + if get_parallel().pp_rank == 0: self.unwrap_pickle_wrapper(recv_reqs) recv_reqs = self._apply_mm_receiver(recv_reqs) @@ -119,7 +119,7 @@ class SchedulerRequestReceiver: return recv_reqs def _pull_raw_reqs(self) -> Optional[List]: - if self.ps.pp_rank == 0: + if get_parallel().pp_rank == 0: if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: recv_reqs = [] @@ -158,10 +158,10 @@ class SchedulerRequestReceiver: ) recv_reqs = point_to_point_pyobj( [], - self.ps.pp_rank * self.ps.tp_size + dp_offset, + get_parallel().pp_rank * self.ps.tp_size + dp_offset, self.world_group.cpu_group, - (self.ps.pp_rank - 1) * self.ps.tp_size + dp_offset, - self.ps.pp_rank * self.ps.tp_size + dp_offset, + (get_parallel().pp_rank - 1) * self.ps.tp_size + dp_offset, + get_parallel().pp_rank * self.ps.tp_size + dp_offset, ) else: recv_reqs = None @@ -232,7 +232,7 @@ class SchedulerRequestReceiver: def _apply_mm_receiver(self, recv_reqs: List) -> List: # Process MM requests under EPD-disaggregation mode if ( - self.ps.pp_rank == 0 + get_parallel().pp_rank == 0 and get_disagg().language_only and get_disagg().encoder_transfer_backend in ["zmq_to_scheduler", "mooncake"] diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 5fdd1fc05..c293589ba 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -103,7 +103,9 @@ class SchedulerPPMixin: for mb_id in range(self.pp_loop_size): self.running_batch = self.running_mbs[mb_id] self.last_batch = self.last_mbs[mb_id] - next_first_rank_mb_id = (mb_id + self.ps.pp_size) % self.pp_loop_size + next_first_rank_mb_id = ( + mb_id + get_parallel().pp_size + ) % self.pp_loop_size next_mb_id = (mb_id + 1) % self.pp_loop_size with torch.profiler.record_function("recv_requests"): recv_reqs = self.ingest_requests() @@ -243,7 +245,9 @@ class SchedulerPPMixin: for mb_id in range(self.pp_loop_size): self.running_batch = self.running_mbs[mb_id] self.last_batch = self.last_mbs[mb_id] - next_first_rank_mb_id = (mb_id + self.ps.pp_size) % self.pp_loop_size + next_first_rank_mb_id = ( + mb_id + get_parallel().pp_size + ) % self.pp_loop_size next_mb_id = (mb_id + 1) % self.pp_loop_size next_pp_outputs = None @@ -391,7 +395,9 @@ class SchedulerPPMixin: for mb_id in range(self.pp_loop_size): self.running_batch = self.running_mbs[mb_id] self.last_batch = self.last_mbs[mb_id] - next_first_rank_mb_id = (mb_id + self.ps.pp_size) % self.pp_loop_size + next_first_rank_mb_id = ( + mb_id + get_parallel().pp_size + ) % self.pp_loop_size next_mb_id = (mb_id + 1) % self.pp_loop_size next_pp_outputs = None @@ -561,7 +567,9 @@ class SchedulerPPMixin: self.on_idle() def init_pp_loop_state(self: Scheduler): - self.pp_loop_size: int = self.ps.pp_size + get_parallel().pp_async_batch_depth + self.pp_loop_size: int = ( + get_parallel().pp_size + get_parallel().pp_async_batch_depth + ) self.mbs = [None] * self.pp_loop_size self.last_mbs = [None] * self.pp_loop_size self.running_mbs = [ @@ -573,7 +581,7 @@ class SchedulerPPMixin: self.last_rank_comm_queue: deque[Tuple[torch.Event, PPProxyTensors]] = deque() self._pp_spec_relay = ( envs.SGLANG_ENABLE_PP_SPEC.get() - and self.ps.pp_size > 1 + and get_parallel().pp_size > 1 and not self.spec_algorithm.is_none() ) @@ -805,31 +813,39 @@ class SchedulerPPMixin: def _pp_send_pyobj_to_next_stage(self: Scheduler, data, async_send: bool = False): p2p_work = [] - if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: + if get_parallel().attn_tp_rank == 0 and get_parallel().attn_cp_rank == 0: dp_offset = ( - self.ps.attn_dp_rank * self.ps.attn_cp_size * self.ps.attn_tp_size + self.ps.attn_dp_rank + * get_parallel().attn_cp_size + * get_parallel().attn_tp_size ) p2p_work = point_to_point_pyobj( data, - self.ps.pp_rank * self.ps.tp_size + dp_offset, + get_parallel().pp_rank * get_parallel().tp_size + dp_offset, self.world_group.cpu_group, - self.ps.pp_rank * self.ps.tp_size + dp_offset, - ((self.ps.pp_rank + 1) % self.ps.pp_size) * self.ps.tp_size + dp_offset, + get_parallel().pp_rank * get_parallel().tp_size + dp_offset, + ((get_parallel().pp_rank + 1) % get_parallel().pp_size) + * get_parallel().tp_size + + dp_offset, async_send=async_send, ) return p2p_work def _pp_recv_pyobj_from_prev_stage(self: Scheduler): - if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: + if get_parallel().attn_tp_rank == 0 and get_parallel().attn_cp_rank == 0: dp_offset = ( - self.ps.attn_dp_rank * self.ps.attn_cp_size * self.ps.attn_tp_size + self.ps.attn_dp_rank + * get_parallel().attn_cp_size + * get_parallel().attn_tp_size ) data = point_to_point_pyobj( [], - self.ps.pp_rank * self.ps.tp_size + dp_offset, + get_parallel().pp_rank * get_parallel().tp_size + dp_offset, self.world_group.cpu_group, - ((self.ps.pp_rank - 1) % self.ps.pp_size) * self.ps.tp_size + dp_offset, - self.ps.pp_rank * self.ps.tp_size + dp_offset, + ((get_parallel().pp_rank - 1) % get_parallel().pp_size) + * get_parallel().tp_size + + dp_offset, + get_parallel().pp_rank * get_parallel().tp_size + dp_offset, ) else: data = None @@ -1470,7 +1486,7 @@ class SchedulerPPMixin: # posted, so the parity-based send-first/recv-first ordering used # for NPU is replaced by batch_isend_irecv which submits all # send/recv operations atomically. - if _is_npu and self.ps.pp_size == 2: + if _is_npu and get_parallel().pp_size == 2: return self._pp2_only_send_recv_output_tensors_npu( next_first_rank_mb_id, next_mb_id, @@ -1498,7 +1514,7 @@ class SchedulerPPMixin: # makes rank 1 post its recv first, which breaks the cycle for any # pp_size > 1. needs_pairing = is_xpu() or self._pp_spec_relay - send_first = (not needs_pairing) or ((self.ps.pp_rank % 2) == 0) + send_first = (not needs_pairing) or ((get_parallel().pp_rank % 2) == 0) def _do_send(): return self._pp_send_output_to_next_stage( diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index db7259755..23f849cd7 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -418,7 +418,7 @@ class TpModelWorker(BaseTpWorker): else: self.random_seed = broadcast_pyobj( [get_device().random_seed], - self.ps.tp_size * self.ps.pp_rank + self.ps.tp_rank, + self.ps.tp_size * get_parallel().pp_rank + self.ps.tp_rank, self.world_group.cpu_group, src=self.world_group.ranks[0], )[0] @@ -451,7 +451,8 @@ class TpModelWorker(BaseTpWorker): assert self.model_runner.max_running_requests > 0, "max_running_request is zero" max_req_len = min( self.model_config.context_len - 1, - self.model_runner.effective_max_total_num_tokens * self.ps.attn_dcp_size + self.model_runner.effective_max_total_num_tokens + * get_parallel().attn_dcp_size - 1, ) assert max_req_len > 0, "Memory pool size is too small" @@ -578,7 +579,8 @@ class TpModelWorker(BaseTpWorker): def get_worker_info(self): max_req_len = min( self.model_config.context_len - 1, - self.model_runner.effective_max_total_num_tokens * self.ps.attn_dcp_size + self.model_runner.effective_max_total_num_tokens + * get_parallel().attn_dcp_size - 1, ) return ( diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index b1732f0c6..6025de05c 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -2420,7 +2420,9 @@ class KVCacheConfigurator: sum(1 for i in all_mamba_layers if start <= i < end) for start, end in ( get_pp_indices( - self.model_config.num_hidden_layers, rank, self.ps.pp_size + self.model_config.num_hidden_layers, + rank, + self.ps.pp_size, ) for rank in range(self.ps.pp_size) ) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index b7856f1a3..fe865d0b8 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -445,7 +445,7 @@ class ModelRunner: self.prefill_shared_read_stager: Optional[Callable[[ForwardBatch], bool]] = None # CPU offload - set_offloader(create_offloader(dp_rank=self.ps.dp_rank)) + set_offloader(create_offloader(dp_rank=get_parallel().dp_rank)) self._weight_checker = WeightChecker(get_model=lambda: self.model, ps=self.ps) @@ -481,7 +481,7 @@ class ModelRunner: "pp_proxy_tensors" in inspect.signature(self.model.forward).parameters ) - if self.ps.pp_size > 1: + if get_parallel().pp_size > 1: if not (envs.SGLANG_ENABLE_PP_SPEC.get() and self.is_draft_worker): assert self.support_pp, ( "Pipeline Parallel is not compatible with this model." @@ -647,8 +647,8 @@ class ModelRunner: from sglang.srt.model_executor.mindspore_runner import init_ms_distributed init_ms_distributed( - world_size=self.ps.tp_size * self.ps.pp_size, - rank=self.ps.tp_size * self.ps.pp_rank + self.ps.tp_rank, + world_size=self.ps.tp_size * get_parallel().pp_size, + rank=self.ps.tp_size * get_parallel().pp_rank + self.ps.tp_rank, local_rank=self.gpu_id, port=self.dist_port, ) @@ -667,8 +667,8 @@ class ModelRunner: prepare_moe_topk( model=self.model, model_config=self.model_config, - moe_ep_size=self.ps.moe_ep_size, - moe_ep_rank=self.ps.moe_ep_rank, + moe_ep_size=get_parallel().moe_ep_size, + moe_ep_rank=get_parallel().moe_ep_rank, ) self.maybe_init_dwdp() @@ -708,7 +708,7 @@ class ModelRunner: def maybe_init_expert_location_metadata(self): if self.is_draft_worker: return - expert_rank = self.ps.moe_ep_rank + ( + expert_rank = get_parallel().moe_ep_rank + ( get_parallel().ep_join_rank_offset if get_exec().moe.is_ep_scale_joiner else 0 @@ -767,8 +767,8 @@ class ModelRunner: self.expert_backup_client = ( ExpertBackupClient( model_config=self.model_config, - moe_ep_size=self.ps.moe_ep_size, - moe_ep_rank=self.ps.moe_ep_rank, + moe_ep_size=get_parallel().moe_ep_size, + moe_ep_rank=get_parallel().moe_ep_rank, get_model=lambda: self.model, ) if ( @@ -803,16 +803,16 @@ class ModelRunner: def get_pp_proxy_topk_size(self) -> Optional[int]: return misc_utils.resolve_pp_proxy_topk_size( model_config=self.model_config, - pp_size=self.ps.pp_size, - pp_rank=self.ps.pp_rank, + pp_size=get_parallel().pp_size, + pp_rank=get_parallel().pp_rank, start_layer=self.layer_info.start_layer, ) def get_pp_proxy_residual_num_blocks(self) -> Optional[int]: return misc_utils.resolve_pp_proxy_residual_num_blocks( model_config=self.model_config, - pp_size=self.ps.pp_size, - pp_rank=self.ps.pp_rank, + pp_size=get_parallel().pp_size, + pp_rank=get_parallel().pp_rank, start_layer=self.layer_info.start_layer, ) @@ -980,7 +980,7 @@ class ModelRunner: swap_in_block_size=hisparse_cfg.swap_in_block_size, shared_index_layers=resolve_shared_index_layers( hf_text_config=self.model_config.hf_text_config, - pp_size=self.ps.pp_size, + pp_size=get_parallel().pp_size, is_speculative=self.spec_algorithm.is_speculative(), ), ) @@ -1157,8 +1157,8 @@ class ModelRunner: check_quantized_moe_compatibility( model_config=self.model_config, tp_size=self.ps.tp_size, - moe_ep_size=self.ps.moe_ep_size, - moe_dp_size=self.ps.moe_dp_size, + moe_ep_size=get_parallel().moe_ep_size, + moe_dp_size=get_parallel().moe_dp_size, ) def init_torch_distributed(self): @@ -1293,7 +1293,7 @@ class ModelRunner: is_draft_worker=self.is_draft_worker, tp_size=self.ps.tp_size, tp_rank=self.ps.tp_rank, - pp_rank=self.ps.pp_rank, + pp_rank=get_parallel().pp_rank, ) if dumper.may_enable: @@ -1580,7 +1580,7 @@ class ModelRunner: dp_size = 1 if get_parallel().enable_dp_attention else self.ps.dp_size self.local_omp_cpuid = numa_utils.init_threads_binding( numa_index=self.gpu_id, - world_size=dp_size * self.ps.tp_size * self.ps.pp_size, + world_size=dp_size * self.ps.tp_size * get_parallel().pp_size, ) def apply_torch_tp(self): diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index d102d94e2..39032b131 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -274,7 +274,7 @@ class BaseRunner(ABC): if ( envs.SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP.get() and deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM - and mr.ps.pp_size > 1 + and get_parallel().pp_size > 1 and not mr.spec_algorithm.is_speculative() ): from sglang.srt.layers.deep_gemm_wrapper.compile_utils import ( @@ -546,7 +546,7 @@ class BaseRunner(ABC): pp_hidden_tokens = num_tokens if ( capture_forward_mode == ForwardMode.EXTEND - and mr.ps.pp_rank != 0 + and get_parallel().pp_rank != 0 and mr.ps.attn_cp_size > 1 ): pp_hidden_tokens = num_tokens // mr.ps.attn_cp_size diff --git a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py index 809b4c593..2f9e89dca 100644 --- a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py +++ b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py @@ -30,6 +30,7 @@ from sglang.srt.runtime_context import ( get_disagg, get_exec, get_model, + get_parallel, get_schedule, get_spec, max_prefill_buffer_tokens, @@ -147,7 +148,7 @@ def flashinfer_autotune_cache_path(model_runner: ModelRunner) -> Path: str(get_model().quantization), str(get_exec().moe.moe_runner_backend), str(mr.ps.tp_size), - str(mr.ps.pp_size), + str(get_parallel().pp_size), str(mr.ps.attn_dp_size), str(mr.ps.moe_ep_size), str(mr.model_config.hf_config.__class__.__name__), @@ -170,7 +171,7 @@ def flashinfer_autotune_cache_path(model_runner: ModelRunner) -> Path: cache_dir.mkdir(parents=True, exist_ok=True) return ( cache_dir - / f"rank_tp{mr.ps.tp_rank}_pp{mr.ps.pp_rank}_dp{mr.ps.dp_rank or 0}.json" + / f"rank_tp{mr.ps.tp_rank}_pp{get_parallel().pp_rank}_dp{mr.ps.dp_rank or 0}.json" ) diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index d58f9bcf1..ede71c989 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -34,6 +34,7 @@ from sglang.srt.runtime_context import ( get_context, get_disagg, get_observability, + get_parallel, get_schedule, get_serving, ) @@ -1118,7 +1119,7 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin): ) enable_kv_cache_events = bool( get_observability().kv_events_config - and ps.pp_rank == 0 + and get_parallel().pp_rank == 0 and ps.attn_tp_rank == 0 and ps.attn_cp_rank == 0 ) diff --git a/python/sglang/srt/ray/scheduler_actor.py b/python/sglang/srt/ray/scheduler_actor.py index 3616d1816..cfd97a58a 100644 --- a/python/sglang/srt/ray/scheduler_actor.py +++ b/python/sglang/srt/ray/scheduler_actor.py @@ -21,7 +21,7 @@ from typing import Any, Dict, Optional import ray from sglang.srt.arg_groups.overrides import declare_resolution -from sglang.srt.runtime_context import publish +from sglang.srt.runtime_context import SpawnRanks, publish, spawn_world_rank from sglang.srt.server_args import PortArgs, ServerArgs logger = logging.getLogger(__name__) @@ -49,7 +49,11 @@ class SchedulerActor: dist_init_addr: Optional[str] = None, ): from sglang.srt.environ import envs - from sglang.srt.managers.scheduler import Scheduler, configure_scheduler_process + from sglang.srt.managers.scheduler import ( + Scheduler, + configure_scheduler_process, + resolve_spawn_dp_rank, + ) from sglang.srt.utils.numa_utils import ( get_numa_node_if_available, numa_bind_to_node, @@ -77,12 +81,23 @@ class SchedulerActor: actual_gpu_id = gpu_id logger.info(f"[TP{tp_rank}] Using passed gpu_id: {gpu_id}") + dp_rank = resolve_spawn_dp_rank(dp_rank) + # This actor takes the place of run_scheduler_process, which is where # a forked scheduler publishes. - publish(server_args, role="scheduler") + publish( + server_args, + role="scheduler", + ranks=SpawnRanks( + world_rank=spawn_world_rank( + server_args, tp_rank=tp_rank, pp_rank=pp_rank + ), + dp_rank=dp_rank, + ), + ) # Configure worker (logging, process title, etc.) - dp_rank = configure_scheduler_process( + configure_scheduler_process( server_args, actual_gpu_id, tp_rank, diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index 1232b480f..f8c424fd5 100644 --- a/python/sglang/srt/runtime_context.py +++ b/python/sglang/srt/runtime_context.py @@ -132,6 +132,10 @@ class Live(msgspec.Struct, frozen=True): source: Any = None doc: str = "" + # For a stamp-only name (`source=None`): what a reader should be told when + # nothing has stamped it. These names have no fallback by construction, so + # the message is the only thing pointing at what did not happen. + unstamped: str = "" _LIVE_READS: dict = { @@ -177,7 +181,35 @@ _LIVE_READS: dict = { "attn_cp_rank": "get_attn_context_model_parallel_rank", "dcp_rank": "get_dcp_rank", "attn_dcp_rank": lambda self: self.dcp_rank if self.dcp_enabled else 0, - "attn_dp_rank": None, + "attn_dp_rank": Live( + source=None, + doc=( + "This process's index in the attention-DP group. Computed from " + "`tp_rank` when the attention topology is initialized, and moved " + "by an elastic scale-up, so no coordinator can answer it." + ), + unstamped=( + "it is computed from this process's `tp_rank` when the attention " + "topology is initialized, so a process that never ran " + "`initialize_dp_attention` has no answer to give" + ), + ), + "dp_rank": Live( + source=None, + doc=( + "Which data-parallel replica this process serves, as the data " + "parallel controller numbered them at spawn. `None` when there is " + "no controller. Unlike `attn_dp_rank` and `moe_dp_rank` it is not " + "a position in any process group -- no group has one member per " + "replica -- which is why nothing can derive it and the spawn " + "states it instead." + ), + unstamped=( + "it is a spawn identity, handed to `publish(..., ranks=...)` by " + "the process entry; a process that published without a rank " + "bundle has no replica index to report" + ), + ), "world_group": "get_world_group", "tp_group": "get_tp_group", "pp_group": "get_pp_group", @@ -227,6 +259,80 @@ def derive_attention_widths( return attn_dp_size, tp_size // attn_dp_size // attn_cp_size +def derive_attention_ranks( + *, tp_rank: int, attn_tp_size: int, attn_cp_size: int, enable_dp_attention: bool +) -> tuple: + """(attn_tp_rank, attn_dp_rank) for a process at `tp_rank`. + + The rank layout is (dp, cp, tp) with tp the fastest-changing dimension:: + + tp_rank = (attn_dp_rank * attn_cp_size + attn_cp_rank) * attn_tp_size + + attn_tp_rank + + Split out beside `derive_attention_widths` because two places need it from + different inputs: `publish` has this process's `tp_rank` from the spawn and + the widths from the configuration, while `initialize_dp_attention` has them + from the groups it just built. They must not carry separate copies of the + arithmetic -- the point of computing it at publish is that the two agree. + """ + attn_tp_rank = tp_rank % attn_tp_size + if not enable_dp_attention: + return attn_tp_rank, 0 + return attn_tp_rank, tp_rank // (attn_tp_size * attn_cp_size) + + +def spawn_world_rank(server_args, *, tp_rank: int, pp_rank: int) -> int: + """This process's place in WORLD, from the ranks its entry was given. + + The inverse of `derive_spawn_ranks`, for the entries that have the pieces + but not the whole: the same expression `bootstrap` hands to + `init_distributed_environment` when it builds the group. + + Reads a resolving view because it runs before `publish`. + """ + from sglang.srt.arg_groups.model_override_base import resolving_view + + cfg = resolving_view(server_args) + return cfg.ep_join_rank_offset + cfg.tp_size * pp_rank + tp_rank + + +def derive_spawn_ranks( + *, + world_rank: int, + tp_size: int, + ep_join_rank_offset: int, + attn_cp_size: int, + attn_tp_size: int, + moe_dp_size: int, + moe_ep_size: int, +) -> dict: + """Every rank a process group would answer, from its place in WORLD. + + WORLD is laid out `rank = ep_join_rank_offset + tp_size * pp_rank + + tp_rank`: `initialize_model_parallel` builds tensor-parallel groups as + contiguous blocks of `tp_size` and pipeline groups strided by it, so the + map is a bijection and this is its inverse. The attention and MoE ranks + are then positions inside the tensor-parallel block, which is what the + launcher computes when it decides what to spawn. + + Pure arithmetic over the published widths: no group is consulted, which is + the point -- this runs at publish, before any of them exist. + """ + local = world_rank - ep_join_rank_offset + tp_rank = local % tp_size + return { + "tp_rank": tp_rank, + "pp_rank": local // tp_size, + "attn_cp_rank": (tp_rank // attn_tp_size) % attn_cp_size, + "moe_dp_rank": tp_rank // (tp_size // moe_dp_size), + "moe_ep_rank": ( + tp_rank + % (tp_size // moe_dp_size) + // (tp_size // moe_dp_size // moe_ep_size) + ), + } + + def derive_parallel_widths( *, tp_size: int, @@ -327,6 +433,31 @@ def dcp_enabled_of(cfg: Any): return parallel_widths_of(cfg)["dcp_enabled"] +class SpawnRanks(msgspec.Struct, frozen=True): + """Where the launcher put this process, in the two numbers only it knows. + + `world_rank` is this process's place in the WORLD group, which fixes every + other rank: the groups are laid out from the published widths, so + `tp_rank`, `pp_rank` and the attention / MoE ranks are functions of it (see + `derive_spawn_ranks`). Passing them separately would be passing the same + fact five more times, with five more ways for an entry to contradict + itself. + + `dp_rank` is the exception, because data-parallel replicas are separate + WORLD groups: with `--dp-size 2` each replica holds ranks `0 .. n-1`, so + 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. + """ + + world_rank: int + dp_rank: Optional[int] = None + + class ParallelContext: """Parallel-topology namespace: one spelling per name. @@ -394,11 +525,10 @@ class ParallelContext: return getattr(_ps(), source)() if source is not None: return source(self) + why = live.unstamped if isinstance(live, Live) else "" raise RuntimeError( - f"parallel rank {name!r} is not available: it is computed from " - "this process's `tp_rank` when the attention topology is " - "initialized, so a process that never ran " - "`initialize_dp_attention` has no answer to give" + f"parallel name {name!r} is not available: " + + (why or "nothing has stamped it in this process") ) if config is None and name in _parallel_config_leaves(): raise ValueError("config namespace 'parallel' not published") @@ -1598,7 +1728,13 @@ def _dump_recorded_namespace_reads() -> None: ) -def publish(server_args, *, role: str, hf_config: Any = None) -> RuntimeContext: +def publish( + server_args, + *, + role: str, + hf_config: Any = None, + ranks: SpawnRanks | None = None, +) -> RuntimeContext: """Install process-wide config for this OS process. Records the process ``role`` — one of the ``ROLE_NAMESPACE_SETS`` keys, @@ -1609,6 +1745,12 @@ def publish(server_args, *, role: str, hf_config: Any = None) -> RuntimeContext: namespace-read enforcement (``record`` audits the reads instead). ``hf_config`` is accepted for forward-compat and currently unused. + ``ranks`` is who this process is, from the entry that spawned it. It is + optional because most roles are not placed in the topology at all -- a + tokenizer has no ``tp_rank`` -- and those processes raise on a rank read + exactly as they do today, with a message naming the missing bundle rather + than an absent process group. + A process holds at most one live config: the bags always describe the engine running now. Re-publish is allowed and is **last-publish-wins** (bags re-projected, provenance reset, role overwritten), which is what @@ -1634,6 +1776,36 @@ def publish(server_args, *, role: str, hf_config: Any = None) -> RuntimeContext: ), ) _CONTEXT._publish_role = role + 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 + # of these then needs no process group, which is the point -- they are + # read long before one exists. The scoped overrides that swap a group + # for a draft worker sit above the record in the read chain, so a + # scope still wins. + parallel = _CONTEXT.parallel + placement = derive_spawn_ranks( + world_rank=ranks.world_rank, + tp_size=parallel.tp_size, + ep_join_rank_offset=parallel.ep_join_rank_offset, + attn_cp_size=parallel.attn_cp_size, + attn_tp_size=parallel.attn_tp_size, + moe_dp_size=parallel.moe_dp_size, + moe_ep_size=parallel.moe_ep_size, + ) + # `moe_dp_rank` is a different quantity when the MoE-DP group is + # aliased to the attention-CP one: the group answers the CP index, + # while this computes the MoE-DP index. Leave it to the group there, so + # one name does not mean two things. + if parallel.moe_dp_size < parallel.attn_cp_size: + placement.pop("moe_dp_rank") + # `dp_rank` is recorded whatever it is, None included: replicas are + # separate WORLD groups, so no rank implies it and `None` is the answer + # "no controller" rather than an absence. + placement["dp_rank"] = ranks.dp_rank + placement["launch_world_rank"] = ranks.world_rank + parallel.override_permanently(**placement) + _stamp_attention_ranks(parallel, placement["tp_rank"]) if _ROLE_NS_MODE == "record": # The '-' marker distinguishes a zero-read role from a process where # recording never ran (signal teardown skips atexit). @@ -1648,6 +1820,27 @@ def publish(server_args, *, role: str, hf_config: Any = None) -> RuntimeContext: return _CONTEXT +def _stamp_attention_ranks(parallel, tp_rank: int) -> None: + """Place this process in the attention topology, from the configuration. + + The widths are already on the bag -- `publish` computed them a moment ago -- + and the rank comes from the spawn, so the position is known here, before any + process group exists. That is the point: a rank read then works in a process + 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. + """ + attn_tp_rank, attn_dp_rank = derive_attention_ranks( + tp_rank=tp_rank, + attn_tp_size=parallel.attn_tp_size, + 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) + + def assert_published(server_args, *, role: str) -> RuntimeContext: """This record, under this role, is already published -- or fail loud. diff --git a/python/sglang/srt/rust_server/server.py b/python/sglang/srt/rust_server/server.py index 803555132..3c71bb9c9 100644 --- a/python/sglang/srt/rust_server/server.py +++ b/python/sglang/srt/rust_server/server.py @@ -124,7 +124,7 @@ class RustServer: # The joining TP group is entirely local to this node. tp_size_per_node = scheduler.ps.tp_size else: - nnodes_per_pp_rank = max(get_parallel().nnodes // scheduler.ps.pp_size, 1) + nnodes_per_pp_rank = max(get_parallel().nnodes // get_parallel().pp_size, 1) tp_size_per_node = scheduler.ps.tp_size // nnodes_per_pp_rank dp_group_width = scheduler.ps.attn_tp_size * scheduler.ps.attn_cp_size # Count DP leaders within this node's TP range. The first leader must diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index fafe1d64d..ec94e1ddb 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -84,6 +84,7 @@ from sglang.srt.speculative.spec_utils import ( GrammarTree, assign_req_to_token_pool_func, build_grammar_vocab_mask, + draft_pp_context, draft_tp_context, ) from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu @@ -402,7 +403,7 @@ class DFlashWorkerV2(BaseSpecWorker): draft_init_ctx = draft_tp_context(get_parallel().attn_tp_group) else: draft_init_ctx = empty_context() - with draft_init_ctx: + with draft_pp_context(), draft_init_ctx: bundle = build_draft_tp_worker( server_args=server_args, gpu_id=gpu_id, @@ -598,7 +599,10 @@ class DFlashWorkerV2(BaseSpecWorker): ) def init_attention_backends(self): - with self.draft_tp_context(self.draft_model_runner.tp_group): + with ( + draft_pp_context(), + self.draft_tp_context(self.draft_model_runner.tp_group), + ): self._draft_worker.init_attention_backends() self._need_mamba_verify_commit = mambaish_config( self.model_runner.model_config @@ -608,7 +612,10 @@ class DFlashWorkerV2(BaseSpecWorker): ) def init_cuda_graphs(self): - with self.draft_tp_context(self.draft_model_runner.tp_group): + with ( + draft_pp_context(), + self.draft_tp_context(self.draft_model_runner.tp_group), + ): capture_decode_cuda_graph = ( get_exec().graph.cuda_graph_config.decode.backend != Backend.DISABLED ) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index 31a7bb82d..1f4038761 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -74,6 +74,7 @@ from sglang.srt.speculative.spec_tp_sync import SpecTpSync, SpecTpSyncSite from sglang.srt.speculative.spec_utils import ( GrammarTree, build_grammar_vocab_mask, + draft_pp_context, draft_tp_context, prepare_mamba_track_for_verify, ) @@ -165,7 +166,7 @@ class DSparkWorkerV2(BaseSpecWorker): "MoE-under-DP all-reduce." ) - with self._draft_context(): + with draft_pp_context(), self._draft_context(): bundle = build_draft_tp_worker( server_args=server_args, gpu_id=gpu_id, @@ -428,7 +429,7 @@ class DSparkWorkerV2(BaseSpecWorker): ) def init_attention_backends(self): - with self._draft_context(): + with draft_pp_context(), self._draft_context(): self._draft_worker.init_attention_backends() self._target_hidden_projection_enabled = _configure_target_hidden_projection( target_model=self.target_worker.model_runner.model, @@ -463,7 +464,7 @@ class DSparkWorkerV2(BaseSpecWorker): "memory is available after target backend initialization.", available_mem, ) - with self._draft_context(): + with draft_pp_context(), self._draft_context(): if capture_decode_cuda_graph: # Keep the draft model graph enabled when folded proposal is # disabled, but do not capture the proposal head as a tail diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 16158434f..b2de353dc 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -338,6 +338,7 @@ class EagleDraftWorker(EagleDraftWorkerBase): def init_attention_backends(self): with ( + draft_pp_context(), self.draft_tp_context(self.draft_runner.tp_group), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(), @@ -347,6 +348,7 @@ class EagleDraftWorker(EagleDraftWorkerBase): def init_cuda_graphs(self): with ( + draft_pp_context(), self.draft_tp_context(self.draft_runner.tp_group), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(), diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index 42a0d212a..4354c9fcb 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -73,6 +73,7 @@ from sglang.srt.speculative.frozen_kv_mtp_utils import ( ) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import ( + draft_pp_context, draft_tp_context, fast_topk, get_plan_stream, @@ -134,7 +135,7 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker): self.hot_token_id = None with ( - empty_context(), + draft_pp_context(), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(), draft_model_build_scope(), @@ -204,6 +205,7 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker): def init_attention_backends(self): with ( + draft_pp_context(), self.draft_tp_context(self.draft_model_runner.tp_group), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(), @@ -214,6 +216,7 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker): def init_cuda_graphs(self): with ( + draft_pp_context(), self.draft_tp_context(self.draft_model_runner.tp_group), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(), diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index b62755c53..08ce08d71 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -82,6 +82,7 @@ from sglang.srt.speculative.multi_layer_eagle_utils import ( ) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import ( + draft_pp_context, draft_tp_context, get_plan_stream, sample_draft_proposal, @@ -163,7 +164,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): # Load draft model weights only. with ( - empty_context(), + draft_pp_context(), speculative_moe_backend_context(), draft_model_build_scope(), ): @@ -221,6 +222,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): def init_attention_backends(self): with ( + draft_pp_context(), self.draft_tp_context(self.draft_runner_list[0].tp_group), speculative_moe_backend_context(), ): @@ -228,6 +230,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): def init_cuda_graphs(self): with ( + draft_pp_context(), self.draft_tp_context(self.draft_runner_list[0].tp_group), speculative_moe_backend_context(), ): diff --git a/python/sglang/srt/speculative/standalone_worker_v2.py b/python/sglang/srt/speculative/standalone_worker_v2.py index ce471794b..4cb7505d8 100644 --- a/python/sglang/srt/speculative/standalone_worker_v2.py +++ b/python/sglang/srt/speculative/standalone_worker_v2.py @@ -27,7 +27,11 @@ from sglang.srt.speculative.base_spec_worker import ( from sglang.srt.speculative.eagle_utils import default_tree_mask_mode from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker, EAGLEWorkerV2 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm -from sglang.srt.speculative.spec_utils import draft_tp_context, get_plan_stream +from sglang.srt.speculative.spec_utils import ( + draft_pp_context, + draft_tp_context, + get_plan_stream, +) from sglang.srt.utils import empty_context, get_bool_env_var logger = logging.getLogger(__name__) @@ -76,7 +80,7 @@ class StandaloneDraftWorker(EagleDraftWorker): # whose MoE gates run during construction; the scope routes their # fusion decision to the speculative leaf (it does not swap # runner_backend — the draft's forwards run outside that context). - with empty_context(), draft_model_build_scope(): + with draft_pp_context(), draft_model_build_scope(): self.draft_worker = TpModelWorker( server_args=server_args, gpu_id=gpu_id, diff --git a/python/sglang/srt/utils/profile_utils.py b/python/sglang/srt/utils/profile_utils.py index bd5f3d810..6db68374c 100644 --- a/python/sglang/srt/utils/profile_utils.py +++ b/python/sglang/srt/utils/profile_utils.py @@ -18,7 +18,7 @@ from sglang.srt.model_executor.step_span_utils import ( set_detailed_annotations_enabled, ) from sglang.srt.platforms import current_platform -from sglang.srt.runtime_context import get_device +from sglang.srt.runtime_context import get_device, get_parallel from sglang.srt.utils import is_npu from sglang.srt.utils.torch_npu_patch_utils import apply_torch_npu_patches @@ -358,15 +358,15 @@ class _ProfilerTorch(_ProfilerConcreteBase): self.torch_profiler.stop() if not _is_npu: # Build filename with only non-zero ranks to maintain backward compatibility - filename_parts = [self.profile_id, f"TP-{self.ps.tp_rank}"] + filename_parts = [self.profile_id, f"TP-{get_parallel().tp_rank}"] # Only add other ranks if parallelism is enabled (size > 1) if self.ps.dp_size > 1: - filename_parts.append(f"DP-{self.ps.dp_rank}") - if self.ps.pp_size > 1: - filename_parts.append(f"PP-{self.ps.pp_rank}") - if self.ps.moe_ep_size > 1: - filename_parts.append(f"EP-{self.ps.moe_ep_rank}") + filename_parts.append(f"DP-{get_parallel().dp_rank}") + if get_parallel().pp_size > 1: + filename_parts.append(f"PP-{get_parallel().pp_rank}") + if get_parallel().moe_ep_size > 1: + filename_parts.append(f"EP-{get_parallel().moe_ep_rank}") filename = ( (self.output_prefix + "-" if self.output_prefix else "") @@ -396,7 +396,7 @@ class _ProfilerMemory(_ProfilerConcreteBase): self.output_dir, (self.output_prefix + "-" if self.output_prefix else "") + str(time.time()) - + f"-TP-{self.ps.tp_rank}-memory" + + f"-TP-{get_parallel().tp_rank}-memory" + self.output_suffix + ".pickle", ) @@ -426,10 +426,13 @@ class _ProfilerRPD(_ProfilerConcreteBase): self.rpd_profile_path = os.path.join( self.output_dir, - "rpd-" + str(time.time()) + f"-TP-{self.ps.tp_rank}" + ".trace.json.gz", + "rpd-" + + str(time.time()) + + f"-TP-{get_parallel().tp_rank}" + + ".trace.json.gz", ) - if self.ps.tp_rank == 0: + if get_parallel().tp_rank == 0: import sqlite3 from rocpd.schema import RocpdSchema @@ -454,7 +457,7 @@ class _ProfilerRPD(_ProfilerConcreteBase): self.rpd_profiler.flush() torch.distributed.barrier(self.cpu_group) - if self.ps.tp_rank == 0: + if get_parallel().tp_rank == 0: from sglang.srt.utils.rpd_utils import rpd_to_chrome_trace rpd_to_chrome_trace("trace.rpd", self.rpd_profile_path) diff --git a/python/sglang/srt/weight_cache/daemon.py b/python/sglang/srt/weight_cache/daemon.py index 3bf39e959..c68af59e2 100644 --- a/python/sglang/srt/weight_cache/daemon.py +++ b/python/sglang/srt/weight_cache/daemon.py @@ -51,9 +51,11 @@ from sglang.srt.arg_groups.overrides import resolving_view from sglang.srt.configs.load_config import LoadConfig from sglang.srt.platforms import current_platform from sglang.srt.runtime_context import ( + SpawnRanks, get_exec, get_parallel, publish, + spawn_world_rank, ) from .protocol import ( @@ -281,7 +283,17 @@ class WeightCacheDaemon: from sglang.srt.model_loader.loader import get_model_loader server_args = self.server_args - publish(server_args, role="weight_cache_daemon") + # The launcher told this daemon where it sits, and it builds the same + # groups a scheduler does, so the same one number places it. + publish( + server_args, + role="weight_cache_daemon", + ranks=SpawnRanks( + world_rank=spawn_world_rank( + server_args, tp_rank=self.tp_rank, pp_rank=self.pp_rank + ) + ), + ) from sglang.srt.layers.moe import initialize_moe_config diff --git a/python/sglang/test/test_utils.py b/python/sglang/test/test_utils.py index eb8f820ce..38bdfb614 100644 --- a/python/sglang/test/test_utils.py +++ b/python/sglang/test/test_utils.py @@ -2,6 +2,7 @@ import argparse import asyncio +import contextlib import copy import doctest import importlib.util @@ -2002,6 +2003,36 @@ def maybe_stub_sgl_kernel(): sys.meta_path.insert(0, _SglKernelFinder()) +@contextlib.contextmanager +def published_topology(role: str = "test", *, ranks=None, **server_args_fields): + """Publish a record describing the parallel topology a test wants. + + Replaces standing a per-process parallel record into the object under + test. The widths arrive the way production gets them -- from published + configuration -- and the per-process ranks the way a spawned process gets + them, so a rank read is answered without building a process group. Stating + the topology through the same door production uses also keeps the derived + widths honest: a hand-built double can claim an `attn_tp_size` the + configuration would never produce. + + `ranks` overrides the spawn identities; by default this process is rank + zero of the world, which fixes every other rank. The context is reset on exit, including when the + test fails. + """ + from sglang.srt.runtime_context import SpawnRanks, publish, reset_context + from sglang.srt.server_args import ServerArgs + + bundle = dict(world_rank=0, dp_rank=None) + bundle.update(ranks or {}) + server_args = ServerArgs(model_path="dummy", **server_args_fields) + reset_context() + publish(server_args, role=role, ranks=SpawnRanks(**bundle)) + try: + yield server_args + finally: + reset_context() + + _GPU_IDLE_TIMEOUT_SECS = 30.0 _GPU_IDLE_POLL_INTERVAL_SECS = 2.0 _GPU_IDLE_USED_MEMORY_THRESHOLD = 2 << 30 # 2 GiB @@ -2253,6 +2284,17 @@ def enter_override(test_case, override): return installed +def enter_scope(test_case, scope): + """Enter a context manager for the length of one test. + + The `with`-statement form of `enter_override` above, and 3.10-safe for the + same reason: `enterContext` arrived in 3.11. + """ + entered = scope.__enter__() + test_case.addCleanup(scope.__exit__, None, None, None) + return entered + + class CustomTestCase(unittest.TestCase): def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) diff --git a/test/registered/unit/constrained/test_grammar_manager.py b/test/registered/unit/constrained/test_grammar_manager.py index 722307b17..55c824b55 100644 --- a/test/registered/unit/constrained/test_grammar_manager.py +++ b/test/registered/unit/constrained/test_grammar_manager.py @@ -26,13 +26,19 @@ from sglang.srt.constrained.base_grammar_backend import ( from sglang.srt.constrained.grammar_manager import GrammarManager from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject from sglang.srt.distributed.communication_tags import P2PTag -from sglang.srt.runtime_context import get_context, publish, reset_context +from sglang.srt.runtime_context import ( + SpawnRanks, + get_context, + get_parallel, + publish, + reset_context, +) from sglang.srt.sampling.sampling_params import ( REQUEST_REASONING_END_TOKEN_IDS_KEY, ) from sglang.srt.server_args import ServerArgs from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import enter_override +from sglang.test.test_utils import enter_override, enter_scope register_cpu_ci(2.0, "base-a-test-cpu") @@ -43,9 +49,10 @@ register_cpu_ci(est_time=5, suite="stage-b-test-cpu-intel") def _make_scheduler(grammar_backend_name="none", skip_tokenizer=False): """Create a mock scheduler with necessary attributes. - The grammar manager reads its config from the bags, so the settings that - used to be hung off the mock are published instead. The caller resets the - context; every test here goes through `_GrammarFixture`. + The grammar manager reads its config and its place in the pipeline from + the context, so the settings that used to be hung off the mock are + published instead. The caller resets the context; every test here goes + through `_GrammarFixture`. """ reset_context() server_args = ServerArgs( @@ -56,7 +63,11 @@ def _make_scheduler(grammar_backend_name="none", skip_tokenizer=False): constrained_json_whitespace_pattern=None, constrained_json_disable_any_whitespace=False, ) - publish(server_args, role="scheduler") + publish( + server_args, + role="scheduler", + ranks=SpawnRanks(world_rank=0), + ) scheduler = MagicMock() scheduler.server_args = server_args scheduler.model_config.request_selectable_think_end_id_sequences = None @@ -66,8 +77,6 @@ def _make_scheduler(grammar_backend_name="none", skip_tokenizer=False): scheduler.dp_tp_group.world_size = 1 scheduler.dp_tp_group.first_rank = 0 scheduler.dp_tp_group.is_first_rank = True - scheduler.ps.pp_rank = 0 - scheduler.ps.pp_size = 1 scheduler.pp_group = None return scheduler @@ -769,8 +778,10 @@ class TestGrammarManagerPPSync(unittest.TestCase): enter_override( self, get_context().override_server_args(skip_tokenizer_init=True) ) - scheduler.ps.pp_rank = pp_rank - scheduler.ps.pp_size = pp_size + # After that override, not before: installing a server-args override + # re-resolves the parallel bag from defaults, which puts `pp_size` + # back to 1 whatever was published. + enter_scope(self, get_parallel().override(pp_size=pp_size, pp_rank=pp_rank)) scheduler.pp_group = pp_group mgr = GrammarManager(scheduler) mgr.grammar_backend = MagicMock(spec=BaseGrammarBackend) diff --git a/test/registered/unit/managers/test_auxiliary_output.py b/test/registered/unit/managers/test_auxiliary_output.py index 6f7fe89bf..049d1343b 100644 --- a/test/registered/unit/managers/test_auxiliary_output.py +++ b/test/registered/unit/managers/test_auxiliary_output.py @@ -27,6 +27,7 @@ from sglang.srt.runtime_context import publish, reset_context from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import published_topology register_cpu_ci(est_time=14, suite="base-a-test-cpu") @@ -453,7 +454,6 @@ def test_active_observer_uses_observer_logits_preprocessing(): def test_scheduler_copies_auxiliary_output_for_non_overlap_results(): event = object() scheduler = object.__new__(Scheduler) - scheduler.ps = SimpleNamespace(pp_size=1) scheduler.device_module = SimpleNamespace(Event=Mock(return_value=event)) result = SimpleNamespace( logits_output=SimpleNamespace(auxiliary_device_output=object()), @@ -463,7 +463,8 @@ def test_scheduler_copies_auxiliary_output_for_non_overlap_results(): ) batch = SimpleNamespace(return_logprob=False, return_hidden_states=False) - Scheduler._copy_auxiliary_output_to_cpu(scheduler, batch, result) + with published_topology(): + Scheduler._copy_auxiliary_output_to_cpu(scheduler, batch, result) assert result.copy_done is event result.copy_to_cpu.assert_called_once_with( @@ -474,7 +475,6 @@ def test_scheduler_copies_auxiliary_output_for_non_overlap_results(): def test_scheduler_preserves_pipeline_parallel_output_for_transport(): scheduler = object.__new__(Scheduler) - scheduler.ps = SimpleNamespace(pp_size=2) scheduler.device_module = SimpleNamespace(Event=Mock()) result = SimpleNamespace( logits_output=SimpleNamespace(auxiliary_device_output=object()), @@ -484,7 +484,8 @@ def test_scheduler_preserves_pipeline_parallel_output_for_transport(): ) batch = SimpleNamespace(return_logprob=False, return_hidden_states=False) - Scheduler._copy_auxiliary_output_to_cpu(scheduler, batch, result) + with published_topology(pp_size=2): + Scheduler._copy_auxiliary_output_to_cpu(scheduler, batch, result) assert result.copy_done is None result.copy_to_cpu.assert_not_called() @@ -514,7 +515,6 @@ def test_pdmux_split_prefill_schedules_auxiliary_output_copy(): scheduler.is_generation = True scheduler.enable_overlap = False scheduler.enable_pdmux = True - scheduler.ps = SimpleNamespace(pp_size=1) scheduler.tp_worker = SimpleNamespace( forward_batch_split_prefill=Mock(return_value=result) ) @@ -535,9 +535,12 @@ def test_pdmux_split_prefill_schedules_auxiliary_output_copy(): return_hidden_states=False, ) - with patch( - "sglang.srt.managers.scheduler.resolve_forward_inputs" - ) as resolve_forward_inputs: + with ( + published_topology(), + patch( + "sglang.srt.managers.scheduler.resolve_forward_inputs" + ) as resolve_forward_inputs, + ): output_result = Scheduler.run_batch(scheduler, batch) resolve_forward_inputs.assert_called_once_with(batch, scheduler.future_map) diff --git a/test/registered/unit/managers/test_output_streamer_customized_info.py b/test/registered/unit/managers/test_output_streamer_customized_info.py index cd103f7f7..3f4094e86 100644 --- a/test/registered/unit/managers/test_output_streamer_customized_info.py +++ b/test/registered/unit/managers/test_output_streamer_customized_info.py @@ -14,6 +14,7 @@ from sglang.srt.utils.weight_versions import ( record_weight_version_events, ) from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import enter_scope, published_topology register_cpu_ci(est_time=10, suite="base-a-test-cpu") @@ -102,6 +103,8 @@ class TestOutputStreamerCustomizedInfo(unittest.TestCase): ) serving_patch.start() observability_patch.start() + # The streamer asks the context which rank it is streaming from. + enter_scope(self, published_topology(ranks={"dp_rank": 0})) self.addCleanup(serving_patch.stop) self.addCleanup(observability_patch.stop) diff --git a/test/registered/unit/managers/test_pp_cp_rank_offsets.py b/test/registered/unit/managers/test_pp_cp_rank_offsets.py index bea857174..ff6cfb8d1 100644 --- a/test/registered/unit/managers/test_pp_cp_rank_offsets.py +++ b/test/registered/unit/managers/test_pp_cp_rank_offsets.py @@ -3,7 +3,11 @@ from types import SimpleNamespace from unittest.mock import patch from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import maybe_stub_sgl_kernel +from sglang.test.test_utils import ( + enter_scope, + maybe_stub_sgl_kernel, + published_topology, +) maybe_stub_sgl_kernel() @@ -32,6 +36,25 @@ def _make_ps(**overrides) -> ParallelState: return ParallelState.trivial(**defaults) +def _published_topology(): + """The topology `_make_ps` describes, published instead of stood in. + + World rank 12 of a `tp=8, pp=2` world is `tp_rank=4` on the second stage, + which puts this process at `attn_dp_rank=1` with `attn_tp_rank=0`: the + context derives all of them from that one number and the widths, where the + record above had to be handed each. + """ + return published_topology( + role="scheduler", + ranks={"world_rank": 12, "dp_rank": 1}, + tp_size=8, + pp_size=2, + dp_size=2, + attn_cp_size=2, + enable_dp_attention=True, + ) + + def _fake_group() -> SimpleNamespace: return SimpleNamespace(rank=0, ranks=[0], cpu_group=object()) @@ -158,6 +181,7 @@ class TestRequestReceiverBroadcast(unittest.TestCase): class TestPPCPRankOffsets(unittest.TestCase): def test_request_receiver_uses_cp_size_for_pp_recv_rank(self): ps = _make_ps() + enter_scope(self, _published_topology()) calls = [] def fake_point_to_point_pyobj(data, rank, group, src, dst, **kwargs): @@ -176,6 +200,7 @@ class TestPPCPRankOffsets(unittest.TestCase): def test_pp_mixin_uses_cp_size_for_pyobj_send_and_recv_rank(self): ps = _make_ps() + enter_scope(self, _published_topology()) scheduler = SchedulerPPMixin() scheduler.ps = ps scheduler.world_group = _fake_group() diff --git a/test/registered/unit/managers/test_scheduler_chunked_abort_race.py b/test/registered/unit/managers/test_scheduler_chunked_abort_race.py index b66ae334e..1f4862e87 100644 --- a/test/registered/unit/managers/test_scheduler_chunked_abort_race.py +++ b/test/registered/unit/managers/test_scheduler_chunked_abort_race.py @@ -5,7 +5,12 @@ from types import SimpleNamespace from unittest.mock import Mock from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel +from sglang.test.test_utils import ( + CustomTestCase, + enter_scope, + maybe_stub_sgl_kernel, + published_topology, +) maybe_stub_sgl_kernel() @@ -38,13 +43,16 @@ def _make_scheduler(pending_req, *, chunked_req, running_reqs) -> Scheduler: sched.disaggregation_mode = None sched.enable_hicache_storage = False sched.mm_receiver = None - sched.ps = SimpleNamespace(pp_size=1) sched.running_batch = SimpleNamespace(reqs=running_reqs) sched.last_batch = None return sched class TestPendingChunkedAbortRace(CustomTestCase): + def setUp(self): + # The abort path asks the context for the pipeline width. + enter_scope(self, published_topology()) + def test_req_left_chunked_slot_is_aborted(self): req = _FakeReq("zombie_rid") sched = _make_scheduler(req, chunked_req=None, running_reqs=[req]) diff --git a/test/registered/unit/managers/test_scheduler_hicache_events.py b/test/registered/unit/managers/test_scheduler_hicache_events.py index 79b764182..8c81c0792 100644 --- a/test/registered/unit/managers/test_scheduler_hicache_events.py +++ b/test/registered/unit/managers/test_scheduler_hicache_events.py @@ -5,7 +5,11 @@ from types import SimpleNamespace from unittest.mock import Mock, call, patch from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import maybe_stub_sgl_kernel +from sglang.test.test_utils import ( + enter_scope, + maybe_stub_sgl_kernel, + published_topology, +) maybe_stub_sgl_kernel() @@ -114,7 +118,7 @@ class TestSchedulerHiCacheEvents(unittest.TestCase): s = self.scheduler s.init_pp_loop_state = Mock() s.pp_loop_size = 1 - s.ps = SimpleNamespace(pp_size=2) + enter_scope(self, published_topology(pp_size=2)) s.pp_group = SimpleNamespace(is_last_rank=True) s.running_mbs = [self.running_batch] s.last_mbs = [None] diff --git a/test/registered/unit/managers/test_scheduler_timeouts.py b/test/registered/unit/managers/test_scheduler_timeouts.py index 518b5a679..242d7bc73 100644 --- a/test/registered/unit/managers/test_scheduler_timeouts.py +++ b/test/registered/unit/managers/test_scheduler_timeouts.py @@ -12,7 +12,12 @@ from unittest.mock import MagicMock, patch from sglang.srt.environ import envs from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel +from sglang.test.test_utils import ( + CustomTestCase, + enter_scope, + maybe_stub_sgl_kernel, + published_topology, +) maybe_stub_sgl_kernel() @@ -62,7 +67,6 @@ def _scheduler(waiting_queue, running_reqs=(), last_batch_reqs=()): s.enable_unified_cache_external_linker = False s.ipc_channels = SimpleNamespace(send_to_tokenizer=MagicMock()) s.beam_coordinator = MagicMock() - s.ps = SimpleNamespace(pp_size=1) s.running_batch = _batch(list(running_reqs)) s.last_batch = _batch(list(last_batch_reqs)) if last_batch_reqs else None return s @@ -126,6 +130,10 @@ class TestWaitingTimeout(CustomTestCase): class TestRunningTimeout(CustomTestCase): + def setUp(self): + # The poll asks the context for the pipeline width. + enter_scope(self, published_topology()) + def test_emits_only_stale_unfinished_reqs_without_marking(self): now = time.perf_counter() stale = _req("stale", forward_entry=now - 10) diff --git a/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py b/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py index 06a017cca..b10cf53ae 100644 --- a/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py +++ b/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py @@ -736,12 +736,15 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase): # publishing a record rather than by standing one in. reset_context() publish( - ServerArgs(model_path="dummy", startup_weight_load_mode=mode), + ServerArgs( + model_path="dummy", + startup_weight_load_mode=mode, + pp_size=pp_size, + ), role="scheduler", ) scheduler = Scheduler.__new__(Scheduler) scheduler.enable_overlap = enable_overlap - scheduler.ps = SimpleNamespace(pp_size=pp_size) scheduler.init_tp_model_worker = lambda: setattr(scheduler, "tp_worker", worker) scheduler.maybe_init_draft_worker = lambda: setattr( scheduler, "draft_worker", draft_worker diff --git a/test/registered/unit/multimodal/rust/shared/test_rust_server_extension.py b/test/registered/unit/multimodal/rust/shared/test_rust_server_extension.py index 96ad091f7..7629e7239 100644 --- a/test/registered/unit/multimodal/rust/shared/test_rust_server_extension.py +++ b/test/registered/unit/multimodal/rust/shared/test_rust_server_extension.py @@ -34,7 +34,6 @@ class TestRustServerExtension(CustomTestCase): attn_dp_rank=1, tp_size=2, tp_rank=1, - pp_size=1, attn_tp_size=1, attn_cp_size=1, ), @@ -49,7 +48,9 @@ class TestRustServerExtension(CustomTestCase): ), ), patch.object( - server_module, "get_parallel", return_value=SimpleNamespace(nnodes=1) + server_module, + "get_parallel", + return_value=SimpleNamespace(nnodes=1, pp_size=1), ), patch.object(ModelServer, "_partition_cores", return_value=(None, None)), patch.object( diff --git a/test/registered/unit/observability/test_forward_pass_metrics.py b/test/registered/unit/observability/test_forward_pass_metrics.py index 808a93c94..a19925d93 100644 --- a/test/registered/unit/observability/test_forward_pass_metrics.py +++ b/test/registered/unit/observability/test_forward_pass_metrics.py @@ -1,4 +1,4 @@ -from sglang.srt.runtime_context import get_context, get_observability +from sglang.srt.runtime_context import get_context, get_observability, get_parallel from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=11, suite="base-a-test-cpu") @@ -15,7 +15,7 @@ from sglang.srt.managers.scheduler_components.metrics_reporter import ( SchedulerMetricsReporter, _CacheHitRateWindow, ) -from sglang.test.test_utils import CustomTestCase +from sglang.test.test_utils import CustomTestCase, enter_scope def _make_ps(**overrides) -> ParallelState: @@ -292,6 +292,8 @@ class TestForwardPassMetrics(unittest.TestCase): kv_events_config=None, ) scheduler.ps = _make_ps(attn_tp_rank=0, dp_rank=2, pp_rank=0, pp_size=1) + # The reporter asks the context whether this is the last stage. + enter_scope(self, get_parallel().override(pp_rank=0, pp_size=1)) scheduler.enable_kv_cache_events = False with patch( @@ -329,6 +331,8 @@ class TestForwardPassMetrics(unittest.TestCase): kv_events_config=None, ) scheduler.ps = _make_ps(attn_tp_rank=0, dp_rank=0, pp_rank=0, pp_size=2) + # The reporter asks the context whether this is the last stage. + enter_scope(self, get_parallel().override(pp_rank=0, pp_size=2)) scheduler.enable_kv_cache_events = False with patch( diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 1c8485908..39675a8a7 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -39,6 +39,7 @@ from sglang.srt.runtime_context import ( Flags, ParallelContext, RuntimeContext, + SpawnRanks, _FlagGroupBase, assert_published, derive_parallel_widths, @@ -200,6 +201,148 @@ class TestTheTwoWorldWidths(_IsolatedOverrides): self.assertEqual(parallel.launch_world_size, 2) +class TestSpawnIdentities(_IsolatedOverrides): + """`dp_rank` and `gpu_id` come from the spawn, because nothing else has them. + + Both vary per process while the record is identical across them, and + neither is a position in any process group -- no group has one member per + data-parallel replica. So the process entry states them at publish. + """ + + def setUp(self): + super().setUp() + parallel = get_parallel() + self._saved_stamp = dict(parallel._stamp) + self.addCleanup( + lambda: ( + parallel.clear_stamp(), + parallel.override_permanently(**self._saved_stamp), + ) + ) + reset_context() + self.addCleanup(reset_context) + + def test_one_rank_fixes_the_rest(self): + """Every other rank is a position in a group laid out from the widths, + so `world_rank` is the whole placement: rank 5 of a `tp=4, pp=2` world + is the second stage's second device.""" + publish( + ServerArgs(model_path="dummy", tp_size=4, pp_size=2), + role="test", + ranks=SpawnRanks(world_rank=5, dp_rank=2), + ) + parallel = get_parallel() + self.assertEqual(parallel.launch_world_rank, 5) + self.assertEqual(parallel.tp_rank, 1) + self.assertEqual(parallel.pp_rank, 1) + self.assertEqual(parallel.dp_rank, 2) + + def test_no_controller_is_an_answer_not_a_failure(self): + """`dp_rank=None` means "not under a data parallel controller", which + is a fact about the deployment, unlike never having been told. The + replicas are separate WORLD groups, so no rank implies it.""" + publish( + ServerArgs(model_path="dummy", tp_size=2), + role="test", + ranks=SpawnRanks(world_rank=0, dp_rank=None), + ) + self.assertIsNone(get_parallel().dp_rank) + + def test_publishing_without_a_bundle_names_what_is_missing(self): + publish(ServerArgs(model_path="dummy", tp_size=2), role="test") + with self.assertRaises(RuntimeError) as caught: + get_parallel().dp_rank + self.assertIn("rank bundle", str(caught.exception)) + + def test_the_attention_rank_keeps_its_own_explanation(self): + """Two stamp-only names, two different reasons to be missing.""" + publish(ServerArgs(model_path="dummy", tp_size=2), role="test") + with self.assertRaises(RuntimeError) as caught: + get_parallel().attn_dp_rank + self.assertIn("initialize_dp_attention", str(caught.exception)) + + +class TestAttentionRanksComeFromPublish(_IsolatedOverrides): + """With a spawn bundle, a rank read works before any group exists. + + This is what `ParallelState` provided by being a plain frozen record, and + what the topology init could not: it needs the groups. Deriving at publish + is what lets a reader ask the context in a process that never initialises + distributed -- every unit test that builds a scheduler component, for one. + """ + + def setUp(self): + super().setUp() + parallel = get_parallel() + self._saved_stamp = dict(parallel._stamp) + self.addCleanup( + lambda: ( + parallel.clear_stamp(), + parallel.override_permanently(**self._saved_stamp), + ) + ) + reset_context() + self.addCleanup(reset_context) + + def test_it_matches_the_topology_init_for_every_shape(self): + """Cross-checked against the function the groups use, not restated. + + Same inputs, two callers: one has them from the configuration and the + spawn, the other from the groups it just built. + """ + from sglang.srt.layers.dp_attention import compute_dp_attention_world_info + + shapes = [ + (8, 1, 1, False), + (8, 2, 1, True), + (8, 4, 1, True), + (8, 2, 2, True), + (16, 4, 2, True), + ] + for tp_size, dp_size, attn_cp_size, dp_attn in shapes: + for tp_rank in range(tp_size): + reset_context() + publish( + ServerArgs( + model_path="dummy", + tp_size=tp_size, + dp_size=dp_size, + attn_cp_size=attn_cp_size, + enable_dp_attention=dp_attn, + ), + role="test", + ranks=SpawnRanks(world_rank=tp_rank), + ) + want_tp, _, want_dp, _ = compute_dp_attention_world_info( + dp_attn, tp_rank, tp_size, dp_size, attn_cp_size + ) + msg = f"tp={tp_size} dp={dp_size} cp={attn_cp_size} rank={tp_rank}" + self.assertEqual(get_parallel().attn_tp_rank, want_tp, msg) + self.assertEqual(get_parallel().attn_dp_rank, want_dp, msg) + + def test_the_rank_reads_without_a_process_group(self): + """No distributed init, no patching of any getter.""" + publish( + ServerArgs( + model_path="dummy", tp_size=8, dp_size=2, enable_dp_attention=True + ), + role="test", + ranks=SpawnRanks(world_rank=5), + ) + with patch( + f"{_PS}.get_attn_tensor_model_parallel_rank", + side_effect=AssertionError("no group must be consulted"), + ): + self.assertEqual(get_parallel().attn_tp_rank, 1) + self.assertEqual(get_parallel().attn_dp_rank, 1) + + def test_without_a_bundle_it_still_asks_the_group(self): + """Unchanged for every process that publishes without a placement.""" + publish(ServerArgs(model_path="dummy", tp_size=8), role="test") + with patch(f"{_PS}.get_attn_tensor_model_parallel_rank", return_value=3): + self.assertEqual(get_parallel().attn_tp_rank, 3) + + class TestStampedRanks(_IsolatedOverrides): """`attn_dp_rank` comes from the stamp, and says so when there is none. @@ -2107,5 +2250,108 @@ class TestTheDerivedHalfIsDeclared(CustomTestCase): self.assertNotIn(name, fields) +class TestAnEntryThatBuildsARunnerHandsOverItsPlacement(CustomTestCase): + """`ModelRunner.__init__` reads a recorded identity, so an entry that + publishes without a bundle and then builds one fails at construction. + + Every such entry is an `__main__`-reachable path, so nothing in the unit + suite exercises it; the benchmark entry was found this way rather than by + a test. This walks the sources instead: a module that publishes and builds + a runner has to pass `ranks=`. + """ + + def test_every_publisher_that_builds_a_runner_passes_a_bundle(self): + import ast as _ast + + root = _pathlib.Path(next(iter(_sglang.__path__))).resolve() + offenders = [] + for path in root.rglob("*.py"): + text = path.read_text(encoding="utf-8-sig") + if "ModelRunner(" not in text or "publish(" not in text: + continue + tree = _ast.parse(text) + builds = any( + isinstance(n, _ast.Call) + and getattr(n.func, "id", getattr(n.func, "attr", None)) + == "ModelRunner" + for n in _ast.walk(tree) + ) + if not builds: + continue + for node in _ast.walk(tree): + if ( + isinstance(node, _ast.Call) + and getattr(node.func, "id", None) == "publish" + and not any(kw.arg == "ranks" for kw in node.keywords) + ): + offenders.append(f"{path.relative_to(root)}:{node.lineno}") + self.assertEqual( + offenders, + [], + "these publish without a spawn bundle and then build a ModelRunner, " + "whose construction reads a recorded identity:\n " + + "\n ".join(offenders), + ) + + +class TestWhoAnswersDuringADraftScope(CustomTestCase): + """A draft worker runs in one process with the target, under a scope. + + Two things have to hold for that to be workable, and neither is visible + from a single read: inside the scope every source agrees on the draft's + shape, and a reader that runs *outside* it still gets the draft's answer + from whatever it carried out. + """ + + def _single_member_group(self): + from sglang.srt.distributed.parallel_state import GroupCoordinator + + group = GroupCoordinator.__new__(GroupCoordinator) + group.world_size = 1 + group.rank_in_group = 0 + return group + + def _two_stage_pipeline(self): + """This process is stage 1 of 2, published the way a spawn states it.""" + reset_context() + self.addCleanup(reset_context) + publish( + ServerArgs(model_path="dummy", pp_size=2), + role="scheduler", + ranks=SpawnRanks(world_rank=1), + ) + + def test_the_pipeline_swap_states_every_member_it_installs(self): + """`pp_size` is a configured leaf: unlike `pp_rank` it does not follow + the group being swapped underneath, so a scope that installs a group + without stating its width reports the target's.""" + from sglang.srt.distributed import parallel_state + + group = self._single_member_group() + self._two_stage_pipeline() + self.assertEqual(get_parallel().pp_size, 2) + with patch.object(parallel_state, "_PP", group): + with parallel_state.patch_pipeline_parallel_group(group): + self.assertEqual(get_parallel().pp_size, 1) + self.assertEqual(get_parallel().pp_rank, 0) + self.assertIs(get_parallel().pp_group, group) + self.assertEqual(get_parallel().pp_size, 2) + self.assertEqual(get_parallel().pp_rank, 1) + + def test_a_report_built_for_a_runner_follows_that_runner(self): + """A weight check is an on-demand request served from the scheduler + loop, so it runs outside the scope that describes a draft runner. Its + report has to name the runner it was built for, which is why it holds + a record instead of asking the context.""" + from sglang.srt.distributed.parallel_state_wrapper import ParallelState + from sglang.srt.utils.weight_checker import WeightChecker + + draft = ParallelState.trivial(pp_rank=0, pp_size=1) + checker = WeightChecker(get_model=lambda: None, ps=draft) + self._two_stage_pipeline() + info = checker._parallelism_info() + self.assertEqual((info.pp_rank, info.pp_size), (0, 1)) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/utils/test_weight_versions.py b/test/registered/unit/utils/test_weight_versions.py index 8a1146b7b..2e2b30b11 100644 --- a/test/registered/unit/utils/test_weight_versions.py +++ b/test/registered/unit/utils/test_weight_versions.py @@ -15,7 +15,11 @@ from sglang.srt.utils.weight_versions import ( truncate_weight_version_events, ) from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase +from sglang.test.test_utils import ( + CustomTestCase, + enter_scope, + published_topology, +) register_cpu_ci(est_time=12, suite="base-a-test-cpu") @@ -300,11 +304,9 @@ class _SchedulerStub: waiting, chunked=None, last_batch=None, - pp_size=1, hisparse=None, ): self.serving = _ServingStub(version) - self.ps = SimpleNamespace(pp_size=pp_size) self.running_batch = SimpleNamespace(reqs=running) self.last_batch = last_batch self.waiting_queue = waiting @@ -313,7 +315,9 @@ class _SchedulerStub: class TestSchedulerRecordWeightVersionChange(CustomTestCase): - def _scheduler(self, *args, **kwargs): + def _scheduler(self, *args, pp_size=1, **kwargs): + # The recording path asks the context for the pipeline width. + enter_scope(self, published_topology(pp_size=pp_size)) scheduler = _SchedulerStub(*args, **kwargs) for name, value in ( ("get_serving", scheduler.serving),