Record a process's placement at publish, not at group build (#40071)
This commit is contained in:
@@ -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 = []
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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(
|
||||||
|
|||||||
@@ -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
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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.
|
||||||
|
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
+5
-2
@@ -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),
|
||||||
|
|||||||
Reference in New Issue
Block a user