minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

A runner and the objects it builds freeze the placement they describe (#40341)

Browse Source
This commit is contained in:
Cheng Wan
2026-09-21 12:24:17 -07:00
committed by GitHub
parent 2d0e94e3a3
commit 65be3fa71a
13 changed files with 289 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/bootstrap.py
+1 -19
View File
@@ -3,14 +3,12 @@ import os
import time
from typing import List, Optional
import msgspec
import torch
import torch.distributed as dist
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.distributed import (
get_default_distributed_backend,
get_pp_group,
get_tp_group,
get_world_group,
init_distributed_environment,
@@ -60,13 +58,6 @@ _is_cpu_arm64 = is_host_cpu_arm64()
_TP_ALL_TO_ALL_WARMUP_BYTES_PER_PEER = 4 << 20
class TorchDistributedResult(msgspec.Struct, frozen=True, kw_only=True):
tp_group: object
pp_group: object
attention_tp_group: object
pre_model_load_memory: float
def init_torch_distributed(
*,
server_args: ServerArgs,
@@ -140,10 +131,6 @@ def init_torch_distributed(
distributed=get_world_group().world_size > 1 and not is_draft_worker,
cpu_group=get_world_group().cpu_group,
)
tp_group = get_tp_group()
pp_group = get_pp_group()
attention_tp_group = get_parallel().attn_tp_group
# Check memory for tensor parallelism
local_gpu_memory = get_available_gpu_memory(device, ps.gpu_id)
if ps.tp_size > 1 and not is_draft_worker:
@@ -156,12 +143,7 @@ def init_torch_distributed(
f"Init torch distributed ends. elapsed={time.perf_counter() - tic:.2f} s, "
f"mem usage={(before_avail_memory - local_gpu_memory):.2f} GB"
)
return TorchDistributedResult(
tp_group=tp_group,
pp_group=pp_group,
attention_tp_group=attention_tp_group,
pre_model_load_memory=pre_model_load_memory,
)
return pre_model_load_memory
def _resolve_backend(*, device: str) -> str:
python/sglang/srt/managers/scheduler.py
-1
View File
@@ -601,7 +601,6 @@ class Scheduler(
and get_parallel().attn_tp_rank == 0
and get_parallel().attn_cp_rank == 0
),
ps=self.ps,
tp_group=self.tp_group,
pp_group=self.pp_group,
enable_hierarchical_cache=self.enable_hierarchical_cache,
python/sglang/srt/mem_cache/kv_cache_builder.py
+9 -8
View File
@@ -64,7 +64,6 @@ if TYPE_CHECKING:
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.distributed.parallel_state import GroupCoordinator
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.managers.tp_worker import BaseTpWorker
from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.base_spec_worker import HiCacheDraftPlan
@@ -263,12 +262,14 @@ def build_kv_cache(
attn_cp_cpu_group: ProcessGroup,
enable_metrics: bool,
enable_kv_cache_events: bool,
ps: ParallelState,
tp_group: GroupCoordinator,
pp_group: GroupCoordinator,
enable_hierarchical_cache: bool,
hicache_draft_plan: Optional[HiCacheDraftPlan] = None,
) -> KVCacheBuildResult:
# Built from the scheduler loop, outside any draft scope, so the context
# answers for the process this cache belongs to.
parallel = get_parallel()
sliding_window_size: Optional[int] = None
full_tokens_per_layer: Optional[int] = None
swa_tokens_per_layer: Optional[int] = None
@@ -369,10 +370,10 @@ def build_kv_cache(
enable_session_radix_cache=get_memory().enable_session_radix_cache,
enable_mamba_extra_buffer=get_exec().mamba.enable_mamba_extra_buffer,
enable_mamba_extra_buffer_lazy=get_exec().mamba.enable_mamba_extra_buffer_lazy,
pp_rank=ps.pp_rank,
pp_size=ps.pp_size,
attn_cp_rank=ps.attn_cp_rank,
attn_cp_size=ps.attn_cp_size,
pp_rank=parallel.pp_rank,
pp_size=parallel.pp_size,
attn_cp_rank=parallel.attn_cp_rank,
attn_cp_size=parallel.attn_cp_size,
chunked_prefill_size=effective_chunked_prefill_size,
sliding_window_size=sliding_window_size,
mtp_draft_device_pools=mtp_draft_device_pools,
@@ -390,8 +391,8 @@ def build_kv_cache(
effective_chunked_prefill_size=effective_chunked_prefill_size,
tp_worker=tp_worker,
model_config=model_config,
tp_size=ps.tp_size,
tp_rank=ps.tp_rank,
tp_size=parallel.tp_size,
tp_rank=parallel.tp_rank,
tp_group=tp_group,
)
with auto_size_hicache(
python/sglang/srt/mem_cache/kv_cache_configurator.py
+13 -11
View File
@@ -206,7 +206,6 @@ def _pp_local_per_request_bytes(
if TYPE_CHECKING:
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.mem_cache.unified_memory_pool import (
UnifiedKVPool,
UnifiedPoolBundle,
@@ -260,7 +259,10 @@ class _PoolSizes(msgspec.Struct, frozen=True, kw_only=True):
class KVCacheConfigurator:
device: str
gpu_id: int
ps: ParallelState
# Frozen at construction, not asked for later: this configurator is built
# inside the scope that describes a draft runner and used outside it.
attn_dp_size: int
pp_size: int
pp_group: Any
model: Any
model_config: ModelConfig
@@ -2324,7 +2326,7 @@ class KVCacheConfigurator:
max_num_reqs = get_schedule().max_running_requests
if max_num_reqs is not None:
requested_per_worker = max_num_reqs // self.ps.attn_dp_size
requested_per_worker = max_num_reqs // self.attn_dp_size
max_num_reqs = min(requested_per_worker, token_capacity // 2)
else:
requested_per_worker = None
@@ -2433,16 +2435,16 @@ class KVCacheConfigurator:
# allocates its own [start_layer, end_layer) slice. Charge the largest
# per-stage share so every rank derives the same pool without a collective.
all_mamba_layers = config.mamba2_cache_params.layers
if self.ps.pp_size > 1 and all_mamba_layers:
if self.pp_size > 1 and all_mamba_layers:
max_stage_mamba_layers = max(
sum(1 for i in all_mamba_layers if start <= i < end)
for start, end in (
get_pp_indices(
self.model_config.num_hidden_layers,
rank,
self.ps.pp_size,
self.pp_size,
)
for rank in range(self.ps.pp_size)
for rank in range(self.pp_size)
)
)
else:
@@ -2474,7 +2476,7 @@ class KVCacheConfigurator:
replayssm_ring_per_req = int(replayssm_ring_per_req * pp_layer_scale)
if replayssm_active and self.hybrid_kda_config is None:
replay_req_slots = (
get_schedule().max_running_requests // self.ps.attn_dp_size + 1
get_schedule().max_running_requests // self.attn_dp_size + 1
)
replayssm_fixed_bytes = replayssm_ring_per_req * replay_req_slots
replayssm_ring_per_slot = 0
@@ -2490,7 +2492,7 @@ class KVCacheConfigurator:
get_context().override(
"mamba_pool.per_dp_shard",
max_mamba_cache_size=get_schedule().max_mamba_cache_size
// self.ps.attn_dp_size,
// self.attn_dp_size,
)
# Reserve intermediate memory based on capped max_num_reqs (+1: the
# pool's padding slot, see memory_pool.py). Skipped under replayssm
@@ -2498,7 +2500,7 @@ class KVCacheConfigurator:
if has_spec_dec and not replayssm_active:
ratio = self._calculate_mamba_ratio()
capped_reqs = min(
get_schedule().max_running_requests // self.ps.attn_dp_size,
get_schedule().max_running_requests // self.attn_dp_size,
get_schedule().max_mamba_cache_size // ratio,
)
intermediate_size = (
@@ -2515,7 +2517,7 @@ class KVCacheConfigurator:
get_context().override(
"mamba_pool.from_max_running_requests",
max_mamba_cache_size=get_schedule().max_running_requests
// self.ps.attn_dp_size,
// self.attn_dp_size,
)
# Reserve intermediate memory based on capped max_num_reqs (+1: the
# pool's padding slot). Skipped under replayssm.
@@ -2555,7 +2557,7 @@ class KVCacheConfigurator:
# Intermediate memory is included in mamba_budget, subtract it
# so the return value only has main_state subtracted from total
capped_reqs = min(
get_schedule().max_running_requests // self.ps.attn_dp_size,
get_schedule().max_running_requests // self.attn_dp_size,
get_schedule().max_mamba_cache_size // ratio,
)
intermediate_size = per_req * (capped_reqs + 1) * D
python/sglang/srt/model_executor/model_runner.py
+59 -42
View File
@@ -421,7 +421,7 @@ class ModelRunner:
import os
logger.warning(
f"Context: {self.device=} {ps.gpu_id=} {os.environ.get('CUDA_VISIBLE_DEVICES')=} {ps.tp_rank=} {ps.tp_size=}"
f"Context: {self.device=} {ps.gpu_id=} {os.environ.get('CUDA_VISIBLE_DEVICES')=} {get_parallel().tp_rank=} {get_parallel().tp_size=}"
)
raise
@@ -447,7 +447,7 @@ class ModelRunner:
# CPU offload
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)
if envs.SGLANG_DETECT_SLOW_RANK.get():
slow_rank_detector.execute()
@@ -503,9 +503,9 @@ class ModelRunner:
):
return
join_effective_ep_size = get_parallel().ep_join_rank_offset + self.ps.tp_size
join_effective_ep_size = get_parallel().ep_join_rank_offset + self.tp_size
dist.barrier(group=self.tp_group.cpu_group)
if self.ps.tp_rank == 0:
if self.tp_rank == 0:
register_scale_cohort(
get_parallel().ep_join_rank_offset,
join_effective_ep_size,
@@ -513,7 +513,7 @@ class ModelRunner:
join_scale_process_group()
get_context().override("elastic_ep.scale_join", ep_size=join_effective_ep_size)
global_ep_rank = self.ps.tp_rank + get_parallel().ep_join_rank_offset
global_ep_rank = self.tp_rank + get_parallel().ep_join_rank_offset
broadcast_global_expert_location_metadata(
model_config=self.model_config,
moe_ep_rank=global_ep_rank,
@@ -563,7 +563,7 @@ class ModelRunner:
def init_weight_updater(self):
self.weight_updater = WeightUpdater(
tp_rank=self.ps.tp_rank,
tp_rank=self.tp_rank,
device=self.device,
gpu_id=self.gpu_id,
model_config=self.model_config,
@@ -586,8 +586,8 @@ class ModelRunner:
def init_weight_exporter(self):
self.weight_exporter = WeightExporter(
tp_rank=self.ps.tp_rank,
tp_size=self.ps.tp_size,
tp_rank=self.tp_rank,
tp_size=self.tp_size,
gpu_id=self.gpu_id,
get_model_path=lambda: self.model_config.model_path,
get_model=lambda: self.model,
@@ -596,7 +596,7 @@ class ModelRunner:
def init_remote_instance_weight_transporter(self):
self.remote_instance_weight_transporter = RemoteInstanceWeightTransporter(
get_model=lambda: self.model,
tp_rank=self.ps.tp_rank,
tp_rank=get_parallel().tp_rank,
gpu_id=self.gpu_id,
)
@@ -610,10 +610,13 @@ class ModelRunner:
)
def init_kv_cache_configurator(self):
# The replica count that shares this KV budget is the deployment's, not
# the one a draft scope reports; the pool is allocated outside any scope.
self.kv_cache_configurator = KVCacheConfigurator(
device=self.device,
gpu_id=self.gpu_id,
ps=self.ps,
attn_dp_size=get_parallel().attn_dp_size,
pp_size=self.pp_size,
pp_group=self.pp_group,
model=self.model,
model_config=self.model_config,
@@ -647,8 +650,8 @@ class ModelRunner:
from sglang.srt.model_executor.mindspore_runner import init_ms_distributed
init_ms_distributed(
world_size=self.ps.tp_size * get_parallel().pp_size,
rank=self.ps.tp_size * get_parallel().pp_rank + self.ps.tp_rank,
world_size=self.tp_size * get_parallel().pp_size,
rank=self.tp_size * get_parallel().pp_rank + self.tp_rank,
local_rank=self.gpu_id,
port=self.dist_port,
)
@@ -718,7 +721,7 @@ class ModelRunner:
moe_ep_rank=expert_rank,
)
)
if self.ps.tp_rank == 0 and envs.SGLANG_LOG_EXPERT_LOCATION_METADATA.get():
if self.tp_rank == 0 and envs.SGLANG_LOG_EXPERT_LOCATION_METADATA.get():
logger.info(
"Initial expert_location_metadata:\n%s",
format_expert_location_layout(get_global_expert_location_metadata()),
@@ -779,7 +782,7 @@ class ModelRunner:
def maybe_apply_post_load_model_transforms(self):
supports_torch_tp = getattr(self.model, "supports_torch_tp", False)
if self.ps.tp_size > 1 and supports_torch_tp:
if self.tp_size > 1 and supports_torch_tp:
self.apply_torch_tp()
def maybe_init_lora_manager(self):
@@ -1015,7 +1018,7 @@ class ModelRunner:
def post_capture_elastic_ep_recover(self):
join_process_groups()
global_ep_rank = self.ps.tp_rank + get_parallel().ep_join_rank_offset
global_ep_rank = self.tp_rank + get_parallel().ep_join_rank_offset
broadcast_global_expert_location_metadata(
model_config=self.model_config,
moe_ep_rank=global_ep_rank,
@@ -1162,13 +1165,13 @@ class ModelRunner:
def check_quantized_moe_compatibility(self):
check_quantized_moe_compatibility(
model_config=self.model_config,
tp_size=self.ps.tp_size,
tp_size=self.tp_size,
moe_ep_size=get_parallel().moe_ep_size,
moe_dp_size=get_parallel().moe_dp_size,
)
def init_torch_distributed(self):
result = bootstrap.init_torch_distributed(
self.pre_model_load_memory = bootstrap.init_torch_distributed(
server_args=self.server_args,
model_config=self.model_config,
device=self.device,
@@ -1177,10 +1180,23 @@ class ModelRunner:
is_draft_worker=self.is_draft_worker,
local_omp_cpuid=self.local_omp_cpuid if self.device == "cpu" else None,
)
self.tp_group = result.tp_group
self.pp_group = result.pp_group
self.attention_tp_group = result.attention_tp_group
self.pre_model_load_memory = result.pre_model_load_memory
# Read once, here: a draft runner is constructed inside the scope that
# states its topology and used outside it, so what it holds has to be
# the placement it was built for rather than whatever the context
# answers later. Groups and widths alike -- a runner asked about its own
# shape after the scope has closed must still describe itself.
parallel = get_parallel()
self.tp_group = parallel.tp_group
self.pp_group = parallel.pp_group
self.attention_tp_group = parallel.attn_tp_group
self.tp_rank = parallel.tp_rank
self.tp_size = parallel.tp_size
self.dp_size = parallel.dp_size
self.attn_dp_size = parallel.attn_dp_size
self.pp_rank = parallel.pp_rank
self.pp_size = parallel.pp_size
self.attn_cp_rank = parallel.attn_cp_rank
self.attn_cp_size = parallel.attn_cp_size
def init_shared_mooncake_transfer_engine(self):
maybe_init_shared_mooncake_transfer_engine(gpu_id=self.gpu_id)
@@ -1204,7 +1220,7 @@ class ModelRunner:
self.load_config = build_load_config(
server_args=self.server_args,
load_format=draft_load_format,
tp_rank=self.ps.tp_rank,
tp_rank=self.tp_rank,
remote_instance_weight_transporter_engine=self.remote_instance_weight_transporter.engine,
remote_instance_weight_transporter_session_id=self.remote_instance_weight_transporter.session_id,
draft_model_idx=self.draft_model_idx,
@@ -1217,11 +1233,11 @@ class ModelRunner:
)
if self.device == "cpu":
self.model_config = adjust_config_with_unaligned_cpu_tp(
self.model_config, self.load_config, self.ps.tp_size
self.model_config, self.load_config, self.tp_size
)
maybe_trigger_remote_instance_nccl_send_group(
tp_rank=self.ps.tp_rank,
tp_rank=self.tp_rank,
load_format=draft_load_format,
)
@@ -1297,8 +1313,8 @@ class ModelRunner:
model=self.model,
spec_algorithm=self.spec_algorithm,
is_draft_worker=self.is_draft_worker,
tp_size=self.ps.tp_size,
tp_rank=self.ps.tp_rank,
tp_size=self.tp_size,
tp_rank=self.tp_rank,
pp_rank=get_parallel().pp_rank,
)
@@ -1316,7 +1332,7 @@ class ModelRunner:
if self.startup_weight_load is None:
dist_barrier_after_load(
elastic_ep_backend=get_exec().moe.elastic_ep_backend,
tp_rank=self.ps.tp_rank,
tp_rank=self.tp_rank,
is_ep_joiner=get_exec().moe.is_ep_joiner,
)
@@ -1339,7 +1355,7 @@ class ModelRunner:
self.startup_weight_load.finalize()
dist_barrier_after_load(
elastic_ep_backend=get_exec().moe.elastic_ep_backend,
tp_rank=self.ps.tp_rank,
tp_rank=self.tp_rank,
is_ep_joiner=get_exec().moe.is_ep_joiner,
)
self.startup_weight_load = None
@@ -1367,8 +1383,8 @@ class ModelRunner:
dtype=self.dtype,
server_args=self.server_args,
lora_backend=get_lora().lora_backend,
tp_size=self.ps.tp_size,
tp_rank=self.ps.tp_rank,
tp_size=self.tp_size,
tp_rank=self.tp_rank,
max_lora_rank=get_lora().max_lora_rank,
target_modules=get_lora().lora_target_modules,
lora_paths=get_lora().lora_paths,
@@ -1583,15 +1599,16 @@ class ModelRunner:
# rather than spawning additional processes, so dp_size must not be
# multiplied into the process count here (unlike regular DP, where
# dp_size * tp_size * pp_size is the true worker count).
dp_size = 1 if get_parallel().enable_dp_attention else self.ps.dp_size
parallel = get_parallel()
dp_size = 1 if parallel.enable_dp_attention else parallel.dp_size
self.local_omp_cpuid = numa_utils.init_threads_binding(
numa_index=self.gpu_id,
world_size=dp_size * self.ps.tp_size * get_parallel().pp_size,
world_size=dp_size * parallel.tp_size * parallel.pp_size,
)
def apply_torch_tp(self):
model_parallel.apply_torch_tp(
model=self.model, device=self.device, tp_size=self.ps.tp_size
model=self.model, device=self.device, tp_size=self.tp_size
)
def update_decode_attn_backend(self, stream_idx: int):
@@ -1717,7 +1734,7 @@ class ModelRunner:
if self.msprobe_debugger is not None:
rank_id = (
self.gpu_id
if self.ps.attn_dp_size is not None and self.ps.attn_dp_size > 1
if self.attn_dp_size is not None and self.attn_dp_size > 1
else None
)
self.msprobe_debugger.start(model=self.model, rank_id=rank_id)
@@ -2098,7 +2115,7 @@ class ModelRunner:
set_global_expert_location_metadata(new_metadata, allow_overwrite=True)
def _elastic_global_rank(self) -> int:
return self.ps.tp_rank + get_parallel().ep_join_rank_offset
return self.tp_rank + get_parallel().ep_join_rank_offset
def _rearm_eplb_after_elastic_scale(self) -> None:
if self.eplb_manager is None:
@@ -2120,7 +2137,7 @@ class ModelRunner:
self._rearm_eplb_after_elastic_scale()
def _report_elastic_scale_failure(self, error: str, effective_size: int) -> None:
if self.ps.tp_rank != 0 or get_exec().moe.is_ep_scale_joiner:
if self.tp_rank != 0 or get_exec().moe.is_ep_scale_joiner:
return
from sglang.srt.managers.io_struct import ElasticScaleUpdateReq
@@ -2131,7 +2148,7 @@ class ModelRunner:
)
def _elastic_scale_ready_barrier(self, target_size: int, log_tag: str) -> None:
if self.ps.tp_rank == 0:
if self.tp_rank == 0:
logger.debug(
"[Elastic EP][scale] %s entering post-scale WORLD barrier "
"(target_ep_size=%d)",
@@ -2139,7 +2156,7 @@ class ModelRunner:
target_size,
)
dist.barrier(group=dist.group.WORLD)
if self.ps.tp_rank == 0:
if self.tp_rank == 0:
logger.debug(
"[Elastic EP][scale] %s passed post-scale WORLD barrier "
"(target_ep_size=%d)",
@@ -2204,7 +2221,7 @@ class ModelRunner:
ElasticEPStateManager.commit_scale()
self._rearm_eplb_after_elastic_scale()
if self.ps.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
if self.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
from sglang.srt.managers.io_struct import ElasticScaleUpdateReq
self._pending_elastic_scale_update = ElasticScaleUpdateReq(
@@ -2237,7 +2254,7 @@ class ModelRunner:
)
ElasticEPStateManager.fail_recovery(error)
self._report_elastic_scale_failure(error, effective_size)
if self.ps.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
if self.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
logger.error("[Elastic EP] %s", error)
return
@@ -2263,7 +2280,7 @@ class ModelRunner:
ElasticEPStateManager.fail_scale(error)
self._reset_eplb_after_elastic_scale_failure()
self._report_elastic_scale_failure(error, effective_size)
if self.ps.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
if self.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
logger.error("[Elastic EP] %s", error)
return
@@ -2279,7 +2296,7 @@ class ModelRunner:
ElasticEPStateManager.fail_scale(error)
self._reset_eplb_after_elastic_scale_failure()
self._report_elastic_scale_failure(error, effective_size)
if self.ps.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
if self.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
logger.error("[Elastic EP] %s", error)
return
if not ElasticEPStateManager.begin_scale():
python/sglang/srt/utils/weight_checker.py
+25 -12
View File
@@ -8,6 +8,7 @@ import torch.distributed as dist
from pydantic import BaseModel, ConfigDict
from sglang.srt.managers.mm_utils import tensor_hash
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils.weight_checker_comparator import (
CHUNK_NUMEL,
ComparableWeight,
@@ -67,9 +68,23 @@ def _is_non_persistent_buffer_name(name: str) -> bool:
class WeightChecker:
def __init__(self, *, get_model: Callable[[], Any], ps: Any):
def __init__(self, *, get_model: Callable[[], Any]):
self._get_model = get_model
self._ps = ps
# A check is served on demand from the scheduler loop, which is outside
# the scope that describes a draft runner. The report has to name the
# runner it was built for, so the placement is read here, at
# construction, rather than asked for when the request arrives.
parallel = get_parallel()
self._placement = ParallelismInfo(
tp_rank=parallel.tp_rank,
tp_size=parallel.tp_size,
dp_rank=parallel.dp_rank if parallel.dp_rank is not None else 0,
dp_size=parallel.attn_dp_size,
pp_rank=parallel.pp_rank,
pp_size=parallel.pp_size,
rank=0,
size=1,
)
self._snapshot_tensors = None
def handle(self, action: str, allow_quant_error: bool = False) -> Optional[Dict]:
@@ -161,16 +176,14 @@ class WeightChecker:
return info.model_dump()
def _parallelism_info(self) -> ParallelismInfo:
ps = self._ps
return ParallelismInfo(
tp_rank=ps.tp_rank,
tp_size=ps.tp_size,
dp_rank=ps.dp_rank if ps.dp_rank is not None else 0,
dp_size=ps.attn_dp_size,
pp_rank=ps.pp_rank,
pp_size=ps.pp_size,
rank=dist.get_rank() if dist.is_initialized() else 0,
size=dist.get_world_size() if dist.is_initialized() else 1,
# The WORLD position is asked for now rather than frozen: unlike the
# runner's placement it is a property of the process, and an elastic
# scale-up moves it.
return self._placement.model_copy(
update={
"rank": dist.get_rank() if dist.is_initialized() else 0,
"size": dist.get_world_size() if dist.is_initialized() else 1,
}
)
def _model_state(self):
test/registered/layers/mamba/test_mamba2_mixer.py
+21 -2
View File
@@ -129,8 +129,27 @@ def mixer2_gated_norm_tensor_parallel(
)
mixer.weight.weight_loader(mixer.weight, weight)
# m2 reads tp via get_parallel().tp_size/rank — force it through the context.
with get_parallel().override(tp_size=1, tp_rank=0):
# m2 reads tp via get_parallel().tp_size/rank — state a single-rank topology
# through the context. Every width that follows from `tp_size` is named:
# narrowing one leaf and leaving the quotients behind describes no layout.
with get_parallel().override(
tp_size=1,
tp_rank=0,
tp_group=None,
attn_tp_size=1,
attn_tp_rank=0,
attn_tp_group=None,
attn_dp_size=1,
attn_dp_rank=0,
attn_cp_size=1,
attn_cp_rank=0,
attn_cp_group=None,
moe_ep_size=1,
moe_ep_rank=0,
moe_ep_group=None,
moe_dp_size=1,
moe_tp_size=1,
):
# create gated-norm without TP to compute reference
mixer_single_gpu = m2.Mixer2RMSNormGated(
full_hidden_size=hidden_size,
test/registered/unit/hardware_backend/mlx/test_max_running_requests.py
+2 -2
View File
@@ -84,7 +84,7 @@ def _stub(
stub._max_mamba_cache_size = max_mamba_cache_size
stub._disable_radix_cache = disable_radix_cache
stub.max_total_num_tokens = max_total_num_tokens
stub.ps = SimpleNamespace(attn_dp_size=dp_size)
stub.attn_dp_size = dp_size
return stub
@@ -94,7 +94,7 @@ def _hybrid_stub_for_initialize(
"""A stub carrying what the real initialize() reads (hybrid path)."""
stub = MlxModelRunnerStub.__new__(MlxModelRunnerStub)
stub._mlx_pool_size = pool
stub.ps = SimpleNamespace(attn_dp_size=1)
stub.attn_dp_size = 1
stub.device = "cpu" # read by init_ngram_embedding_manager
# Evaluated as a call argument in init_ngram_embedding_manager before
# the use_ngram_embedding short-circuit; never read.
test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py
+2 -1
View File
@@ -91,7 +91,8 @@ def _make_prefill_aware_swa_runner(
kv_cache_dtype=torch.float16,
kv_cache_dtype_str="auto",
page_size=1,
ps=SimpleNamespace(attn_cp_size=1, tp_size=1),
attn_cp_size=1,
tp_size=1,
is_draft_worker=False,
server_args=server_args,
attention_chunk_size=None,
test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py
+1 -1
View File
@@ -671,7 +671,7 @@ class TestModelRunnerStartupWeightLoadOwnership(CustomTestCase):
elastic_ep_backend=None,
is_ep_joiner=False,
)
runner.ps = SimpleNamespace(tp_rank=0)
runner.tp_rank = 0
return runner
def test_start_delegates_to_the_manager(self):
test/registered/unit/spec/test_dflash_logits.py
+2 -2
View File
@@ -248,7 +248,7 @@ def test_worker_folds_a_gate_admitted_quantized_selector_head(monkeypatch):
worker = SimpleNamespace(
block_size=8,
selector=object(),
ps=SimpleNamespace(tp_rank=0),
model_runner=SimpleNamespace(tp_rank=0),
draft_model=SimpleNamespace(lm_head=None),
device="cpu",
_selector_sampling_enabled=True,
@@ -282,7 +282,7 @@ def test_worker_warns_once_when_selector_sampling_is_disabled(monkeypatch):
selector=object(),
_selector_sampling_enabled=False,
_warned_sampling_fallback=False,
ps=SimpleNamespace(tp_rank=0),
model_runner=SimpleNamespace(tp_rank=0),
)
batch = SimpleNamespace(sampling_info=SimpleNamespace(is_all_greedy=False))
test/registered/unit/test_runtime_context.py
+142 -5
View File
@@ -2781,17 +2781,154 @@ class TestWhoAnswersDuringADraftScope(CustomTestCase):
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
report has to name the runner it was built for, which is why it reads
the placement once, where it is built, instead of asking again when the
request arrives."""
from sglang.srt.distributed import parallel_state
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()
self.assertEqual(get_parallel().pp_size, 2)
group = self._single_member_group()
with patch.object(parallel_state, "_PP", group):
with parallel_state.patch_pipeline_parallel_group(group):
checker = WeightChecker(get_model=lambda: None)
# The scope has closed and the context answers the target's shape again.
self.assertEqual(get_parallel().pp_size, 2)
info = checker._parallelism_info()
self.assertEqual((info.pp_rank, info.pp_size), (0, 1))
class TestNothingReadsThePlacementBeforeItIsFrozen(CustomTestCase):
"""`ModelRunner.__init__` freezes its placement partway through.
A method called before that point reads an attribute that does not exist
yet, and only on the configuration that reaches it -- a remote weight
transporter, a NUMA binding -- so the suites say nothing and a GPU job is
where it surfaces. The order is what makes it wrong, so the order is what
is checked.
"""
def _model_runner(self):
import ast as _ast
source = (_SRT / "model_executor" / "model_runner.py").read_text()
for node in _ast.parse(source).body:
if isinstance(node, _ast.ClassDef) and node.name == "ModelRunner":
return {m.name: m for m in node.body if isinstance(m, _ast.FunctionDef)}
raise AssertionError("ModelRunner not found")
def _frozen_names(self, methods):
import ast as _ast
return {
target.attr
for node in _ast.walk(methods["init_torch_distributed"])
if isinstance(node, _ast.Assign)
for target in node.targets
if isinstance(target, _ast.Attribute)
and isinstance(target.value, _ast.Name)
and target.value.id == "self"
}
def _reads(self, methods, fn, frozen, depth=0):
import ast as _ast
found = set()
for node in _ast.walk(fn):
if (
isinstance(node, _ast.Attribute)
and isinstance(node.value, _ast.Name)
and node.value.id == "self"
and isinstance(node.ctx, _ast.Load)
and node.attr in frozen
):
found.add(node.attr)
if (
depth < 2
and isinstance(node, _ast.Call)
and isinstance(node.func, _ast.Attribute)
and isinstance(node.func.value, _ast.Name)
and node.func.value.id == "self"
and node.func.attr in methods
and node.func.attr != "init_torch_distributed"
):
found |= self._reads(
methods, methods[node.func.attr], frozen, depth + 1
)
return found
def _calls_before_the_freeze(self, methods):
import ast as _ast
calls = []
for statement in methods["__init__"].body:
for node in _ast.walk(statement):
if (
isinstance(node, _ast.Call)
and isinstance(node.func, _ast.Attribute)
and isinstance(node.func.value, _ast.Name)
and node.func.value.id == "self"
):
calls.append((node.lineno, node.func.attr))
calls.sort()
names = [name for _, name in calls]
self.assertIn(
"init_torch_distributed",
names,
"the freeze moved; this census is keyed on where it happens",
)
return calls[: names.index("init_torch_distributed")]
def test_no_method_called_before_the_freeze_reads_what_it_freezes(self):
methods = self._model_runner()
frozen = self._frozen_names(methods)
self.assertGreater(len(frozen), 5, "found no frozen names; census is broken")
offenders = []
for lineno, name in self._calls_before_the_freeze(methods):
fn = methods.get(name)
if fn is None:
continue
read = self._reads(methods, fn, frozen)
if read:
offenders.append(
f"__init__:{lineno} self.{name}() reads {sorted(read)}"
)
self.assertEqual(
offenders,
[],
"these run before init_torch_distributed and read what it sets; "
"ask get_parallel() there, or move the call after the freeze:\n "
+ "\n ".join(offenders),
)
def test_the_census_would_notice_one(self):
"""The positive control: the walk has to find a read that is there."""
import ast as _ast
import textwrap
methods = {
m.name: m
for m in _ast.parse(
textwrap.dedent(
"""
class R:
def init_torch_distributed(self):
self.tp_rank = 0
def early(self):
return self.tp_rank
"""
)
)
.body[0]
.body
}
frozen = self._frozen_names(methods)
self.assertEqual(frozen, {"tp_rank"})
self.assertEqual(self._reads(methods, methods["early"], frozen), {"tp_rank"})
if __name__ == "__main__":
unittest.main()
test/registered/unit/utils/test_weight_checker.py
+12 -4
View File
@@ -44,7 +44,7 @@ from sglang.srt.utils.weight_checker_comparator import (
RawComparable,
)
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import CustomTestCase
from sglang.test.test_utils import CustomTestCase, enter_scope, published_topology
register_amd_ci(est_time=30, suite="stage-b-test-1-gpu-small-amd")
register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
@@ -551,7 +551,8 @@ class _WeightCheckerTestBase(CustomTestCase):
torch.manual_seed(0)
self.model = _TinyModel().cuda()
runner = _FakeModelRunner(self.model)
self.checker = WeightChecker(get_model=lambda: runner.model, ps=runner.ps)
enter_scope(self, published_topology())
self.checker = WeightChecker(get_model=lambda: runner.model)
class TestSnapshot(_WeightCheckerTestBase):
@@ -760,9 +761,16 @@ class _ChecksumTestBase(CustomTestCase):
pp_rank=0,
pp_size=1,
)
self.checker = WeightChecker(
get_model=lambda: self.runner.model, ps=self.runner.ps
enter_scope(
self,
published_topology(
tp_size=4,
dp_size=2,
enable_dp_attention=True,
ranks={"world_rank": 2, "dp_rank": 1},
),
)
self.checker = WeightChecker(get_model=lambda: self.runner.model)
class TestComputeChecksum(_ChecksumTestBase):