From fea2d5211d174012f04d0da96546bd9c34da6868 Mon Sep 17 00:00:00 2001 From: Xuchun Shang Date: Thu, 18 Dec 2025 16:59:56 +0800 Subject: [PATCH] [bug fix][pp] fix inconsistent latency between tp (#15379) Signed-off-by: Xuchun Shang Co-authored-by: Shangming Cai --- python/sglang/srt/managers/scheduler_pp_mixin.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index ea15cfd8b..c7d0ea51b 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -640,6 +640,16 @@ class SchedulerPPMixin: f"seq_lens={seq_lens}, latencies_ms={latencies}" ) + if self.attn_tp_size > 1: + data_to_sync_tp = [seq_lens, latencies] + data_to_sync_tp = broadcast_pyobj( + data_to_sync_tp, + self.attn_tp_group.rank, + self.attn_tp_cpu_group, + src=self.attn_tp_group.ranks[0], + ) + seq_lens, latencies = data_to_sync_tp + # Broadcast data to all ranks if torch.distributed.is_available() and torch.distributed.is_initialized(): data_to_sync = [seq_lens, latencies] @@ -869,7 +879,7 @@ class SchedulerPPMixin: else: data = None - if self.attn_tp_size != 1: + if self.attn_tp_size > 1: data = broadcast_pyobj( data, self.attn_tp_group.rank,