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:
Cheng Wan
2026-08-27 12:56:42 -07:00
committed by GitHub
co-authored by Claude Opus 5
parent ca1d7ed8e6
commit fd40a331bf
62 changed files with 439 additions and 1313 deletions
+5 -2
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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 = (
+3 -3
View File
@@ -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
+1 -1
View File
@@ -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
+4 -4
View File
@@ -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)
+2 -2
View File
@@ -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,
),
+1 -1
View File
@@ -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
+9 -5
View File
@@ -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
+4 -6
View File
@@ -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
)
)
+1 -1
View File
@@ -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
)
+23 -18
View File
@@ -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(
+5 -3
View File
@@ -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,
+1 -1
View File
@@ -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()
+1 -1
View File
@@ -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
+11 -6
View File
@@ -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
+36 -70
View File
@@ -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
)
+8 -12
View File
@@ -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()
+2 -2
View File
@@ -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:
+2 -1
View File
@@ -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