EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

[Bugfix] Avoid unnecessary reduce-scatter call in prepare_mlp (#9169)

Browse Source
This commit is contained in:
Huaixin Chang
2025-08-14 12:04:41 +08:00
committed by GitHub
parent 0fc8bf2cd4
commit 98457c0453
1 changed files with 4 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/layers/dp_attention.py
View File

@@ -292,6 +292,10 @@ def _dp_gather_via_all_gather(
forward_batch: ForwardBatch, forward_batch: ForwardBatch,
is_partial: bool, is_partial: bool,
): ):
if get_attention_tp_size() == 1:
get_tp_group().all_gather_into_tensor(global_tokens, local_tokens)
return
if not is_partial: if not is_partial:
if get_attention_tp_rank() != 0: if get_attention_tp_rank() != 0:
local_tokens.fill_(0) local_tokens.fill_(0)