minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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 import time
from typing import List, Optional from typing import List, Optional
import msgspec
import torch import torch
import torch.distributed as dist import torch.distributed as dist
from sglang.srt.configs.model_config import ModelConfig from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.distributed import ( from sglang.srt.distributed import (
get_default_distributed_backend, get_default_distributed_backend,
get_pp_group,
get_tp_group, get_tp_group,
get_world_group, get_world_group,
init_distributed_environment, init_distributed_environment,
@@ -60,13 +58,6 @@ _is_cpu_arm64 = is_host_cpu_arm64()
_TP_ALL_TO_ALL_WARMUP_BYTES_PER_PEER = 4 << 20 _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( def init_torch_distributed(
*, *,
server_args: ServerArgs, server_args: ServerArgs,
@@ -140,10 +131,6 @@ def init_torch_distributed(
distributed=get_world_group().world_size > 1 and not is_draft_worker, distributed=get_world_group().world_size > 1 and not is_draft_worker,
cpu_group=get_world_group().cpu_group, 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 # Check memory for tensor parallelism
local_gpu_memory = get_available_gpu_memory(device, ps.gpu_id) local_gpu_memory = get_available_gpu_memory(device, ps.gpu_id)
if ps.tp_size > 1 and not is_draft_worker: 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"Init torch distributed ends. elapsed={time.perf_counter() - tic:.2f} s, "
f"mem usage={(before_avail_memory - local_gpu_memory):.2f} GB" f"mem usage={(before_avail_memory - local_gpu_memory):.2f} GB"
) )
return TorchDistributedResult( return pre_model_load_memory
tp_group=tp_group,
pp_group=pp_group,
attention_tp_group=attention_tp_group,
pre_model_load_memory=pre_model_load_memory,
)
def _resolve_backend(*, device: str) -> str: 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_tp_rank == 0
and get_parallel().attn_cp_rank == 0 and get_parallel().attn_cp_rank == 0
), ),
ps=self.ps,
tp_group=self.tp_group, tp_group=self.tp_group,
pp_group=self.pp_group, pp_group=self.pp_group,
enable_hierarchical_cache=self.enable_hierarchical_cache, 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.configs.model_config import ModelConfig
from sglang.srt.distributed.parallel_state import GroupCoordinator 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.managers.tp_worker import BaseTpWorker
from sglang.srt.server_args import ServerArgs from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.base_spec_worker import HiCacheDraftPlan from sglang.srt.speculative.base_spec_worker import HiCacheDraftPlan
@@ -263,12 +262,14 @@ def build_kv_cache(
attn_cp_cpu_group: ProcessGroup, attn_cp_cpu_group: ProcessGroup,
enable_metrics: bool, enable_metrics: bool,
enable_kv_cache_events: bool, enable_kv_cache_events: bool,
ps: ParallelState,
tp_group: GroupCoordinator, tp_group: GroupCoordinator,
pp_group: GroupCoordinator, pp_group: GroupCoordinator,
enable_hierarchical_cache: bool, enable_hierarchical_cache: bool,
hicache_draft_plan: Optional[HiCacheDraftPlan] = None, hicache_draft_plan: Optional[HiCacheDraftPlan] = None,
) -> KVCacheBuildResult: ) -> 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 sliding_window_size: Optional[int] = None
full_tokens_per_layer: Optional[int] = None full_tokens_per_layer: Optional[int] = None
swa_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_session_radix_cache=get_memory().enable_session_radix_cache,
enable_mamba_extra_buffer=get_exec().mamba.enable_mamba_extra_buffer, enable_mamba_extra_buffer=get_exec().mamba.enable_mamba_extra_buffer,
enable_mamba_extra_buffer_lazy=get_exec().mamba.enable_mamba_extra_buffer_lazy, enable_mamba_extra_buffer_lazy=get_exec().mamba.enable_mamba_extra_buffer_lazy,
pp_rank=ps.pp_rank, pp_rank=parallel.pp_rank,
pp_size=ps.pp_size, pp_size=parallel.pp_size,
attn_cp_rank=ps.attn_cp_rank, attn_cp_rank=parallel.attn_cp_rank,
attn_cp_size=ps.attn_cp_size, attn_cp_size=parallel.attn_cp_size,
chunked_prefill_size=effective_chunked_prefill_size, chunked_prefill_size=effective_chunked_prefill_size,
sliding_window_size=sliding_window_size, sliding_window_size=sliding_window_size,
mtp_draft_device_pools=mtp_draft_device_pools, mtp_draft_device_pools=mtp_draft_device_pools,
@@ -390,8 +391,8 @@ def build_kv_cache(
effective_chunked_prefill_size=effective_chunked_prefill_size, effective_chunked_prefill_size=effective_chunked_prefill_size,
tp_worker=tp_worker, tp_worker=tp_worker,
model_config=model_config, model_config=model_config,
tp_size=ps.tp_size, tp_size=parallel.tp_size,
tp_rank=ps.tp_rank, tp_rank=parallel.tp_rank,
tp_group=tp_group, tp_group=tp_group,
) )
with auto_size_hicache( 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: if TYPE_CHECKING:
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.mem_cache.unified_memory_pool import ( from sglang.srt.mem_cache.unified_memory_pool import (
UnifiedKVPool, UnifiedKVPool,
UnifiedPoolBundle, UnifiedPoolBundle,
@@ -260,7 +259,10 @@ class _PoolSizes(msgspec.Struct, frozen=True, kw_only=True):
class KVCacheConfigurator: class KVCacheConfigurator:
device: str device: str
gpu_id: int 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 pp_group: Any
model: Any model: Any
model_config: ModelConfig model_config: ModelConfig
@@ -2324,7 +2326,7 @@ class KVCacheConfigurator:
max_num_reqs = get_schedule().max_running_requests max_num_reqs = get_schedule().max_running_requests
if max_num_reqs is not None: 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) max_num_reqs = min(requested_per_worker, token_capacity // 2)
else: else:
requested_per_worker = None requested_per_worker = None
@@ -2433,16 +2435,16 @@ class KVCacheConfigurator:
# allocates its own [start_layer, end_layer) slice. Charge the largest # allocates its own [start_layer, end_layer) slice. Charge the largest
# per-stage share so every rank derives the same pool without a collective. # per-stage share so every rank derives the same pool without a collective.
all_mamba_layers = config.mamba2_cache_params.layers 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( max_stage_mamba_layers = max(
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, self.model_config.num_hidden_layers,
rank, rank,
self.ps.pp_size, self.pp_size,
) )
for rank in range(self.ps.pp_size) for rank in range(self.pp_size)
) )
) )
else: else:
@@ -2474,7 +2476,7 @@ class KVCacheConfigurator:
replayssm_ring_per_req = int(replayssm_ring_per_req * pp_layer_scale) replayssm_ring_per_req = int(replayssm_ring_per_req * pp_layer_scale)
if replayssm_active and self.hybrid_kda_config is None: if replayssm_active and self.hybrid_kda_config is None:
replay_req_slots = ( 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_fixed_bytes = replayssm_ring_per_req * replay_req_slots
replayssm_ring_per_slot = 0 replayssm_ring_per_slot = 0
@@ -2490,7 +2492,7 @@ class KVCacheConfigurator:
get_context().override( get_context().override(
"mamba_pool.per_dp_shard", "mamba_pool.per_dp_shard",
max_mamba_cache_size=get_schedule().max_mamba_cache_size 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 # Reserve intermediate memory based on capped max_num_reqs (+1: the
# pool's padding slot, see memory_pool.py). Skipped under replayssm # 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: if has_spec_dec and not replayssm_active:
ratio = self._calculate_mamba_ratio() ratio = self._calculate_mamba_ratio()
capped_reqs = min( 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, get_schedule().max_mamba_cache_size // ratio,
) )
intermediate_size = ( intermediate_size = (
@@ -2515,7 +2517,7 @@ class KVCacheConfigurator:
get_context().override( get_context().override(
"mamba_pool.from_max_running_requests", "mamba_pool.from_max_running_requests",
max_mamba_cache_size=get_schedule().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 # Reserve intermediate memory based on capped max_num_reqs (+1: the
# pool's padding slot). Skipped under replayssm. # pool's padding slot). Skipped under replayssm.
@@ -2555,7 +2557,7 @@ class KVCacheConfigurator:
# Intermediate memory is included in mamba_budget, subtract it # Intermediate memory is included in mamba_budget, subtract it
# so the return value only has main_state subtracted from total # so the return value only has main_state subtracted from total
capped_reqs = min( 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, get_schedule().max_mamba_cache_size // ratio,
) )
intermediate_size = per_req * (capped_reqs + 1) * D 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 import os
logger.warning( 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 raise
@@ -447,7 +447,7 @@ class ModelRunner:
# CPU offload # CPU offload
set_offloader(create_offloader(dp_rank=get_parallel().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)
if envs.SGLANG_DETECT_SLOW_RANK.get(): if envs.SGLANG_DETECT_SLOW_RANK.get():
slow_rank_detector.execute() slow_rank_detector.execute()
@@ -503,9 +503,9 @@ class ModelRunner:
): ):
return 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) dist.barrier(group=self.tp_group.cpu_group)
if self.ps.tp_rank == 0: if self.tp_rank == 0:
register_scale_cohort( register_scale_cohort(
get_parallel().ep_join_rank_offset, get_parallel().ep_join_rank_offset,
join_effective_ep_size, join_effective_ep_size,
@@ -513,7 +513,7 @@ class ModelRunner:
join_scale_process_group() join_scale_process_group()
get_context().override("elastic_ep.scale_join", ep_size=join_effective_ep_size) 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( broadcast_global_expert_location_metadata(
model_config=self.model_config, model_config=self.model_config,
moe_ep_rank=global_ep_rank, moe_ep_rank=global_ep_rank,
@@ -563,7 +563,7 @@ class ModelRunner:
def init_weight_updater(self): def init_weight_updater(self):
self.weight_updater = WeightUpdater( self.weight_updater = WeightUpdater(
tp_rank=self.ps.tp_rank, tp_rank=self.tp_rank,
device=self.device, device=self.device,
gpu_id=self.gpu_id, gpu_id=self.gpu_id,
model_config=self.model_config, model_config=self.model_config,
@@ -586,8 +586,8 @@ class ModelRunner:
def init_weight_exporter(self): def init_weight_exporter(self):
self.weight_exporter = WeightExporter( self.weight_exporter = WeightExporter(
tp_rank=self.ps.tp_rank, tp_rank=self.tp_rank,
tp_size=self.ps.tp_size, tp_size=self.tp_size,
gpu_id=self.gpu_id, gpu_id=self.gpu_id,
get_model_path=lambda: self.model_config.model_path, get_model_path=lambda: self.model_config.model_path,
get_model=lambda: self.model, get_model=lambda: self.model,
@@ -596,7 +596,7 @@ class ModelRunner:
def init_remote_instance_weight_transporter(self): def init_remote_instance_weight_transporter(self):
self.remote_instance_weight_transporter = RemoteInstanceWeightTransporter( self.remote_instance_weight_transporter = RemoteInstanceWeightTransporter(
get_model=lambda: self.model, get_model=lambda: self.model,
tp_rank=self.ps.tp_rank, tp_rank=get_parallel().tp_rank,
gpu_id=self.gpu_id, gpu_id=self.gpu_id,
) )
@@ -610,10 +610,13 @@ class ModelRunner:
) )
def init_kv_cache_configurator(self): 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( self.kv_cache_configurator = KVCacheConfigurator(
device=self.device, device=self.device,
gpu_id=self.gpu_id, 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, pp_group=self.pp_group,
model=self.model, model=self.model,
model_config=self.model_config, model_config=self.model_config,
@@ -647,8 +650,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 * get_parallel().pp_size, world_size=self.tp_size * get_parallel().pp_size,
rank=self.ps.tp_size * get_parallel().pp_rank + self.ps.tp_rank, rank=self.tp_size * get_parallel().pp_rank + self.tp_rank,
local_rank=self.gpu_id, local_rank=self.gpu_id,
port=self.dist_port, port=self.dist_port,
) )
@@ -718,7 +721,7 @@ class ModelRunner:
moe_ep_rank=expert_rank, 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( logger.info(
"Initial expert_location_metadata:\n%s", "Initial expert_location_metadata:\n%s",
format_expert_location_layout(get_global_expert_location_metadata()), format_expert_location_layout(get_global_expert_location_metadata()),
@@ -779,7 +782,7 @@ class ModelRunner:
def maybe_apply_post_load_model_transforms(self): def maybe_apply_post_load_model_transforms(self):
supports_torch_tp = getattr(self.model, "supports_torch_tp", False) 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() self.apply_torch_tp()
def maybe_init_lora_manager(self): def maybe_init_lora_manager(self):
@@ -1015,7 +1018,7 @@ class ModelRunner:
def post_capture_elastic_ep_recover(self): def post_capture_elastic_ep_recover(self):
join_process_groups() 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( broadcast_global_expert_location_metadata(
model_config=self.model_config, model_config=self.model_config,
moe_ep_rank=global_ep_rank, moe_ep_rank=global_ep_rank,
@@ -1162,13 +1165,13 @@ class ModelRunner:
def check_quantized_moe_compatibility(self): def check_quantized_moe_compatibility(self):
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.tp_size,
moe_ep_size=get_parallel().moe_ep_size, moe_ep_size=get_parallel().moe_ep_size,
moe_dp_size=get_parallel().moe_dp_size, moe_dp_size=get_parallel().moe_dp_size,
) )
def init_torch_distributed(self): def init_torch_distributed(self):
result = bootstrap.init_torch_distributed( self.pre_model_load_memory = bootstrap.init_torch_distributed(
server_args=self.server_args, server_args=self.server_args,
model_config=self.model_config, model_config=self.model_config,
device=self.device, device=self.device,
@@ -1177,10 +1180,23 @@ class ModelRunner:
is_draft_worker=self.is_draft_worker, is_draft_worker=self.is_draft_worker,
local_omp_cpuid=self.local_omp_cpuid if self.device == "cpu" else None, local_omp_cpuid=self.local_omp_cpuid if self.device == "cpu" else None,
) )
self.tp_group = result.tp_group # Read once, here: a draft runner is constructed inside the scope that
self.pp_group = result.pp_group # states its topology and used outside it, so what it holds has to be
self.attention_tp_group = result.attention_tp_group # the placement it was built for rather than whatever the context
self.pre_model_load_memory = result.pre_model_load_memory # 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): def init_shared_mooncake_transfer_engine(self):
maybe_init_shared_mooncake_transfer_engine(gpu_id=self.gpu_id) maybe_init_shared_mooncake_transfer_engine(gpu_id=self.gpu_id)
@@ -1204,7 +1220,7 @@ class ModelRunner:
self.load_config = build_load_config( self.load_config = build_load_config(
server_args=self.server_args, server_args=self.server_args,
load_format=draft_load_format, 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_engine=self.remote_instance_weight_transporter.engine,
remote_instance_weight_transporter_session_id=self.remote_instance_weight_transporter.session_id, remote_instance_weight_transporter_session_id=self.remote_instance_weight_transporter.session_id,
draft_model_idx=self.draft_model_idx, draft_model_idx=self.draft_model_idx,
@@ -1217,11 +1233,11 @@ class ModelRunner:
) )
if self.device == "cpu": if self.device == "cpu":
self.model_config = adjust_config_with_unaligned_cpu_tp( 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( maybe_trigger_remote_instance_nccl_send_group(
tp_rank=self.ps.tp_rank, tp_rank=self.tp_rank,
load_format=draft_load_format, load_format=draft_load_format,
) )
@@ -1297,8 +1313,8 @@ class ModelRunner:
model=self.model, model=self.model,
spec_algorithm=self.spec_algorithm, spec_algorithm=self.spec_algorithm,
is_draft_worker=self.is_draft_worker, is_draft_worker=self.is_draft_worker,
tp_size=self.ps.tp_size, tp_size=self.tp_size,
tp_rank=self.ps.tp_rank, tp_rank=self.tp_rank,
pp_rank=get_parallel().pp_rank, pp_rank=get_parallel().pp_rank,
) )
@@ -1316,7 +1332,7 @@ class ModelRunner:
if self.startup_weight_load is None: if self.startup_weight_load is None:
dist_barrier_after_load( dist_barrier_after_load(
elastic_ep_backend=get_exec().moe.elastic_ep_backend, 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, is_ep_joiner=get_exec().moe.is_ep_joiner,
) )
@@ -1339,7 +1355,7 @@ class ModelRunner:
self.startup_weight_load.finalize() self.startup_weight_load.finalize()
dist_barrier_after_load( dist_barrier_after_load(
elastic_ep_backend=get_exec().moe.elastic_ep_backend, 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, is_ep_joiner=get_exec().moe.is_ep_joiner,
) )
self.startup_weight_load = None self.startup_weight_load = None
@@ -1367,8 +1383,8 @@ class ModelRunner:
dtype=self.dtype, dtype=self.dtype,
server_args=self.server_args, server_args=self.server_args,
lora_backend=get_lora().lora_backend, lora_backend=get_lora().lora_backend,
tp_size=self.ps.tp_size, tp_size=self.tp_size,
tp_rank=self.ps.tp_rank, tp_rank=self.tp_rank,
max_lora_rank=get_lora().max_lora_rank, max_lora_rank=get_lora().max_lora_rank,
target_modules=get_lora().lora_target_modules, target_modules=get_lora().lora_target_modules,
lora_paths=get_lora().lora_paths, lora_paths=get_lora().lora_paths,
@@ -1583,15 +1599,16 @@ class ModelRunner:
# rather than spawning additional processes, so dp_size must not be # rather than spawning additional processes, so dp_size must not be
# multiplied into the process count here (unlike regular DP, where # multiplied into the process count here (unlike regular DP, where
# dp_size * tp_size * pp_size is the true worker count). # 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( 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 * get_parallel().pp_size, world_size=dp_size * parallel.tp_size * parallel.pp_size,
) )
def apply_torch_tp(self): def apply_torch_tp(self):
model_parallel.apply_torch_tp( 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): def update_decode_attn_backend(self, stream_idx: int):
@@ -1717,7 +1734,7 @@ class ModelRunner:
if self.msprobe_debugger is not None: if self.msprobe_debugger is not None:
rank_id = ( rank_id = (
self.gpu_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 else None
) )
self.msprobe_debugger.start(model=self.model, rank_id=rank_id) 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) set_global_expert_location_metadata(new_metadata, allow_overwrite=True)
def _elastic_global_rank(self) -> int: 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: def _rearm_eplb_after_elastic_scale(self) -> None:
if self.eplb_manager is None: if self.eplb_manager is None:
@@ -2120,7 +2137,7 @@ class ModelRunner:
self._rearm_eplb_after_elastic_scale() self._rearm_eplb_after_elastic_scale()
def _report_elastic_scale_failure(self, error: str, effective_size: int) -> None: 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 return
from sglang.srt.managers.io_struct import ElasticScaleUpdateReq 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: 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( logger.debug(
"[Elastic EP][scale] %s entering post-scale WORLD barrier " "[Elastic EP][scale] %s entering post-scale WORLD barrier "
"(target_ep_size=%d)", "(target_ep_size=%d)",
@@ -2139,7 +2156,7 @@ class ModelRunner:
target_size, target_size,
) )
dist.barrier(group=dist.group.WORLD) dist.barrier(group=dist.group.WORLD)
if self.ps.tp_rank == 0: if self.tp_rank == 0:
logger.debug( logger.debug(
"[Elastic EP][scale] %s passed post-scale WORLD barrier " "[Elastic EP][scale] %s passed post-scale WORLD barrier "
"(target_ep_size=%d)", "(target_ep_size=%d)",
@@ -2204,7 +2221,7 @@ class ModelRunner:
ElasticEPStateManager.commit_scale() ElasticEPStateManager.commit_scale()
self._rearm_eplb_after_elastic_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 from sglang.srt.managers.io_struct import ElasticScaleUpdateReq
self._pending_elastic_scale_update = ElasticScaleUpdateReq( self._pending_elastic_scale_update = ElasticScaleUpdateReq(
@@ -2237,7 +2254,7 @@ class ModelRunner:
) )
ElasticEPStateManager.fail_recovery(error) ElasticEPStateManager.fail_recovery(error)
self._report_elastic_scale_failure(error, effective_size) 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) logger.error("[Elastic EP] %s", error)
return return
@@ -2263,7 +2280,7 @@ class ModelRunner:
ElasticEPStateManager.fail_scale(error) ElasticEPStateManager.fail_scale(error)
self._reset_eplb_after_elastic_scale_failure() self._reset_eplb_after_elastic_scale_failure()
self._report_elastic_scale_failure(error, effective_size) 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) logger.error("[Elastic EP] %s", error)
return return
@@ -2279,7 +2296,7 @@ class ModelRunner:
ElasticEPStateManager.fail_scale(error) ElasticEPStateManager.fail_scale(error)
self._reset_eplb_after_elastic_scale_failure() self._reset_eplb_after_elastic_scale_failure()
self._report_elastic_scale_failure(error, effective_size) 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) logger.error("[Elastic EP] %s", error)
return return
if not ElasticEPStateManager.begin_scale(): 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 pydantic import BaseModel, ConfigDict
from sglang.srt.managers.mm_utils import tensor_hash 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 ( from sglang.srt.utils.weight_checker_comparator import (
CHUNK_NUMEL, CHUNK_NUMEL,
ComparableWeight, ComparableWeight,
@@ -67,9 +68,23 @@ def _is_non_persistent_buffer_name(name: str) -> bool:
class WeightChecker: class WeightChecker:
def __init__(self, *, get_model: Callable[[], Any], ps: Any): def __init__(self, *, get_model: Callable[[], Any]):
self._get_model = get_model 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 self._snapshot_tensors = None
def handle(self, action: str, allow_quant_error: bool = False) -> Optional[Dict]: def handle(self, action: str, allow_quant_error: bool = False) -> Optional[Dict]:
@@ -161,16 +176,14 @@ class WeightChecker:
return info.model_dump() return info.model_dump()
def _parallelism_info(self) -> ParallelismInfo: def _parallelism_info(self) -> ParallelismInfo:
ps = self._ps # The WORLD position is asked for now rather than frozen: unlike the
return ParallelismInfo( # runner's placement it is a property of the process, and an elastic
tp_rank=ps.tp_rank, # scale-up moves it.
tp_size=ps.tp_size, return self._placement.model_copy(
dp_rank=ps.dp_rank if ps.dp_rank is not None else 0, update={
dp_size=ps.attn_dp_size, "rank": dist.get_rank() if dist.is_initialized() else 0,
pp_rank=ps.pp_rank, "size": dist.get_world_size() if dist.is_initialized() else 1,
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,
) )
def _model_state(self): 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) mixer.weight.weight_loader(mixer.weight, weight)
# m2 reads tp via get_parallel().tp_size/rank — force it through the context. # m2 reads tp via get_parallel().tp_size/rank — state a single-rank topology
with get_parallel().override(tp_size=1, tp_rank=0): # 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 # create gated-norm without TP to compute reference
mixer_single_gpu = m2.Mixer2RMSNormGated( mixer_single_gpu = m2.Mixer2RMSNormGated(
full_hidden_size=hidden_size, 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._max_mamba_cache_size = max_mamba_cache_size
stub._disable_radix_cache = disable_radix_cache stub._disable_radix_cache = disable_radix_cache
stub.max_total_num_tokens = max_total_num_tokens 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 return stub
@@ -94,7 +94,7 @@ def _hybrid_stub_for_initialize(
"""A stub carrying what the real initialize() reads (hybrid path).""" """A stub carrying what the real initialize() reads (hybrid path)."""
stub = MlxModelRunnerStub.__new__(MlxModelRunnerStub) stub = MlxModelRunnerStub.__new__(MlxModelRunnerStub)
stub._mlx_pool_size = pool 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 stub.device = "cpu" # read by init_ngram_embedding_manager
# Evaluated as a call argument in init_ngram_embedding_manager before # Evaluated as a call argument in init_ngram_embedding_manager before
# the use_ngram_embedding short-circuit; never read. # 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=torch.float16,
kv_cache_dtype_str="auto", kv_cache_dtype_str="auto",
page_size=1, page_size=1,
ps=SimpleNamespace(attn_cp_size=1, tp_size=1), attn_cp_size=1,
tp_size=1,
is_draft_worker=False, is_draft_worker=False,
server_args=server_args, server_args=server_args,
attention_chunk_size=None, 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, elastic_ep_backend=None,
is_ep_joiner=False, is_ep_joiner=False,
) )
runner.ps = SimpleNamespace(tp_rank=0) runner.tp_rank = 0
return runner return runner
def test_start_delegates_to_the_manager(self): 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( worker = SimpleNamespace(
block_size=8, block_size=8,
selector=object(), selector=object(),
ps=SimpleNamespace(tp_rank=0), model_runner=SimpleNamespace(tp_rank=0),
draft_model=SimpleNamespace(lm_head=None), draft_model=SimpleNamespace(lm_head=None),
device="cpu", device="cpu",
_selector_sampling_enabled=True, _selector_sampling_enabled=True,
@@ -282,7 +282,7 @@ def test_worker_warns_once_when_selector_sampling_is_disabled(monkeypatch):
selector=object(), selector=object(),
_selector_sampling_enabled=False, _selector_sampling_enabled=False,
_warned_sampling_fallback=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)) 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): def test_a_report_built_for_a_runner_follows_that_runner(self):
"""A weight check is an on-demand request served from the scheduler """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 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 report has to name the runner it was built for, which is why it reads
a record instead of asking the context.""" the placement once, where it is built, instead of asking again when the
from sglang.srt.distributed.parallel_state_wrapper import ParallelState request arrives."""
from sglang.srt.distributed import parallel_state
from sglang.srt.utils.weight_checker import WeightChecker 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._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() info = checker._parallelism_info()
self.assertEqual((info.pp_rank, info.pp_size), (0, 1)) 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__": if __name__ == "__main__":
unittest.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, RawComparable,
) )
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci 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_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") 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) torch.manual_seed(0)
self.model = _TinyModel().cuda() self.model = _TinyModel().cuda()
runner = _FakeModelRunner(self.model) 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): class TestSnapshot(_WeightCheckerTestBase):
@@ -760,9 +761,16 @@ class _ChecksumTestBase(CustomTestCase):
pp_rank=0, pp_rank=0,
pp_size=1, pp_size=1,
) )
self.checker = WeightChecker( enter_scope(
get_model=lambda: self.runner.model, ps=self.runner.ps 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): class TestComputeChecksum(_ChecksumTestBase):