MSCCL++ Integration (#22734)

Co-authored-by: Caio Rocha <caiorocha@microsof.com>
Co-authored-by: empyreus <rjsouza1995@gmail.com>

python/sglang/bench_one_batch.py

@@ -65,7 +65,10 @@ import torch
 import torch.distributed as dist
 
 from sglang.srt.configs.model_config import ModelConfig
-from sglang.srt.distributed.parallel_state import destroy_distributed_environment
+from sglang.srt.distributed.parallel_state import (
+    destroy_distributed_environment,
+    destroy_model_parallel,
+)
 from sglang.srt.entrypoints.engine import _set_envs_and_config
 from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.layers.moe import initialize_moe_config
@@ -936,6 +939,7 @@ def latency_test(
                 fout.write(json.dumps(result) + "\n")
 
     if server_args.tp_size > 1:
+        destroy_model_parallel()
         destroy_distributed_environment()
 
 
@@ -957,12 +961,17 @@ def main(server_args, bench_args):
 
     port_args = PortArgs.init_new(server_args)
 
+    # Calculate local ranks for multi-node setup
+    nranks_per_node = server_args.tp_size // server_args.nnodes
+    local_rank_start = server_args.node_rank * nranks_per_node
+    local_rank_end = local_rank_start + nranks_per_node
+
     if server_args.tp_size == 1:
         work_func(server_args, port_args, bench_args, 0, 0)
     else:
         workers = []
-        for tp_rank in range(server_args.tp_size):
-            with maybe_reindex_device_id(tp_rank) as gpu_id:
+        for tp_rank in range(local_rank_start, local_rank_end):
+            with maybe_reindex_device_id(tp_rank - local_rank_start) as gpu_id:
                 proc = multiprocessing.Process(
                     target=work_func,
                     args=(

python/sglang/srt/distributed/device_communicators/custom_all_reduce.py

@@ -328,7 +328,8 @@ class CustomAllreduce:
 
     def close(self):
         if not self.disabled and self._ptr:
-            ops.dispose(self._ptr)
+            if ops is not None:
+                ops.dispose(self._ptr)
             if _is_cuda:
                 self.free_shared_buffer(self.meta_ptrs)
                 self.free_shared_buffer(self.buffer_ptrs)

python/sglang/srt/distributed/device_communicators/custom_all_reduce_ops.py

@@ -16,8 +16,6 @@ _is_musa = is_musa()
 
 IS_CUSTOM_AR_AVAILABLE = _is_cuda or _is_hip or _is_musa
 IS_QUICK_AR_AVAILABLE = _is_hip
-# TODO(zyksir): mscclpp is untested on AMD and therefore disabled.
-IS_MSCCLPP_AR_AVAILABLE = _is_cuda
 
 try:
     import sgl_kernel.allreduce as _custom_ar
@@ -26,7 +24,6 @@ except ImportError as e:
         logger.warning("Failed to import from custom_ar with %r", e)
     IS_CUSTOM_AR_AVAILABLE = False
     IS_QUICK_AR_AVAILABLE = False
-    IS_MSCCLPP_AR_AVAILABLE = False
 
 # region IS_CUSTOM_AR_AVAILABLE
 
@@ -167,44 +164,3 @@ elif _is_hip:
 
 
 # endregion
-
-# region IS_MSCCLPP_AR_AVAILABLE
-
-if not IS_MSCCLPP_AR_AVAILABLE:
-    pass
-
-elif _is_cuda:
-
-    def mscclpp_generate_unique_id() -> bytes:
-        return _custom_ar.mscclpp_generate_unique_id()
-
-    def mscclpp_init_context(
-        unique_id: bytes,
-        rank: int,
-        world_size: int,
-        scratch: torch.Tensor,
-        put_buffer: torch.Tensor,
-        nranks_per_node: int,
-        rank_to_node: List[int],
-        rank_to_ib: List[int],
-        context_selection: int,
-    ) -> int:
-        return _custom_ar.mscclpp_init_context(
-            unique_id,
-            rank,
-            world_size,
-            scratch,
-            put_buffer,
-            nranks_per_node,
-            rank_to_node,
-            rank_to_ib,
-            context_selection,
-        )
-
-    def mscclpp_allreduce(
-        context: int, inp: torch.Tensor, out: torch.Tensor, nthreads: int, nblocks: int
-    ) -> None:
-        return _custom_ar.mscclpp_allreduce(context, inp, out, nthreads, nblocks)
-
-
-# endregion

python/sglang/srt/distributed/device_communicators/pymscclpp.py

@@ -1,108 +1,247 @@
-import bisect
+import importlib
 import logging
-import math
-import os
 from contextlib import contextmanager
-from enum import IntEnum
 from typing import Optional, Union
 
 import torch
 import torch.distributed as dist
 from torch.distributed import ProcessGroup, ReduceOp
 
-import sglang.srt.distributed.device_communicators.custom_all_reduce_ops as ops
-from sglang.srt.utils import is_hip
+from sglang.srt.compilation.piecewise_context_manager import (
+    get_pcg_capture_stream,
+    is_in_pcg_torch_compile,
+    is_in_piecewise_cuda_graph,
+)
+from sglang.srt.server_args import get_global_server_args
 
 logger = logging.getLogger(__name__)
 
-_is_hip = is_hip()
-
-
-class MscclContextSelection(IntEnum):
-    MSCCL1SHOT1NODELL = 1
-    MSCCL1SHOT2NODELL = 2
-
-
-def mscclpp_is_weak_contiguous(inp: torch.Tensor):
-    return inp.is_contiguous() or (
-        inp.storage().nbytes() - inp.storage_offset() * inp.element_size()
-        == inp.numel() * inp.element_size()
-    )
-
-
-def mscclpp_convert_to_bytes(size_str):
-    """
-    Converts a human-readable size string (e.g., "1MB", "2.5kb", "3 GB")
-    into the equivalent number of bytes using binary units.
-
-    Args:
-        size_str (str): A string representing size with unit (KB, MB, GB).
-
-    Returns:
-        int: Number of bytes.
-    """
-    size_str = size_str.strip().lower()
-
-    if not size_str:
-        raise ValueError("Empty input string")
-
-    # Extract numeric part and unit
-    for i in range(len(size_str)):
-        if not size_str[i].isdigit() and size_str[i] != ".":
-            break
-    num_str = size_str[:i]
-    unit = size_str[i:].strip()
-
-    try:
-        num = float(num_str)
-    except ValueError:
-        raise ValueError(f"Invalid numeric value in '{size_str}'")
-
-    # Conversion factors
-    if unit == "b":
-        return int(num)
-    elif unit == "kb":
-        return int(num * 1024)
-    elif unit == "mb":
-        return int(num * 1024 * 1024)
-    elif unit == "gb":
-        return int(num * 1024 * 1024 * 1024)
-    else:
-        raise ValueError(f"Unsupported unit: {unit}, support B, KB, MB, GB only")
-
-
-def mscclpp_bench_time(func, test_niter: int = 10, warmup_niter: int = 2):
-    # warmup
-    for _ in range(warmup_niter):
-        func()
-    start_event = torch.cuda.Event(enable_timing=True)
-    end_event = torch.cuda.Event(enable_timing=True)
-    torch.cuda.synchronize()
-    dist.barrier()
-    start_event.record()
-    for _ in range(test_niter):
-        func()
-    end_event.record()
-    end_event.synchronize()
-    func_cost_us = start_event.elapsed_time(end_event) / test_niter * 1000
-    return func_cost_us
-
 
 class PyMscclppCommunicator:
-    _SUPPORTED_WORLD_SIZES = [8, 16]
-    _MAX_BYTES = mscclpp_convert_to_bytes(os.getenv("SGLANG_MSCCLPP_MAX_BYTES", "1MB"))
+    _SUPPORTED_WORLD_SIZES = [8, 16, 32]
     _SUPPORTED_DTYPE = [torch.float, torch.float16, torch.bfloat16]
 
-    # max_bytes: max supported mscclpp allreduce size
-    # in A100 mscclpp is faster than nccl only under condition of msg size smaller than1MB
+    def _is_symm_mem_enabled(self) -> bool:
+        try:
+            return get_global_server_args().enable_symm_mem
+        except ValueError:
+            return False
+
+    def _is_weak_contiguous(self, inp: torch.Tensor):
+        return inp.is_contiguous() or (
+            inp.storage().nbytes() - inp.storage_offset() * inp.element_size()
+            == inp.numel() * inp.element_size()
+        )
+
+    def _get_tuned_config(self, size):
+        if size <= 512:
+            target_size = 512
+        elif size > 256 * 1024 * 1024:
+            target_size = 256 * 1024 * 1024
+        else:
+            target_size = 1 << (size - 1).bit_length()
+        return self.best_configs.get(target_size)
+
+    def _create_dsl_algorithms(self):
+        dsl_algos_config = []
+        n_nodes = self.world_size // self.nranks_per_node
+        if n_nodes == 2 or n_nodes == 4:
+            for tbg in [1, 2, 4, 8]:
+                for num_threads_per_block in [256, 512, 768, 1024]:
+                    spec = self.mscclpp.language.AlgoSpec(
+                        name=f"allreduce_{n_nodes}node_{tbg}TBG_{num_threads_per_block}TPB",
+                        collective=self.mscclpp.language.collectives.AllReduce(
+                            self.world_size, 1, True
+                        ),
+                        nranks_per_node=self.nranks_per_node,
+                        world_size=self.world_size,
+                        in_place=True,
+                        instances=1,
+                        protocol="LL",
+                        auto_sync=False,
+                        num_threads_per_block=num_threads_per_block,
+                        reuse_resources=True,
+                        use_double_scratch_buffer=True,
+                        min_message_size=tbg * (1 << 10),
+                        max_message_size=8 << 20,
+                        tags={"default": 1},
+                    )
+                    algo = self.mscclpp.compile(
+                        self.def_algo.allreduce_multi_nodes,
+                        spec,
+                        self.rank,
+                        thread_block_group_size=tbg,
+                    )
+                    dsl_algos_config.append((algo, [0], [0]))
+        return dsl_algos_config
+
+    def _create_native_algorithms(self):
+        navitve_algorithms_config = []
+        dlpack = self.mscclpp.RawGpuBuffer(1 << 27).to_dlpack(
+            data_type=str(torch.float16)
+        )
+        self.scratch_buffer = torch.utils.dlpack.from_dlpack(dlpack)
+        self.flag_buffer = torch.ones(128, dtype=torch.uint32, device="cuda")
+        algos = self.mscclpp_ext.AlgorithmCollectionBuilder().build_default_algorithms(
+            scratch_buffer=self.scratch_buffer.data_ptr(),
+            scratch_buffer_size=self.scratch_buffer.nbytes,
+            rank=self.rank,
+        )
+
+        for algo in algos:
+            if algo.name == "default_allreduce_nvls_packet":
+                algo.set_message_size_range(0, 512 << 10)
+                navitve_algorithms_config.append(
+                    (algo, [4, 8, 12, 16], [256, 512, 768, 1024])
+                )
+            if algo.name == "default_allreduce_packet":
+                algo.set_message_size_range(0, 2 << 20)
+                navitve_algorithms_config.append(
+                    (algo, [14, 21, 28, 42, 56], [256, 512, 768, 1024])
+                )
+            if algo.name == "default_allreduce_rsag_zero_copy":
+                algo.set_message_size_range(512 << 10, 4 << 30)
+                navitve_algorithms_config.append(
+                    (algo, [32, 48, 64, 128], [256, 512, 768, 1024])
+                )
+            if (
+                self.symm_mem_enabled
+                and algo.name == "default_allreduce_nvls_zero_copy"
+            ):
+                algo.set_message_size_range(512 << 10, 4 << 30)
+                navitve_algorithms_config.append(
+                    (algo, [4, 8, 12, 16, 32], [256, 512, 768, 1024])
+                )
+
+        return navitve_algorithms_config
+
+    def _create_algorithms(self):
+        if self.world_size == 8:
+            self.algos_config = self._create_native_algorithms()
+            self._tune(5, 10, 20, self.algos_config)
+        elif self.world_size == 16 or self.world_size == 32:
+            self.dsl_algos_config = self._create_dsl_algorithms()
+            self._tune(5, 10, 20, self.dsl_algos_config)
+
+    def _get_time(
+        self,
+        algo,
+        tune_tensor,
+        size,
+        nb,
+        nt,
+        n_warmup,
+        n_graph_launches,
+        n_ops_per_graph,
+    ):
+        # Check if the algorithm can run with the given configuration
+        if self._run_algo(algo, tune_tensor, size, nb, nt, True) != 0:
+            return float("inf")
+
+        # Warmup iterations to stabilize performance
+        for _ in range(n_warmup):
+            self._run_algo(algo, tune_tensor, size, nb, nt, True)
+
+        # Warmup on capture stream
+        capture_stream = torch.cuda.Stream()
+        capture_stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(capture_stream):
+            self._run_algo(algo, tune_tensor, size, nb, nt, True)
+        capture_stream.synchronize()
+
+        # Capture the algorithm execution in a CUDA graph
+        g = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(g, stream=capture_stream):
+            for _ in range(n_ops_per_graph):
+                self._run_algo(algo, tune_tensor, size, nb, nt, True)
+
+        # Measure the execution time of the captured graph
+        start_event = torch.cuda.Event(enable_timing=True)
+        end_event = torch.cuda.Event(enable_timing=True)
+        start_event.record(capture_stream)
+        with torch.cuda.stream(capture_stream):
+            for _ in range(n_graph_launches):
+                g.replay()
+        end_event.record(capture_stream)
+        end_event.synchronize()
+        elapsed = start_event.elapsed_time(end_event)
+
+        # Synchronize timing results across all ranks to ensure consistent algorithm selection
+        # replicate n times such due to algo limitations
+        time_tensor = torch.full(
+            (self.world_size,), elapsed, dtype=torch.float64, device="cuda"
+        ).to(dtype=torch.float32)
+        torch.cuda.current_stream().wait_stream(capture_stream)
+        if self.rank == 0:
+            avg_time = time_tensor[self.rank].item() / self.world_size
+            tensor = torch.tensor([avg_time])
+        else:
+            tensor = torch.empty(1)
+        dist.broadcast(tensor, src=0, group=self.group)
+        avg_time = tensor.item()
+
+        return avg_time
+
+    def _tune(self, n_warmup, n_graph_launches, n_ops_per_graph, algos_config):
+        sizes = [1 << i for i in range(9, 24)]
+        dlpack = self.mscclpp.RawGpuBuffer(1 << 27).to_dlpack(
+            data_type=str(torch.float16)
+        )
+        tune_tensor = torch.utils.dlpack.from_dlpack(dlpack)
+
+        for size in sizes:
+            best_time = float("inf")
+            best_config = None
+            for i in range(len(algos_config)):
+                algo, candidates_nblocks, candidates_nthreads = algos_config[i]
+                if (
+                    size >= algo.message_size_range[0]
+                    and size <= algo.message_size_range[1]
+                ):
+                    for nb in candidates_nblocks:
+                        for nt in candidates_nthreads:
+                            avg_time = self._get_time(
+                                algo,
+                                tune_tensor,
+                                size,
+                                nb,
+                                nt,
+                                n_warmup,
+                                n_graph_launches,
+                                n_ops_per_graph,
+                            )
+                            if avg_time < best_time:
+                                best_time = avg_time
+                                best_config = (algo, nb, nt)
+            if best_config:
+                self.best_configs[size] = best_config
+
+        torch.cuda.synchronize()
+        for algo, _, _ in algos_config:
+            algo.reset()
+
+    def _run_algo(self, algo, tensor, size, nblocks, nthreads, sym_mem_enabled=False):
+        return algo.execute(
+            comm=self.comm.communicator,
+            executor=self.executor,
+            input_buffer=tensor.data_ptr(),
+            output_buffer=tensor.data_ptr(),
+            input_size=size,
+            output_size=size,
+            dtype=self.dtype_to_mscclpp_dtype(tensor.dtype),
+            op=self.mscclpp.ReduceOp.SUM,
+            stream=torch.cuda.current_stream().cuda_stream,
+            nblocks=nblocks,
+            nthreads_per_block=nthreads,
+            symmetric_memory=sym_mem_enabled,
+        )
+
     def __init__(
         self,
         group: ProcessGroup,
         device: Union[int, str, torch.device],
-        max_bytes=_MAX_BYTES,
     ) -> None:
-        """
-        Args:
+        """Args:
             group: the process group to work on. If None, it will use the
                 default process group.
             device: the device to bind the CustomAllreduce to. If None,
@@ -114,11 +253,16 @@ class PyMscclppCommunicator:
         self._IS_CAPTURING = False
         self.disabled = True
 
-        if not ops.IS_MSCCLPP_AR_AVAILABLE:
-            # disable because of missing mscclpp library
-            # e.g. in a non-cuda environment
+        try:
+            self.mscclpp = importlib.import_module("mscclpp")
+            self.mscclpp_ext = importlib.import_module("mscclpp.ext")
+            self.def_algo = importlib.import_module("mscclpp.default_algos")
+        except ImportError:
+            self.available = False
+            self.mscclpp = None
             return
 
+        self.available = True
         self.group = group
 
         assert (
@@ -161,137 +305,83 @@ class PyMscclppCommunicator:
         assert isinstance(device, torch.device)
         self.device = device
 
-        self.max_bytes = max_bytes
         self.rank = rank
         self.world_size = world_size
-
-        if dist.get_rank(group) == 0:
-            unique_id = [ops.mscclpp_generate_unique_id()]
-        else:
-            unique_id = [None]
-        dist.broadcast_object_list(unique_id, src=self.ranks[0], group=self.group)
-        self.unique_id = unique_id[0]
-        self.rank_to_node, self.rank_to_ib = list(range(world_size)), list(
-            range(world_size)
+        self.comm = self.mscclpp.CommGroup(
+            torch_group=self.group, rank=rank, size=world_size
         )
-        for r in range(world_size):
-            self.rank_to_node[r] = r // 8
-            self.rank_to_ib[r] = self.rank % 8
+        self.executor = self.mscclpp.Executor(self.comm.communicator)
+        self.symm_mem_enabled = self._is_symm_mem_enabled()
+        self.best_configs = {}
+        self._create_algorithms()
 
-        self._context = None
-        self.context_selection = None
-        self.msg_size_for_finetune = [
-            2**i for i in range(10, math.floor(math.log2(self.max_bytes)) + 1)
-        ]
-        self.msg_size2best_config = {}
-        if world_size == 8:
-            self.context_selection = MscclContextSelection.MSCCL1SHOT1NODELL
-        elif world_size == 16:
-            self.context_selection = MscclContextSelection.MSCCL1SHOT2NODELL
-        if not _is_hip:
-            self.scratch = torch.empty(
-                self.max_bytes * 8,
-                dtype=torch.uint8,
-                device=self.device,
-            )
-            self.put_buffer = torch.empty(
-                self.max_bytes * 8 // self.nranks_per_node,
-                dtype=torch.uint8,
-                device=self.device,
-            )
-            self._context = ops.mscclpp_init_context(
-                self.unique_id,
-                self.rank,
-                self.world_size,
-                self.scratch,
-                self.put_buffer,
-                self.nranks_per_node,
-                self.rank_to_node,
-                self.rank_to_ib,
-                int(self.context_selection),
-            )
-        else:
-            raise NotImplementedError("HIP Mscclpp is not supported yet.")
-
-        self.msg_size2best_config = {}
-        self.pre_tune_config()
-        if dist.get_rank(group) == 0:
-            msg_size2best_config = [self.msg_size2best_config]
-        else:
-            msg_size2best_config = [None]
-        dist.broadcast_object_list(
-            msg_size2best_config, src=self.ranks[0], group=self.group
-        )
-        self.msg_size2best_config = msg_size2best_config[0]
-
-        # PyMscclpp is enabled only in cuda graph
-        self.disabled = True
-
-    def pre_tune_config(self, dtype=torch.bfloat16) -> bool:
-        logger.debug(f"start to pre-tune configs for rank {self.rank}")
-        nthreads_to_try = [256, 512, 1024]
-        nblocks_to_try = [21, 42, 84]
-        inp_randn = torch.ones(
-            self.msg_size_for_finetune[-1] // dtype.itemsize, dtype=dtype, device="cuda"
-        )
-        oup_randn = torch.empty_like(inp_randn)
-        for msg_size in self.msg_size_for_finetune:
-            mock_inp, mock_outp = (
-                inp_randn[: msg_size // dtype.itemsize],
-                oup_randn[: msg_size // dtype.itemsize],
-            )
-            best_config, best_time = None, None
-            for nthreads in nthreads_to_try:
-                for nblocks in nblocks_to_try:
-                    cur_cost = mscclpp_bench_time(
-                        lambda: ops.mscclpp_allreduce(
-                            self._context, mock_inp, mock_outp, nthreads, nblocks
-                        )
-                    )
-                    if best_time is None or cur_cost < best_time:
-                        best_config = (nthreads, nblocks)
-                        best_time = cur_cost
-            self.msg_size2best_config[msg_size] = best_config
-            if self.rank == 0:
-                logger.debug(
-                    f"for msg_size {msg_size}, best_config: {best_config}, best_time: {best_time}us"
-                )
+    def destroy(self):
+        self.algos_config = None
+        self.best_configs = None
+        self.executor = None
+        self.scratch_buffer = None
+        self.flag_buffer = None
+        self.comm = None
 
     def should_mscclpp_allreduce(
         self, inp: torch.Tensor, op: ReduceOp = ReduceOp.SUM
     ) -> bool:
-        if self.disabled or self._context is None:
+        if (
+            self.disabled
+            or self.world_size not in PyMscclppCommunicator._SUPPORTED_WORLD_SIZES
+        ):
             return False
         if inp.dtype not in PyMscclppCommunicator._SUPPORTED_DTYPE:
             return False
-        if not mscclpp_is_weak_contiguous(inp):
+        if not self._is_weak_contiguous(inp):
             return False
-        # only support sum op
-        if op != ReduceOp.SUM:
+        if op is not ReduceOp.SUM:
             return False
-        if inp.numel() * inp.element_size() > self.max_bytes:
+        if self._get_tuned_config(inp.numel() * inp.element_size()) is None:
+            return False
+        # mscclpp must not be used during any piecewise CUDA graph phase
+        # (compile, capture, or replay) as it changes the allreduce dispatch
+        # path and triggers recompilation.
+        if (
+            is_in_piecewise_cuda_graph()
+            or is_in_pcg_torch_compile()
+            or get_pcg_capture_stream() is not None
+        ):
             return False
         return True
 
-    def all_reduce(self, tensor: torch.Tensor, op: ReduceOp = ReduceOp.SUM):
-        if self._IS_CAPTURING:
-            if torch.cuda.is_current_stream_capturing():
-                self.graph_input_set.add((tensor.dtype, tensor.numel()))
-        msg_size = tensor.numel() * tensor.itemsize
-        index = bisect.bisect_left(self.msg_size_for_finetune, msg_size)
-        msg_size_finetune = self.msg_size_for_finetune[index]
-        nthreads, nblocks = self.msg_size2best_config[msg_size_finetune]
-        result = torch.empty_like(tensor)
-        ops.mscclpp_allreduce(self._context, tensor, result, nthreads, nblocks)
-        return result
+    def dtype_to_mscclpp_dtype(self, dtype: torch.dtype):
+        if dtype == torch.float16:
+            return self.mscclpp.DataType.float16
+        elif dtype == torch.float32:
+            return self.mscclpp.DataType.float32
+        elif dtype == torch.int32:
+            return self.mscclpp.DataType.int32
+        elif dtype == torch.bfloat16:
+            return self.mscclpp.DataType.bfloat16
+        else:
+            raise ValueError(f"Unknown data type: {dtype}")
+
+    def all_reduce(
+        self,
+        tensor: torch.Tensor,
+        op: ReduceOp = ReduceOp.SUM,
+        stream: torch.cuda.Stream = None,
+    ):
+        assert op == torch.distributed.ReduceOp.SUM
+        nbytes = tensor.numel() * tensor.element_size()
+        algo, nblocks, nthreads = self._get_tuned_config(nbytes)
+        self._run_algo(algo, tensor, nbytes, nblocks, nthreads, self.symm_mem_enabled)
+        return tensor
 
     @contextmanager
     def change_state(
         self,
         enable: Optional[bool] = None,
     ):
-        if enable is None:
+        if enable is None or self.available is False:
             # guess a default value when not specified
+            # DO: Decided if raise an exception here or not
             enable = self.available
 
         old_disable = self.disabled

python/sglang/srt/distributed/parallel_state.py

@@ -605,7 +605,15 @@ class GroupCoordinator:
         if self.npu_communicator is not None and not self.npu_communicator.disabled:
             return self.npu_communicator.all_reduce(input_)
 
-        if self.pynccl_comm is not None and self.is_symmetric_memory_enabled():
+        should_use_pymscclpp_allreduce = (
+            self.pymscclpp_comm is not None
+            and self.pymscclpp_comm.should_mscclpp_allreduce(input_)
+        )
+        if (
+            self.pynccl_comm is not None
+            and self.is_symmetric_memory_enabled()
+            and not should_use_pymscclpp_allreduce
+        ):
             self.debug_check_symmetric_mempool(self, {"input": input_}, "all_reduce")
             with self.pynccl_comm.change_state(enable=True):
                 self.pynccl_comm.all_reduce(input_)
@@ -615,6 +623,7 @@ class GroupCoordinator:
         if (
             self.ca_comm is not None
             and not self.ca_comm.disabled
+            and not should_use_pymscclpp_allreduce
             and self.ca_comm.should_custom_ar(input_)
         ):
             outplace_all_reduce_method = "ca"
@@ -624,11 +633,7 @@ class GroupCoordinator:
             and self.qr_comm.should_quick_allreduce(input_)
         ):
             outplace_all_reduce_method = "qr"
-        elif (
-            self.pymscclpp_comm is not None
-            and not self.pymscclpp_comm.disabled
-            and self.pymscclpp_comm.should_mscclpp_allreduce(input_)
-        ):
+        elif self.pymscclpp_comm is not None and should_use_pymscclpp_allreduce:
             outplace_all_reduce_method = "pymscclpp"
         elif (
             self.torch_symm_mem_comm is not None
@@ -1473,6 +1478,8 @@ class GroupCoordinator:
             self.cpu_group = None
         if self.pynccl_comm is not None:
             self.pynccl_comm = None
+        if self.pymscclpp_comm is not None:
+            self.pymscclpp_comm.destroy()
         if self.ca_comm is not None:
             self.ca_comm = None
         if self.mq_broadcaster is not None: