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

refactor context parallel state (#17213)

Browse Source 
Co-authored-by: Shunkang <182541032+Shunkangz@users.noreply.github.co>
This commit is contained in:
dongjiyingdjy
2026-02-13 23:18:17 +08:00
committed by GitHub
parent 0012d6a4eb
commit 8b4c364960
27 changed files with 847 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
python/sglang/srt/managers/data_parallel_controller.py
View File

@@ -418,6 +418,8 @@ class DataParallelController:
tp_size_per_node * (server_args.node_rank % nnodes_per_tp_group + 1),
)
attn_cp_rank = 0
moe_dp_rank = 0
for pp_rank in pp_rank_range:
for tp_rank in tp_rank_range:
rank_port_args = port_args
@@ -429,6 +431,7 @@ class DataParallelController:
tp_rank,
server_args.tp_size,
server_args.dp_size,
server_args.attn_cp_size,
)
# compute zmq ports for this dp rank
rank_port_args = PortArgs.init_new(
@@ -445,7 +448,30 @@ class DataParallelController:
+ ((pp_rank % pp_size_per_node) * tp_size_per_node)
+ (tp_rank % tp_size_per_node) * server_args.gpu_id_step
)
moe_ep_rank = tp_rank // (server_args.tp_size // server_args.ep_size)
attn_dp_size = (
server_args.dp_size if server_args.enable_dp_attention else 1
)
# Parallelism hierarchy (outermost to innermost):
# - 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 // server_args.attn_cp_size
)
attn_cp_rank = (tp_rank // attn_tp_size) % server_args.attn_cp_size
moe_dp_rank = tp_rank // (
server_args.tp_size // server_args.moe_dp_size
)
moe_ep_rank = (
tp_rank
% (server_args.tp_size // server_args.moe_dp_size)
// (
server_args.tp_size
// server_args.moe_dp_size
// server_args.ep_size
)
)
with self.env_lock, maybe_reindex_device_id(gpu_id) as gpu_id:
proc = mp.Process(
target=self.run_scheduler_process_func,
@@ -454,6 +480,8 @@ class DataParallelController:
rank_port_args,
gpu_id,
tp_rank,
attn_cp_rank,
moe_dp_rank,
moe_ep_rank,
pp_rank,
dp_rank,