Record a process's placement at publish, not at group build (#40071)

This commit is contained in:
Cheng Wan
2026-09-18 23:52:37 -07:00
committed by GitHub
parent 36aa8479ef
commit 3a5f52e144
47 changed files with 917 additions and 231 deletions
+18 -2
View File
@@ -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.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import ( from sglang.srt.runtime_context import (
SpawnRanks,
get_model, get_model,
get_parallel, get_parallel,
get_schedule, get_schedule,
publish, publish,
spawn_world_rank,
) )
from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.sampling.sampling_params import SamplingParams
from sglang.srt.server_args import PortArgs, ServerArgs from sglang.srt.server_args import PortArgs, ServerArgs
@@ -707,7 +709,15 @@ def correctness_test(
gpu_id, gpu_id,
tp_rank, 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 the logger
configure_logger(server_args, prefix=f" TP{tp_rank}") configure_logger(server_args, prefix=f" TP{tp_rank}")
@@ -912,7 +922,13 @@ def latency_test(
cfg = resolving_view(server_args) cfg = resolving_view(server_args)
# `main` runs this inline for tp_size == 1 and spawns it per rank otherwise; # `main` runs this inline for tp_size == 1 and spawns it per rank otherwise;
# a spawned child arrives with nothing published. # 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_moe_config()
initialize_fp8_gemm_config() initialize_fp8_gemm_config()
initialize_fp4_gemm_config() initialize_fp4_gemm_config()
@@ -14,7 +14,7 @@ from sglang.srt.constrained.base_grammar_backend import (
from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject
from sglang.srt.distributed.communication_tags import P2PTag from sglang.srt.distributed.communication_tags import P2PTag
from sglang.srt.environ import envs 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 ( from sglang.srt.sampling.sampling_params import (
get_request_reasoning_end_token_ids, 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_size = scheduler.dp_tp_group.world_size
self.grammar_sync_entry = scheduler.dp_tp_group.first_rank self.grammar_sync_entry = scheduler.dp_tp_group.first_rank
self.is_grammar_sync_entry = scheduler.dp_tp_group.is_first_rank self.is_grammar_sync_entry = scheduler.dp_tp_group.is_first_rank
self.pp_rank = scheduler.ps.pp_rank self.pp_rank = get_parallel().pp_rank
self.pp_size = scheduler.ps.pp_size self.pp_size = get_parallel().pp_size
self.pp_group = scheduler.pp_group self.pp_group = scheduler.pp_group
self.grammar_pp_sync_work_list = [] self.grammar_pp_sync_work_list = []
+1 -1
View File
@@ -393,7 +393,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
self.bootstrap_port = bootstrap_port self.bootstrap_port = bootstrap_port
self.max_total_num_tokens = max_total_num_tokens self.max_total_num_tokens = max_total_num_tokens
self.pp_rank = pp_rank 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.num_reserved_decode_tokens = num_reserved_decode_tokens
self.transfer_backend = transfer_backend self.transfer_backend = transfer_backend
# Queue for requests pending pre-allocation # Queue for requests pending pre-allocation
+3 -3
View File
@@ -997,7 +997,7 @@ class SchedulerDisaggregationPrefillMixin:
KVPoll.Failed, KVPoll.Failed,
): ):
logger.warning_once( 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", f"from consensus; treating as undone",
) )
undone_reqs.append(req) undone_reqs.append(req)
@@ -1077,7 +1077,7 @@ class SchedulerDisaggregationPrefillMixin:
) -> Optional[Exception]: ) -> Optional[Exception]:
"""Conclude an inflight request whose KV transfer failed.""" """Conclude an inflight request whose KV transfer failed."""
error_message = ( 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=}" f"{req.rid=} {req.bootstrap_room=}"
) )
exc: Optional[Exception] = None exc: Optional[Exception] = None
@@ -1142,7 +1142,7 @@ class SchedulerDisaggregationPrefillMixin:
def handle_bootstrap_failure(self: Scheduler, req: Req) -> None: def handle_bootstrap_failure(self: Scheduler, req: Req) -> None:
self.clear_pending_chunk_send(req) self.clear_pending_chunk_send(req)
error_message = ( 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=}" f"{req.rid=} {req.bootstrap_room=}"
) )
is_propagated = False is_propagated = False
@@ -99,6 +99,7 @@ def init_torch_distributed(
) )
# Only initialize the distributed environment on the target model worker. # 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( _init_parallel_groups(
backend=backend, backend=backend,
dist_init_method=dist_init_method, dist_init_method=dist_init_method,
@@ -3029,6 +3029,13 @@ def patch_pipeline_parallel_group(pp_group: GroupCoordinator):
global _PP global _PP
_PP = pp_group _PP = pp_group
try: try:
# `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 yield
finally: finally:
_PP_STATE_PATCHED = False _PP_STATE_PATCHED = False
@@ -509,7 +509,7 @@ def pp_parallel_deep_gemm_warmup(runner) -> None:
logger.info( logger.info(
"PP-parallel DeepGEMM warmup start " "PP-parallel DeepGEMM warmup start "
"(pp_rank=%d, tp_rank=%d, batch_sizes=%s, disagg=%s).", "(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, model_runner.ps.tp_rank,
batch_sizes, batch_sizes,
disagg_mode, disagg_mode,
@@ -538,5 +538,5 @@ def pp_parallel_deep_gemm_warmup(runner) -> None:
logger.info( logger.info(
"PP-parallel DeepGEMM warmup done in %.2fs (pp_rank=%d).", "PP-parallel DeepGEMM warmup done in %.2fs (pp_rank=%d).",
time.perf_counter() - t0, time.perf_counter() - t0,
model_runner.ps.pp_rank, get_parallel().pp_rank,
) )
+7 -9
View File
@@ -27,6 +27,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
) )
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.runtime_context import ( from sglang.srt.runtime_context import (
derive_attention_ranks,
derive_attention_widths, derive_attention_widths,
get_device, get_device,
get_exec, get_exec,
@@ -349,15 +350,12 @@ def compute_dp_attention_world_info(
dp_size=dp_size, dp_size=dp_size,
enable_dp_attention=enable_dp_attention, enable_dp_attention=enable_dp_attention,
) )
attn_tp_rank = tp_rank % attn_tp_size attn_tp_rank, attn_dp_rank = derive_attention_ranks(
tp_rank=tp_rank,
if not enable_dp_attention: attn_tp_size=attn_tp_size,
attn_dp_rank = 0 attn_cp_size=attn_cp_size,
else: enable_dp_attention=enable_dp_attention,
# 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)
return attn_tp_rank, attn_tp_size, attn_dp_rank, attn_dp_size return attn_tp_rank, attn_tp_size, attn_dp_rank, attn_dp_size
+98 -67
View File
@@ -29,6 +29,7 @@ from http import HTTPStatus
from typing import TYPE_CHECKING, Any, Deque, Dict, List, Optional, Set, Tuple, Union from typing import TYPE_CHECKING, Any, Deque, Dict, List, Optional, Set, Tuple, Union
from sglang.srt.runtime_context import ( from sglang.srt.runtime_context import (
SpawnRanks,
attention_backends, attention_backends,
get_context, get_context,
get_device, get_device,
@@ -44,6 +45,7 @@ from sglang.srt.runtime_context import (
get_serving, get_serving,
get_spec, get_spec,
publish, publish,
spawn_world_rank,
) )
from sglang.srt.utils.common import suppress_noisy_warnings # isort: skip from sglang.srt.utils.common import suppress_noisy_warnings # isort: skip
@@ -595,9 +597,9 @@ class Scheduler(
enable_metrics=get_observability().enable_metrics, enable_metrics=get_observability().enable_metrics,
enable_kv_cache_events=bool( enable_kv_cache_events=bool(
get_observability().kv_events_config get_observability().kv_events_config
and self.ps.pp_rank == 0 and get_parallel().pp_rank == 0
and self.ps.attn_tp_rank == 0 and get_parallel().attn_tp_rank == 0
and self.ps.attn_cp_rank == 0 and get_parallel().attn_cp_rank == 0
), ),
ps=self.ps, ps=self.ps,
tp_group=self.tp_group, tp_group=self.tp_group,
@@ -759,8 +761,8 @@ class Scheduler(
return return
rank = ( rank = (
self.ps.dp_rank get_parallel().dp_rank
if self.ps.dp_rank is not None if get_parallel().dp_rank is not None
else self.tp_group.rank_in_group else self.tp_group.rank_in_group
) )
logger.info("HCCL DP prewarm start: rank=%s", rank) logger.info("HCCL DP prewarm start: rank=%s", rank)
@@ -775,10 +777,10 @@ class Scheduler(
if _is_npu: if _is_npu:
from sglang.srt.hardware_backend.npu.utils import init_zbal 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!") logger.error("only zbal mix mode support pp_size > 1!")
init_zbal( 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 ) # only switch allocator if is mix mode
def init_model_config(self): def init_model_config(self):
@@ -811,9 +813,9 @@ class Scheduler(
def init_ipc_channels(self, port_args: PortArgs): def init_ipc_channels(self, port_args: PortArgs):
is_rank_zero = ( is_rank_zero = (
self.ps.pp_rank == 0 get_parallel().pp_rank == 0
and self.ps.attn_tp_rank == 0 and get_parallel().attn_tp_rank == 0
and self.ps.attn_cp_rank == 0 and get_parallel().attn_cp_rank == 0
) )
self.ipc_channels = SchedulerIpcChannels.create( self.ipc_channels = SchedulerIpcChannels.create(
port_args=port_args, port_args=port_args,
@@ -824,7 +826,7 @@ class Scheduler(
skip_tokenizer_init=self.skip_tokenizer_init, skip_tokenizer_init=self.skip_tokenizer_init,
metrics_enabled=get_observability().enable_metrics metrics_enabled=get_observability().enable_metrics
and ( and (
self.ps.attn_tp_rank == 0 get_parallel().attn_tp_rank == 0
or get_observability().enable_metrics_for_all_schedulers or get_observability().enable_metrics_for_all_schedulers
), ),
enable_scripted_runtime=envs.SGLANG_TEST_SCRIPTED_RUNTIME.get(), enable_scripted_runtime=envs.SGLANG_TEST_SCRIPTED_RUNTIME.get(),
@@ -837,7 +839,7 @@ class Scheduler(
return return
self.recv_from_tokenizer = self.ipc_channels.recv_from_tokenizer 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: try:
self.load_snapshot_writer = create_load_snapshot_writer( self.load_snapshot_writer = create_load_snapshot_writer(
port_args, port_args,
@@ -850,9 +852,9 @@ class Scheduler(
def init_idle_sleeper(self) -> None: def init_idle_sleeper(self) -> None:
if ( if (
self.ps.pp_rank == 0 get_parallel().pp_rank == 0
and self.ps.attn_tp_rank == 0 and get_parallel().attn_tp_rank == 0
and self.ps.attn_cp_rank == 0 and get_parallel().attn_cp_rank == 0
and get_device().sleep_on_idle and get_device().sleep_on_idle
): ):
self.idle_sleeper = IdleSleeper( self.idle_sleeper = IdleSleeper(
@@ -1032,8 +1034,8 @@ class Scheduler(
if ( if (
envs.SGLANG_ENABLE_PP_SPEC.get() envs.SGLANG_ENABLE_PP_SPEC.get()
and self.ps.pp_size > 1 and get_parallel().pp_size > 1
and self.ps.pp_rank != self.ps.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 # 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. # 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. # Match run_batch / _pp_launch_batch so warmup allocations stay reusable.
forward_stream = ( forward_stream = (
model_runner.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 else self.schedule_stream
) )
with device_module.stream(forward_stream): with device_module.stream(forward_stream):
@@ -1208,7 +1210,7 @@ class Scheduler(
get_context().override( get_context().override(
"scheduler.pp_max_micro_batch_size_default", "scheduler.pp_max_micro_batch_size_default",
pp_max_micro_batch_size=max( 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.startup_available_gpu_memory_gb = get_available_gpu_memory(
self.device, self.ps.gpu_id, empty_cache=False self.device, self.ps.gpu_id, empty_cache=False
) )
if self.ps.tp_rank == 0: if get_parallel().tp_rank == 0:
logger.info( logger.info(
f"max_total_num_tokens={self.max_total_num_tokens}, " f"max_total_num_tokens={self.max_total_num_tokens}, "
f"chunked_prefill_size={get_schedule().chunked_prefill_size}, " f"chunked_prefill_size={get_schedule().chunked_prefill_size}, "
@@ -1344,7 +1346,7 @@ class Scheduler(
def maybe_init_dynamic_chunk_sizer(self) -> None: def maybe_init_dynamic_chunk_sizer(self) -> None:
"""Profile a PP prefill latency model that sizes chunks per stage.""" """Profile a PP prefill latency model that sizes chunks per stage."""
self.dynamic_chunk_sizer: Optional[DynamicChunkSizer] = None 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 return
sizer = DynamicChunkSizer( sizer = DynamicChunkSizer(
model_runner=self.tp_worker.model_runner, model_runner=self.tp_worker.model_runner,
@@ -1359,7 +1361,7 @@ class Scheduler(
device=self.device, device=self.device,
pp_group=self.pp_group, pp_group=self.pp_group,
world_group=self.world_group, world_group=self.world_group,
pp_rank=self.ps.pp_rank, pp_rank=get_parallel().pp_rank,
) )
if sizer.profile_and_fit(): if sizer.profile_and_fit():
self.dynamic_chunk_sizer = sizer self.dynamic_chunk_sizer = sizer
@@ -1422,7 +1424,7 @@ class Scheduler(
else: else:
self.prefill_delayer = PrefillDelayer( self.prefill_delayer = PrefillDelayer(
dp_size=self.ps.dp_size, 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, cpu_group=self.tp_cpu_group,
device_group=self.tp_group.device_group, device_group=self.tp_group.device_group,
metrics_collector=( metrics_collector=(
@@ -1433,7 +1435,7 @@ class Scheduler(
max_delay_passes=get_schedule().prefill_delayer_max_delay_passes, max_delay_passes=get_schedule().prefill_delayer_max_delay_passes,
token_usage_low_watermark=get_schedule().prefill_delayer_token_usage_low_watermark, token_usage_low_watermark=get_schedule().prefill_delayer_token_usage_low_watermark,
device=self.tp_group.device, 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. # NOTE: preemption is enabled by default for priority scheduling.
@@ -1464,7 +1466,7 @@ class Scheduler(
# Init recv skipper and input blocker # Init recv skipper and input blocker
self.recv_skipper = SchedulerRecvSkipper.maybe_create() self.recv_skipper = SchedulerRecvSkipper.maybe_create()
self.input_blocker = ( 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") if get_bool_env_var("SGLANG_ENABLE_COLOCATED_BATCH_GEN")
else None else None
) )
@@ -1553,7 +1555,7 @@ class Scheduler(
self.disagg_decode_transfer_queue = DecodeTransferQueue( self.disagg_decode_transfer_queue = DecodeTransferQueue(
gloo_group=self.attn_tp_cpu_group, gloo_group=self.attn_tp_cpu_group,
req_to_metadata_buffer_idx_allocator=self.req_to_metadata_buffer_idx_allocator, 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, metadata_buffers=self.disagg_metadata_buffers,
scheduler=self, scheduler=self,
tree_cache=self.tree_cache, tree_cache=self.tree_cache,
@@ -1570,13 +1572,13 @@ class Scheduler(
transfer_queue=self.disagg_decode_transfer_queue, transfer_queue=self.disagg_decode_transfer_queue,
tree_cache=self.tree_cache, tree_cache=self.tree_cache,
gloo_group=self.attn_tp_cpu_group, gloo_group=self.attn_tp_cpu_group,
tp_rank=self.ps.tp_rank, tp_rank=get_parallel().tp_rank,
tp_size=self.ps.tp_size, tp_size=get_parallel().tp_size,
dp_size=get_parallel().dp_size, dp_size=get_parallel().dp_size,
gpu_id=self.ps.gpu_id, gpu_id=self.ps.gpu_id,
bootstrap_port=get_disagg().disaggregation_bootstrap_port, bootstrap_port=get_disagg().disaggregation_bootstrap_port,
max_total_num_tokens=self.max_total_num_tokens, 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, num_reserved_decode_tokens=get_disagg().num_reserved_decode_tokens,
transfer_backend=self.transfer_backend, transfer_backend=self.transfer_backend,
) )
@@ -1602,16 +1604,16 @@ class Scheduler(
draft_token_to_kv_pool=draft_token_to_kv_pool, draft_token_to_kv_pool=draft_token_to_kv_pool,
req_to_metadata_buffer_idx_allocator=self.req_to_metadata_buffer_idx_allocator, req_to_metadata_buffer_idx_allocator=self.req_to_metadata_buffer_idx_allocator,
metadata_buffers=self.disagg_metadata_buffers, metadata_buffers=self.disagg_metadata_buffers,
tp_rank=self.ps.tp_rank, tp_rank=get_parallel().tp_rank,
tp_size=self.ps.tp_size, tp_size=get_parallel().tp_size,
gpu_id=self.ps.gpu_id, gpu_id=self.ps.gpu_id,
bootstrap_port=get_disagg().disaggregation_bootstrap_port, bootstrap_port=get_disagg().disaggregation_bootstrap_port,
gloo_group=self.attn_tp_cpu_group, gloo_group=self.attn_tp_cpu_group,
max_total_num_tokens=self.max_total_num_tokens, max_total_num_tokens=self.max_total_num_tokens,
scheduler=self, scheduler=self,
scheduler_stage_metrics=self.scheduler_stage_metrics, scheduler_stage_metrics=self.scheduler_stage_metrics,
pp_rank=self.ps.pp_rank, pp_rank=get_parallel().pp_rank,
pp_size=self.ps.pp_size, pp_size=get_parallel().pp_size,
transfer_backend=self.transfer_backend, transfer_backend=self.transfer_backend,
) )
# The prefill requests that are in the middle of kv sending # The prefill requests that are in the middle of kv sending
@@ -1638,8 +1640,8 @@ class Scheduler(
self.server_args, self.server_args,
dtype=self.model_config.dtype, dtype=self.model_config.dtype,
hf_config=self.model_config.hf_config, hf_config=self.model_config.hf_config,
pp_rank=self.ps.pp_rank, pp_rank=get_parallel().pp_rank,
tp_rank=self.ps.tp_rank, tp_rank=get_parallel().tp_rank,
tp_group=self.tp_group, tp_group=self.tp_group,
scheduler=self, scheduler=self,
) )
@@ -1895,7 +1897,9 @@ class Scheduler(
if self.device == "cpu": if self.device == "cpu":
self.schedule_stream.synchronize = lambda: None # No-op for 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 # CUDA/HIP streams come from a fixed round-robin pool. Redraw if this
# stream aliases forward_stream, which would eliminate scheduler # stream aliases forward_stream, which would eliminate scheduler
# overlap. Only CUDA/HIP streams expose a ``cuda_stream`` handle; # overlap. Only CUDA/HIP streams expose a ``cuda_stream`` handle;
@@ -2090,9 +2094,9 @@ class Scheduler(
""" """
local_reqs = [] local_reqs = []
if ( if (
self.ps.pp_rank == 0 get_parallel().pp_rank == 0
and self.ps.attn_tp_rank == 0 and get_parallel().attn_tp_rank == 0
and self.ps.attn_cp_rank == 0 and get_parallel().attn_cp_rank == 0
): ):
local_reqs = self._poll_timeout_aborts() local_reqs = self._poll_timeout_aborts()
recv_reqs = self.request_receiver.recv_requests(local_reqs=local_reqs) 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`` and with it the server-process duties a Python ``TokenizerManager``
would otherwise own (e.g. serving the PD KV bootstrap registry).""" would otherwise own (e.g. serving the PD KV bootstrap registry)."""
return envs.SGLANG_RUST_SERVER.get() and ( return envs.SGLANG_RUST_SERVER.get() and (
self.ps.pp_rank == 0 get_parallel().pp_rank == 0
and self.ps.attn_tp_rank == 0 and get_parallel().attn_tp_rank == 0
and self.ps.attn_cp_rank == 0 and get_parallel().attn_cp_rank == 0
) )
def maybe_init_rust_server(self) -> None: def maybe_init_rust_server(self) -> None:
@@ -2417,10 +2421,10 @@ class Scheduler(
self.kv_events_publisher = SchedulerKvEventsPublisher( self.kv_events_publisher = SchedulerKvEventsPublisher(
kv_events_config=get_observability().kv_events_config, kv_events_config=get_observability().kv_events_config,
ps=self.ps, ps=self.ps,
attn_tp_rank=self.ps.attn_tp_rank, attn_tp_rank=get_parallel().attn_tp_rank,
attn_cp_rank=self.ps.attn_cp_rank, attn_cp_rank=get_parallel().attn_cp_rank,
attn_dp_rank=self.ps.attn_dp_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, tree_cache=self.tree_cache,
send_metrics_from_scheduler=self.ipc_channels.send_metrics_from_scheduler, send_metrics_from_scheduler=self.ipc_channels.send_metrics_from_scheduler,
max_running_requests=self.max_running_requests, max_running_requests=self.max_running_requests,
@@ -3398,7 +3402,7 @@ class Scheduler(
if (timeout_s := envs.SGLANG_REQ_RUNNING_TIMEOUT.get()) > 0: if (timeout_s := envs.SGLANG_REQ_RUNNING_TIMEOUT.get()) > 0:
deadline = time.perf_counter() - timeout_s 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] inflight_batches = [self.running_batch, self.last_batch]
else: else:
inflight_batches = [*self.running_mbs, *self.mbs] inflight_batches = [*self.running_mbs, *self.mbs]
@@ -4460,7 +4464,7 @@ class Scheduler(
batch.input_ids = None batch.input_ids = None
self._copy_auxiliary_output_to_cpu(batch, batch_result) self._copy_auxiliary_output_to_cpu(batch, batch_result)
elif not batch.spec_algorithm.is_none(): 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 batch.forward_mode.is_extend() or batch.is_extend_in_batch
) )
# The relayed tree is what the requests carry between rounds; # 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 # Only the last PP rank owns real results requiring D2H; other ranks
# consume device tensors rebuilt from the output ring. # consume device tensors rebuilt from the output ring.
batch_result.copy_done = self.device_module.Event() 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( batch_result.copy_to_cpu(
return_logprob=batch.return_logprob, return_logprob=batch.return_logprob,
return_hidden_states=batch.return_hidden_states, return_hidden_states=batch.return_hidden_states,
@@ -4685,7 +4692,7 @@ class Scheduler(
): ):
return return
# PP transports the device output to the first rank before copying it. # 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 return
if result.copy_done is not None: if result.copy_done is not None:
raise RuntimeError( raise RuntimeError(
@@ -5024,7 +5031,7 @@ class Scheduler(
return idle return idle
def _pp_microbatches_drained(self) -> bool: def _pp_microbatches_drained(self) -> bool:
if self.ps.pp_size == 1: if get_parallel().pp_size == 1:
return True return True
return all(x.is_empty() for x in self.running_mbs) and all( 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 mb is None or mb.is_empty() for mb in self.mbs
@@ -5265,10 +5272,10 @@ class Scheduler(
if_success = False if_success = False
break break
elif k == "pp_max_micro_batch_size" and ( 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( 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 if_success = False
break break
@@ -5382,7 +5389,7 @@ class Scheduler(
) )
def collect_inflight_reqs(self) -> Set[Req]: 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] inflight_batches = [self.running_batch, self.last_batch]
else: else:
inflight_batches = [*self.running_mbs, *self.mbs] inflight_batches = [*self.running_mbs, *self.mbs]
@@ -5472,7 +5479,7 @@ class Scheduler(
if hasattr(req.disagg_kv_sender, "abort"): if hasattr(req.disagg_kv_sender, "abort"):
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.") prepare_abort(req, "Aborted by AbortReq.")
# Abort in-flight requests # Abort in-flight requests
@@ -5488,7 +5495,7 @@ class Scheduler(
if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid): 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=}") logger.debug(f"Abort prealloc queue request. {decode_req.req.rid=}")
decode_req.kv_receiver.abort() 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.") prepare_abort(decode_req.req, "Aborted by AbortReq.")
# Abort requests waiting for kvcache to release tree cache # Abort requests waiting for kvcache to release tree cache
@@ -5821,7 +5828,11 @@ class Scheduler(
output = self.session_controller.open(recv_req) output = self.session_controller.open(recv_req)
if output.success and self.enable_session_radix_cache: if output.success and self.enable_session_radix_cache:
self.tree_cache.open_radix_session(recv_req.session_id) 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 output
return None return None
@@ -5922,6 +5933,20 @@ def _dispatch_event_loop_once(scheduler: Scheduler):
scheduler.event_loop_normal_disagg_decode() 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( def configure_scheduler_process(
server_args: ServerArgs, server_args: ServerArgs,
gpu_id: int, gpu_id: int,
@@ -5934,18 +5959,15 @@ def configure_scheduler_process(
display_tp_rank: Optional[int] = None, display_tp_rank: Optional[int] = None,
display_dp_rank: Optional[int] = None, display_dp_rank: Optional[int] = None,
display_moe_ep_rank: Optional[int] = None, display_moe_ep_rank: Optional[int] = None,
) -> Optional[int]: ) -> None:
"""Configure scheduler worker logging and process title. """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() kill_itself_when_parent_died()
# Generate the logger prefix # 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_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_tp = display_tp_rank if display_tp_rank is not None else tp_rank
shown_moe_ep = ( shown_moe_ep = (
@@ -5987,8 +6009,6 @@ def configure_scheduler_process(
if numa_node is not None: if numa_node is not None:
numa_bind_to_node(numa_node) numa_bind_to_node(numa_node)
return dp_rank
def run_scheduler_process( def run_scheduler_process(
server_args: ServerArgs, server_args: ServerArgs,
@@ -6007,9 +6027,20 @@ def run_scheduler_process(
): ):
# Load plugins so hooks can override Scheduler and its dependencies. # Load plugins so hooks can override Scheduler and its dependencies.
load_plugins() load_plugins()
# Publish before anything in this process reads configuration. dp_rank = resolve_spawn_dp_rank(dp_rank)
publish(server_args, role="scheduler") # Publish before anything in this process reads configuration, with the
dp_rank = configure_scheduler_process( # 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, server_args,
gpu_id, gpu_id,
tp_rank, tp_rank,
@@ -548,8 +548,8 @@ class SchedulerDPAttnAdapter:
local_batch, local_batch,
model_runner=self.model_runner, model_runner=self.model_runner,
dp_size=get_parallel().dp_size, dp_size=get_parallel().dp_size,
attn_tp_size=self.ps.attn_tp_size, attn_tp_size=get_parallel().attn_tp_size,
attn_cp_size=self.ps.attn_cp_size, attn_cp_size=get_parallel().attn_cp_size,
tp_group=self.tp_group, tp_group=self.tp_group,
get_idle_batch=self.get_idle_batch, get_idle_batch=self.get_idle_batch,
disable_cuda_graph=cuda_graph_fully_disabled(), disable_cuda_graph=cuda_graph_fully_disabled(),
@@ -19,6 +19,7 @@ from sglang.srt.disaggregation.kv_events import (
select_kv_publisher_dp_rank, select_kv_publisher_dp_rank,
) )
from sglang.srt.managers.io_struct import hook_custom_types, sock_send from sglang.srt.managers.io_struct import hook_custom_types, sock_send
from sglang.srt.runtime_context import get_parallel
if TYPE_CHECKING: if TYPE_CHECKING:
from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.distributed.parallel_state_wrapper import ParallelState
@@ -68,7 +69,7 @@ class SchedulerKvEventsPublisher:
self.kv_event_publisher = EventPublisherFactory.create( self.kv_event_publisher = EventPublisherFactory.create(
kv_events_config, kv_events_config,
select_kv_publisher_dp_rank( 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_cache_usage_perc = self.get_stats().token_usage
kv_metrics.gpu_prefix_cache_hit_rate = self.get_stats().cache_hit_rate kv_metrics.gpu_prefix_cache_hit_rate = self.get_stats().cache_hit_rate
kv_metrics.data_parallel_rank = ( 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: if not self.send_metrics_from_scheduler.closed:
@@ -14,7 +14,7 @@ from sglang.srt.managers.load_snapshot import (
QueueMetrics, QueueMetrics,
SpeculativeMetrics, SpeculativeMetrics,
) )
from sglang.srt.runtime_context import get_lora from sglang.srt.runtime_context import get_lora, get_parallel
if TYPE_CHECKING: if TYPE_CHECKING:
from sglang.srt.distributed.parallel_state_wrapper import ParallelState 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 decode_moments = list(totals) if totals[0] > 0 else None
return LoadSnapshot( 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(), timestamp=time.time(),
num_running_reqs=num_running_reqs, num_running_reqs=num_running_reqs,
num_waiting_reqs=num_waiting_reqs, num_waiting_reqs=num_waiting_reqs,
@@ -309,7 +309,7 @@ class SchedulerMetricsReporter:
if ( if (
get_observability().enable_forward_pass_metrics get_observability().enable_forward_pass_metrics
and self.scheduler.ps.attn_tp_rank == 0 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 ( from sglang.srt.observability.forward_pass_metrics import (
_FpmPublisherThread, _FpmPublisherThread,
@@ -32,7 +32,7 @@ from sglang.srt.managers.schedule_batch import (
Req, Req,
) )
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache 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.server_args import ServerArgs
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.utils.weight_versions import compute_weight_version_spans from sglang.srt.utils.weight_versions import compute_weight_version_spans
@@ -205,7 +205,7 @@ class SchedulerOutputStreamer:
# Send to detokenizer # Send to detokenizer
payload = acc.to_payload( payload = acc.to_payload(
dp_rank=self.ps.dp_rank, dp_rank=get_parallel().dp_rank,
is_idle_batch=is_idle_batch, is_idle_batch=is_idle_batch,
) )
if payload is not None: if payload is not None:
@@ -232,7 +232,7 @@ class SchedulerOutputStreamer:
def _maybe_log_time_stats(self, *, req: Req) -> None: def _maybe_log_time_stats(self, *, req: Req) -> None:
if ( if (
req.finished() req.finished()
and self.ps.attn_tp_rank == 0 and get_parallel().attn_tp_rank == 0
and get_observability().enable_request_time_stats_logging and get_observability().enable_request_time_stats_logging
): ):
req.log_time_stats() req.log_time_stats()
@@ -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.forward_batch_info import ForwardMode
from sglang.srt.model_executor.step_span_utils import set_detailed_annotations_enabled from sglang.srt.model_executor.step_span_utils import set_detailed_annotations_enabled
from sglang.srt.platforms import current_platform 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 import is_mps, is_npu
from sglang.srt.utils.profile_merger import ProfileMerger from sglang.srt.utils.profile_merger import ProfileMerger
from sglang.srt.utils.profile_utils import ProfileManager from sglang.srt.utils.profile_utils import ProfileManager
@@ -207,10 +207,13 @@ class SchedulerProfilerManager:
self.rpd_profile_path = os.path.join( self.rpd_profile_path = os.path.join(
self.torch_profiler_output_dir, 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 import sqlite3
from rocpd.schema import RocpdSchema from rocpd.schema import RocpdSchema
@@ -282,13 +285,13 @@ class SchedulerProfilerManager:
if not self.merge_profiles: if not self.merge_profiles:
return "" return ""
if self.ps.tp_rank != 0: if get_parallel().tp_rank != 0:
return "" 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 "" 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 "" 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 "" return ""
try: try:
@@ -336,15 +339,15 @@ class SchedulerProfilerManager:
self.torch_profiler.stop() self.torch_profiler.stop()
if not _is_npu: if not _is_npu:
# Build filename with only non-zero ranks to maintain backward compatibility # 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) # Only add other ranks if parallelism is enabled (size > 1)
if self.ps.dp_size > 1: if self.ps.dp_size > 1:
filename_parts.append(f"DP-{self.ps.dp_rank}") filename_parts.append(f"DP-{get_parallel().dp_rank}")
if self.ps.pp_size > 1: if get_parallel().pp_size > 1:
filename_parts.append(f"PP-{self.ps.pp_rank}") filename_parts.append(f"PP-{get_parallel().pp_rank}")
if self.ps.moe_ep_size > 1: if get_parallel().moe_ep_size > 1:
filename_parts.append(f"EP-{self.ps.moe_ep_rank}") filename_parts.append(f"EP-{get_parallel().moe_ep_rank}")
filename = ( filename = (
stage_prefix stage_prefix
@@ -364,7 +367,7 @@ class SchedulerProfilerManager:
self.rpd_profiler.flush() self.rpd_profiler.flush()
torch.distributed.barrier(self.dp_tp_cpu_group) 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 from sglang.srt.utils.rpd_utils import rpd_to_chrome_trace
rpd_to_chrome_trace("trace.rpd", self.rpd_profile_path) rpd_to_chrome_trace("trace.rpd", self.rpd_profile_path)
@@ -376,7 +379,7 @@ class SchedulerProfilerManager:
self.torch_profiler_output_dir, self.torch_profiler_output_dir,
stage_prefix stage_prefix
+ str(time.time()) + str(time.time())
+ f"-TP-{self.ps.tp_rank}-memory" + f"-TP-{get_parallel().tp_rank}-memory"
+ stage_suffix + stage_suffix
+ ".pickle", + ".pickle",
) )
@@ -109,7 +109,7 @@ class SchedulerRequestReceiver:
recv_reqs = self._broadcast_reqs_across_ranks(recv_reqs, local_reqs) 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) self.unwrap_pickle_wrapper(recv_reqs)
recv_reqs = self._apply_mm_receiver(recv_reqs) recv_reqs = self._apply_mm_receiver(recv_reqs)
@@ -119,7 +119,7 @@ class SchedulerRequestReceiver:
return recv_reqs return recv_reqs
def _pull_raw_reqs(self) -> Optional[List]: 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: if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0:
recv_reqs = [] recv_reqs = []
@@ -158,10 +158,10 @@ class SchedulerRequestReceiver:
) )
recv_reqs = point_to_point_pyobj( 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.world_group.cpu_group,
(self.ps.pp_rank - 1) * self.ps.tp_size + dp_offset, (get_parallel().pp_rank - 1) * self.ps.tp_size + dp_offset,
self.ps.pp_rank * self.ps.tp_size + dp_offset, get_parallel().pp_rank * self.ps.tp_size + dp_offset,
) )
else: else:
recv_reqs = None recv_reqs = None
@@ -232,7 +232,7 @@ class SchedulerRequestReceiver:
def _apply_mm_receiver(self, recv_reqs: List) -> List: def _apply_mm_receiver(self, recv_reqs: List) -> List:
# Process MM requests under EPD-disaggregation mode # Process MM requests under EPD-disaggregation mode
if ( if (
self.ps.pp_rank == 0 get_parallel().pp_rank == 0
and get_disagg().language_only and get_disagg().language_only
and get_disagg().encoder_transfer_backend and get_disagg().encoder_transfer_backend
in ["zmq_to_scheduler", "mooncake"] in ["zmq_to_scheduler", "mooncake"]
@@ -103,7 +103,9 @@ class SchedulerPPMixin:
for mb_id in range(self.pp_loop_size): for mb_id in range(self.pp_loop_size):
self.running_batch = self.running_mbs[mb_id] self.running_batch = self.running_mbs[mb_id]
self.last_batch = self.last_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_mb_id = (mb_id + 1) % self.pp_loop_size
with torch.profiler.record_function("recv_requests"): with torch.profiler.record_function("recv_requests"):
recv_reqs = self.ingest_requests() recv_reqs = self.ingest_requests()
@@ -243,7 +245,9 @@ class SchedulerPPMixin:
for mb_id in range(self.pp_loop_size): for mb_id in range(self.pp_loop_size):
self.running_batch = self.running_mbs[mb_id] self.running_batch = self.running_mbs[mb_id]
self.last_batch = self.last_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_mb_id = (mb_id + 1) % self.pp_loop_size
next_pp_outputs = None next_pp_outputs = None
@@ -391,7 +395,9 @@ class SchedulerPPMixin:
for mb_id in range(self.pp_loop_size): for mb_id in range(self.pp_loop_size):
self.running_batch = self.running_mbs[mb_id] self.running_batch = self.running_mbs[mb_id]
self.last_batch = self.last_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_mb_id = (mb_id + 1) % self.pp_loop_size
next_pp_outputs = None next_pp_outputs = None
@@ -561,7 +567,9 @@ class SchedulerPPMixin:
self.on_idle() self.on_idle()
def init_pp_loop_state(self: Scheduler): 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.mbs = [None] * self.pp_loop_size
self.last_mbs = [None] * self.pp_loop_size self.last_mbs = [None] * self.pp_loop_size
self.running_mbs = [ self.running_mbs = [
@@ -573,7 +581,7 @@ class SchedulerPPMixin:
self.last_rank_comm_queue: deque[Tuple[torch.Event, PPProxyTensors]] = deque() self.last_rank_comm_queue: deque[Tuple[torch.Event, PPProxyTensors]] = deque()
self._pp_spec_relay = ( self._pp_spec_relay = (
envs.SGLANG_ENABLE_PP_SPEC.get() 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() 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): def _pp_send_pyobj_to_next_stage(self: Scheduler, data, async_send: bool = False):
p2p_work = [] 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 = ( 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( p2p_work = point_to_point_pyobj(
data, 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.world_group.cpu_group,
self.ps.pp_rank * self.ps.tp_size + dp_offset, get_parallel().pp_rank * get_parallel().tp_size + dp_offset,
((self.ps.pp_rank + 1) % self.ps.pp_size) * self.ps.tp_size + dp_offset, ((get_parallel().pp_rank + 1) % get_parallel().pp_size)
* get_parallel().tp_size
+ dp_offset,
async_send=async_send, async_send=async_send,
) )
return p2p_work return p2p_work
def _pp_recv_pyobj_from_prev_stage(self: Scheduler): 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 = ( 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( 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.world_group.cpu_group,
((self.ps.pp_rank - 1) % self.ps.pp_size) * self.ps.tp_size + dp_offset, ((get_parallel().pp_rank - 1) % get_parallel().pp_size)
self.ps.pp_rank * self.ps.tp_size + dp_offset, * get_parallel().tp_size
+ dp_offset,
get_parallel().pp_rank * get_parallel().tp_size + dp_offset,
) )
else: else:
data = None data = None
@@ -1470,7 +1486,7 @@ class SchedulerPPMixin:
# posted, so the parity-based send-first/recv-first ordering used # posted, so the parity-based send-first/recv-first ordering used
# for NPU is replaced by batch_isend_irecv which submits all # for NPU is replaced by batch_isend_irecv which submits all
# send/recv operations atomically. # 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( return self._pp2_only_send_recv_output_tensors_npu(
next_first_rank_mb_id, next_first_rank_mb_id,
next_mb_id, next_mb_id,
@@ -1498,7 +1514,7 @@ class SchedulerPPMixin:
# makes rank 1 post its recv first, which breaks the cycle for any # makes rank 1 post its recv first, which breaks the cycle for any
# pp_size > 1. # pp_size > 1.
needs_pairing = is_xpu() or self._pp_spec_relay 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(): def _do_send():
return self._pp_send_output_to_next_stage( return self._pp_send_output_to_next_stage(
+5 -3
View File
@@ -418,7 +418,7 @@ class TpModelWorker(BaseTpWorker):
else: else:
self.random_seed = broadcast_pyobj( self.random_seed = broadcast_pyobj(
[get_device().random_seed], [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, self.world_group.cpu_group,
src=self.world_group.ranks[0], src=self.world_group.ranks[0],
)[0] )[0]
@@ -451,7 +451,8 @@ class TpModelWorker(BaseTpWorker):
assert self.model_runner.max_running_requests > 0, "max_running_request is zero" assert self.model_runner.max_running_requests > 0, "max_running_request is zero"
max_req_len = min( max_req_len = min(
self.model_config.context_len - 1, 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, - 1,
) )
assert max_req_len > 0, "Memory pool size is too small" assert max_req_len > 0, "Memory pool size is too small"
@@ -578,7 +579,8 @@ class TpModelWorker(BaseTpWorker):
def get_worker_info(self): def get_worker_info(self):
max_req_len = min( max_req_len = min(
self.model_config.context_len - 1, 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, - 1,
) )
return ( return (
@@ -2420,7 +2420,9 @@ class KVCacheConfigurator:
sum(1 for i in all_mamba_layers if start <= i < end) sum(1 for i in all_mamba_layers if start <= i < end)
for start, end in ( for start, end in (
get_pp_indices( 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) for rank in range(self.ps.pp_size)
) )
@@ -445,7 +445,7 @@ class ModelRunner:
self.prefill_shared_read_stager: Optional[Callable[[ForwardBatch], bool]] = None self.prefill_shared_read_stager: Optional[Callable[[ForwardBatch], bool]] = None
# CPU offload # 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) 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 "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): if not (envs.SGLANG_ENABLE_PP_SPEC.get() and self.is_draft_worker):
assert self.support_pp, ( assert self.support_pp, (
"Pipeline Parallel is not compatible with this model." "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 from sglang.srt.model_executor.mindspore_runner import init_ms_distributed
init_ms_distributed( init_ms_distributed(
world_size=self.ps.tp_size * self.ps.pp_size, world_size=self.ps.tp_size * get_parallel().pp_size,
rank=self.ps.tp_size * self.ps.pp_rank + self.ps.tp_rank, rank=self.ps.tp_size * get_parallel().pp_rank + self.ps.tp_rank,
local_rank=self.gpu_id, local_rank=self.gpu_id,
port=self.dist_port, port=self.dist_port,
) )
@@ -667,8 +667,8 @@ class ModelRunner:
prepare_moe_topk( prepare_moe_topk(
model=self.model, model=self.model,
model_config=self.model_config, model_config=self.model_config,
moe_ep_size=self.ps.moe_ep_size, moe_ep_size=get_parallel().moe_ep_size,
moe_ep_rank=self.ps.moe_ep_rank, moe_ep_rank=get_parallel().moe_ep_rank,
) )
self.maybe_init_dwdp() self.maybe_init_dwdp()
@@ -708,7 +708,7 @@ class ModelRunner:
def maybe_init_expert_location_metadata(self): def maybe_init_expert_location_metadata(self):
if self.is_draft_worker: if self.is_draft_worker:
return return
expert_rank = self.ps.moe_ep_rank + ( expert_rank = get_parallel().moe_ep_rank + (
get_parallel().ep_join_rank_offset get_parallel().ep_join_rank_offset
if get_exec().moe.is_ep_scale_joiner if get_exec().moe.is_ep_scale_joiner
else 0 else 0
@@ -767,8 +767,8 @@ class ModelRunner:
self.expert_backup_client = ( self.expert_backup_client = (
ExpertBackupClient( ExpertBackupClient(
model_config=self.model_config, model_config=self.model_config,
moe_ep_size=self.ps.moe_ep_size, moe_ep_size=get_parallel().moe_ep_size,
moe_ep_rank=self.ps.moe_ep_rank, moe_ep_rank=get_parallel().moe_ep_rank,
get_model=lambda: self.model, get_model=lambda: self.model,
) )
if ( if (
@@ -803,16 +803,16 @@ class ModelRunner:
def get_pp_proxy_topk_size(self) -> Optional[int]: def get_pp_proxy_topk_size(self) -> Optional[int]:
return misc_utils.resolve_pp_proxy_topk_size( return misc_utils.resolve_pp_proxy_topk_size(
model_config=self.model_config, model_config=self.model_config,
pp_size=self.ps.pp_size, pp_size=get_parallel().pp_size,
pp_rank=self.ps.pp_rank, pp_rank=get_parallel().pp_rank,
start_layer=self.layer_info.start_layer, start_layer=self.layer_info.start_layer,
) )
def get_pp_proxy_residual_num_blocks(self) -> Optional[int]: def get_pp_proxy_residual_num_blocks(self) -> Optional[int]:
return misc_utils.resolve_pp_proxy_residual_num_blocks( return misc_utils.resolve_pp_proxy_residual_num_blocks(
model_config=self.model_config, model_config=self.model_config,
pp_size=self.ps.pp_size, pp_size=get_parallel().pp_size,
pp_rank=self.ps.pp_rank, pp_rank=get_parallel().pp_rank,
start_layer=self.layer_info.start_layer, start_layer=self.layer_info.start_layer,
) )
@@ -980,7 +980,7 @@ class ModelRunner:
swap_in_block_size=hisparse_cfg.swap_in_block_size, swap_in_block_size=hisparse_cfg.swap_in_block_size,
shared_index_layers=resolve_shared_index_layers( shared_index_layers=resolve_shared_index_layers(
hf_text_config=self.model_config.hf_text_config, 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(), is_speculative=self.spec_algorithm.is_speculative(),
), ),
) )
@@ -1157,8 +1157,8 @@ class ModelRunner:
check_quantized_moe_compatibility( check_quantized_moe_compatibility(
model_config=self.model_config, model_config=self.model_config,
tp_size=self.ps.tp_size, tp_size=self.ps.tp_size,
moe_ep_size=self.ps.moe_ep_size, moe_ep_size=get_parallel().moe_ep_size,
moe_dp_size=self.ps.moe_dp_size, moe_dp_size=get_parallel().moe_dp_size,
) )
def init_torch_distributed(self): def init_torch_distributed(self):
@@ -1293,7 +1293,7 @@ class ModelRunner:
is_draft_worker=self.is_draft_worker, is_draft_worker=self.is_draft_worker,
tp_size=self.ps.tp_size, tp_size=self.ps.tp_size,
tp_rank=self.ps.tp_rank, tp_rank=self.ps.tp_rank,
pp_rank=self.ps.pp_rank, pp_rank=get_parallel().pp_rank,
) )
if dumper.may_enable: if dumper.may_enable:
@@ -1580,7 +1580,7 @@ class ModelRunner:
dp_size = 1 if get_parallel().enable_dp_attention else self.ps.dp_size dp_size = 1 if get_parallel().enable_dp_attention else self.ps.dp_size
self.local_omp_cpuid = numa_utils.init_threads_binding( self.local_omp_cpuid = numa_utils.init_threads_binding(
numa_index=self.gpu_id, 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): def apply_torch_tp(self):
@@ -274,7 +274,7 @@ class BaseRunner(ABC):
if ( if (
envs.SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP.get() envs.SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP.get()
and deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM 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() and not mr.spec_algorithm.is_speculative()
): ):
from sglang.srt.layers.deep_gemm_wrapper.compile_utils import ( from sglang.srt.layers.deep_gemm_wrapper.compile_utils import (
@@ -546,7 +546,7 @@ class BaseRunner(ABC):
pp_hidden_tokens = num_tokens pp_hidden_tokens = num_tokens
if ( if (
capture_forward_mode == ForwardMode.EXTEND capture_forward_mode == ForwardMode.EXTEND
and mr.ps.pp_rank != 0 and get_parallel().pp_rank != 0
and mr.ps.attn_cp_size > 1 and mr.ps.attn_cp_size > 1
): ):
pp_hidden_tokens = num_tokens // mr.ps.attn_cp_size pp_hidden_tokens = num_tokens // mr.ps.attn_cp_size
@@ -30,6 +30,7 @@ from sglang.srt.runtime_context import (
get_disagg, get_disagg,
get_exec, get_exec,
get_model, get_model,
get_parallel,
get_schedule, get_schedule,
get_spec, get_spec,
max_prefill_buffer_tokens, max_prefill_buffer_tokens,
@@ -147,7 +148,7 @@ def flashinfer_autotune_cache_path(model_runner: ModelRunner) -> Path:
str(get_model().quantization), str(get_model().quantization),
str(get_exec().moe.moe_runner_backend), str(get_exec().moe.moe_runner_backend),
str(mr.ps.tp_size), str(mr.ps.tp_size),
str(mr.ps.pp_size), str(get_parallel().pp_size),
str(mr.ps.attn_dp_size), str(mr.ps.attn_dp_size),
str(mr.ps.moe_ep_size), str(mr.ps.moe_ep_size),
str(mr.model_config.hf_config.__class__.__name__), 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) cache_dir.mkdir(parents=True, exist_ok=True)
return ( return (
cache_dir 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"
) )
@@ -34,6 +34,7 @@ from sglang.srt.runtime_context import (
get_context, get_context,
get_disagg, get_disagg,
get_observability, get_observability,
get_parallel,
get_schedule, get_schedule,
get_serving, get_serving,
) )
@@ -1118,7 +1119,7 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin):
) )
enable_kv_cache_events = bool( enable_kv_cache_events = bool(
get_observability().kv_events_config 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_tp_rank == 0
and ps.attn_cp_rank == 0 and ps.attn_cp_rank == 0
) )
+19 -4
View File
@@ -21,7 +21,7 @@ from typing import Any, Dict, Optional
import ray import ray
from sglang.srt.arg_groups.overrides import declare_resolution 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 from sglang.srt.server_args import PortArgs, ServerArgs
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -49,7 +49,11 @@ class SchedulerActor:
dist_init_addr: Optional[str] = None, dist_init_addr: Optional[str] = None,
): ):
from sglang.srt.environ import envs 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 ( from sglang.srt.utils.numa_utils import (
get_numa_node_if_available, get_numa_node_if_available,
numa_bind_to_node, numa_bind_to_node,
@@ -77,12 +81,23 @@ class SchedulerActor:
actual_gpu_id = gpu_id actual_gpu_id = gpu_id
logger.info(f"[TP{tp_rank}] Using passed 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 # This actor takes the place of run_scheduler_process, which is where
# a forked scheduler publishes. # 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.) # Configure worker (logging, process title, etc.)
dp_rank = configure_scheduler_process( configure_scheduler_process(
server_args, server_args,
actual_gpu_id, actual_gpu_id,
tp_rank, tp_rank,
+199 -6
View File
@@ -132,6 +132,10 @@ class Live(msgspec.Struct, frozen=True):
source: Any = None source: Any = None
doc: str = "" 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 = { _LIVE_READS: dict = {
@@ -177,7 +181,35 @@ _LIVE_READS: dict = {
"attn_cp_rank": "get_attn_context_model_parallel_rank", "attn_cp_rank": "get_attn_context_model_parallel_rank",
"dcp_rank": "get_dcp_rank", "dcp_rank": "get_dcp_rank",
"attn_dcp_rank": lambda self: self.dcp_rank if self.dcp_enabled else 0, "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", "world_group": "get_world_group",
"tp_group": "get_tp_group", "tp_group": "get_tp_group",
"pp_group": "get_pp_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 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( def derive_parallel_widths(
*, *,
tp_size: int, tp_size: int,
@@ -327,6 +433,31 @@ def dcp_enabled_of(cfg: Any):
return parallel_widths_of(cfg)["dcp_enabled"] 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: class ParallelContext:
"""Parallel-topology namespace: one spelling per name. """Parallel-topology namespace: one spelling per name.
@@ -394,11 +525,10 @@ class ParallelContext:
return getattr(_ps(), source)() return getattr(_ps(), source)()
if source is not None: if source is not None:
return source(self) return source(self)
why = live.unstamped if isinstance(live, Live) else ""
raise RuntimeError( raise RuntimeError(
f"parallel rank {name!r} is not available: it is computed from " f"parallel name {name!r} is not available: "
"this process's `tp_rank` when the attention topology is " + (why or "nothing has stamped it in this process")
"initialized, so a process that never ran "
"`initialize_dp_attention` has no answer to give"
) )
if config is None and name in _parallel_config_leaves(): if config is None and name in _parallel_config_leaves():
raise ValueError("config namespace 'parallel' not published") 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. """Install process-wide config for this OS process.
Records the process ``role`` — one of the ``ROLE_NAMESPACE_SETS`` keys, 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). namespace-read enforcement (``record`` audits the reads instead).
``hf_config`` is accepted for forward-compat and currently unused. ``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 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** engine running now. Re-publish is allowed and is **last-publish-wins**
(bags re-projected, provenance reset, role overwritten), which is what (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 _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": if _ROLE_NS_MODE == "record":
# The '-' marker distinguishes a zero-read role from a process where # The '-' marker distinguishes a zero-read role from a process where
# recording never ran (signal teardown skips atexit). # recording never ran (signal teardown skips atexit).
@@ -1648,6 +1820,27 @@ def publish(server_args, *, role: str, hf_config: Any = None) -> RuntimeContext:
return _CONTEXT 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: def assert_published(server_args, *, role: str) -> RuntimeContext:
"""This record, under this role, is already published -- or fail loud. """This record, under this role, is already published -- or fail loud.
+1 -1
View File
@@ -124,7 +124,7 @@ class RustServer:
# The joining TP group is entirely local to this node. # The joining TP group is entirely local to this node.
tp_size_per_node = scheduler.ps.tp_size tp_size_per_node = scheduler.ps.tp_size
else: 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 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 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 # Count DP leaders within this node's TP range. The first leader must
@@ -84,6 +84,7 @@ from sglang.srt.speculative.spec_utils import (
GrammarTree, GrammarTree,
assign_req_to_token_pool_func, assign_req_to_token_pool_func,
build_grammar_vocab_mask, build_grammar_vocab_mask,
draft_pp_context,
draft_tp_context, draft_tp_context,
) )
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu 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) draft_init_ctx = draft_tp_context(get_parallel().attn_tp_group)
else: else:
draft_init_ctx = empty_context() draft_init_ctx = empty_context()
with draft_init_ctx: with draft_pp_context(), draft_init_ctx:
bundle = build_draft_tp_worker( bundle = build_draft_tp_worker(
server_args=server_args, server_args=server_args,
gpu_id=gpu_id, gpu_id=gpu_id,
@@ -598,7 +599,10 @@ class DFlashWorkerV2(BaseSpecWorker):
) )
def init_attention_backends(self): 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._draft_worker.init_attention_backends()
self._need_mamba_verify_commit = mambaish_config( self._need_mamba_verify_commit = mambaish_config(
self.model_runner.model_config self.model_runner.model_config
@@ -608,7 +612,10 @@ class DFlashWorkerV2(BaseSpecWorker):
) )
def init_cuda_graphs(self): 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 = ( capture_decode_cuda_graph = (
get_exec().graph.cuda_graph_config.decode.backend != Backend.DISABLED get_exec().graph.cuda_graph_config.decode.backend != Backend.DISABLED
) )
@@ -74,6 +74,7 @@ from sglang.srt.speculative.spec_tp_sync import SpecTpSync, SpecTpSyncSite
from sglang.srt.speculative.spec_utils import ( from sglang.srt.speculative.spec_utils import (
GrammarTree, GrammarTree,
build_grammar_vocab_mask, build_grammar_vocab_mask,
draft_pp_context,
draft_tp_context, draft_tp_context,
prepare_mamba_track_for_verify, prepare_mamba_track_for_verify,
) )
@@ -165,7 +166,7 @@ class DSparkWorkerV2(BaseSpecWorker):
"MoE-under-DP all-reduce." "MoE-under-DP all-reduce."
) )
with self._draft_context(): with draft_pp_context(), self._draft_context():
bundle = build_draft_tp_worker( bundle = build_draft_tp_worker(
server_args=server_args, server_args=server_args,
gpu_id=gpu_id, gpu_id=gpu_id,
@@ -428,7 +429,7 @@ class DSparkWorkerV2(BaseSpecWorker):
) )
def init_attention_backends(self): def init_attention_backends(self):
with self._draft_context(): with draft_pp_context(), self._draft_context():
self._draft_worker.init_attention_backends() self._draft_worker.init_attention_backends()
self._target_hidden_projection_enabled = _configure_target_hidden_projection( self._target_hidden_projection_enabled = _configure_target_hidden_projection(
target_model=self.target_worker.model_runner.model, target_model=self.target_worker.model_runner.model,
@@ -463,7 +464,7 @@ class DSparkWorkerV2(BaseSpecWorker):
"memory is available after target backend initialization.", "memory is available after target backend initialization.",
available_mem, available_mem,
) )
with self._draft_context(): with draft_pp_context(), self._draft_context():
if capture_decode_cuda_graph: if capture_decode_cuda_graph:
# Keep the draft model graph enabled when folded proposal is # Keep the draft model graph enabled when folded proposal is
# disabled, but do not capture the proposal head as a tail # disabled, but do not capture the proposal head as a tail
@@ -338,6 +338,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
def init_attention_backends(self): def init_attention_backends(self):
with ( with (
draft_pp_context(),
self.draft_tp_context(self.draft_runner.tp_group), self.draft_tp_context(self.draft_runner.tp_group),
speculative_moe_backend_context(), speculative_moe_backend_context(),
speculative_moe_a2a_backend_context(), speculative_moe_a2a_backend_context(),
@@ -347,6 +348,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
def init_cuda_graphs(self): def init_cuda_graphs(self):
with ( with (
draft_pp_context(),
self.draft_tp_context(self.draft_runner.tp_group), self.draft_tp_context(self.draft_runner.tp_group),
speculative_moe_backend_context(), speculative_moe_backend_context(),
speculative_moe_a2a_backend_context(), speculative_moe_a2a_backend_context(),
@@ -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_info import SpeculativeAlgorithm
from sglang.srt.speculative.spec_utils import ( from sglang.srt.speculative.spec_utils import (
draft_pp_context,
draft_tp_context, draft_tp_context,
fast_topk, fast_topk,
get_plan_stream, get_plan_stream,
@@ -134,7 +135,7 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
self.hot_token_id = None self.hot_token_id = None
with ( with (
empty_context(), draft_pp_context(),
speculative_moe_backend_context(), speculative_moe_backend_context(),
speculative_moe_a2a_backend_context(), speculative_moe_a2a_backend_context(),
draft_model_build_scope(), draft_model_build_scope(),
@@ -204,6 +205,7 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
def init_attention_backends(self): def init_attention_backends(self):
with ( with (
draft_pp_context(),
self.draft_tp_context(self.draft_model_runner.tp_group), self.draft_tp_context(self.draft_model_runner.tp_group),
speculative_moe_backend_context(), speculative_moe_backend_context(),
speculative_moe_a2a_backend_context(), speculative_moe_a2a_backend_context(),
@@ -214,6 +216,7 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
def init_cuda_graphs(self): def init_cuda_graphs(self):
with ( with (
draft_pp_context(),
self.draft_tp_context(self.draft_model_runner.tp_group), self.draft_tp_context(self.draft_model_runner.tp_group),
speculative_moe_backend_context(), speculative_moe_backend_context(),
speculative_moe_a2a_backend_context(), speculative_moe_a2a_backend_context(),
@@ -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_info import SpeculativeAlgorithm
from sglang.srt.speculative.spec_utils import ( from sglang.srt.speculative.spec_utils import (
draft_pp_context,
draft_tp_context, draft_tp_context,
get_plan_stream, get_plan_stream,
sample_draft_proposal, sample_draft_proposal,
@@ -163,7 +164,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
# Load draft model weights only. # Load draft model weights only.
with ( with (
empty_context(), draft_pp_context(),
speculative_moe_backend_context(), speculative_moe_backend_context(),
draft_model_build_scope(), draft_model_build_scope(),
): ):
@@ -221,6 +222,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
def init_attention_backends(self): def init_attention_backends(self):
with ( with (
draft_pp_context(),
self.draft_tp_context(self.draft_runner_list[0].tp_group), self.draft_tp_context(self.draft_runner_list[0].tp_group),
speculative_moe_backend_context(), speculative_moe_backend_context(),
): ):
@@ -228,6 +230,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
def init_cuda_graphs(self): def init_cuda_graphs(self):
with ( with (
draft_pp_context(),
self.draft_tp_context(self.draft_runner_list[0].tp_group), self.draft_tp_context(self.draft_runner_list[0].tp_group),
speculative_moe_backend_context(), speculative_moe_backend_context(),
): ):
@@ -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_utils import default_tree_mask_mode
from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker, EAGLEWorkerV2 from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker, EAGLEWorkerV2
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm 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 from sglang.srt.utils import empty_context, get_bool_env_var
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -76,7 +80,7 @@ class StandaloneDraftWorker(EagleDraftWorker):
# whose MoE gates run during construction; the scope routes their # whose MoE gates run during construction; the scope routes their
# fusion decision to the speculative leaf (it does not swap # fusion decision to the speculative leaf (it does not swap
# runner_backend — the draft's forwards run outside that context). # 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( self.draft_worker = TpModelWorker(
server_args=server_args, server_args=server_args,
gpu_id=gpu_id, gpu_id=gpu_id,
+14 -11
View File
@@ -18,7 +18,7 @@ from sglang.srt.model_executor.step_span_utils import (
set_detailed_annotations_enabled, set_detailed_annotations_enabled,
) )
from sglang.srt.platforms import current_platform 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 import is_npu
from sglang.srt.utils.torch_npu_patch_utils import apply_torch_npu_patches from sglang.srt.utils.torch_npu_patch_utils import apply_torch_npu_patches
@@ -358,15 +358,15 @@ class _ProfilerTorch(_ProfilerConcreteBase):
self.torch_profiler.stop() self.torch_profiler.stop()
if not _is_npu: if not _is_npu:
# Build filename with only non-zero ranks to maintain backward compatibility # 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) # Only add other ranks if parallelism is enabled (size > 1)
if self.ps.dp_size > 1: if self.ps.dp_size > 1:
filename_parts.append(f"DP-{self.ps.dp_rank}") filename_parts.append(f"DP-{get_parallel().dp_rank}")
if self.ps.pp_size > 1: if get_parallel().pp_size > 1:
filename_parts.append(f"PP-{self.ps.pp_rank}") filename_parts.append(f"PP-{get_parallel().pp_rank}")
if self.ps.moe_ep_size > 1: if get_parallel().moe_ep_size > 1:
filename_parts.append(f"EP-{self.ps.moe_ep_rank}") filename_parts.append(f"EP-{get_parallel().moe_ep_rank}")
filename = ( filename = (
(self.output_prefix + "-" if self.output_prefix else "") (self.output_prefix + "-" if self.output_prefix else "")
@@ -396,7 +396,7 @@ class _ProfilerMemory(_ProfilerConcreteBase):
self.output_dir, self.output_dir,
(self.output_prefix + "-" if self.output_prefix else "") (self.output_prefix + "-" if self.output_prefix else "")
+ str(time.time()) + str(time.time())
+ f"-TP-{self.ps.tp_rank}-memory" + f"-TP-{get_parallel().tp_rank}-memory"
+ self.output_suffix + self.output_suffix
+ ".pickle", + ".pickle",
) )
@@ -426,10 +426,13 @@ class _ProfilerRPD(_ProfilerConcreteBase):
self.rpd_profile_path = os.path.join( self.rpd_profile_path = os.path.join(
self.output_dir, 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 import sqlite3
from rocpd.schema import RocpdSchema from rocpd.schema import RocpdSchema
@@ -454,7 +457,7 @@ class _ProfilerRPD(_ProfilerConcreteBase):
self.rpd_profiler.flush() self.rpd_profiler.flush()
torch.distributed.barrier(self.cpu_group) 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 from sglang.srt.utils.rpd_utils import rpd_to_chrome_trace
rpd_to_chrome_trace("trace.rpd", self.rpd_profile_path) rpd_to_chrome_trace("trace.rpd", self.rpd_profile_path)
+13 -1
View File
@@ -51,9 +51,11 @@ from sglang.srt.arg_groups.overrides import resolving_view
from sglang.srt.configs.load_config import LoadConfig from sglang.srt.configs.load_config import LoadConfig
from sglang.srt.platforms import current_platform from sglang.srt.platforms import current_platform
from sglang.srt.runtime_context import ( from sglang.srt.runtime_context import (
SpawnRanks,
get_exec, get_exec,
get_parallel, get_parallel,
publish, publish,
spawn_world_rank,
) )
from .protocol import ( from .protocol import (
@@ -281,7 +283,17 @@ class WeightCacheDaemon:
from sglang.srt.model_loader.loader import get_model_loader from sglang.srt.model_loader.loader import get_model_loader
server_args = self.server_args 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 from sglang.srt.layers.moe import initialize_moe_config
+42
View File
@@ -2,6 +2,7 @@
import argparse import argparse
import asyncio import asyncio
import contextlib
import copy import copy
import doctest import doctest
import importlib.util import importlib.util
@@ -2002,6 +2003,36 @@ def maybe_stub_sgl_kernel():
sys.meta_path.insert(0, _SglKernelFinder()) 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_TIMEOUT_SECS = 30.0
_GPU_IDLE_POLL_INTERVAL_SECS = 2.0 _GPU_IDLE_POLL_INTERVAL_SECS = 2.0
_GPU_IDLE_USED_MEMORY_THRESHOLD = 2 << 30 # 2 GiB _GPU_IDLE_USED_MEMORY_THRESHOLD = 2 << 30 # 2 GiB
@@ -2253,6 +2284,17 @@ def enter_override(test_case, override):
return installed 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): class CustomTestCase(unittest.TestCase):
def __init_subclass__(cls, **kwargs): def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs) super().__init_subclass__(**kwargs)
@@ -26,13 +26,19 @@ from sglang.srt.constrained.base_grammar_backend import (
from sglang.srt.constrained.grammar_manager import GrammarManager from sglang.srt.constrained.grammar_manager import GrammarManager
from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject
from sglang.srt.distributed.communication_tags import P2PTag 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 ( from sglang.srt.sampling.sampling_params import (
REQUEST_REASONING_END_TOKEN_IDS_KEY, REQUEST_REASONING_END_TOKEN_IDS_KEY,
) )
from sglang.srt.server_args import ServerArgs from sglang.srt.server_args import ServerArgs
from sglang.test.ci.ci_register import register_cpu_ci 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") 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): def _make_scheduler(grammar_backend_name="none", skip_tokenizer=False):
"""Create a mock scheduler with necessary attributes. """Create a mock scheduler with necessary attributes.
The grammar manager reads its config from the bags, so the settings that The grammar manager reads its config and its place in the pipeline from
used to be hung off the mock are published instead. The caller resets the the context, so the settings that used to be hung off the mock are
context; every test here goes through `_GrammarFixture`. published instead. The caller resets the context; every test here goes
through `_GrammarFixture`.
""" """
reset_context() reset_context()
server_args = ServerArgs( server_args = ServerArgs(
@@ -56,7 +63,11 @@ def _make_scheduler(grammar_backend_name="none", skip_tokenizer=False):
constrained_json_whitespace_pattern=None, constrained_json_whitespace_pattern=None,
constrained_json_disable_any_whitespace=False, constrained_json_disable_any_whitespace=False,
) )
publish(server_args, role="scheduler") publish(
server_args,
role="scheduler",
ranks=SpawnRanks(world_rank=0),
)
scheduler = MagicMock() scheduler = MagicMock()
scheduler.server_args = server_args scheduler.server_args = server_args
scheduler.model_config.request_selectable_think_end_id_sequences = None 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.world_size = 1
scheduler.dp_tp_group.first_rank = 0 scheduler.dp_tp_group.first_rank = 0
scheduler.dp_tp_group.is_first_rank = True scheduler.dp_tp_group.is_first_rank = True
scheduler.ps.pp_rank = 0
scheduler.ps.pp_size = 1
scheduler.pp_group = None scheduler.pp_group = None
return scheduler return scheduler
@@ -769,8 +778,10 @@ class TestGrammarManagerPPSync(unittest.TestCase):
enter_override( enter_override(
self, get_context().override_server_args(skip_tokenizer_init=True) self, get_context().override_server_args(skip_tokenizer_init=True)
) )
scheduler.ps.pp_rank = pp_rank # After that override, not before: installing a server-args override
scheduler.ps.pp_size = pp_size # 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 scheduler.pp_group = pp_group
mgr = GrammarManager(scheduler) mgr = GrammarManager(scheduler)
mgr.grammar_backend = MagicMock(spec=BaseGrammarBackend) mgr.grammar_backend = MagicMock(spec=BaseGrammarBackend)
@@ -27,6 +27,7 @@ from sglang.srt.runtime_context import publish, reset_context
from sglang.srt.server_args import ServerArgs from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.test.ci.ci_register import register_cpu_ci 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") 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(): def test_scheduler_copies_auxiliary_output_for_non_overlap_results():
event = object() event = object()
scheduler = object.__new__(Scheduler) scheduler = object.__new__(Scheduler)
scheduler.ps = SimpleNamespace(pp_size=1)
scheduler.device_module = SimpleNamespace(Event=Mock(return_value=event)) scheduler.device_module = SimpleNamespace(Event=Mock(return_value=event))
result = SimpleNamespace( result = SimpleNamespace(
logits_output=SimpleNamespace(auxiliary_device_output=object()), logits_output=SimpleNamespace(auxiliary_device_output=object()),
@@ -463,6 +463,7 @@ def test_scheduler_copies_auxiliary_output_for_non_overlap_results():
) )
batch = SimpleNamespace(return_logprob=False, return_hidden_states=False) batch = SimpleNamespace(return_logprob=False, return_hidden_states=False)
with published_topology():
Scheduler._copy_auxiliary_output_to_cpu(scheduler, batch, result) Scheduler._copy_auxiliary_output_to_cpu(scheduler, batch, result)
assert result.copy_done is event assert result.copy_done is event
@@ -474,7 +475,6 @@ def test_scheduler_copies_auxiliary_output_for_non_overlap_results():
def test_scheduler_preserves_pipeline_parallel_output_for_transport(): def test_scheduler_preserves_pipeline_parallel_output_for_transport():
scheduler = object.__new__(Scheduler) scheduler = object.__new__(Scheduler)
scheduler.ps = SimpleNamespace(pp_size=2)
scheduler.device_module = SimpleNamespace(Event=Mock()) scheduler.device_module = SimpleNamespace(Event=Mock())
result = SimpleNamespace( result = SimpleNamespace(
logits_output=SimpleNamespace(auxiliary_device_output=object()), logits_output=SimpleNamespace(auxiliary_device_output=object()),
@@ -484,6 +484,7 @@ def test_scheduler_preserves_pipeline_parallel_output_for_transport():
) )
batch = SimpleNamespace(return_logprob=False, return_hidden_states=False) batch = SimpleNamespace(return_logprob=False, return_hidden_states=False)
with published_topology(pp_size=2):
Scheduler._copy_auxiliary_output_to_cpu(scheduler, batch, result) Scheduler._copy_auxiliary_output_to_cpu(scheduler, batch, result)
assert result.copy_done is None assert result.copy_done is None
@@ -514,7 +515,6 @@ def test_pdmux_split_prefill_schedules_auxiliary_output_copy():
scheduler.is_generation = True scheduler.is_generation = True
scheduler.enable_overlap = False scheduler.enable_overlap = False
scheduler.enable_pdmux = True scheduler.enable_pdmux = True
scheduler.ps = SimpleNamespace(pp_size=1)
scheduler.tp_worker = SimpleNamespace( scheduler.tp_worker = SimpleNamespace(
forward_batch_split_prefill=Mock(return_value=result) 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, return_hidden_states=False,
) )
with patch( with (
published_topology(),
patch(
"sglang.srt.managers.scheduler.resolve_forward_inputs" "sglang.srt.managers.scheduler.resolve_forward_inputs"
) as resolve_forward_inputs: ) as resolve_forward_inputs,
):
output_result = Scheduler.run_batch(scheduler, batch) output_result = Scheduler.run_batch(scheduler, batch)
resolve_forward_inputs.assert_called_once_with(batch, scheduler.future_map) resolve_forward_inputs.assert_called_once_with(batch, scheduler.future_map)
@@ -14,6 +14,7 @@ from sglang.srt.utils.weight_versions import (
record_weight_version_events, record_weight_version_events,
) )
from sglang.test.ci.ci_register import register_cpu_ci 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") register_cpu_ci(est_time=10, suite="base-a-test-cpu")
@@ -102,6 +103,8 @@ class TestOutputStreamerCustomizedInfo(unittest.TestCase):
) )
serving_patch.start() serving_patch.start()
observability_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(serving_patch.stop)
self.addCleanup(observability_patch.stop) self.addCleanup(observability_patch.stop)
@@ -3,7 +3,11 @@ from types import SimpleNamespace
from unittest.mock import patch from unittest.mock import patch
from sglang.test.ci.ci_register import register_cpu_ci 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() maybe_stub_sgl_kernel()
@@ -32,6 +36,25 @@ def _make_ps(**overrides) -> ParallelState:
return ParallelState.trivial(**defaults) 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: def _fake_group() -> SimpleNamespace:
return SimpleNamespace(rank=0, ranks=[0], cpu_group=object()) return SimpleNamespace(rank=0, ranks=[0], cpu_group=object())
@@ -158,6 +181,7 @@ class TestRequestReceiverBroadcast(unittest.TestCase):
class TestPPCPRankOffsets(unittest.TestCase): class TestPPCPRankOffsets(unittest.TestCase):
def test_request_receiver_uses_cp_size_for_pp_recv_rank(self): def test_request_receiver_uses_cp_size_for_pp_recv_rank(self):
ps = _make_ps() ps = _make_ps()
enter_scope(self, _published_topology())
calls = [] calls = []
def fake_point_to_point_pyobj(data, rank, group, src, dst, **kwargs): 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): def test_pp_mixin_uses_cp_size_for_pyobj_send_and_recv_rank(self):
ps = _make_ps() ps = _make_ps()
enter_scope(self, _published_topology())
scheduler = SchedulerPPMixin() scheduler = SchedulerPPMixin()
scheduler.ps = ps scheduler.ps = ps
scheduler.world_group = _fake_group() scheduler.world_group = _fake_group()
@@ -5,7 +5,12 @@ from types import SimpleNamespace
from unittest.mock import Mock from unittest.mock import Mock
from sglang.test.ci.ci_register import register_cpu_ci 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() maybe_stub_sgl_kernel()
@@ -38,13 +43,16 @@ def _make_scheduler(pending_req, *, chunked_req, running_reqs) -> Scheduler:
sched.disaggregation_mode = None sched.disaggregation_mode = None
sched.enable_hicache_storage = False sched.enable_hicache_storage = False
sched.mm_receiver = None sched.mm_receiver = None
sched.ps = SimpleNamespace(pp_size=1)
sched.running_batch = SimpleNamespace(reqs=running_reqs) sched.running_batch = SimpleNamespace(reqs=running_reqs)
sched.last_batch = None sched.last_batch = None
return sched return sched
class TestPendingChunkedAbortRace(CustomTestCase): 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): def test_req_left_chunked_slot_is_aborted(self):
req = _FakeReq("zombie_rid") req = _FakeReq("zombie_rid")
sched = _make_scheduler(req, chunked_req=None, running_reqs=[req]) sched = _make_scheduler(req, chunked_req=None, running_reqs=[req])
@@ -5,7 +5,11 @@ from types import SimpleNamespace
from unittest.mock import Mock, call, patch from unittest.mock import Mock, call, patch
from sglang.test.ci.ci_register import register_cpu_ci 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() maybe_stub_sgl_kernel()
@@ -114,7 +118,7 @@ class TestSchedulerHiCacheEvents(unittest.TestCase):
s = self.scheduler s = self.scheduler
s.init_pp_loop_state = Mock() s.init_pp_loop_state = Mock()
s.pp_loop_size = 1 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.pp_group = SimpleNamespace(is_last_rank=True)
s.running_mbs = [self.running_batch] s.running_mbs = [self.running_batch]
s.last_mbs = [None] s.last_mbs = [None]
@@ -12,7 +12,12 @@ from unittest.mock import MagicMock, patch
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.test.ci.ci_register import register_cpu_ci 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() 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.enable_unified_cache_external_linker = False
s.ipc_channels = SimpleNamespace(send_to_tokenizer=MagicMock()) s.ipc_channels = SimpleNamespace(send_to_tokenizer=MagicMock())
s.beam_coordinator = MagicMock() s.beam_coordinator = MagicMock()
s.ps = SimpleNamespace(pp_size=1)
s.running_batch = _batch(list(running_reqs)) s.running_batch = _batch(list(running_reqs))
s.last_batch = _batch(list(last_batch_reqs)) if last_batch_reqs else None s.last_batch = _batch(list(last_batch_reqs)) if last_batch_reqs else None
return s return s
@@ -126,6 +130,10 @@ class TestWaitingTimeout(CustomTestCase):
class TestRunningTimeout(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): def test_emits_only_stale_unfinished_reqs_without_marking(self):
now = time.perf_counter() now = time.perf_counter()
stale = _req("stale", forward_entry=now - 10) stale = _req("stale", forward_entry=now - 10)
@@ -736,12 +736,15 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase):
# publishing a record rather than by standing one in. # publishing a record rather than by standing one in.
reset_context() reset_context()
publish( 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", role="scheduler",
) )
scheduler = Scheduler.__new__(Scheduler) scheduler = Scheduler.__new__(Scheduler)
scheduler.enable_overlap = enable_overlap scheduler.enable_overlap = enable_overlap
scheduler.ps = SimpleNamespace(pp_size=pp_size)
scheduler.init_tp_model_worker = lambda: setattr(scheduler, "tp_worker", worker) scheduler.init_tp_model_worker = lambda: setattr(scheduler, "tp_worker", worker)
scheduler.maybe_init_draft_worker = lambda: setattr( scheduler.maybe_init_draft_worker = lambda: setattr(
scheduler, "draft_worker", draft_worker scheduler, "draft_worker", draft_worker
@@ -34,7 +34,6 @@ class TestRustServerExtension(CustomTestCase):
attn_dp_rank=1, attn_dp_rank=1,
tp_size=2, tp_size=2,
tp_rank=1, tp_rank=1,
pp_size=1,
attn_tp_size=1, attn_tp_size=1,
attn_cp_size=1, attn_cp_size=1,
), ),
@@ -49,7 +48,9 @@ class TestRustServerExtension(CustomTestCase):
), ),
), ),
patch.object( 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(ModelServer, "_partition_cores", return_value=(None, None)),
patch.object( patch.object(
@@ -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 from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=11, suite="base-a-test-cpu") 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, SchedulerMetricsReporter,
_CacheHitRateWindow, _CacheHitRateWindow,
) )
from sglang.test.test_utils import CustomTestCase from sglang.test.test_utils import CustomTestCase, enter_scope
def _make_ps(**overrides) -> ParallelState: def _make_ps(**overrides) -> ParallelState:
@@ -292,6 +292,8 @@ class TestForwardPassMetrics(unittest.TestCase):
kv_events_config=None, kv_events_config=None,
) )
scheduler.ps = _make_ps(attn_tp_rank=0, dp_rank=2, pp_rank=0, pp_size=1) 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 scheduler.enable_kv_cache_events = False
with patch( with patch(
@@ -329,6 +331,8 @@ class TestForwardPassMetrics(unittest.TestCase):
kv_events_config=None, kv_events_config=None,
) )
scheduler.ps = _make_ps(attn_tp_rank=0, dp_rank=0, pp_rank=0, pp_size=2) 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 scheduler.enable_kv_cache_events = False
with patch( with patch(
@@ -39,6 +39,7 @@ from sglang.srt.runtime_context import (
Flags, Flags,
ParallelContext, ParallelContext,
RuntimeContext, RuntimeContext,
SpawnRanks,
_FlagGroupBase, _FlagGroupBase,
assert_published, assert_published,
derive_parallel_widths, derive_parallel_widths,
@@ -200,6 +201,148 @@ class TestTheTwoWorldWidths(_IsolatedOverrides):
self.assertEqual(parallel.launch_world_size, 2) 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): class TestStampedRanks(_IsolatedOverrides):
"""`attn_dp_rank` comes from the stamp, and says so when there is none. """`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) 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__": if __name__ == "__main__":
unittest.main() unittest.main()
@@ -15,7 +15,11 @@ from sglang.srt.utils.weight_versions import (
truncate_weight_version_events, truncate_weight_version_events,
) )
from sglang.test.ci.ci_register import register_cpu_ci 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") register_cpu_ci(est_time=12, suite="base-a-test-cpu")
@@ -300,11 +304,9 @@ class _SchedulerStub:
waiting, waiting,
chunked=None, chunked=None,
last_batch=None, last_batch=None,
pp_size=1,
hisparse=None, hisparse=None,
): ):
self.serving = _ServingStub(version) self.serving = _ServingStub(version)
self.ps = SimpleNamespace(pp_size=pp_size)
self.running_batch = SimpleNamespace(reqs=running) self.running_batch = SimpleNamespace(reqs=running)
self.last_batch = last_batch self.last_batch = last_batch
self.waiting_queue = waiting self.waiting_queue = waiting
@@ -313,7 +315,9 @@ class _SchedulerStub:
class TestSchedulerRecordWeightVersionChange(CustomTestCase): 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) scheduler = _SchedulerStub(*args, **kwargs)
for name, value in ( for name, value in (
("get_serving", scheduler.serving), ("get_serving", scheduler.serving),