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

Avoid logits multimem all-gather on cross-node TP groups (#29881)

Browse Source
This commit is contained in:
san-tian
2026-07-04 00:24:58 -07:00
committed by GitHub
parent 88db9e033a
commit 576dc31e33
1 changed files with 14 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/device_communicators/triton_symm_mem_ag.py
+14
View File
@@ -462,6 +462,20 @@ class MultimemAllGatherer:
self._skip_entry_sync = skip_entry_sync self._skip_entry_sync = skip_entry_sync
# None => always NCCL; _UNINIT => build on first eager call. # None => always NCCL; _UNINIT => build on first eager call.
self._state = self._UNINIT if enabled else None self._state = self._UNINIT if enabled else None
if self._state is self._UNINIT:
# Lazy import avoids a module-load dependency on the distributed facade.
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.parallel_state import in_the_same_node_as
tp_group = get_tp_group()
if tp_group.world_size > 1 and not all(
in_the_same_node_as(tp_group.cpu_group, source_rank=0)
):
logger.warning(
"multimem all-gather disabled because the TP group spans "
"across nodes."
)
self._state = None
def __call__(self, x: torch.Tensor) -> torch.Tensor: def __call__(self, x: torch.Tensor) -> torch.Tensor:
state = self._state state = self._state