Disable extra NCCL CUDA event synchronization with symm mem (#27089)

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.

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",

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