From e1f2f9d1fa84cd1b8d9020377fdd707b3a485687 Mon Sep 17 00:00:00 2001 From: Nicolas Castet <26874160+nvcastet@users.noreply.github.com> Date: Wed, 29 Jul 2026 09:39:36 -0500 Subject: [PATCH] Disable extra NCCL CUDA event synchronization with symm mem (#27089) --- .../device_communicators/pynccl.py | 18 +++- .../device_communicators/pynccl_wrapper.py | 102 ++++++++++++++++++ .../sglang/srt/distributed/parallel_state.py | 1 + 3 files changed, 118 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/distributed/device_communicators/pynccl.py b/python/sglang/srt/distributed/device_communicators/pynccl.py index cb74d4d35..9072ec096 100644 --- a/python/sglang/srt/distributed/device_communicators/pynccl.py +++ b/python/sglang/srt/distributed/device_communicators/pynccl.py @@ -16,6 +16,7 @@ from sglang.srt.distributed.device_communicators.pynccl_wrapper import ( buffer_type, cudaStream_t, ncclComm_t, + ncclConfig_t, ncclDataTypeEnum, ncclRedOpTypeEnum, ncclUniqueId, @@ -33,6 +34,7 @@ class PyNcclCommunicator: group: Union[ProcessGroup, StatelessProcessGroup], device: Union[int, str, torch.device], library_path: Optional[str] = None, + is_symmetric_memory_enabled: bool = False, ): """ Args: @@ -42,6 +44,7 @@ class PyNcclCommunicator: it will be bind to f"cuda:{local_rank}". library_path: the path to the NCCL library. If None, it will use the default library path. + is_symmetric_memory_enabled: whether symmetric memory is enabled. It is the caller's responsibility to make sure each communicator is bind to a unique device. """ @@ -108,9 +111,18 @@ class PyNcclCommunicator: # `torch.cuda.device` is a context manager that changes the # current cuda device to the specified one with torch.cuda.device(device): - self.comm: ncclComm_t = self.nccl.ncclCommInitRank( - self.world_size, self.unique_id, self.rank - ) + if is_symmetric_memory_enabled: + # When symmetric memory is enabled, disable internal + # NCCL cuda event synchronizations to improve performance. + config = ncclConfig_t.create() + config.graphUsageMode = 1 + self.comm: ncclComm_t = self.nccl.ncclCommInitRankConfig( + self.world_size, self.unique_id, self.rank, config + ) + else: + self.comm: ncclComm_t = self.nccl.ncclCommInitRank( + self.world_size, self.unique_id, self.rank + ) warmup_stream = torch.cuda.Stream() # A small all_reduce for warmup. diff --git a/python/sglang/srt/distributed/device_communicators/pynccl_wrapper.py b/python/sglang/srt/distributed/device_communicators/pynccl_wrapper.py index adea26090..d9afda433 100644 --- a/python/sglang/srt/distributed/device_communicators/pynccl_wrapper.py +++ b/python/sglang/srt/distributed/device_communicators/pynccl_wrapper.py @@ -74,6 +74,75 @@ ncclComm_t = ctypes.c_void_p ncclWindow_t = ctypes.c_void_p +# Sentinels NCCL uses to mark "use the default value", see +# NCCL_CONFIG_UNDEF_INT / NCCL_CONFIG_UNDEF_PTR / NCCL_API_MAGIC in nccl.h.in. +NCCL_CONFIG_UNDEF_INT = -(2**31) # INT_MIN +NCCL_CONFIG_UNDEF_PTR = None +NCCL_API_MAGIC = 0xCAFEBEEF + +# Version code of the header layout the struct below was extracted from, i.e. +# the equivalent of NCCL_VERSION_CODE baked into NCCL_CONFIG_INITIALIZER at +# compile time. It declares which fields this initializer knows about, so a +# (possibly older/newer) runtime libnccl.so honors exactly the fields this +# struct exposes. Bump this together with the struct layout when re-extracting +# from nccl.h.in. +NCCL_CONFIG_VERSION = 23000 # NCCL_VERSION(2, 30, 0) + + +class ncclConfig_t(ctypes.Structure): + """Mirror of ``ncclConfig_t`` (``ncclConfig_v22800``) from nccl.h.in. + + Fields are append-only across NCCL versions, and NCCL only copies + ``min(config.size, sizeof(its own struct))`` bytes, so declaring the full + (newest) layout stays compatible with both older and newer libnccl.so. + """ + + _fields_ = [ + # attributes that users should never touch. + ("size", ctypes.c_size_t), + ("magic", ctypes.c_uint), + ("version", ctypes.c_uint), + # attributes that users are able to customize. + ("blocking", ctypes.c_int), + ("cgaClusterSize", ctypes.c_int), + ("minCTAs", ctypes.c_int), + ("maxCTAs", ctypes.c_int), + ("netName", ctypes.c_char_p), + ("splitShare", ctypes.c_int), + ("trafficClass", ctypes.c_int), + ("commName", ctypes.c_char_p), + ("collnetEnable", ctypes.c_int), + ("CTAPolicy", ctypes.c_int), + ("shrinkShare", ctypes.c_int), + ("nvlsCTAs", ctypes.c_int), + ("nChannelsPerNetPeer", ctypes.c_int), + ("nvlinkCentricSched", ctypes.c_int), + ("graphUsageMode", ctypes.c_int), + ("numRmaCtx", ctypes.c_int), + ("maxP2pPeers", ctypes.c_int), + ] + + @classmethod + def create(cls) -> "ncclConfig_t": + """Build a config equivalent to ``NCCL_CONFIG_INITIALIZER``. + + All customizable fields are set to the UNDEF sentinels; callers may + override individual fields afterwards. + """ + config = cls() + config.size = ctypes.sizeof(cls) + config.magic = NCCL_API_MAGIC + config.version = NCCL_CONFIG_VERSION + # Must explicitly set sentinels: ctypes zero-initializes, and 0 is a + # valid value for many fields (e.g. blocking=0 means non-blocking). + for name, ctype in cls._fields_[3:]: + if ctype is ctypes.c_char_p: + setattr(config, name, NCCL_CONFIG_UNDEF_PTR) + else: + setattr(config, name, NCCL_CONFIG_UNDEF_INT) + return config + + class ncclUniqueId(ctypes.Structure): _fields_ = [("internal", ctypes.c_byte * 128)] @@ -173,6 +242,20 @@ class NCCLLibrary: ncclResult_t, [ctypes.POINTER(ncclComm_t), ctypes.c_int, ncclUniqueId, ctypes.c_int], ), + # ncclResult_t ncclCommInitRankConfig( + # ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank, + # ncclConfig_t* config); + Function( + "ncclCommInitRankConfig", + ncclResult_t, + [ + ctypes.POINTER(ncclComm_t), + ctypes.c_int, + ncclUniqueId, + ctypes.c_int, + ctypes.POINTER(ncclConfig_t), + ], + ), # ncclResult_t ncclAllReduce( # const void* sendbuff, void* recvbuff, size_t count, # ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm, @@ -396,6 +479,24 @@ class NCCLLibrary: self.NCCL_CHECK(self._funcs["ncclGetUniqueId"](ctypes.byref(unique_id))) return unique_id + def ncclCommInitRankConfig( + self, + world_size: int, + unique_id: ncclUniqueId, + rank: int, + config: Optional[ncclConfig_t] = None, + ) -> ncclComm_t: + comm = ncclComm_t() + if config is None: + # Equivalent to NCCL_CONFIG_INITIALIZER (header-layout version). + config = ncclConfig_t.create() + self.NCCL_CHECK( + self._funcs["ncclCommInitRankConfig"]( + ctypes.byref(comm), world_size, unique_id, rank, ctypes.byref(config) + ) + ) + return comm + def ncclCommInitRank( self, world_size: int, unique_id: ncclUniqueId, rank: int ) -> ncclComm_t: @@ -561,6 +662,7 @@ __all__ = [ "ncclDataTypeEnum", "ncclRedOpTypeEnum", "ncclUniqueId", + "ncclConfig_t", "ncclComm_t", "cudaStream_t", "buffer_type", diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 2d3f8932a..9347cde21 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -429,6 +429,7 @@ class GroupCoordinator: self.pynccl_comm = PyNcclCommunicator( group=self.cpu_group, device=self.device, + is_symmetric_memory_enabled=self.is_symmetric_memory_enabled(), ) self.pymscclpp_comm: Optional[PyMscclppCommunicator] = None