From e47afa0237756e0968a7e0ff8840c677dc298548 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Wed, 31 Dec 2025 17:58:18 +0800 Subject: [PATCH] [DP]Fix sync bubble in adjust_num_token_non_padded_for_attn_tp (#16178) --- .../sglang/srt/model_executor/forward_batch_info.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index e71b71b83..f43d690a3 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -216,6 +216,10 @@ def compute_local_num_token_non_padded( attn_tp_size = get_attention_tp_size() tokens_per_rank = num_tokens_per_dp // attn_tp_size + # Make sure global_num_token_non_padded is tensor so torch.clamp doesn't break + if isinstance(global_num_token_non_padded, int): + global_num_token_non_padded = torch.tensor(global_num_token_non_padded) + return torch.clamp( global_num_token_non_padded - tokens_per_rank * attn_tp_rank, 0, @@ -545,14 +549,15 @@ class ForwardBatch: from sglang.srt.utils.common import require_mlp_tp_gather dp_rank = get_attention_dp_rank() + assert self.global_num_tokens_cpu is not None if require_mlp_tp_gather(server_args): - num_tokens_per_dp = self.global_num_tokens_gpu[dp_rank] + num_tokens_per_dp = self.global_num_tokens_cpu[dp_rank] else: - num_tokens_per_dp = self.global_num_tokens_gpu[0] + num_tokens_per_dp = self.global_num_tokens_cpu[0] self.num_token_non_padded = compute_local_num_token_non_padded( - global_num_token_non_padded=self.num_token_non_padded, + global_num_token_non_padded=self.num_token_non_padded_cpu, num_tokens_per_dp=num_tokens_per_dp, )