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

[DP]Fix sync bubble in adjust_num_token_non_padded_for_attn_tp (#16178)

Browse Source
This commit is contained in:
Baizhou Zhang
2025-12-31 17:58:18 +08:00
committed by GitHub
parent 5ed384d07b
commit e47afa0237
1 changed files with 8 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/sglang/srt/model_executor/forward_batch_info.py
View File

@@ -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,
)