config: a parallel size has one spelling; a patched scope declares its own (#36621)
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
ca1d7ed8e6
commit
fd40a331bf
@@ -899,9 +899,12 @@ def latency_test(
|
||||
initialize_fp4_gemm_config()
|
||||
|
||||
if get_bool_env_var("SGLANG_SET_CPU_AFFINITY"):
|
||||
parallel = get_parallel().config
|
||||
parallel = get_parallel()
|
||||
set_gpu_proc_affinity(
|
||||
parallel.pp_size, parallel.tp_size, parallel.nnodes, tp_rank
|
||||
parallel.pp_size,
|
||||
parallel.tp_size,
|
||||
parallel.nnodes,
|
||||
tp_rank,
|
||||
)
|
||||
|
||||
# Configure the logger
|
||||
|
||||
@@ -81,7 +81,7 @@ async def warm_up_compile(
|
||||
)
|
||||
generate_req_input.bootstrap_host = [FAKE_BOOTSTRAP_HOST] * dp_size
|
||||
generate_req_input.bootstrap_room = [
|
||||
i * (2**63 // dp_size) + (i % get_parallel().config.tp_size)
|
||||
i * (2**63 // dp_size) + (i % get_parallel().tp_size)
|
||||
for i in range(dp_size)
|
||||
]
|
||||
else:
|
||||
|
||||
@@ -701,7 +701,16 @@ def model_parallel_is_initialized() -> bool:
|
||||
|
||||
@contextmanager
|
||||
def use_tensor_parallel_group(tp_group: GroupCoordinator):
|
||||
"""Use one TP group consistently across diffusion and reused SRT modules."""
|
||||
"""Use one TP group consistently across diffusion and reused SRT modules.
|
||||
|
||||
The scope replaces the module globals that ``get_tp_group()`` and srt's
|
||||
``get_tp_group()`` / ``get_attention_tp_group()`` read, and — like srt's
|
||||
``patch_tensor_parallel_group`` — the three members the runtime context
|
||||
answers with, so that a size read from the published bag cannot disagree
|
||||
with a rank read from the swapped group.
|
||||
"""
|
||||
from sglang.srt.runtime_context import get_parallel
|
||||
|
||||
old_tp_group = get_tp_group()
|
||||
import sglang.srt.distributed.parallel_state as srt_parallel_state
|
||||
|
||||
@@ -712,7 +721,12 @@ def use_tensor_parallel_group(tp_group: GroupCoordinator):
|
||||
srt_parallel_state._TP = tp_group
|
||||
srt_parallel_state._ATTN_TP = tp_group
|
||||
try:
|
||||
yield
|
||||
with get_parallel().override(
|
||||
tp_size=tp_group.world_size,
|
||||
tp_rank=tp_group.rank_in_group,
|
||||
tp_group=tp_group,
|
||||
):
|
||||
yield
|
||||
finally:
|
||||
_TP = old_tp_group
|
||||
srt_parallel_state._TP = old_srt_tp_group
|
||||
|
||||
@@ -291,7 +291,12 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
|
||||
from sglang.srt.server_args import ServerArgs as SrtServerArgs
|
||||
|
||||
if get_context()._server_args is None:
|
||||
publish(SrtServerArgs(model_path="dummy"), role="diffusion_gpu_worker")
|
||||
# srt reads the size from the configuration and the rank from the
|
||||
# live group, so the dummy carries the width just installed.
|
||||
publish(
|
||||
SrtServerArgs(model_path="dummy", tp_size=self.server_args.tp_size),
|
||||
role="diffusion_gpu_worker",
|
||||
)
|
||||
|
||||
# set proc title
|
||||
if model_parallel_is_initialized():
|
||||
|
||||
@@ -17,10 +17,19 @@ from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import
|
||||
initialize_parallel_runtime,
|
||||
)
|
||||
from sglang.srt.distributed import parallel_state as srt_parallel_state
|
||||
from sglang.srt.runtime_context import get_parallel
|
||||
|
||||
_UTILS = "sglang.multimodal_gen.test.single_test_file.component_accuracy.utils"
|
||||
|
||||
|
||||
def _tp_group(world_size: int = 1, rank_in_group: int = 0) -> SimpleNamespace:
|
||||
"""A TP group handle carrying the two members the scope declares to the
|
||||
runtime context (`use_tensor_parallel_group` overrides `tp_size` /
|
||||
`tp_rank` / `tp_group` for its duration); the scope otherwise only stores
|
||||
the handle and compares it by identity."""
|
||||
return SimpleNamespace(world_size=world_size, rank_in_group=rank_in_group)
|
||||
|
||||
|
||||
def _server_args(*, ulysses_degree: int, ring_degree: int) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
tp_size=1,
|
||||
@@ -162,7 +171,7 @@ def test_srt_tp_groups_follow_encoder_folding_context():
|
||||
original_diffusion_tp_group = object()
|
||||
original_srt_tp_group = object()
|
||||
original_srt_attention_tp_group = object()
|
||||
folding_tp_group = object()
|
||||
folding_tp_group = _tp_group(world_size=2, rank_in_group=1)
|
||||
|
||||
with (
|
||||
patch.object(parallel_state, "_TP", original_diffusion_tp_group),
|
||||
@@ -177,6 +186,9 @@ def test_srt_tp_groups_follow_encoder_folding_context():
|
||||
assert parallel_state._TP is folding_tp_group
|
||||
assert srt_parallel_state._TP is folding_tp_group
|
||||
assert srt_parallel_state._ATTN_TP is folding_tp_group
|
||||
assert get_parallel().tp_size == 2
|
||||
assert get_parallel().tp_rank == 1
|
||||
assert get_parallel().tp_group is folding_tp_group
|
||||
|
||||
assert parallel_state._TP is original_diffusion_tp_group
|
||||
assert srt_parallel_state._TP is original_srt_tp_group
|
||||
@@ -185,8 +197,8 @@ def test_srt_tp_groups_follow_encoder_folding_context():
|
||||
|
||||
def test_encoder_folding_context_is_nested_and_restores_each_group():
|
||||
original_tp_group = object()
|
||||
outer_tp_group = object()
|
||||
inner_tp_group = object()
|
||||
outer_tp_group = _tp_group(world_size=4, rank_in_group=3)
|
||||
inner_tp_group = _tp_group(world_size=2, rank_in_group=1)
|
||||
|
||||
with (
|
||||
patch.object(parallel_state, "_TP", original_tp_group),
|
||||
@@ -194,14 +206,19 @@ def test_encoder_folding_context_is_nested_and_restores_each_group():
|
||||
patch.object(srt_parallel_state, "_ATTN_TP", original_tp_group),
|
||||
):
|
||||
with parallel_state.use_tensor_parallel_group(outer_tp_group):
|
||||
assert get_parallel().tp_size == 4
|
||||
with parallel_state.use_tensor_parallel_group(inner_tp_group):
|
||||
assert parallel_state._TP is inner_tp_group
|
||||
assert srt_parallel_state._TP is inner_tp_group
|
||||
assert srt_parallel_state._ATTN_TP is inner_tp_group
|
||||
assert get_parallel().tp_size == 2
|
||||
assert get_parallel().tp_rank == 1
|
||||
|
||||
assert parallel_state._TP is outer_tp_group
|
||||
assert srt_parallel_state._TP is outer_tp_group
|
||||
assert srt_parallel_state._ATTN_TP is outer_tp_group
|
||||
assert get_parallel().tp_size == 4
|
||||
assert get_parallel().tp_rank == 3
|
||||
|
||||
assert parallel_state._TP is original_tp_group
|
||||
assert srt_parallel_state._TP is original_tp_group
|
||||
|
||||
@@ -268,7 +268,7 @@ class ZayaConfig(PretrainedConfig):
|
||||
try:
|
||||
|
||||
tp_size = get_parallel().tp_size
|
||||
except (AssertionError, RuntimeError):
|
||||
except (AssertionError, RuntimeError, ValueError):
|
||||
tp_size = 1
|
||||
|
||||
in_out_ch_full = (
|
||||
|
||||
@@ -1735,8 +1735,8 @@ class _SGLangPlugin(_FrameworkPlugin):
|
||||
info["moe_tp_rank"] = parallel.moe_tp_rank
|
||||
info["moe_tp_size"] = parallel.moe_tp_size
|
||||
info["moe_dp_rank"] = parallel.moe_dp_rank
|
||||
info["moe_dp_size"] = parallel.moe_dp_size
|
||||
except (AttributeError, AssertionError):
|
||||
info["moe_dp_size"] = self._dp_attn.get_moe_cp_size()
|
||||
except (AttributeError, AssertionError, ValueError):
|
||||
info["distributed_error"] = True
|
||||
|
||||
try:
|
||||
@@ -1748,7 +1748,7 @@ class _SGLangPlugin(_FrameworkPlugin):
|
||||
info["attn_dp_size"] = self._dp_attn.get_attention_dp_size()
|
||||
info["attn_cp_rank"] = parallel.attn_cp_rank
|
||||
info["attn_cp_size"] = parallel.attn_cp_size
|
||||
except (AttributeError, AssertionError):
|
||||
except (AttributeError, AssertionError, ValueError):
|
||||
info["dp_attention_error"] = True
|
||||
|
||||
return info
|
||||
|
||||
@@ -186,7 +186,7 @@ class CommonKVManager(BaseKVManager):
|
||||
self.system_dp_rank = (
|
||||
self.kv_args.system_dp_rank if self.kv_args.system_dp_rank else 0
|
||||
)
|
||||
self.pp_size = get_parallel().config.pp_size
|
||||
self.pp_size = get_parallel().pp_size
|
||||
self.pp_rank = self.kv_args.pp_rank
|
||||
self.local_ip = get_local_ip_auto()
|
||||
cp_sharded_prefill = self.attn_cp_size > 1 and (
|
||||
|
||||
@@ -221,7 +221,7 @@ async def serve_grpc_encoder(server_args: ServerArgs):
|
||||
).to_tcp()
|
||||
|
||||
send_sockets: List[zmq.Socket] = []
|
||||
for rank in range(1, get_parallel().config.tp_size):
|
||||
for rank in range(1, get_parallel().tp_size):
|
||||
schedule_path = f"ipc:///tmp/{ipc_path_prefix}_schedule_{rank}"
|
||||
send_sockets.append(
|
||||
get_zmq_socket(zmq_ctx, zmq.PUSH, schedule_path, bind=False)
|
||||
|
||||
@@ -1734,7 +1734,7 @@ class MMReceiverBase(ABC):
|
||||
self.host = get_local_ip_auto(get_serving().host)
|
||||
self.pp_rank = pp_rank
|
||||
self.tp_rank = tp_rank
|
||||
self.tp_size = get_parallel().config.tp_size
|
||||
self.tp_size = get_parallel().tp_size
|
||||
self.tp_group = tp_group
|
||||
self.nnodes = server_args.nnodes
|
||||
self.hostname = get_local_ip_auto()
|
||||
|
||||
@@ -1530,7 +1530,7 @@ def launch_local_runtime(server_args: ServerArgs) -> EncoderRuntime:
|
||||
|
||||
send_sockets: List[zmq.Socket] = []
|
||||
tp_processes: List[mp.Process] = []
|
||||
for rank in range(1, get_parallel().config.tp_size):
|
||||
for rank in range(1, get_parallel().tp_size):
|
||||
schedule_path = f"ipc:///tmp/{ipc_path_prefix}_schedule_{rank}"
|
||||
send_sockets.append(
|
||||
get_zmq_socket(zmq_context, zmq.PUSH, schedule_path, bind=False)
|
||||
@@ -1570,10 +1570,10 @@ def launch_dp_runtime(server_args: ServerArgs) -> DPDispatcher:
|
||||
HTTP uses this entry point today. gRPC can reuse it later without
|
||||
importing HTTP application state or Uvicorn.
|
||||
"""
|
||||
if get_parallel().dp_size <= 1 or get_parallel().config.tp_size != 1:
|
||||
if get_parallel().dp_size <= 1 or get_parallel().tp_size != 1:
|
||||
raise ValueError(
|
||||
"Encoder DP mode requires --dp-size > 1 and --tp-size 1; got "
|
||||
f"dp_size={get_parallel().dp_size}, tp_size={get_parallel().config.tp_size}."
|
||||
f"dp_size={get_parallel().dp_size}, tp_size={get_parallel().tp_size}."
|
||||
)
|
||||
dp_size = get_parallel().dp_size
|
||||
logger.info(f"Launching encoder in DP mode: dp_size={dp_size}")
|
||||
|
||||
@@ -451,7 +451,7 @@ class MMEncoder:
|
||||
this instance's value, not a config change, so it travels as an
|
||||
argument."""
|
||||
assert_published(server_args, role="encoder")
|
||||
logger.info(f"init MMEncoder {rank}/{get_parallel().config.tp_size}")
|
||||
logger.info(f"init MMEncoder {rank}/{get_parallel().tp_size}")
|
||||
self.server_args = server_args
|
||||
configure_media_url_security(
|
||||
get_mm().allowed_media_domains,
|
||||
@@ -492,14 +492,12 @@ class MMEncoder:
|
||||
|
||||
init_distributed_environment(
|
||||
backend=get_default_distributed_backend(self.device),
|
||||
world_size=get_parallel().config.tp_size,
|
||||
world_size=get_parallel().tp_size,
|
||||
rank=rank,
|
||||
distributed_init_method=dist_init_method,
|
||||
local_rank=rank,
|
||||
)
|
||||
initialize_model_parallel(
|
||||
tensor_model_parallel_size=get_parallel().config.tp_size
|
||||
)
|
||||
initialize_model_parallel(tensor_model_parallel_size=get_parallel().tp_size)
|
||||
initialize_dp_attention(server_args, self.model_config)
|
||||
|
||||
self.model = load_model(
|
||||
@@ -557,7 +555,7 @@ class MMEncoder:
|
||||
)
|
||||
self.mm_global_cache = EmbeddingCacheController(
|
||||
rank,
|
||||
get_parallel().config.tp_size,
|
||||
get_parallel().tp_size,
|
||||
embedding_store=embedding_store,
|
||||
hidden_dims=self._embedding_dims,
|
||||
tp_group=get_tp_group().cpu_group,
|
||||
@@ -1035,7 +1033,7 @@ class MMEncoder:
|
||||
)
|
||||
|
||||
def _broadcast_global_cache_mask(self, mask_tensor: torch.Tensor):
|
||||
if get_parallel().config.tp_size > 1:
|
||||
if get_parallel().tp_size > 1:
|
||||
torch.distributed.broadcast(
|
||||
mask_tensor,
|
||||
src=0,
|
||||
|
||||
@@ -405,8 +405,7 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager):
|
||||
):
|
||||
super().__init__(args, disaggregation_mode, server_args, is_mla_backend)
|
||||
self.transfer_source_rank = (
|
||||
self.kv_args.pp_rank * get_parallel().config.tp_size
|
||||
+ self.kv_args.engine_rank
|
||||
self.kv_args.pp_rank * get_parallel().tp_size + self.kv_args.engine_rank
|
||||
)
|
||||
self.kv_args.kv_data_mem_kinds = _normalize_kv_mem_kinds(
|
||||
getattr(self.kv_args, "kv_data_mem_kinds", None),
|
||||
|
||||
@@ -2803,14 +2803,19 @@ _TP_STATE_PATCHED = False
|
||||
|
||||
@contextmanager
|
||||
def patch_tensor_parallel_group(tp_group: GroupCoordinator):
|
||||
"""Patch the tp group temporarily until this function ends.
|
||||
"""Run under a different tensor-parallel group until this scope ends.
|
||||
|
||||
This method is for draft workers of speculative decoding to run draft model
|
||||
with different tp degree from that of target model workers.
|
||||
This is for draft workers of speculative decoding, which run the draft model
|
||||
at the target's attention-TP width rather than its global TP width.
|
||||
|
||||
The scope replaces both the module global that ``get_tp_group()`` reads and
|
||||
the three members the runtime context answers with.
|
||||
|
||||
Args:
|
||||
tp_group (GroupCoordinator): the tp group coordinator
|
||||
"""
|
||||
from sglang.srt.runtime_context import get_parallel
|
||||
|
||||
global _TP_STATE_PATCHED
|
||||
assert not _TP_STATE_PATCHED, "Should not call when it's already patched"
|
||||
|
||||
@@ -2819,9 +2824,13 @@ def patch_tensor_parallel_group(tp_group: GroupCoordinator):
|
||||
global _TP
|
||||
_TP = tp_group
|
||||
try:
|
||||
yield
|
||||
with get_parallel().override(
|
||||
tp_size=tp_group.world_size,
|
||||
tp_rank=tp_group.rank_in_group,
|
||||
tp_group=tp_group,
|
||||
):
|
||||
yield
|
||||
finally:
|
||||
# restore the original state
|
||||
_TP_STATE_PATCHED = False
|
||||
_TP = old_tp_group
|
||||
|
||||
|
||||
@@ -127,7 +127,7 @@ class ElasticEPStateManager:
|
||||
|
||||
if get_exec().moe.ep_join_mode == "scale":
|
||||
inst.effective_ep_size = (
|
||||
get_parallel().ep_join_rank_offset + get_parallel().config.tp_size
|
||||
get_parallel().ep_join_rank_offset + get_parallel().tp_size
|
||||
)
|
||||
inst.original_ep_size = (
|
||||
get_parallel().elastic_ep_initial_size
|
||||
|
||||
@@ -72,7 +72,7 @@ class ExpertBackupManager:
|
||||
# losing the initial PUB message due to slow joiners.
|
||||
num_ready_clients = 0
|
||||
|
||||
while num_ready_clients < get_parallel().config.tp_size:
|
||||
while num_ready_clients < get_parallel().tp_size:
|
||||
sock_recv(self.recv_from_expert_backup_client)
|
||||
num_ready_clients += 1
|
||||
|
||||
|
||||
@@ -683,7 +683,7 @@ class Engine(EngineScoreMixin, EngineBase):
|
||||
pp_rank_range, tp_rank_range, pp_size_per_node, tp_size_per_node = (
|
||||
_calculate_rank_ranges(
|
||||
server_args.nnodes,
|
||||
get_parallel().config.pp_size,
|
||||
get_parallel().pp_size,
|
||||
tp_size,
|
||||
server_args.node_rank,
|
||||
)
|
||||
@@ -844,7 +844,7 @@ class Engine(EngineScoreMixin, EngineBase):
|
||||
pp_rank_range, tp_rank_range, pp_size_per_node, tp_size_per_node = (
|
||||
_calculate_rank_ranges(
|
||||
server_args.nnodes,
|
||||
get_parallel().config.pp_size,
|
||||
get_parallel().pp_size,
|
||||
server_args.tp_size,
|
||||
server_args.node_rank,
|
||||
)
|
||||
@@ -1843,8 +1843,8 @@ def _compute_parallelism_ranks(
|
||||
"""
|
||||
attn_dp_size = get_parallel().dp_size if get_parallel().enable_dp_attention else 1
|
||||
tp_size = server_args.tp_size
|
||||
attn_cp_size = get_parallel().config.attn_cp_size
|
||||
moe_dp_size = get_parallel().config.moe_dp_size
|
||||
attn_cp_size = get_parallel().attn_cp_size
|
||||
moe_dp_size = get_parallel().moe_dp_size
|
||||
|
||||
# Parallelism hierarchy (outermost to innermost):
|
||||
# - Attention: Global(TP) -> DP -> ATTN_CP -> ATTN_TP (innermost)
|
||||
|
||||
@@ -146,8 +146,8 @@ async def get_loads(
|
||||
"version": __version__,
|
||||
"accelerator": _accelerator_name(),
|
||||
"num_accelerators": _num_accelerators_per_dp_rank(
|
||||
get_parallel().config.tp_size,
|
||||
get_parallel().config.pp_size,
|
||||
get_parallel().tp_size,
|
||||
get_parallel().pp_size,
|
||||
get_parallel().dp_size,
|
||||
get_parallel().enable_dp_attention,
|
||||
),
|
||||
|
||||
@@ -245,7 +245,7 @@ class ExpertLocationMetadata:
|
||||
if get_exec().moe.ep_join_mode == "scale":
|
||||
ep_size = max(
|
||||
ep_size,
|
||||
get_parallel().ep_join_rank_offset + get_parallel().config.tp_size,
|
||||
get_parallel().ep_join_rank_offset + get_parallel().tp_size,
|
||||
)
|
||||
num_physical_experts, num_local_physical_experts = (
|
||||
_compute_elastic_expert_layout(
|
||||
|
||||
@@ -248,7 +248,7 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp):
|
||||
if _is_cuda:
|
||||
self.sm_count = deep_gemm.get_num_sms()
|
||||
self.half_device_sm_count = ceil_align(self.sm_count // 2, 8)
|
||||
pp_size = get_parallel().config.pp_size
|
||||
pp_size = get_parallel().pp_size
|
||||
self.logits_with_pp_recv = pp_size > 1 and not get_pp_group().is_last_rank
|
||||
else:
|
||||
self.logits_with_pp_recv = False
|
||||
|
||||
@@ -282,15 +282,19 @@ def get_cp_strategy() -> Optional[ContextParallelStrategy]:
|
||||
global _STRATEGY
|
||||
|
||||
if _STRATEGY is None:
|
||||
# The reads are what raise, so they sit inside the guard.
|
||||
try:
|
||||
parallel = get_parallel().config
|
||||
except ValueError:
|
||||
parallel = get_parallel()
|
||||
enable_prefill_cp = parallel.enable_prefill_cp
|
||||
cp_size = parallel.attn_cp_size
|
||||
cp_strategy = parallel.cp_strategy
|
||||
except (AssertionError, AttributeError, RuntimeError, ValueError):
|
||||
return None
|
||||
if parallel.enable_prefill_cp:
|
||||
if enable_prefill_cp:
|
||||
init_cp_strategy(
|
||||
enable_prefill_cp=True,
|
||||
cp_size=parallel.attn_cp_size,
|
||||
cp_strategy=parallel.cp_strategy,
|
||||
cp_size=cp_size,
|
||||
cp_strategy=cp_strategy,
|
||||
)
|
||||
return _STRATEGY
|
||||
|
||||
|
||||
@@ -349,7 +349,7 @@ def initialize_dp_attention(
|
||||
)
|
||||
enable_dp_attention = get_parallel().enable_dp_attention
|
||||
dp_size = get_parallel().dp_size
|
||||
attn_cp_size = get_parallel().config.attn_cp_size
|
||||
attn_cp_size = get_parallel().attn_cp_size
|
||||
|
||||
dp.enabled = enable_dp_attention
|
||||
|
||||
@@ -1026,12 +1026,10 @@ def get_moe_cp_size() -> int:
|
||||
def is_enable_moe_cp_allgather() -> bool:
|
||||
"""True when moe_dp_size < attn_cp_size, requiring allgather across CP ranks before MoE.
|
||||
|
||||
Reads the configured sizes, not the live groups: that very configuration makes
|
||||
``initialize_model_parallel`` alias ``_MOE_DP`` to ``_ATTN_CP``
|
||||
(``parallel_state.py``), so the live sizes are equal and the comparison would
|
||||
always be false.
|
||||
In that configuration ``initialize_model_parallel`` aliases ``_MOE_DP`` to
|
||||
``_ATTN_CP``, so the two groups report equal widths.
|
||||
"""
|
||||
return get_parallel().config.attn_cp_size > get_parallel().config.moe_dp_size
|
||||
return get_parallel().attn_cp_size > get_parallel().moe_dp_size
|
||||
|
||||
|
||||
def moe_cp_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor):
|
||||
|
||||
@@ -2120,7 +2120,9 @@ def validate_fp8_block_shape(
|
||||
) -> None:
|
||||
"""Validate block quantization shapes for tensor parallelism."""
|
||||
|
||||
tp_size = getattr(layer, "tp_size", get_parallel().tp_size)
|
||||
# Lazy: a ``getattr`` default would read the published bag even for a
|
||||
# layer that carries its own tp_size.
|
||||
tp_size = layer.tp_size if hasattr(layer, "tp_size") else get_parallel().tp_size
|
||||
block_n, block_k = block_size[0], block_size[1]
|
||||
|
||||
# Required by row parallel
|
||||
|
||||
@@ -392,9 +392,7 @@ class DataParallelController:
|
||||
)
|
||||
threads.append(thread)
|
||||
base_gpu_id += (
|
||||
server_args.tp_size
|
||||
* get_parallel().config.pp_size
|
||||
* server_args.gpu_id_step
|
||||
server_args.tp_size * get_parallel().pp_size * server_args.gpu_id_step
|
||||
)
|
||||
|
||||
if server_args.node_rank == 0:
|
||||
@@ -615,8 +613,8 @@ class DataParallelController:
|
||||
|
||||
scheduler_pipe_readers = []
|
||||
|
||||
pp_size_per_node = max(get_parallel().config.pp_size // server_args.nnodes, 1)
|
||||
nnodes_per_pp_rank = max(server_args.nnodes // get_parallel().config.pp_size, 1)
|
||||
pp_size_per_node = max(get_parallel().pp_size // server_args.nnodes, 1)
|
||||
nnodes_per_pp_rank = max(server_args.nnodes // get_parallel().pp_size, 1)
|
||||
pp_rank_range = range(
|
||||
pp_size_per_node * (server_args.node_rank // nnodes_per_pp_rank),
|
||||
pp_size_per_node * (server_args.node_rank // nnodes_per_pp_rank + 1),
|
||||
@@ -647,7 +645,7 @@ class DataParallelController:
|
||||
tp_rank,
|
||||
server_args.tp_size,
|
||||
get_parallel().dp_size,
|
||||
get_parallel().config.attn_cp_size,
|
||||
get_parallel().attn_cp_size,
|
||||
)
|
||||
# compute zmq ports for this dp rank
|
||||
rank_port_args = PortArgs.init_new(
|
||||
@@ -683,22 +681,18 @@ class DataParallelController:
|
||||
# - Attention: Global(TP) -> DP -> ATTN_CP -> ATTN_TP (innermost)
|
||||
# - MoE: Global(TP) -> MOE_DP -> EP -> MOE_TP (innermost)
|
||||
attn_tp_size = (
|
||||
server_args.tp_size
|
||||
// attn_dp_size
|
||||
// get_parallel().config.attn_cp_size
|
||||
server_args.tp_size // attn_dp_size // get_parallel().attn_cp_size
|
||||
)
|
||||
attn_cp_rank = (
|
||||
tp_rank // attn_tp_size
|
||||
) % get_parallel().config.attn_cp_size
|
||||
attn_cp_rank = (tp_rank // attn_tp_size) % get_parallel().attn_cp_size
|
||||
moe_dp_rank = tp_rank // (
|
||||
server_args.tp_size // get_parallel().config.moe_dp_size
|
||||
server_args.tp_size // get_parallel().moe_dp_size
|
||||
)
|
||||
moe_ep_rank = (
|
||||
tp_rank
|
||||
% (server_args.tp_size // get_parallel().config.moe_dp_size)
|
||||
% (server_args.tp_size // get_parallel().moe_dp_size)
|
||||
// (
|
||||
server_args.tp_size
|
||||
// get_parallel().config.moe_dp_size
|
||||
// get_parallel().moe_dp_size
|
||||
// get_parallel().ep_size
|
||||
)
|
||||
)
|
||||
|
||||
@@ -275,7 +275,7 @@ class NativeMmHost:
|
||||
)
|
||||
|
||||
return (
|
||||
get_parallel().config.tp_size > 1
|
||||
get_parallel().tp_size > 1
|
||||
and determine_tensor_transport_mode() != "default"
|
||||
and not self.server_args.skip_tokenizer_init
|
||||
)
|
||||
|
||||
@@ -478,30 +478,30 @@ class Scheduler(
|
||||
compute_dp_attention_world_info(
|
||||
get_parallel().enable_dp_attention,
|
||||
tp_rank,
|
||||
get_parallel().config.tp_size,
|
||||
get_parallel().tp_size,
|
||||
get_parallel().dp_size,
|
||||
get_parallel().config.attn_cp_size,
|
||||
get_parallel().attn_cp_size,
|
||||
)
|
||||
)
|
||||
self.ps = ParallelState(
|
||||
tp_rank=tp_rank,
|
||||
tp_size=get_parallel().config.tp_size,
|
||||
tp_size=get_parallel().tp_size,
|
||||
pp_rank=pp_rank,
|
||||
pp_size=get_parallel().config.pp_size,
|
||||
pp_size=get_parallel().pp_size,
|
||||
dp_rank=dp_rank,
|
||||
dp_size=get_parallel().dp_size,
|
||||
attn_tp_rank=attn_tp_rank,
|
||||
attn_tp_size=attn_tp_size,
|
||||
attn_cp_rank=attn_cp_rank,
|
||||
attn_cp_size=get_parallel().config.attn_cp_size,
|
||||
attn_dcp_rank=tp_rank % get_parallel().config.dcp_size,
|
||||
attn_dcp_size=get_parallel().config.dcp_size,
|
||||
attn_cp_size=get_parallel().attn_cp_size,
|
||||
attn_dcp_rank=tp_rank % get_parallel().dcp_size,
|
||||
attn_dcp_size=get_parallel().dcp_size,
|
||||
attn_dp_rank=attn_dp_rank,
|
||||
attn_dp_size=attn_dp_size,
|
||||
moe_ep_rank=moe_ep_rank,
|
||||
moe_ep_size=get_parallel().ep_size,
|
||||
moe_dp_rank=moe_dp_rank,
|
||||
moe_dp_size=get_parallel().config.moe_dp_size,
|
||||
moe_dp_size=get_parallel().moe_dp_size,
|
||||
gpu_id=gpu_id,
|
||||
)
|
||||
|
||||
@@ -4544,7 +4544,12 @@ class Scheduler(
|
||||
# Resolved config (pristine server_args + post-publish overrides) so a
|
||||
# readback reflects values changed via /set_internal_state, not startup.
|
||||
ret = get_context().resolved_server_args_dict()
|
||||
ret["world_size"] = compute_world_size(get_parallel().config)
|
||||
ret["world_size"] = compute_world_size(
|
||||
enable_dp_attention=get_parallel().enable_dp_attention,
|
||||
dp_size=get_parallel().dp_size,
|
||||
tp_size=get_parallel().tp_size,
|
||||
pp_size=get_parallel().pp_size,
|
||||
)
|
||||
ret["last_gen_throughput"] = self.metrics_reporter.last_gen_throughput
|
||||
draft_graph_memory_usage = (
|
||||
None if self.draft_worker is None else self.draft_worker.graph_memory_usage
|
||||
@@ -5220,7 +5225,7 @@ def dispatch_event_loop(scheduler: Scheduler):
|
||||
if disaggregation_mode == DisaggregationMode.NULL:
|
||||
if scheduler.enable_pdmux:
|
||||
scheduler.event_loop_pdmux()
|
||||
elif get_parallel().config.pp_size > 1:
|
||||
elif get_parallel().pp_size > 1:
|
||||
scheduler.event_loop_pp()
|
||||
elif scheduler.enable_overlap_mlx:
|
||||
scheduler.event_loop_overlap_mlx()
|
||||
@@ -5229,14 +5234,14 @@ def dispatch_event_loop(scheduler: Scheduler):
|
||||
else:
|
||||
scheduler.event_loop_normal()
|
||||
elif disaggregation_mode == DisaggregationMode.PREFILL:
|
||||
if get_parallel().config.pp_size > 1:
|
||||
if get_parallel().pp_size > 1:
|
||||
scheduler.event_loop_pp_disagg_prefill()
|
||||
elif scheduler.enable_overlap:
|
||||
scheduler.event_loop_overlap_disagg_prefill()
|
||||
else:
|
||||
scheduler.event_loop_normal_disagg_prefill()
|
||||
elif disaggregation_mode == DisaggregationMode.DECODE:
|
||||
if get_parallel().config.pp_size > 1:
|
||||
if get_parallel().pp_size > 1:
|
||||
scheduler.event_loop_pp_disagg_decode()
|
||||
elif scheduler.enable_overlap:
|
||||
scheduler.event_loop_overlap_disagg_decode()
|
||||
@@ -5277,13 +5282,13 @@ def configure_scheduler_process(
|
||||
prefix = ""
|
||||
if shown_dp is not None:
|
||||
prefix += f" DP{shown_dp}"
|
||||
if get_parallel().config.pp_size > 1:
|
||||
if get_parallel().pp_size > 1:
|
||||
prefix += f" PP{pp_rank}"
|
||||
if get_parallel().config.attn_cp_size > 1:
|
||||
if get_parallel().attn_cp_size > 1:
|
||||
prefix += f" ATTN_CP{attn_cp_rank}"
|
||||
if get_parallel().config.moe_dp_size > 1:
|
||||
if get_parallel().moe_dp_size > 1:
|
||||
prefix += f" MOE_DP{moe_dp_rank}"
|
||||
if get_parallel().config.tp_size > 1:
|
||||
if get_parallel().tp_size > 1:
|
||||
prefix += f" TP{shown_tp}"
|
||||
if get_parallel().ep_size > 1:
|
||||
prefix += f" EP{shown_moe_ep}"
|
||||
@@ -5299,8 +5304,8 @@ def configure_scheduler_process(
|
||||
# Set cpu affinity to this gpu process
|
||||
if envs.SGLANG_SET_CPU_AFFINITY.get():
|
||||
set_gpu_proc_affinity(
|
||||
get_parallel().config.pp_size,
|
||||
get_parallel().config.tp_size,
|
||||
get_parallel().pp_size,
|
||||
get_parallel().tp_size,
|
||||
get_parallel().nnodes,
|
||||
gpu_id,
|
||||
)
|
||||
|
||||
@@ -1117,7 +1117,7 @@ class SchedulerMetricsReporter:
|
||||
active_lora_ids = set()
|
||||
|
||||
# For PP mode, check all running micro batches
|
||||
if get_parallel().config.pp_size > 1:
|
||||
if get_parallel().pp_size > 1:
|
||||
for batch in self.scheduler.running_mbs:
|
||||
if batch and hasattr(batch, "reqs"):
|
||||
for req in batch.reqs:
|
||||
|
||||
@@ -179,8 +179,8 @@ class TokenizerControlMixin:
|
||||
)
|
||||
if primary_group_control:
|
||||
control_fan_out = (
|
||||
worker_count + get_parallel().config.tp_size - 1
|
||||
) // get_parallel().config.tp_size
|
||||
worker_count + get_parallel().tp_size - 1
|
||||
) // get_parallel().tp_size
|
||||
else:
|
||||
control_fan_out = worker_count
|
||||
|
||||
|
||||
@@ -1906,7 +1906,7 @@ class KVCacheConfigurator:
|
||||
token_capacity = min(token_capacity, user_limit)
|
||||
|
||||
# Sync across PP ranks (each may have different layer counts)
|
||||
if get_parallel().config.pp_size > 1:
|
||||
if get_parallel().pp_size > 1:
|
||||
tensor = torch.tensor(token_capacity, dtype=torch.int64)
|
||||
torch.distributed.all_reduce(
|
||||
tensor,
|
||||
|
||||
@@ -609,9 +609,9 @@ class CPUGraphRunner:
|
||||
self.enable_profile_cuda_graph = (
|
||||
model_runner.server_args.enable_profile_cuda_graph
|
||||
)
|
||||
self.tp_size = get_parallel().config.tp_size
|
||||
self.tp_size = get_parallel().tp_size
|
||||
self.dp_size = get_parallel().dp_size
|
||||
self.pp_size = get_parallel().config.pp_size
|
||||
self.pp_size = get_parallel().pp_size
|
||||
|
||||
self.capture_forward_mode = ForwardMode.DECODE
|
||||
self.capture_hidden_mode = self.return_hidden_states_mode
|
||||
|
||||
@@ -240,7 +240,7 @@ def _resolve_dflash_draft_cell_size(
|
||||
draft_model_config=draft_model_config,
|
||||
draft_num_layers=draft_num_layers,
|
||||
draft_kv_cache_dtype=draft_kv_cache_dtype,
|
||||
tp_size=get_parallel().config.tp_size,
|
||||
tp_size=get_parallel().tp_size,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning(
|
||||
|
||||
@@ -119,10 +119,10 @@ class StartupWeightLoadOptions:
|
||||
prefill_cuda_graph_backend=cuda_graph_config.prefill.backend,
|
||||
is_draft_worker=is_draft_worker,
|
||||
speculative_algorithm=get_spec().speculative_algorithm,
|
||||
tp_size=get_parallel().config.tp_size,
|
||||
attn_cp_size=get_parallel().config.attn_cp_size,
|
||||
dcp_size=get_parallel().config.dcp_size,
|
||||
pp_size=get_parallel().config.pp_size,
|
||||
tp_size=get_parallel().tp_size,
|
||||
attn_cp_size=get_parallel().attn_cp_size,
|
||||
dcp_size=get_parallel().dcp_size,
|
||||
pp_size=get_parallel().pp_size,
|
||||
dp_size=get_parallel().dp_size,
|
||||
ep_size=get_parallel().ep_size,
|
||||
cpu_offload_gb=get_exec().offload.cpu_offload_gb,
|
||||
|
||||
@@ -216,10 +216,10 @@ class BaseRunner(ABC):
|
||||
self.model_runner = model_runner
|
||||
self.device = model_runner.device
|
||||
self.device_module = torch.get_device_module(self.device)
|
||||
self.tp_size = get_parallel().config.tp_size
|
||||
self.tp_size = get_parallel().tp_size
|
||||
# elastic-EP scale-up rewrites dp_size on the published config
|
||||
self.dp_size = get_parallel().dp_size
|
||||
self.pp_size = get_parallel().config.pp_size
|
||||
self.pp_size = get_parallel().pp_size
|
||||
self.enable_pdmux = model_runner.server_args.enable_pdmux
|
||||
self.return_hidden_states_mode = (
|
||||
CaptureHiddenMode.NULL
|
||||
@@ -349,7 +349,7 @@ class BaseRunner(ABC):
|
||||
vocab_size=mr.model_config.vocab_size,
|
||||
dtype=mr.model_config.dtype,
|
||||
dp_size=get_parallel().dp_size,
|
||||
pp_size=get_parallel().config.pp_size,
|
||||
pp_size=get_parallel().pp_size,
|
||||
is_encoder_decoder=mr.model_config.is_encoder_decoder,
|
||||
require_mlp_tp_gather=require_mlp_tp_gather(),
|
||||
seq_len_fill_value=mr.attn_backend.get_cuda_graph_seq_len_fill_value(),
|
||||
@@ -521,7 +521,7 @@ class BaseRunner(ABC):
|
||||
extend_prefix_lens = None
|
||||
extend_start_loc = None
|
||||
|
||||
if get_parallel().config.pp_size > 1:
|
||||
if get_parallel().pp_size > 1:
|
||||
# PP0 already cp-split hidden_states before send.
|
||||
pp_hidden_tokens = num_tokens
|
||||
if (
|
||||
@@ -645,7 +645,7 @@ class BaseRunner(ABC):
|
||||
|
||||
kwargs = {}
|
||||
if (
|
||||
get_parallel().config.pp_size > 1
|
||||
get_parallel().pp_size > 1
|
||||
and "pp_proxy_tensors" in inspect.signature(mr.model.forward).parameters
|
||||
):
|
||||
kwargs["pp_proxy_tensors"] = PPProxyTensors(
|
||||
|
||||
@@ -1884,17 +1884,19 @@ class PreshardedModelLoader(DefaultModelLoader):
|
||||
def _safe(fn) -> int:
|
||||
try:
|
||||
return fn()
|
||||
except (AssertionError, AttributeError, RuntimeError):
|
||||
except (AssertionError, AttributeError, RuntimeError, ValueError):
|
||||
return 1
|
||||
|
||||
from sglang.srt.layers.dp_attention import get_moe_cp_size
|
||||
|
||||
parallel = get_parallel()
|
||||
return {
|
||||
"tp": _safe(lambda: parallel.tp_size),
|
||||
"dp": _safe(lambda: parallel.moe_dp_size),
|
||||
"dp": _safe(get_moe_cp_size),
|
||||
"ep": _safe(lambda: parallel.moe_ep_size),
|
||||
"pp": _safe(lambda: parallel.pp_size),
|
||||
"moe_dense_tp_size": parallel.moe_dense_tp_size,
|
||||
"moe_dp_size": get_parallel().config.moe_dp_size,
|
||||
"moe_dp_size": get_parallel().moe_dp_size,
|
||||
"enable_dp_lm_head": parallel.enable_dp_lm_head,
|
||||
"enable_fp32_lm_head": get_exec().features.enable_fp32_lm_head,
|
||||
"quantization": model_config.quantization,
|
||||
|
||||
@@ -733,7 +733,7 @@ class KimiK25ForConditionalGeneration(nn.Module):
|
||||
# Match the configured TP consumer count captured when the
|
||||
# tokenizer creates MmItemMemoryPool. A live attention subgroup
|
||||
# size could leave acknowledgements missing and strand the lease.
|
||||
ipc_consumer_count = max(get_parallel().config.tp_size, 1)
|
||||
ipc_consumer_count = max(get_parallel().tp_size, 1)
|
||||
device_index = device.index
|
||||
if device.type == "cuda" and device_index is None:
|
||||
device_index = torch.cuda.current_device()
|
||||
|
||||
@@ -3382,7 +3382,7 @@ class KimiK3ForConditionalGeneration(nn.Module):
|
||||
# Match the configured TP consumer count captured when the
|
||||
# tokenizer creates MmItemMemoryPool. A live attention subgroup
|
||||
# size could leave acknowledgements missing and strand the lease.
|
||||
ipc_consumer_count = max(get_parallel().config.tp_size, 1)
|
||||
ipc_consumer_count = max(get_parallel().tp_size, 1)
|
||||
device_index = device.index
|
||||
if device.type == "cuda" and device_index is None:
|
||||
device_index = torch.cuda.current_device()
|
||||
|
||||
@@ -147,8 +147,8 @@ class RayDataParallelController(DataParallelController):
|
||||
bundle_idx = self.bundle_for_node[node_idx]
|
||||
pp_range, tp_range, pp_per_node, tp_per_node = _calculate_rank_ranges(
|
||||
nnodes,
|
||||
get_parallel().config.pp_size,
|
||||
get_parallel().config.tp_size,
|
||||
get_parallel().pp_size,
|
||||
get_parallel().tp_size,
|
||||
node_rank=node_idx,
|
||||
)
|
||||
for pp_rank in pp_range:
|
||||
@@ -160,7 +160,7 @@ class RayDataParallelController(DataParallelController):
|
||||
tp_rank % tp_per_node
|
||||
)
|
||||
|
||||
parallel = get_parallel().config
|
||||
parallel = get_parallel()
|
||||
if parallel.enable_dp_attention:
|
||||
_, _, actual_dp_rank, _ = compute_dp_attention_world_info(
|
||||
parallel.enable_dp_attention,
|
||||
@@ -208,7 +208,7 @@ class RayDataParallelController(DataParallelController):
|
||||
world_size = _compute_world_size()
|
||||
bundle_indices = _resolve_bundle_indices(self.pg, world_size)
|
||||
|
||||
parallel = get_parallel().config
|
||||
parallel = get_parallel()
|
||||
ranks_per_tp_group = parallel.tp_size * parallel.pp_size
|
||||
if dp_rank is not None:
|
||||
start_rank = dp_rank * ranks_per_tp_group
|
||||
@@ -237,9 +237,9 @@ class RayDataParallelController(DataParallelController):
|
||||
_, _, actual_dp_rank, _ = compute_dp_attention_world_info(
|
||||
get_parallel().enable_dp_attention,
|
||||
tp_rank,
|
||||
get_parallel().config.tp_size,
|
||||
get_parallel().tp_size,
|
||||
get_parallel().dp_size,
|
||||
get_parallel().config.attn_cp_size,
|
||||
get_parallel().attn_cp_size,
|
||||
)
|
||||
rank_port_args = PortArgs.init_new(
|
||||
server_args, actual_dp_rank, worker_ports
|
||||
|
||||
@@ -111,7 +111,12 @@ def _compute_world_size() -> int:
|
||||
Reads the published parallel leaves: the driver is sizing the actors that
|
||||
will hold the process groups, so there is nothing live to ask.
|
||||
"""
|
||||
return compute_world_size(get_parallel().config)
|
||||
return compute_world_size(
|
||||
enable_dp_attention=get_parallel().enable_dp_attention,
|
||||
dp_size=get_parallel().dp_size,
|
||||
tp_size=get_parallel().tp_size,
|
||||
pp_size=get_parallel().pp_size,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_bundle_indices(pg: PlacementGroup, world_size: int) -> List[int]:
|
||||
@@ -269,7 +274,7 @@ class RayEngine(Engine):
|
||||
placement_group as create_placement_group,
|
||||
)
|
||||
|
||||
parallel = get_parallel().config
|
||||
parallel = get_parallel()
|
||||
if parallel.enable_dp_attention:
|
||||
total_gpus = parallel.tp_size * parallel.pp_size
|
||||
else:
|
||||
@@ -332,8 +337,8 @@ class RayEngine(Engine):
|
||||
pp_range, tp_range, pp_per_node, tp_per_node = (
|
||||
_calculate_rank_ranges(
|
||||
nnodes,
|
||||
get_parallel().config.pp_size,
|
||||
get_parallel().config.tp_size,
|
||||
get_parallel().pp_size,
|
||||
get_parallel().tp_size,
|
||||
node_rank=node_idx,
|
||||
)
|
||||
)
|
||||
@@ -369,7 +374,7 @@ class RayEngine(Engine):
|
||||
f"bundle_indices={bundle_indices}"
|
||||
)
|
||||
|
||||
tp_size = get_parallel().config.tp_size
|
||||
tp_size = get_parallel().tp_size
|
||||
for rank in range(world_size):
|
||||
pp_rank = rank // tp_size
|
||||
tp_rank = rank % tp_size
|
||||
@@ -448,7 +453,7 @@ class RayEngine(Engine):
|
||||
RayDataParallelController,
|
||||
)
|
||||
|
||||
parallel = get_parallel().config
|
||||
parallel = get_parallel()
|
||||
if parallel.enable_dp_attention:
|
||||
# DP attention folds DP into TP — total GPUs = tp_size * pp_size
|
||||
total_gpus = parallel.tp_size * parallel.pp_size
|
||||
|
||||
@@ -13,16 +13,13 @@
|
||||
# ==============================================================================
|
||||
"""A single structured accessor for process-static runtime state.
|
||||
|
||||
``get_parallel()`` returns a ``ParallelContext`` whose bare attributes — tp / dcp
|
||||
/ pp / moe / attn size and rank, plus the process-group handles — each delegate
|
||||
live to the canonical getter in ``distributed.parallel_state`` /
|
||||
``layers.dp_attention``. Returned values are exactly what those getters return;
|
||||
this is a read-through wrapper, not a cache. It gives call-sites one import and
|
||||
one naming scheme in place of a dozen free functions, plus a test-only
|
||||
``override()`` hook to force a topology without monkeypatching the underlying
|
||||
getters. The resolved parallel **configuration** is the same object's ``config``
|
||||
hop (``get_parallel().config.tp_size``), which reads the published ``parallel``
|
||||
bag: bare is the live group, ``config`` is what was configured.
|
||||
``get_parallel()`` returns a ``ParallelContext``. Ranks and process-group handles
|
||||
read through **live** to the canonical getter in ``distributed.parallel_state`` /
|
||||
``layers.dp_attention`` — exactly what those getters return, a read-through
|
||||
wrapper and not a cache. Every other name, the sizes included, is a leaf of the
|
||||
published ``parallel`` bag. It gives call-sites one import and one naming scheme
|
||||
in place of a dozen free functions, plus an ``override()`` hook to force a
|
||||
topology without monkeypatching the underlying getters.
|
||||
|
||||
``get_server_args()`` returns the process-wide ``ServerArgs``. This is the
|
||||
user's raw input, kept **read-only** for debug and reproduction; what
|
||||
@@ -135,21 +132,26 @@ _PARALLEL_FIELDS = frozenset(
|
||||
|
||||
|
||||
class ParallelContext:
|
||||
"""Parallel-topology namespace: the live groups bare, configuration under
|
||||
``config``.
|
||||
"""Parallel-topology namespace: one spelling per name.
|
||||
|
||||
``get_parallel().tp_size`` and its size / rank / group siblings are
|
||||
read-through ``@property`` over the canonical getters, so they answer with
|
||||
the **live** process groups and raise before distributed init. The resolved
|
||||
parallel **configuration** is one hop away, on the published bag:
|
||||
``get_parallel().config.tp_size``, ``.config.nccl_port``. It answers in any
|
||||
process at any point after publish, and follows a post-publish ``override``.
|
||||
Ranks and group handles are read-through ``@property`` over the canonical
|
||||
getters, so they answer with the **live** process groups and raise before
|
||||
distributed init. Every other name — ``tp_size`` and its size siblings
|
||||
included, alongside config-only leaves such as ``nccl_port`` — is answered
|
||||
from the published ``parallel`` bag, in any process at any point after
|
||||
publish.
|
||||
|
||||
The two disagree by design, so which one a call site wants is spelled at the
|
||||
call site — no ``config`` means live. Elastic EP scales the live world away
|
||||
from the configured one, and ``initialize_model_parallel`` aliases
|
||||
``_MOE_DP`` to ``_ATTN_CP`` when ``attn_cp_size > moe_dp_size``, which makes
|
||||
a live comparison of that pair degenerate.
|
||||
A size is read from the configuration because the groups are built at
|
||||
exactly the configured widths. Two things do not follow that rule and are
|
||||
asked of the group itself: ``initialize_model_parallel`` aliases ``_MOE_DP``
|
||||
to ``_ATTN_CP`` when ``attn_cp_size > moe_dp_size``, so a reader that means
|
||||
the MoE communicator's width calls ``get_moe_cp_size()``; and
|
||||
``patch_tensor_parallel_group`` runs a scope under a different TP group,
|
||||
which it declares by overriding ``tp_size``, ``tp_rank`` and ``tp_group``
|
||||
for its duration. Elastic EP is a third case, and it needs no rule here: it
|
||||
scales ``ep_size`` / ``dp_size`` on the published bag while the group
|
||||
coordinators keep the width they were constructed with, so the two are
|
||||
different names rather than two answers to one name.
|
||||
"""
|
||||
|
||||
__slots__ = ("_overrides", "_config")
|
||||
@@ -158,28 +160,15 @@ class ParallelContext:
|
||||
self._overrides = {}
|
||||
self._config = None # parallel config bag, wired at publish
|
||||
|
||||
@property
|
||||
def config(self) -> _ConfigBag:
|
||||
"""The published ``parallel`` config bag.
|
||||
|
||||
Reads the slot directly: ``parallel`` sits outside the per-role
|
||||
namespace table (every process reads topology config), so no role check
|
||||
applies here. The body stays
|
||||
dynamo-traceable — ``get_parallel().config.moe_dense_tp_size`` and the
|
||||
gate helpers over it run inside compiled model forwards.
|
||||
"""
|
||||
config = self._config
|
||||
if config is None:
|
||||
raise ValueError("config namespace 'parallel' not published")
|
||||
return config
|
||||
|
||||
def __getattr__(self, name):
|
||||
# Reached only for names with no live @property: the bare config leaves.
|
||||
if name.startswith("_"):
|
||||
# This also breaks the recursion when the ``_config`` slot itself is
|
||||
# still unset (pickle/copy protocols probe attributes before
|
||||
# __init__ runs).
|
||||
raise AttributeError(name)
|
||||
overrides = self._overrides
|
||||
if name in overrides:
|
||||
return overrides[name]
|
||||
config = self._config
|
||||
if config is not None:
|
||||
if name in config._fields:
|
||||
@@ -214,18 +203,10 @@ class ParallelContext:
|
||||
def world_rank(self) -> int:
|
||||
return self._v("world_rank", _ps().get_world_rank)
|
||||
|
||||
@property
|
||||
def tp_size(self) -> int:
|
||||
return self._v("tp_size", _ps().get_tensor_model_parallel_world_size)
|
||||
|
||||
@property
|
||||
def tp_rank(self) -> int:
|
||||
return self._v("tp_rank", _ps().get_tensor_model_parallel_rank)
|
||||
|
||||
@property
|
||||
def pp_size(self) -> int:
|
||||
return self._v("pp_size", _ps().get_pipeline_model_parallel_world_size)
|
||||
|
||||
@property
|
||||
def pp_rank(self) -> int:
|
||||
return self._v("pp_rank", _ps().get_pipeline_model_parallel_rank)
|
||||
@@ -238,10 +219,6 @@ class ParallelContext:
|
||||
def moe_ep_rank(self) -> int:
|
||||
return self._v("moe_ep_rank", _ps().get_moe_expert_parallel_rank)
|
||||
|
||||
@property
|
||||
def moe_dp_size(self) -> int:
|
||||
return self._v("moe_dp_size", _ps().get_moe_data_parallel_world_size)
|
||||
|
||||
@property
|
||||
def moe_dp_rank(self) -> int:
|
||||
return self._v("moe_dp_rank", _ps().get_moe_data_parallel_rank)
|
||||
@@ -262,18 +239,10 @@ class ParallelContext:
|
||||
def attn_tp_rank(self) -> int:
|
||||
return self._v("attn_tp_rank", _ps().get_attn_tensor_model_parallel_rank)
|
||||
|
||||
@property
|
||||
def attn_cp_size(self) -> int:
|
||||
return self._v("attn_cp_size", _ps().get_attn_context_model_parallel_world_size)
|
||||
|
||||
@property
|
||||
def attn_cp_rank(self) -> int:
|
||||
return self._v("attn_cp_rank", _ps().get_attn_context_model_parallel_rank)
|
||||
|
||||
@property
|
||||
def dcp_size(self) -> int:
|
||||
return self._v("dcp_size", _ps().get_dcp_world_size)
|
||||
|
||||
@property
|
||||
def dcp_rank(self) -> int:
|
||||
return self._v("dcp_rank", _ps().get_dcp_rank)
|
||||
@@ -283,14 +252,15 @@ class ParallelContext:
|
||||
def getter():
|
||||
if _ps().get_dcp_group_no_assert() is None:
|
||||
return False
|
||||
return self.dcp_size > 1
|
||||
return _ps().get_dcp_world_size() > 1
|
||||
|
||||
return self._v("dcp_enabled", getter)
|
||||
|
||||
@property
|
||||
def attn_dcp_size(self) -> int:
|
||||
return self._v(
|
||||
"attn_dcp_size", lambda: self.dcp_size if self.dcp_enabled else 1
|
||||
"attn_dcp_size",
|
||||
lambda: _ps().get_dcp_world_size() if self.dcp_enabled else 1,
|
||||
)
|
||||
|
||||
@property
|
||||
@@ -1163,8 +1133,8 @@ def get_forward() -> ForwardFlags:
|
||||
# --- Resolved config namespaces -------------------------
|
||||
# Each returns the top-level snapshot bag; reads are `get_exec().moe.field` etc.
|
||||
# All fail with ValueError("... not published") until publish has projected them.
|
||||
# ``parallel`` has no getter of its own: its bag is reached as
|
||||
# ``get_parallel().config``, alongside the live topology it belongs to.
|
||||
# ``parallel`` has no bag getter: ``get_parallel()`` answers its leaves
|
||||
# directly, alongside the live topology they belong to.
|
||||
def get_device() -> _ConfigBag:
|
||||
return _CONTEXT.config_bag("device")
|
||||
|
||||
@@ -1215,7 +1185,7 @@ def get_observability() -> _ConfigBag:
|
||||
# table declares which top-level config namespaces each role reads. ``None``
|
||||
# means the full tree — either the role genuinely needs everything (scheduler)
|
||||
# or its deployment shape has not been audited yet (restrict only what smoke
|
||||
# coverage can verify). ``parallel`` is served by ``get_parallel().config`` and
|
||||
# coverage can verify). ``parallel`` is served by ``get_parallel()`` and
|
||||
# every process legitimately reads topology config, so it is not in this table.
|
||||
#
|
||||
# ``SGLANG_ROLE_NAMESPACES`` selects the mode (read once at import):
|
||||
@@ -1611,11 +1581,7 @@ def max_prefill_buffer_tokens() -> int:
|
||||
else 0
|
||||
)
|
||||
tokens = chunked
|
||||
if (
|
||||
schedule.enable_dynamic_chunking
|
||||
and get_parallel().config.pp_size > 1
|
||||
and chunked
|
||||
):
|
||||
if schedule.enable_dynamic_chunking and get_parallel().pp_size > 1 and chunked:
|
||||
tokens = max(
|
||||
tokens, schedule.max_prefill_tokens or 0, math.ceil(chunked * 1.25)
|
||||
)
|
||||
@@ -1645,7 +1611,7 @@ def pre_capture_activation_reserve_mb(gpu_mem: float | None) -> float:
|
||||
activation_tokens = max(schedule.chunked_prefill_size, 2048)
|
||||
else:
|
||||
activation_tokens = max(schedule.max_prefill_tokens, 2048)
|
||||
parallel = get_parallel().config
|
||||
parallel = get_parallel()
|
||||
reserved_mem = (
|
||||
512 + activation_tokens * 1.5 + parallel.tp_size * parallel.pp_size / 8 * 1024
|
||||
)
|
||||
|
||||
@@ -11024,20 +11024,16 @@ def resolve_encoder_transfer_backend(
|
||||
return "zmq_to_scheduler"
|
||||
|
||||
|
||||
def compute_world_size(config) -> int:
|
||||
"""Return the total GPU count across all data-parallel replicas.
|
||||
def compute_world_size(
|
||||
*, enable_dp_attention: bool, dp_size: int, tp_size: int, pp_size: int
|
||||
) -> int:
|
||||
"""Total GPU count across all data-parallel replicas.
|
||||
|
||||
Takes the resolved topology -- the published `parallel` bag, or a view over
|
||||
the declarations. `enable_dp_attention` and `dp_size` are both resolution's
|
||||
answers (`_handle_dwdp` fills the pair, DeepSeek MLA context parallelism
|
||||
turns DP attention on), so a raw-record read would size the world from what
|
||||
the operator typed.
|
||||
Takes the values rather than a config object: the two sizes are the widths
|
||||
the launch asked for, which the Ray driver needs before any process group
|
||||
exists, and passing a context would hand it the live groups instead.
|
||||
"""
|
||||
return (
|
||||
(1 if config.enable_dp_attention else config.dp_size)
|
||||
* config.tp_size
|
||||
* config.pp_size
|
||||
)
|
||||
return (1 if enable_dp_attention else dp_size) * tp_size * pp_size
|
||||
|
||||
|
||||
def m3_fp8_attn_gemm_enabled(args) -> bool:
|
||||
|
||||
@@ -113,7 +113,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
|
||||
self.device_module = torch.get_device_module(self.device)
|
||||
self.tp_size = model_runner.ps.tp_size
|
||||
self.attn_dp_size = model_runner.ps.attn_dp_size
|
||||
self.pp_size = get_parallel().config.pp_size
|
||||
self.pp_size = get_parallel().pp_size
|
||||
self.enable_torch_compile = get_flags().capture.enable_torch_compile
|
||||
self.disable_padding = model_runner.server_args.disable_cuda_graph_padding
|
||||
self.require_gathered_buffer = require_gathered_buffer()
|
||||
|
||||
@@ -99,7 +99,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
|
||||
self.device_module = torch.get_device_module(self.device)
|
||||
self.tp_size = model_runner.ps.tp_size
|
||||
self.attn_dp_size = model_runner.ps.attn_dp_size
|
||||
self.pp_size = get_parallel().config.pp_size
|
||||
self.pp_size = get_parallel().pp_size
|
||||
self.enable_torch_compile = get_flags().capture.enable_torch_compile
|
||||
self.disable_padding = model_runner.server_args.disable_cuda_graph_padding
|
||||
self.require_gathered_buffer = require_gathered_buffer()
|
||||
|
||||
@@ -99,7 +99,7 @@ class FrozenKVMTPCudaGraphRunner(DecodeCudaGraphRunner):
|
||||
self.require_attn_tp_gather = require_attn_tp_gather()
|
||||
self.tp_size = self.model_runner.ps.tp_size
|
||||
self.attn_dp_size = self.model_runner.ps.attn_dp_size
|
||||
self.pp_size = get_parallel().config.pp_size
|
||||
self.pp_size = get_parallel().pp_size
|
||||
self.speculative_num_steps = get_spec().speculative_num_steps
|
||||
self.topk = get_spec().speculative_eagle_topk
|
||||
self.draft_attn_backend = frozen_kv_mtp_worker.draft_attn_backend
|
||||
|
||||
@@ -155,7 +155,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
|
||||
self.device_module = torch.get_device_module(self.device)
|
||||
self.tp_size = model_runner.ps.tp_size
|
||||
self.dp_size = get_parallel().dp_size
|
||||
self.pp_size = get_parallel().config.pp_size
|
||||
self.pp_size = get_parallel().pp_size
|
||||
self.enable_torch_compile = get_flags().capture.enable_torch_compile
|
||||
self.disable_padding = model_runner.server_args.disable_cuda_graph_padding
|
||||
self.require_gathered_buffer = require_gathered_buffer()
|
||||
|
||||
@@ -3756,7 +3756,7 @@ def require_mlp_tp_gather():
|
||||
else:
|
||||
return (
|
||||
get_parallel().moe_dense_tp_size
|
||||
> get_parallel().config.tp_size // get_parallel().dp_size
|
||||
> get_parallel().tp_size // get_parallel().dp_size
|
||||
)
|
||||
else:
|
||||
return False
|
||||
@@ -3782,7 +3782,7 @@ def require_attn_tp_gather():
|
||||
or get_parallel().moe_dense_tp_size is not None
|
||||
):
|
||||
if get_parallel().enable_dp_attention:
|
||||
return get_parallel().dp_size < get_parallel().config.tp_size
|
||||
return get_parallel().dp_size < get_parallel().tp_size
|
||||
else:
|
||||
return True
|
||||
else:
|
||||
|
||||
@@ -163,8 +163,8 @@ def _contains_tensor_container(value) -> bool:
|
||||
|
||||
def get_vmm_feature_consumer_count() -> int:
|
||||
if get_parallel().enable_dp_attention:
|
||||
return get_parallel().config.tp_size // get_parallel().dp_size
|
||||
return get_parallel().config.tp_size
|
||||
return get_parallel().tp_size // get_parallel().dp_size
|
||||
return get_parallel().tp_size
|
||||
|
||||
|
||||
class CudaVmmMemoryPool:
|
||||
|
||||
@@ -494,6 +494,7 @@ class IpcModelLoader(BaseModelLoader):
|
||||
|
||||
try:
|
||||
# Build engine's config fingerprint
|
||||
from sglang.srt.layers.dp_attention import get_moe_cp_size
|
||||
from sglang.srt.runtime_context import get_exec, get_parallel
|
||||
|
||||
ps = get_parallel()
|
||||
@@ -504,7 +505,7 @@ class IpcModelLoader(BaseModelLoader):
|
||||
pp_rank = ps.pp_rank
|
||||
|
||||
ep_size = ps.moe_ep_size
|
||||
moe_dp_size = ps.moe_dp_size
|
||||
moe_dp_size = get_moe_cp_size()
|
||||
moe_dp_rank = ps.moe_dp_rank
|
||||
moe_ep_rank = ps.moe_ep_rank
|
||||
|
||||
|
||||
Reference in New Issue
Block a user