minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[PP] bugfix: include CP size in PP rank offset (#29571)

Browse Source
This commit is contained in:
SSSunzt
2026-06-29 16:01:41 +08:00
committed by GitHub
parent c7b9b92d9a
commit 91cf159696
3 changed files with 143 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler_components/request_receiver.py
+3 -1
View File
@@ -124,7 +124,9 @@ class SchedulerRequestReceiver:
recv_reqs = None
else:
if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0:
dp_offset = self.ps.attn_dp_rank * self.ps.attn_tp_size
dp_offset = (
self.ps.attn_dp_rank * self.ps.attn_cp_size * self.ps.attn_tp_size
)
recv_reqs = point_to_point_pyobj(
[],
self.ps.pp_rank * self.ps.tp_size + dp_offset,
python/sglang/srt/managers/scheduler_pp_mixin.py
+6 -2
View File
@@ -927,7 +927,9 @@ class SchedulerPPMixin:
def _pp_send_pyobj_to_next_stage(self: Scheduler, data, async_send: bool = False):
p2p_work = []
if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0:
dp_offset = self.ps.attn_dp_rank * self.ps.attn_tp_size
dp_offset = (
self.ps.attn_dp_rank * self.ps.attn_cp_size * self.ps.attn_tp_size
)
p2p_work = point_to_point_pyobj(
data,
self.ps.pp_rank * self.ps.tp_size + dp_offset,
@@ -940,7 +942,9 @@ class SchedulerPPMixin:
def _pp_recv_pyobj_from_prev_stage(self: Scheduler):
if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0:
dp_offset = self.ps.attn_dp_rank * self.ps.attn_tp_size
dp_offset = (
self.ps.attn_dp_rank * self.ps.attn_cp_size * self.ps.attn_tp_size
)
data = point_to_point_pyobj(
[],
self.ps.pp_rank * self.ps.tp_size + dp_offset,