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

fix: support PP2+CP8+TP8 (PP with context parallelism) (#19548)

Browse Source
This commit is contained in:
ybyang
2026-03-17 00:51:47 +08:00
committed by GitHub
parent 6489f77733
commit 289cbcf482
2 changed files with 11 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
python/sglang/srt/managers/scheduler_pp_mixin.py
View File

@@ -865,7 +865,7 @@ class SchedulerPPMixin:
def _pp_send_pyobj_to_next_stage(self: Scheduler, data, async_send: bool = False):
p2p_work = []
if self.attn_tp_rank == 0:
if self.attn_tp_rank == 0 and self.attn_cp_rank == 0:
dp_offset = self.attn_dp_rank * self.attn_tp_size
p2p_work = point_to_point_pyobj(
data,
@@ -878,7 +878,7 @@ class SchedulerPPMixin:
return p2p_work
def _pp_recv_pyobj_from_prev_stage(self: Scheduler):
if self.attn_tp_rank == 0:
if self.attn_tp_rank == 0 and self.attn_cp_rank == 0:
dp_offset = self.attn_dp_rank * self.attn_tp_size
data = point_to_point_pyobj(
[],
@@ -898,6 +898,14 @@ class SchedulerPPMixin:
src=self.attn_tp_group.ranks[0],
)
if self.attn_cp_size > 1:
data = broadcast_pyobj(
data,
self.attn_cp_group.rank,
self.attn_cp_cpu_group,
src=self.attn_cp_group.ranks[0],
)
return data
def _pp_prepare_tensor_dict(

2
python/sglang/srt/server_args.py
View File

@@ -1474,6 +1474,7 @@ class ServerArgs:
assert (
self.tp_size == 8
), "Current multi-machine CP support suffers from precision issues. So context parallel only support Single machine(tp_size == 8)"
self.attn_cp_size = self.tp_size
logger.warning(
f"Enable Context Parallel opt for deeeseekv3.2-DSA, Setting dp_size == {self.dp_size} and moe_dense_tp_size == {self.moe_dense_tp_size}, ep_size == {self.ep_size}, tp_size == {self.tp_size}, kv_cache_dtype == {self.kv_cache_dtype}, moe_a2a_backend {self.moe_a2a_backend} "
@@ -2438,7 +2439,6 @@ class ServerArgs:
assert (
self.tp_size % (self.dp_size * self.attn_cp_size) == 0
), "tp_size must be divisible by dp_size * attn_cp_size"
assert self.pp_size == 1, "PP is not supported with context parallelism"
if self.moe_dp_size > 1:
# The tp_size is the world size, not the real tensor parallel size