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

2025-08-14 12:04:41 +08:00
parent 0fc8bf2cd4
commit 98457c0453
1 changed files with 4 additions and 0 deletions
--- a/python/sglang/srt/layers/dp_attention.py
+++ b/python/sglang/srt/layers/dp_attention.py
@@ -292,6 +292,10 @@ def _dp_gather_via_all_gather(
    forward_batch: ForwardBatch,
    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 get_attention_tp_rank() != 0:
            local_tokens.fill_(0)