MSCCL++ Integration (#22734)

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

benchmark/kernels/all_reduce/README.md

@@ -0,0 +1,55 @@
+## MSCCL++ All-Reduce Benchmark
+
+[MSCCL++](https://github.com/microsoft/mscclpp) is a GPU-driven communication library that can replace NCCL for all-reduce operations. It supports CUDA graph capture and is optimized for small-to-medium message sizes commonly seen in tensor-parallel inference.
+
+Currently supported configurations: **TP=8** (single-node) and **TP=16** (two-node).
+
+### Prerequisites
+
+1. If you use the default SGLang Docker image build from `docker/Dockerfile`, [MSCCL++](https://github.com/microsoft/mscclpp) is already installed by default.
+2. If you are not using that Docker image (or want to install manually), install [MSCCL++](https://github.com/microsoft/mscclpp) from source (requires CMake and a CUDA toolkit):
+    ```bash
+    git clone https://github.com/microsoft/mscclpp.git
+    cd mscclpp && mkdir build && cd build
+    cmake .. && make -j && pip install ..
+    ```
+3. Ensure `mscclpp` is importable in your Python environment before running the benchmark or using MSCCL++ for inference.
+
+### Running the Benchmark
+
+The benchmark compares all-reduce latency across torch/NCCL (eager), MSCCL++ (eager and graph), and PyNccl (graph) for power-of-two message sizes.
+
+```bash
+torchrun --nproc_per_node 8 \
+    --nnodes 1 \
+    --node_rank 0 \
+    benchmark/kernels/all_reduce/benchmark_mscclpp.py
+```
+
+For multi-node (TP=16):
+```bash
+export WORLD_SIZE=2
+export MASTER_ADDR=<master-ip>
+export MASTER_PORT=12345
+
+# Run on each node with the appropriate RANK (0 or 1):
+torchrun --nproc_per_node 8 \
+    --nnodes $WORLD_SIZE \
+    --node_rank $RANK \
+    --master_addr $MASTER_ADDR \
+    --master_port $MASTER_PORT \
+    benchmark/kernels/all_reduce/benchmark_mscclpp.py
+```
+
+### Inference with MSCCL++
+
+Use the `--enable-mscclpp` flag to select MSCCL++ as the all-reduce backend during CUDA-graph-captured inference:
+
+```bash
+python -m sglang.launch_server \
+    --model-path Qwen/Qwen3-8B \
+    --tp-size 8 \
+    --enable-mscclpp
+```
+
+> **Note:** MSCCL++ performs auto-tuning on first initialization, which may add a few seconds to startup time. The tuned configurations are cached for the lifetime of the process.

benchmark/kernels/all_reduce/benchmark_mscclpp.py

@@ -24,6 +24,7 @@ from sglang.srt.distributed import init_distributed_environment
 from sglang.srt.distributed.device_communicators.pymscclpp import PyMscclppCommunicator
 from sglang.srt.distributed.device_communicators.pynccl import PyNcclCommunicator
 from sglang.srt.distributed.parallel_state import (
+    cleanup_dist_env_and_memory,
     get_tensor_model_parallel_group,
     graph_capture,
     initialize_model_parallel,
@@ -51,10 +52,12 @@ def pynccl_allreduce(
 
 def _bench_graph_time(func, inp_randn, warmup_loop=2, graph_loop=10, test_loop=10):
     graph_input = inp_randn.clone()
+    graph_input_snapshot = inp_randn.clone()
     with graph_capture() as graph_capture_context:
         graph = torch.cuda.CUDAGraph()
         with torch.cuda.graph(graph, stream=graph_capture_context.stream):
             for _ in range(graph_loop):
+                graph_input.copy_(graph_input_snapshot)
                 graph_out = func(graph_input)
 
     graph.replay()
@@ -222,3 +225,7 @@ if __name__ == "__main__":
         prof_dir = f"prof/msccl"
         os.makedirs(prof_dir, exist_ok=True)
         ctx.export_chrome_trace(f"{prof_dir}/trace_rank{dist.get_rank()}.json.gz")
+
+    pymscclpp_comm.destroy()
+    dist.barrier()
+    cleanup_dist_env_and_memory()

docker/Dockerfile

@@ -21,6 +21,7 @@ ARG GITHUB_ARTIFACTORY=github.com
 ARG INSTALL_FLASHINFER_JIT_CACHE=0
 ARG FLASHINFER_VERSION=0.6.12
 ARG MOONCAKE_VERSION=0.3.11.post1
+ARG MSCCLPP_VERSION=sglang-v0.9.1
 #if need other arg please add in MOONCAKE_COMPILE_ARG
 ARG MOONCAKE_COMPILE_ARG="-DUSE_HTTP=ON -DUSE_MNNVL=ON -DUSE_CUDA=ON -DWITH_EP=ON"
 
@@ -443,6 +444,7 @@ ARG USE_LATEST_SGLANG
 ARG GITHUB_ARTIFACTORY
 ARG MOONCAKE_VERSION
 ARG MOONCAKE_COMPILE_ARG
+ARG MSCCLPP_VERSION
 
 WORKDIR /sgl-workspace
 
@@ -523,6 +525,24 @@ RUN --mount=type=cache,target=/root/.cache/pip \
         python3 -m pip install mooncake-transfer-engine==${MOONCAKE_VERSION}; \
     fi
 
+# Install MSCCL++ Python dependencies and package (builds extension via CMake through pip)
+RUN --mount=type=cache,target=/root/.cache/pip \
+    git clone --depth=1 --branch ${MSCCLPP_VERSION} https://${GITHUB_ARTIFACTORY}/microsoft/mscclpp.git /tmp/mscclpp \
+    && case "${CUDA_VERSION}" in \
+           12.*) \
+               CMAKE_ARGS="-DMSCCLPP_BYPASS_GPU_CHECK=ON -DMSCCLPP_USE_CUDA=ON -DMSCCLPP_GPU_ARCHS=80,90,100,100a,103,103a" \
+               python3 -m pip install "/tmp/mscclpp[cuda12]"; \
+               ;; \
+           13.*) \
+               CMAKE_ARGS="-DMSCCLPP_BYPASS_GPU_CHECK=ON -DMSCCLPP_USE_CUDA=ON -DMSCCLPP_GPU_ARCHS=80,90,100,100a,103,103a" \
+               python3 -m pip install "/tmp/mscclpp[cuda13]"; \
+               ;; \
+           *) \
+               echo "Unsupported CUDA version for MSCCL++: ${CUDA_VERSION}" && exit 1; \
+               ;; \
+       esac \
+    && rm -rf /tmp/mscclpp
+
 # Install essential Python packages (use constraints to prevent conflicts)
 RUN --mount=type=cache,target=/root/.cache/pip \
     python3 -m pip install -c /sgl-workspace/constraints.txt \

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):
+        for tp_rank in range(local_rank_start, local_rank_end):
-            with maybe_reindex_device_id(tp_rank) as gpu_id:
+            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.compilation.piecewise_context_manager import (
-from sglang.srt.utils import is_hip
+    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]
+    _SUPPORTED_WORLD_SIZES = [8, 16, 32]
-    _MAX_BYTES = mscclpp_convert_to_bytes(os.getenv("SGLANG_MSCCLPP_MAX_BYTES", "1MB"))
     _SUPPORTED_DTYPE = [torch.float, torch.float16, torch.bfloat16]
 
-    # max_bytes: max supported mscclpp allreduce size
+    def _is_symm_mem_enabled(self) -> bool:
-    # in A100 mscclpp is faster than nccl only under condition of msg size smaller than1MB
+        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:
+        try:
-            # disable because of missing mscclpp library
+            self.mscclpp = importlib.import_module("mscclpp")
-            # e.g. in a non-cuda environment
+            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
-
+        self.comm = self.mscclpp.CommGroup(
-        if dist.get_rank(group) == 0:
+            torch_group=self.group, rank=rank, size=world_size
-            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)
         )
-        for r in range(world_size):
+        self.executor = self.mscclpp.Executor(self.comm.communicator)
-            self.rank_to_node[r] = r // 8
+        self.symm_mem_enabled = self._is_symm_mem_enabled()
-            self.rank_to_ib[r] = self.rank % 8
+        self.best_configs = {}
+        self._create_algorithms()
 
-        self._context = None
+    def destroy(self):
-        self.context_selection = None
+        self.algos_config = None
-        self.msg_size_for_finetune = [
+        self.best_configs = None
-            2**i for i in range(10, math.floor(math.log2(self.max_bytes)) + 1)
+        self.executor = None
-        ]
+        self.scratch_buffer = None
-        self.msg_size2best_config = {}
+        self.flag_buffer = None
-        if world_size == 8:
+        self.comm = None
-            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 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 is not ReduceOp.SUM:
-        if op != 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):
+    def dtype_to_mscclpp_dtype(self, dtype: torch.dtype):
-        if self._IS_CAPTURING:
+        if dtype == torch.float16:
-            if torch.cuda.is_current_stream_capturing():
+            return self.mscclpp.DataType.float16
-                self.graph_input_set.add((tensor.dtype, tensor.numel()))
+        elif dtype == torch.float32:
-        msg_size = tensor.numel() * tensor.itemsize
+            return self.mscclpp.DataType.float32
-        index = bisect.bisect_left(self.msg_size_for_finetune, msg_size)
+        elif dtype == torch.int32:
-        msg_size_finetune = self.msg_size_for_finetune[index]
+            return self.mscclpp.DataType.int32
-        nthreads, nblocks = self.msg_size2best_config[msg_size_finetune]
+        elif dtype == torch.bfloat16:
-        result = torch.empty_like(tensor)
+            return self.mscclpp.DataType.bfloat16
-        ops.mscclpp_allreduce(self._context, tensor, result, nthreads, nblocks)
+        else:
-        return result
+            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 (
+        elif self.pymscclpp_comm is not None and should_use_pymscclpp_allreduce:
-            self.pymscclpp_comm is not None
-            and not self.pymscclpp_comm.disabled
-            and self.pymscclpp_comm.should_mscclpp_allreduce(input_)
-        ):
             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:

sgl-kernel/CMakeLists.txt

@@ -86,24 +86,6 @@ FetchContent_Declare(
 )
 FetchContent_Populate(repo-flash-attention)
 
-# mscclpp
-FetchContent_Declare(
-    repo-mscclpp
-    URL      https://${GITHUB_ARTIFACTORY}/microsoft/mscclpp/archive/51eca89d20f0cfb3764ccd764338d7b22cd486a6.tar.gz
-    URL_HASH SHA256=b064de701da5253e32f4031d16b01a61d78c369327d585ff7bfb521bbe742677
-)
-FetchContent_Populate(repo-mscclpp)
-
-# mscclpp's own CMakeLists.txt hardcodes a github.com FetchContent for nlohmann/json.
-# Patch it to route through GITHUB_ARTIFACTORY so it follows the same mirror as our deps.
-file(READ "${repo-mscclpp_SOURCE_DIR}/CMakeLists.txt" _mscclpp_cmakelists)
-string(REPLACE
-    "https://github.com/nlohmann/json"
-    "https://${GITHUB_ARTIFACTORY}/nlohmann/json"
-    _mscclpp_cmakelists "${_mscclpp_cmakelists}")
-file(WRITE "${repo-mscclpp_SOURCE_DIR}/CMakeLists.txt" "${_mscclpp_cmakelists}")
-unset(_mscclpp_cmakelists)
-
 # ccache option
 option(ENABLE_CCACHE "Whether to use ccache" ON)
 find_program(CCACHE_FOUND ccache)
@@ -257,7 +239,6 @@ endif()
 # NOTE: Please sort the filenames alphabetically
 set(SOURCES
     "csrc/allreduce/custom_all_reduce.cu"
-    "csrc/allreduce/mscclpp_allreduce.cu"
     "csrc/attention/cutlass_mla_kernel.cu"
     "csrc/attention/merge_attn_states.cu"
     "csrc/attention/vertical_slash_index.cu"
@@ -329,7 +310,6 @@ set(INCLUDES
     ${repo-cutlass_SOURCE_DIR}/tools/util/include
     ${repo-flashinfer_SOURCE_DIR}/include
     ${repo-flashinfer_SOURCE_DIR}/csrc
-    ${repo-mscclpp_SOURCE_DIR}/include
     ${repo-cutlass_SOURCE_DIR}/examples/77_blackwell_fmha
     ${repo-cutlass_SOURCE_DIR}/examples/common
     ${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src
@@ -379,19 +359,8 @@ else()
     set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -D_GLIBCXX_USE_CXX11_ABI=1")
 endif()
 
-# mscclpp option
+target_link_libraries(common_ops_sm90_build PRIVATE ${TORCH_LIBRARIES} c10 cuda cublas cublasLt)
-set(MSCCLPP_USE_CUDA ON)
+target_link_libraries(common_ops_sm100_build PRIVATE ${TORCH_LIBRARIES} c10 cuda cublas cublasLt)
-set(MSCCLPP_BYPASS_GPU_CHECK ON)
-set(MSCCLPP_BUILD_TESTS OFF)
-set(MSCCLPP_BUILD_PYTHON_BINDINGS OFF)
-set(MSCCLPP_BUILD_APPS_NCCL OFF)
-add_subdirectory(
-    ${repo-mscclpp_SOURCE_DIR}
-    ${CMAKE_CURRENT_BINARY_DIR}/mscclpp-build
-)
-
-target_link_libraries(common_ops_sm90_build PRIVATE ${TORCH_LIBRARIES} c10 cuda cublas cublasLt mscclpp_static)
-target_link_libraries(common_ops_sm100_build PRIVATE ${TORCH_LIBRARIES} c10 cuda cublas cublasLt mscclpp_static)
 
 # sparse flash attention
 target_compile_definitions(common_ops_sm90_build PRIVATE

sgl-kernel/csrc/allreduce/mscclpp_allreduce.cu

@@ -1,140 +0,0 @@
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAStream.h>
-#include <torch/all.h>
-#include <torch/library.h>
-
-#include "mscclpp_allreduce.cuh"
-
-enum MscclContextSelection {
-  MSCCL1NODELL = 1,
-  MSCCL2NODELL = 2,
-};
-
-class MscclContext {
- public:
-  MscclContextSelection selection_;
-  std::shared_ptr<sglang::Msccl1NodeLLcontext> msccl_1nodeLL_context;
-  std::shared_ptr<sglang::Msccl2NodeLLcontext> msccl_2nodeLL_context;
-  MscclContext(MscclContextSelection selection) : selection_(selection) {}
-  template <typename T>
-  void allreduce(
-      cudaStream_t stream, T* input, T* output, const size_t input_numel, int threads = 512, int block_limit = 21) {
-    if (selection_ == MSCCL1NODELL) {
-      msccl_1nodeLL_context->allreduce<T>(stream, input, output, input_numel, threads, block_limit);
-    } else if (selection_ == MSCCL2NODELL) {
-      msccl_2nodeLL_context->allreduce<T>(stream, input, output, input_numel, threads, block_limit);
-    }
-  }
-};
-
-using fptr_t = int64_t;
-static_assert(sizeof(void*) == sizeof(fptr_t));
-
-torch::Tensor _unique_id2tensor(const mscclpp::UniqueId& unique_id) {
-  auto options = torch::TensorOptions().dtype(torch::kByte).device(torch::kCPU);
-  auto tensor = torch::empty({static_cast<int64_t>(unique_id.size())}, options);
-  std::memcpy(tensor.data_ptr<uint8_t>(), unique_id.data(), unique_id.size());
-  return tensor;
-}
-
-// Function to convert vector of int32_t back to array of uint8_t
-mscclpp::UniqueId _tensor2unique_id(const torch::Tensor& tensor) {
-  mscclpp::UniqueId unique_id;
-  std::memcpy(unique_id.data(), tensor.data_ptr<uint8_t>(), unique_id.size());
-  return unique_id;
-}
-
-torch::Tensor mscclpp_generate_unique_id() {
-  mscclpp::UniqueId unique_id = mscclpp::TcpBootstrap::createUniqueId();
-  return _unique_id2tensor(unique_id);
-}
-
-fptr_t mscclpp_init_context(
-    const torch::Tensor& unique_id,
-    const int64_t rank,
-    const int64_t world_size,
-    torch::Tensor& scratch,
-    torch::Tensor& put_buffer,
-    const int64_t nranks_per_node,
-    const std::vector<int64_t>& rank_to_node,
-    const std::vector<int64_t>& rank_to_ib,
-    const int64_t context_selection) {
-  MscclContext* context_ptr = new MscclContext(static_cast<MscclContextSelection>(context_selection));
-  mscclpp::UniqueId uid = _tensor2unique_id(unique_id);
-  if (context_selection == MSCCL1NODELL) {
-    void* scratch_ptr = reinterpret_cast<void*>(scratch.data_ptr());
-    const size_t scratch_bytes = scratch.numel() * scratch.element_size();
-    context_ptr->msccl_1nodeLL_context = std::make_shared<sglang::Msccl1NodeLLcontext>(
-        uid, rank, world_size, scratch_ptr, scratch_bytes, nranks_per_node, rank_to_node, rank_to_ib);
-  } else if (context_selection == MSCCL2NODELL) {
-    void* scratch_ptr = reinterpret_cast<void*>(scratch.data_ptr());
-    const size_t scratch_bytes = scratch.numel() * scratch.element_size();
-    void* put_buffer_ptr = reinterpret_cast<void*>(put_buffer.data_ptr());
-    const size_t put_buffer_bytes = put_buffer.numel() * put_buffer.element_size();
-    context_ptr->msccl_2nodeLL_context = std::make_shared<sglang::Msccl2NodeLLcontext>(
-        uid,
-        rank,
-        world_size,
-        scratch_ptr,
-        scratch_bytes,
-        put_buffer_ptr,
-        put_buffer_bytes,
-        nranks_per_node,
-        rank_to_node,
-        rank_to_ib);
-  } else {
-    throw std::runtime_error("invalid context selection");
-  }
-  return (fptr_t)context_ptr;
-}
-
-bool _mscclpp_is_weak_contiguous(torch::Tensor& t) {
-  return t.is_contiguous() ||
-         (t.storage().nbytes() - t.storage_offset() * t.element_size() == t.numel() * t.element_size());
-}
-void mscclpp_allreduce(fptr_t _context, torch::Tensor& inp, torch::Tensor& out, int64_t nthreads, int64_t nblocks) {
-  MscclContext* context = reinterpret_cast<MscclContext*>(_context);
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(inp));
-  auto stream = c10::cuda::getCurrentCUDAStream().stream();
-
-  TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
-  TORCH_CHECK_EQ(inp.numel(), out.numel());
-  TORCH_CHECK(_mscclpp_is_weak_contiguous(out));
-  TORCH_CHECK(_mscclpp_is_weak_contiguous(inp));
-  switch (out.scalar_type()) {
-    case at::ScalarType::Float: {
-      context->allreduce<float>(
-          stream,
-          reinterpret_cast<float*>(inp.data_ptr()),
-          reinterpret_cast<float*>(out.data_ptr()),
-          inp.numel(),
-          nthreads,
-          nblocks);
-      break;
-    }
-    case at::ScalarType::Half: {
-      context->allreduce<half>(
-          stream,
-          reinterpret_cast<half*>(inp.data_ptr()),
-          reinterpret_cast<half*>(out.data_ptr()),
-          inp.numel(),
-          nthreads,
-          nblocks);
-      break;
-    }
-#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
-    case at::ScalarType::BFloat16: {
-      context->allreduce<__nv_bfloat16>(
-          stream,
-          reinterpret_cast<__nv_bfloat16*>(inp.data_ptr()),
-          reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
-          inp.numel(),
-          nthreads,
-          nblocks);
-      break;
-    }
-#endif
-    default:
-      throw std::runtime_error("custom allreduce only supports float32, float16 and bfloat16");
-  }
-}

sgl-kernel/csrc/allreduce/mscclpp_allreduce.cuh

@@ -1,779 +0,0 @@
-// Copyright (c) Microsoft Corporation.
-// Licensed under the MIT license.
-#pragma once
-#ifdef USE_ROCM
-#include <hip/hip_fp16.h>
-#else
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
-#endif
-
-#include <mscclpp/concurrency_device.hpp>
-#include <mscclpp/core.hpp>
-#include <mscclpp/memory_channel.hpp>
-#include <mscclpp/memory_channel_device.hpp>
-#include <mscclpp/nvls_device.hpp>
-#include <mscclpp/port_channel.hpp>
-#include <mscclpp/port_channel_device.hpp>
-
-// comment this for test_mscclpp_allreduce.cu
-#include "utils.h"
-
-namespace sglang {
-
-__device__ mscclpp::DeviceSyncer deviceSyncer;
-__device__ mscclpp::DeviceSyncer allGatherDeviceSyncer;
-__device__ mscclpp::DeviceSyncer reduceScatterDeviceSyncer;
-__device__ mscclpp::DeviceSyncer ibDeviceSyncer;
-
-template <typename To, typename From>
-__forceinline__ __device__ To bit_cast(const From& src) {
-  static_assert(sizeof(To) == sizeof(From), "Size mismatch for bit_cast");
-
-  union {
-    From f;
-    To t;
-  } u;
-  u.f = src;
-  return u.t;
-}
-
-template <typename T>
-__forceinline__ __device__ T add_elements(T a, T b) {
-  return a + b;
-}
-
-template <>
-__forceinline__ __device__ __half2 add_elements(__half2 a, __half2 b) {
-  return __hadd2(a, b);
-}
-
-#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
-template <>
-__forceinline__ __device__ __nv_bfloat162 add_elements(__nv_bfloat162 a, __nv_bfloat162 b) {
-  return __hadd2(a, b);
-}
-#endif
-
-template <typename T>
-__forceinline__ __device__ int4 add_vectors_helper(int4 a, int4 b) {
-  int4 ret;
-  ret.w = bit_cast<int, T>(add_elements(bit_cast<T, int>(a.w), bit_cast<T, int>(b.w)));
-  ret.x = bit_cast<int, T>(add_elements(bit_cast<T, int>(a.x), bit_cast<T, int>(b.x)));
-  ret.y = bit_cast<int, T>(add_elements(bit_cast<T, int>(a.y), bit_cast<T, int>(b.y)));
-  ret.z = bit_cast<int, T>(add_elements(bit_cast<T, int>(a.z), bit_cast<T, int>(b.z)));
-  return ret;
-}
-
-template <typename T>
-__forceinline__ __device__ int4 add_vectors(int4 a, int4 b) {
-  return add_vectors_helper<T>(a, b);
-}
-
-#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
-template <>
-__forceinline__ __device__ int4 add_vectors<__nv_bfloat16>(int4 a, int4 b) {
-  return add_vectors_helper<__nv_bfloat162>(a, b);
-}
-#endif
-
-template <>
-__forceinline__ __device__ int4 add_vectors<__half>(int4 a, int4 b) {
-  return add_vectors_helper<__half2>(a, b);
-}
-
-template <typename T>
-__forceinline__ __device__ uint2 add_vectors_helper(uint2 a, uint2 b) {
-  uint2 ret;
-  ret.x = bit_cast<int, T>(add_elements(bit_cast<T, int>(a.x), bit_cast<T, int>(b.x)));
-  ret.y = bit_cast<int, T>(add_elements(bit_cast<T, int>(a.y), bit_cast<T, int>(b.y)));
-  return ret;
-}
-
-template <typename T>
-__forceinline__ __device__ uint2 add_vectors(uint2 a, uint2 b) {
-  return add_vectors_helper<T>(a, b);
-}
-
-#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
-template <>
-__forceinline__ __device__ uint2 add_vectors<__nv_bfloat16>(uint2 a, uint2 b) {
-  return add_vectors_helper<__nv_bfloat162>(a, b);
-}
-#endif
-
-template <>
-__forceinline__ __device__ uint2 add_vectors<__half>(uint2 a, uint2 b) {
-  return add_vectors_helper<__half2>(a, b);
-}
-
-template <typename T>
-__forceinline__ __device__ int add_vectors_helper(int a, int b) {
-  return bit_cast<int, T>(add_elements(bit_cast<T, int>(a), bit_cast<T, int>(b)));
-}
-
-template <typename T>
-__forceinline__ __device__ int add_vectors(int a, int b) {
-  return add_vectors_helper<T>(a, b);
-}
-
-#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
-template <>
-__forceinline__ __device__ int add_vectors<__nv_bfloat16>(int a, int b) {
-  return add_vectors_helper<__nv_bfloat162>(a, b);
-}
-#endif
-
-template <>
-__forceinline__ __device__ int add_vectors<__half>(int a, int b) {
-  return add_vectors_helper<__half2>(a, b);
-}
-
-// -------------------------------------------------------
-// allreduce_LL_1node using LLPacket, origin allreduce2
-// -------------------------------------------------------
-
-__device__ uint64_t globalFlag = 1;
-
-template <typename TYPE>
-__global__ void __launch_bounds__(1024, 1) allreduce_LL_1node(
-    mscclpp::MemoryChannelDeviceHandle* memChans,
-    TYPE* buff,
-    TYPE* scratch,
-    void* resultBuff,
-    int rank,
-    int worldSize,
-    size_t nelems) {
-  nelems = nelems / (sizeof(int) / sizeof(TYPE));
-  // This version of allreduce only works for single nodes
-  const int nPeers = worldSize - 1;
-  const size_t nPkts = nelems / 2;
-  const int nelemsPerRank = nelems / worldSize;
-  const int nPktsPerRank = nelemsPerRank / 2;
-  // flag for packets. Initially 1
-  const uint32_t flag = (uint32_t)globalFlag;
-  // thread block & channel info
-  const int nBlocksPerPeer = gridDim.x / nPeers;
-  const int localBlockIdx = blockIdx.x % nBlocksPerPeer;
-  const int peerIdx = blockIdx.x / nBlocksPerPeer;
-  const int remoteRank = peerIdx < rank ? peerIdx : peerIdx + 1;
-  mscclpp::MemoryChannelDeviceHandle memChan = memChans[peerIdx];
-  const int tid = threadIdx.x + localBlockIdx * blockDim.x;
-  // double buffering
-  size_t scratchBaseOffset = (flag & 1) ? 0 : nPkts * sizeof(mscclpp::LLPacket);
-  void* scratchBuff = (void*)((char*)scratch + scratchBaseOffset);
-  size_t scratchOffset = scratchBaseOffset + rank * nPktsPerRank * sizeof(mscclpp::LLPacket);
-  size_t scratchResultOffset =
-      (flag & 1) ? 2 * nPkts * sizeof(mscclpp::LLPacket) : 3 * nPkts * sizeof(mscclpp::LLPacket);
-  size_t srcOffset = remoteRank * nelemsPerRank * sizeof(int);
-  uint2* src = (uint2*)((char*)buff + rank * nelemsPerRank * sizeof(int));
-  uint2* dst = (uint2*)((char*)resultBuff + rank * nelemsPerRank * sizeof(int));
-
-  // step 1: write to scratch buffer
-  memChan.putPackets(scratchOffset, srcOffset, nelemsPerRank * sizeof(int), tid, blockDim.x * nBlocksPerPeer, flag);
-  // step 2: get data from scratch buffer, reduce data and write result to remote scratch buffer
-  for (int idx = threadIdx.x + blockIdx.x * blockDim.x; idx < nPktsPerRank; idx += blockDim.x * gridDim.x) {
-    uint2 data = make_uint2(0, 0);
-    for (int index = 0; index < nPeers; index++) {
-      const int remoteRank = index < rank ? index : index + 1;
-      mscclpp::LLPacket* dstPkt = (mscclpp::LLPacket*)scratchBuff + remoteRank * nPktsPerRank;
-      uint2 val = dstPkt[idx].read(flag);
-      data = add_vectors<TYPE>(val, data);
-    }
-    data = add_vectors<TYPE>(data, src[idx]);
-    dst[idx] = data;
-
-    mscclpp::LLPacket packet;
-    packet.data1 = data.x;
-    packet.flag1 = flag;
-    packet.data2 = data.y;
-    packet.flag2 = flag;
-    size_t offset = scratchResultOffset / sizeof(mscclpp::LLPacket) + (idx + rank * nPktsPerRank);
-    for (int index = 0; index < nPeers; index++) {
-      memChans[index].write(offset, packet);
-    }
-  }
-  // step 3: get data result from scratch buffer
-  mscclpp::LLPacket* dstPkt = (mscclpp::LLPacket*)((char*)scratch + scratchResultOffset);
-  const int dstOffset = remoteRank * nPktsPerRank;
-  uint2* result = (uint2*)((char*)resultBuff + remoteRank * nelemsPerRank * sizeof(int));
-  for (int idx = threadIdx.x + localBlockIdx * blockDim.x; idx < nPktsPerRank; idx += blockDim.x * nBlocksPerPeer) {
-    uint2 data = dstPkt[idx + dstOffset].read(flag);
-    result[idx].x = data.x;
-    result[idx].y = data.y;
-  }
-  if (threadIdx.x == 0 && blockIdx.x == 0) {
-    globalFlag += 1;
-  }
-}
-
-// -------------------------------------------------------
-// allreduce_LL_2node using LLPacket, origin allreduce5
-// -------------------------------------------------------
-
-template <typename TYPE>
-__global__ void __launch_bounds__(1024, 1) allreduce_LL_2node(
-    mscclpp::MemoryChannelDeviceHandle* memChans,
-    mscclpp::PortChannelDeviceHandle* portChans,
-    TYPE* buff,
-    TYPE* scratch,
-    TYPE* putBuff,
-    TYPE* resultBuff,
-    int rank,
-    int nRanksPerNode,
-    int worldSize,
-    size_t nelems) {
-  nelems = nelems / (sizeof(int) / sizeof(TYPE));
-  // This version of allreduce only works for single nodes
-  const int nPeersInNode = nRanksPerNode - 1;
-  const int nPkts = nelems / 2;
-  const int nelemsPerLocalRank = nelems / nRanksPerNode;
-  const int nPktsPerLocalRank = nelemsPerLocalRank / 2;
-  const int localRankId = rank % nRanksPerNode;
-  // flag for packets. Initially 1
-  const uint32_t flag = (uint32_t)globalFlag;
-  // thread block & channel info
-  const int nBlocksPerPeer = gridDim.x / nPeersInNode;
-  const int localBlockIdx = blockIdx.x % nBlocksPerPeer;
-  const int peerIdx = blockIdx.x / nBlocksPerPeer;
-  const int remoteRankIdx = peerIdx < localRankId ? peerIdx : peerIdx + 1;
-  mscclpp::MemoryChannelDeviceHandle memChan = memChans[peerIdx];
-  mscclpp::PortChannelDeviceHandle portChan = portChans[localRankId];
-  const int tid = threadIdx.x + localBlockIdx * blockDim.x;
-  // double buffering
-  size_t scratchBaseOffset = (flag & 1) ? 0 : nPkts * sizeof(mscclpp::LLPacket);
-  size_t putBaseOffset = (flag & 1) ? 0 : nPktsPerLocalRank * sizeof(mscclpp::LLPacket);
-  void* scratchBuff = (void*)((char*)scratch + scratchBaseOffset);
-  size_t scratchOffset = scratchBaseOffset + localRankId * nPktsPerLocalRank * sizeof(mscclpp::LLPacket);
-  size_t scratchResultOffset =
-      (flag & 1) ? 2 * nPkts * sizeof(mscclpp::LLPacket) : 3 * nPkts * sizeof(mscclpp::LLPacket);
-  size_t srcOffset = remoteRankIdx * nelemsPerLocalRank * sizeof(int);
-  uint2* src = (uint2*)((char*)buff + localRankId * nelemsPerLocalRank * sizeof(int));
-  uint2* dst = (uint2*)((char*)resultBuff + localRankId * nelemsPerLocalRank * sizeof(int));
-
-  // step 1: write to scratch buffer
-  if (nRanksPerNode > 1) {
-    memChan.putPackets(
-        scratchOffset, srcOffset, nelemsPerLocalRank * sizeof(int), tid, blockDim.x * nBlocksPerPeer, flag);
-  }
-  // step 2: get data from scratch buffer, do local reduce-scatter in each node.
-  mscclpp::LLPacket* putPkt = (mscclpp::LLPacket*)((char*)putBuff + putBaseOffset);
-  for (int idx = threadIdx.x + blockIdx.x * blockDim.x; idx < nPktsPerLocalRank; idx += blockDim.x * gridDim.x) {
-    uint2 data = make_uint2(0, 0);
-    for (int index = 0; index < nPeersInNode; index++) {
-      const int remoteRank = index < localRankId ? index : index + 1;
-      mscclpp::LLPacket* dstPkt = (mscclpp::LLPacket*)scratchBuff + remoteRank * nPktsPerLocalRank;
-      uint2 val = dstPkt[idx].read(flag);
-      data = add_vectors<TYPE>(val, data);
-    }
-    data = add_vectors<TYPE>(data, src[idx]);
-    putPkt[idx].write(data.x, data.y, flag);
-    dst[idx] = data;
-  }
-  deviceSyncer.sync(gridDim.x);
-  // step 3. send local reduced data to remote node.
-  if (threadIdx.x == 0 && blockIdx.x == 0) {
-    portChan.put(scratchOffset, putBaseOffset, nPktsPerLocalRank * sizeof(mscclpp::LLPacket));
-    if ((flag & 63) == 0) {
-      portChan.flush();
-    }
-  }
-  // step 4. try to read the data from scratch buffer and write to local peers
-  mscclpp::LLPacket* dstPkt = (mscclpp::LLPacket*)scratchBuff + localRankId * nPktsPerLocalRank;
-  for (int idx = threadIdx.x + blockIdx.x * blockDim.x; idx < nPktsPerLocalRank; idx += blockDim.x * gridDim.x) {
-    uint2 res = dst[idx];
-    uint2 val = dstPkt[idx].read(flag);
-    res = add_vectors<TYPE>(res, val);
-
-    mscclpp::LLPacket packet;
-    packet.data1 = res.x;
-    packet.flag1 = flag;
-    packet.data2 = res.y;
-    packet.flag2 = flag;
-    size_t offset = scratchResultOffset / sizeof(mscclpp::LLPacket) + (idx + localRankId * nPktsPerLocalRank);
-    for (int index = 0; index < nPeersInNode; index++) {
-      memChans[index].write(offset, packet);
-    }
-    dst[idx] = res;
-  }
-
-  // step 5: get data result from scratch buffer
-  dstPkt = (mscclpp::LLPacket*)((char*)scratch + scratchResultOffset);
-  const int dstOffset = remoteRankIdx * nPktsPerLocalRank;
-  uint2* result = (uint2*)((char*)resultBuff + remoteRankIdx * nelemsPerLocalRank * sizeof(int));
-  if (nRanksPerNode > 1) {
-    for (int idx = threadIdx.x + localBlockIdx * blockDim.x; idx < nPktsPerLocalRank;
-         idx += blockDim.x * nBlocksPerPeer) {
-      uint2 data = dstPkt[idx + dstOffset].read(flag);
-      result[idx] = data;
-    }
-  }
-  if (threadIdx.x == 0 && blockIdx.x == 0) {
-    globalFlag += 1;
-  }
-}
-
-static const mscclpp::Transport IBs[] = {
-    mscclpp::Transport::IB0,
-    mscclpp::Transport::IB1,
-    mscclpp::Transport::IB2,
-    mscclpp::Transport::IB3,
-    mscclpp::Transport::IB4,
-    mscclpp::Transport::IB5,
-    mscclpp::Transport::IB6,
-    mscclpp::Transport::IB7};
-
-class MscclCommGroup {
- public:
-  std::shared_ptr<mscclpp::Communicator> comm_;
-  const size_t rank_;
-  const size_t world_size_;
-  const std::vector<int64_t> rank_to_node_;
-  const std::vector<int64_t> rank_to_ib_;
-  MscclCommGroup(
-      mscclpp::UniqueId unique_id,
-      const size_t rank,
-      const size_t world_size,
-      const std::vector<int64_t>& rank_to_node,
-      const std::vector<int64_t>& rank_to_ib)
-      : rank_(rank), world_size_(world_size), rank_to_node_(rank_to_node), rank_to_ib_(rank_to_ib) {
-    auto bootstrap = std::make_shared<mscclpp::TcpBootstrap>(rank, world_size);
-    bootstrap->initialize(unique_id);
-    comm_ = std::make_shared<mscclpp::Communicator>(bootstrap);
-  }
-  template <typename T>
-  void allreduce(cudaStream_t stream, T* output, size_t input_numel, int threads = 512, int block_limit = 21) {
-    throw std::runtime_error("you should not call allreduce of a base context");
-  }
-  bool is_same_node(int r1, int r2) {
-    return rank_to_node_[r1] == rank_to_node_[r2];
-  }
-
-  void make_connection(
-      std::unordered_map<int, std::shared_ptr<mscclpp::Connection>>& same_node_connections,
-      std::unordered_map<int, std::shared_ptr<mscclpp::Connection>>& cross_node_connections) {
-    same_node_connections.clear();
-    cross_node_connections.clear();
-    std::unordered_map<int, mscclpp::NonblockingFuture<std::shared_ptr<mscclpp::Connection>>> conn_futures;
-    for (int r = 0; r < world_size_; ++r) {
-      if (r == rank_) continue;
-      mscclpp::Transport transport = is_same_node(r, rank_) ? mscclpp::Transport::CudaIpc : IBs[rank_to_ib_[r]];
-      conn_futures.emplace(r, comm_->connectOnSetup(r, 0, transport));
-    }
-    comm_->setup();
-    for (int r = 0; r < world_size_; ++r) {
-      if (r == rank_) continue;
-      if (is_same_node(r, rank_)) {
-        same_node_connections.emplace(r, conn_futures[r].get());
-      } else {
-        cross_node_connections.emplace(r, conn_futures[r].get());
-      }
-    }
-  }
-
-  void make_memory_channels_with_scratch(
-      void* tensor_ptr,
-      const size_t tensor_bytes,
-      void* scratch_ptr,
-      const size_t scratch_bytes,
-      const std::unordered_map<int, std::shared_ptr<mscclpp::Connection>>& connections,
-      std::unordered_map<int, std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>>& semaphores,
-      std::unordered_map<int, mscclpp::RegisteredMemory>& registered_memories,
-      std::unordered_map<int, mscclpp::MemoryChannel>& channels) {
-    channels.clear();
-    make_semaphores<mscclpp::MemoryDevice2DeviceSemaphore>(connections, semaphores);
-    register_tensor_with_connections(scratch_ptr, scratch_bytes, connections, registered_memories);
-    for (const auto& [peer, _] : connections) {
-      channels.emplace(
-          peer, mscclpp::MemoryChannel(semaphores[peer], registered_memories[peer], tensor_ptr, scratch_ptr));
-    }
-  }
-  void make_port_channels_with_scratch(
-      std::shared_ptr<mscclpp::ProxyService> proxyService,
-      void* tensor_ptr,
-      const size_t tensor_bytes,
-      void* scratch_ptr,
-      const size_t scratch_bytes,
-      const std::unordered_map<int, std::shared_ptr<mscclpp::Connection>>& connections,
-      std::unordered_map<int, std::shared_ptr<mscclpp::Host2DeviceSemaphore>>& semaphores,
-      std::unordered_map<int, mscclpp::RegisteredMemory>& registered_memories,
-      std::unordered_map<int, mscclpp::PortChannel>& channels) {
-    channels.clear();
-    make_semaphores<mscclpp::Host2DeviceSemaphore>(connections, semaphores);
-
-    mscclpp::TransportFlags flags;
-    for (const auto& [_, conn] : connections) {
-      flags |= conn->transport();
-    }
-    auto local_reg_memory = comm_->registerMemory(tensor_ptr, tensor_bytes, flags);
-
-    register_tensor_with_connections(scratch_ptr, scratch_bytes, connections, registered_memories);
-    std::unordered_map<int, mscclpp::SemaphoreId> semaphore_ids;
-    std::unordered_map<int, size_t> memory_ids;
-    memory_ids[rank_] = proxyService->addMemory(local_reg_memory);
-    for (const auto& [peer, memory] : registered_memories) {
-      if (peer == rank_) continue;
-      memory_ids[peer] = proxyService->addMemory(memory);
-    }
-    for (const auto& [peer, semaphore] : semaphores) {
-      semaphore_ids[peer] = proxyService->addSemaphore(semaphore);
-    }
-
-    for (const auto& [peer, _] : connections) {
-      channels.emplace(peer, proxyService->portChannel(semaphore_ids[peer], memory_ids[peer], memory_ids[rank_]));
-    }
-  }
-
-  template <typename SemaphoreType>
-  void make_semaphores(
-      const std::unordered_map<int, std::shared_ptr<mscclpp::Connection>>& connections,
-      std::unordered_map<int, std::shared_ptr<SemaphoreType>>& semaphores) {
-    semaphores.clear();
-    for (const auto& [peer, conn] : connections) {
-      semaphores[peer] = std::make_shared<SemaphoreType>(*comm_, conn);
-    }
-    comm_->setup();
-  }
-
-  void register_tensor_with_connections(
-      void* tensor_ptr,
-      size_t tensor_bytes,
-      const std::unordered_map<int, std::shared_ptr<mscclpp::Connection>>& connections,
-      std::unordered_map<int, mscclpp::RegisteredMemory>& registered_memories) {
-    registered_memories.clear();
-    mscclpp::TransportFlags all_transports;
-    for (const auto& [_, connection] : connections) {
-      all_transports |= connection->transport();
-    }
-    mscclpp::RegisteredMemory buf_reg_mem = comm_->registerMemory(tensor_ptr, tensor_bytes, all_transports);
-    registered_memories[rank_] = buf_reg_mem;
-
-    std::unordered_map<int, mscclpp::NonblockingFuture<mscclpp::RegisteredMemory>> remote_mem_futures;
-    for (const auto& [r, connection] : connections) {
-      comm_->sendMemoryOnSetup(buf_reg_mem, r, 0);
-      auto remoteMemory = comm_->recvMemoryOnSetup(r, 0);
-      remote_mem_futures.emplace(r, remoteMemory);
-    }
-    comm_->setup();
-    for (auto& [r, mem_feature] : remote_mem_futures) {
-      registered_memories.emplace(r, mem_feature.get());
-    }
-  }
-
-  void make_device_memory_handle_base_on_new_ptr(
-      const std::unordered_map<int, mscclpp::MemoryChannel>& old_memory_channels,
-      std::unordered_map<int, mscclpp::RegisteredMemory>& registered_sm_memories,
-      std::unordered_map<int, std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>>& memory_semaphores,
-      std::unordered_map<int, mscclpp::MemoryChannel>& memory_channels,
-      mscclpp::GpuBuffer<mscclpp::MemoryChannelDeviceHandle>& device_memory_handle,
-      void* input,
-      void* scratch,
-      const cudaStream_t stream) {
-    memory_channels.clear();
-    for (const auto& [peer, channel] : old_memory_channels) {
-      memory_channels.emplace(
-          peer, mscclpp::MemoryChannel(memory_semaphores[peer], registered_sm_memories[peer], input, scratch));
-    }
-    std::vector<mscclpp::MemoryChannel> memory_channels_list;
-    for (int r = 0; r < world_size_; r++) {
-      if (r == rank_) continue;
-      if (is_same_node(r, rank_)) {
-        memory_channels_list.push_back(memory_channels[r]);
-      }
-    }
-    std::vector<mscclpp::MemoryChannelDeviceHandle> memory_channel_handlers(memory_channels_list.size());
-    std::transform(
-        memory_channels_list.begin(),
-        memory_channels_list.end(),
-        memory_channel_handlers.begin(),
-        [](const mscclpp::MemoryChannel& channel) { return channel.deviceHandle(); });
-    mscclpp::gpuMemcpyAsync<mscclpp::MemoryChannelDeviceHandle>(
-        device_memory_handle.data(),
-        memory_channel_handlers.data(),
-        memory_channel_handlers.size(),
-        stream,
-        cudaMemcpyHostToDevice);
-  }
-};
-
-class Msccl1NodeLLcontext {
- private:
-  std::shared_ptr<MscclCommGroup> comm_group_ = nullptr;
-  void* scratch_;
-  const size_t scratch_bytes_;
-  std::unordered_map<int, std::shared_ptr<mscclpp::Connection>> same_node_connections_;
-  std::unordered_map<int, std::shared_ptr<mscclpp::Connection>> cross_node_connections_;
-
-  std::unordered_map<int, mscclpp::RegisteredMemory> registered_sm_memories_;
-  std::unordered_map<int, std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>> memory_semaphores_;
-  std::unordered_map<int, mscclpp::MemoryChannel> memory_channels_;
-  mscclpp::GpuBuffer<mscclpp::MemoryChannelDeviceHandle> d_memHandles_;
-  std::unordered_map<void*, std::unordered_map<int, mscclpp::MemoryChannel>> input_ptr2memory_channels_;
-  std::unordered_map<void*, mscclpp::GpuBuffer<mscclpp::MemoryChannelDeviceHandle>> input_ptr2d_memHandles_;
-  cudaStream_t h2d_stream;
-  const size_t nranks_per_node_;
-
- public:
-  Msccl1NodeLLcontext(
-      mscclpp::UniqueId unique_id,
-      const size_t rank,
-      const size_t world_size,
-      void* scratch,
-      const size_t scratch_bytes,
-      const size_t nranks_per_node,
-      const std::vector<int64_t>& rank_to_node,
-      const std::vector<int64_t>& rank_to_ib)
-      : scratch_(scratch),
-        scratch_bytes_(scratch_bytes),
-        nranks_per_node_(nranks_per_node),
-        d_memHandles_(nranks_per_node - 1) {
-    CHECK_CUDA_SUCCESS(cudaStreamCreateWithFlags(&h2d_stream, cudaStreamNonBlocking));
-    comm_group_ = std::make_shared<MscclCommGroup>(unique_id, rank, world_size, rank_to_node, rank_to_ib);
-    comm_group_->make_connection(same_node_connections_, cross_node_connections_);
-    comm_group_->make_memory_channels_with_scratch(
-        scratch_,
-        scratch_bytes_,
-        scratch_,
-        scratch_bytes_,
-        same_node_connections_,
-        memory_semaphores_,
-        registered_sm_memories_,
-        memory_channels_);
-    std::vector<mscclpp::MemoryChannel> memory_channels_list;
-    for (int r = 0; r < comm_group_->world_size_; r++) {
-      if (r == comm_group_->rank_) continue;
-      memory_channels_list.push_back(memory_channels_[r]);
-    }
-    std::vector<mscclpp::MemoryChannelDeviceHandle> memory_channel_handlers(memory_channels_list.size());
-    std::transform(
-        memory_channels_list.begin(),
-        memory_channels_list.end(),
-        memory_channel_handlers.begin(),
-        [](const mscclpp::MemoryChannel& channel) { return channel.deviceHandle(); });
-    mscclpp::gpuMemcpy<mscclpp::MemoryChannelDeviceHandle>(
-        d_memHandles_.data(), memory_channel_handlers.data(), memory_channel_handlers.size(), cudaMemcpyHostToDevice);
-  }
-
-  ~Msccl1NodeLLcontext() {
-    CHECK_CUDA_SUCCESS(cudaStreamDestroy(h2d_stream));
-  }
-
-  template <typename T>
-  void allreduce(cudaStream_t stream, T* input, T* output, size_t input_numel, int nthreads = 512, int nblocks = 21) {
-    dim3 nthrs(nthreads);
-    dim3 nblks(nblocks);
-    cudaStreamCaptureStatus capturing_status;
-    CHECK_CUDA_SUCCESS(cudaStreamIsCapturing(stream, &capturing_status));
-    mscclpp::MemoryChannelDeviceHandle* memChans;
-    if (capturing_status != cudaStreamCaptureStatusActive) {
-      std::unordered_map<int, mscclpp::MemoryChannel> memory_channels;
-      comm_group_->make_device_memory_handle_base_on_new_ptr(
-          memory_channels_,
-          registered_sm_memories_,
-          memory_semaphores_,
-          memory_channels,
-          d_memHandles_,
-          input,
-          scratch_,
-          h2d_stream);
-      CHECK_CUDA_SUCCESS(cudaStreamSynchronize(h2d_stream));
-      memChans = d_memHandles_.data();
-    } else {
-      void* input_void_ptr = reinterpret_cast<void*>(input);
-      if (input_ptr2d_memHandles_.find(input_void_ptr) == input_ptr2d_memHandles_.end()) {
-        std::unordered_map<int, mscclpp::MemoryChannel> memory_channels;
-        mscclpp::GpuBuffer<mscclpp::MemoryChannelDeviceHandle> device_memory_handle(comm_group_->world_size_ - 1);
-        comm_group_->make_device_memory_handle_base_on_new_ptr(
-            memory_channels_,
-            registered_sm_memories_,
-            memory_semaphores_,
-            memory_channels,
-            device_memory_handle,
-            input,
-            scratch_,
-            h2d_stream);
-        input_ptr2memory_channels_.emplace(input_void_ptr, memory_channels);
-        input_ptr2d_memHandles_.emplace(input_void_ptr, device_memory_handle);
-      }
-      auto it = input_ptr2d_memHandles_.find(input_void_ptr);
-      memChans = it->second.data();
-    }
-    allreduce_LL_1node<T><<<nblks, nthrs, 0, stream>>>(
-        memChans, (T*)input, (T*)scratch_, output, comm_group_->rank_, comm_group_->world_size_, input_numel);
-
-    cudaError_t status = cudaGetLastError();
-    if (status != cudaSuccess) {
-      printf("rank: %lu failed to launch allreduce_LL_1node: %s\n", comm_group_->rank_, cudaGetErrorString(status));
-    }
-  }
-};
-
-class Msccl2NodeLLcontext {
- private:
-  std::shared_ptr<MscclCommGroup> comm_group_ = nullptr;
-  void* scratch_;
-  const size_t scratch_bytes_;
-  void* put_buffer_;
-  const size_t put_buffer_bytes_;
-  std::unordered_map<int, std::shared_ptr<mscclpp::Connection>> same_node_connections_;
-  std::unordered_map<int, std::shared_ptr<mscclpp::Connection>> cross_node_connections_;
-
-  std::unordered_map<int, mscclpp::RegisteredMemory> registered_sm_memories_;
-  std::unordered_map<int, mscclpp::RegisteredMemory> registered_port_memories_;
-
-  std::unordered_map<int, std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>> memory_semaphores_;
-  std::unordered_map<int, std::shared_ptr<mscclpp::Host2DeviceSemaphore>> port_semaphores_;
-
-  std::unordered_map<int, mscclpp::MemoryChannel> memory_channels_;
-  std::unordered_map<int, mscclpp::PortChannel> port_channels_;
-
-  mscclpp::GpuBuffer<mscclpp::MemoryChannelDeviceHandle> d_memHandles_;
-  mscclpp::GpuBuffer<mscclpp::PortChannelDeviceHandle> d_portHandles_;
-
-  std::shared_ptr<mscclpp::ProxyService> proxyService;
-  cudaStream_t h2d_stream;
-  const size_t nranks_per_node_;
-
-  std::unordered_map<void*, std::unordered_map<int, mscclpp::MemoryChannel>> input_ptr2memory_channels_;
-  std::unordered_map<void*, mscclpp::GpuBuffer<mscclpp::MemoryChannelDeviceHandle>> input_ptr2d_memHandles_;
-
- public:
-  Msccl2NodeLLcontext(
-      mscclpp::UniqueId unique_id,
-      const size_t rank,
-      const size_t world_size,
-      void* scratch,
-      const size_t scratch_bytes,
-      void* put_buffer,
-      const size_t put_buffer_bytes,
-      const size_t nranks_per_node,
-      const std::vector<int64_t>& rank_to_node,
-      const std::vector<int64_t>& rank_to_ib)
-      : scratch_(scratch),
-        scratch_bytes_(scratch_bytes),
-        put_buffer_(put_buffer),
-        put_buffer_bytes_(put_buffer_bytes),
-        nranks_per_node_(nranks_per_node),
-        d_memHandles_(nranks_per_node - 1),
-        d_portHandles_(world_size - nranks_per_node) {
-    CHECK_CUDA_SUCCESS(cudaStreamCreateWithFlags(&h2d_stream, cudaStreamNonBlocking));
-    comm_group_ = std::make_shared<MscclCommGroup>(unique_id, rank, world_size, rank_to_node, rank_to_ib);
-    proxyService = std::make_shared<mscclpp::ProxyService>();
-    proxyService->startProxy();
-    comm_group_->make_connection(same_node_connections_, cross_node_connections_);
-    comm_group_->make_memory_channels_with_scratch(
-        scratch_,
-        scratch_bytes_,
-        scratch_,
-        scratch_bytes_,
-        same_node_connections_,
-        memory_semaphores_,
-        registered_sm_memories_,
-        memory_channels_);
-    comm_group_->make_port_channels_with_scratch(
-        proxyService,
-        put_buffer_,
-        put_buffer_bytes_,
-        scratch_,
-        scratch_bytes_,
-        cross_node_connections_,
-        port_semaphores_,
-        registered_port_memories_,
-        port_channels_);
-    std::vector<mscclpp::MemoryChannel> memory_channels_list;
-    std::vector<mscclpp::PortChannel> port_channels_list;
-    for (int r = 0; r < comm_group_->world_size_; r++) {
-      if (r == comm_group_->rank_) continue;
-      if (comm_group_->is_same_node(r, comm_group_->rank_)) {
-        memory_channels_list.push_back(memory_channels_[r]);
-      } else {
-        port_channels_list.push_back(port_channels_[r]);
-      }
-    }
-    std::vector<mscclpp::MemoryChannelDeviceHandle> memory_channel_handlers(memory_channels_list.size());
-    std::transform(
-        memory_channels_list.begin(),
-        memory_channels_list.end(),
-        memory_channel_handlers.begin(),
-        [](const mscclpp::MemoryChannel& channel) { return channel.deviceHandle(); });
-    mscclpp::gpuMemcpy<mscclpp::MemoryChannelDeviceHandle>(
-        d_memHandles_.data(), memory_channel_handlers.data(), memory_channel_handlers.size(), cudaMemcpyHostToDevice);
-
-    std::vector<mscclpp::PortChannelDeviceHandle> port_channel_handlers(port_channels_list.size());
-    std::transform(
-        port_channels_list.begin(),
-        port_channels_list.end(),
-        port_channel_handlers.begin(),
-        [](const mscclpp::PortChannel& channel) { return channel.deviceHandle(); });
-    mscclpp::gpuMemcpy<mscclpp::PortChannelDeviceHandle>(
-        d_portHandles_.data(), port_channel_handlers.data(), port_channel_handlers.size(), cudaMemcpyHostToDevice);
-  }
-
-  ~Msccl2NodeLLcontext() {
-    CHECK_CUDA_SUCCESS(cudaStreamDestroy(h2d_stream));
-    if (proxyService) {
-      proxyService->stopProxy();
-    }
-  }
-
-  template <typename T>
-  void
-  allreduce(cudaStream_t stream, T* input, T* output, const size_t input_numel, int nthreads = 512, int nblocks = 21) {
-    dim3 nthrs(nthreads);
-    dim3 nblks(nblocks);
-    cudaStreamCaptureStatus capturing_status;
-    CHECK_CUDA_SUCCESS(cudaStreamIsCapturing(stream, &capturing_status));
-    mscclpp::MemoryChannelDeviceHandle* memChans;
-    if (capturing_status != cudaStreamCaptureStatusActive) {
-      std::unordered_map<int, mscclpp::MemoryChannel> memory_channels;
-      comm_group_->make_device_memory_handle_base_on_new_ptr(
-          memory_channels_,
-          registered_sm_memories_,
-          memory_semaphores_,
-          memory_channels,
-          d_memHandles_,
-          input,
-          scratch_,
-          h2d_stream);
-      CHECK_CUDA_SUCCESS(cudaStreamSynchronize(h2d_stream));
-      memChans = d_memHandles_.data();
-    } else {
-      void* input_void_ptr = reinterpret_cast<void*>(input);
-      if (input_ptr2d_memHandles_.find(input_void_ptr) == input_ptr2d_memHandles_.end()) {
-        std::unordered_map<int, mscclpp::MemoryChannel> memory_channels;
-        mscclpp::GpuBuffer<mscclpp::MemoryChannelDeviceHandle> device_memory_handle(7);
-        comm_group_->make_device_memory_handle_base_on_new_ptr(
-            memory_channels_,
-            registered_sm_memories_,
-            memory_semaphores_,
-            memory_channels,
-            device_memory_handle,
-            input,
-            scratch_,
-            h2d_stream);
-        input_ptr2memory_channels_.emplace(input_void_ptr, memory_channels);
-        input_ptr2d_memHandles_.emplace(input_void_ptr, device_memory_handle);
-      }
-      auto it = input_ptr2d_memHandles_.find(input_void_ptr);
-      memChans = it->second.data();
-    }
-    allreduce_LL_2node<T><<<nblks, nthrs, 0, stream>>>(
-        memChans,
-        d_portHandles_.data(),
-        (T*)input,
-        (T*)scratch_,
-        (T*)put_buffer_,
-        output,
-        comm_group_->rank_,
-        nranks_per_node_,
-        comm_group_->world_size_,
-        input_numel);
-
-    cudaError_t status = cudaGetLastError();
-    if (status != cudaSuccess) {
-      printf("rank: %lu failed to launch allreduce_LL_2node: %s\n", comm_group_->rank_, cudaGetErrorString(status));
-    }
-  }
-};
-
-}  // namespace sglang

sgl-kernel/csrc/allreduce/test_mscclpp_allreduce.cu

@@ -1,153 +0,0 @@
-/*
- * this file is used to test mscclpp_allreduce.cu using mpirun
- * this file is adapted from https://github.com/flashinfer-ai/flashinfer/blob/v0.2.5/src/test_sum_all_reduce.cu
-usage:
-cd PATH-TO-THIS-FILE
-export MPI_HOME=/usr/local/mpi
-# export MPI_HOME=/opt/hpcx/ompi/
-export MSCCLPP_HOME=/workspace/test/mscclpp
-nvcc -O2 -arch=native -std=c++17 test_mscclpp_allreduce.cu \
-  -o test_mscclpp_allreduce -D_GLIBCXX_USE_CXX11_ABI=0 \
-  -I${MSCCLPP_HOME}/include -L${MSCCLPP_HOME}/build -lmscclpp \
-  -lnccl -I${MPI_HOME}/include -L${MPI_HOME}/lib -lmpi
-
-/opt/hpcx/ompi/bin/
-mpirun --allow-run-as-root -H 127.0.0.1:8 -np 8 \
-  --map-by ppr:8:node \
-  --mca btl_openib_warn_no_device_params_found 0 \
-  --mca btl_tcp_if_include bond0 \
-  --allow-run-as-root -np 8 \
-  -x NCCL_RUNTIME_CONNECT=0 -x NCCL_IB_GID_INDEX=3 -x NCCL_DEBUG=WARN \
-  -x LD_PRELOAD=${MSCCLPP_HOME}/build/libmscclpp.so ./test_mscclpp_allreduce
- */
-#include <mpi.h>
-#include <thrust/detail/raw_pointer_cast.h>
-#include <thrust/device_vector.h>
-#include <thrust/host_vector.h>
-
-#ifndef CHECK_CUDA_SUCCESS
-#define CHECK_CUDA_SUCCESS(cmd)                                                             \
-  do {                                                                                      \
-    cudaError_t e = cmd;                                                                    \
-    if (e != cudaSuccess) {                                                                 \
-      printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, cudaGetErrorString(e)); \
-      exit(EXIT_FAILURE);                                                                   \
-    }                                                                                       \
-  } while (0)
-#endif
-
-#include <cstdint>
-
-#include "mscclpp_allreduce.cuh"
-
-template <typename T>
-bool isclose(T a, T b, float rtol = 1e-5, float atol = 1e-8) {
-  return fabs(a - b) <= (atol + rtol * fabs(b));
-}
-
-int main(int argc, char* argv[]) {
-  // init mpi
-  MPI_Init(&argc, &argv);
-  printf("MPI Initialized.\n");
-  int nranks, rank;
-
-  // get work size and rank id
-  MPI_Comm_size(MPI_COMM_WORLD, &nranks);
-  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
-  cudaSetDevice(rank);
-  printf("nranks: %d, rank: %d\n", nranks, rank);
-
-  // init host and device buffers
-  using T = float;
-  using ReduceT = float;
-  const size_t num_elems = 2 * 1024 * 1024;
-  std::vector<T> host_buf(num_elems);
-  for (uint32_t i = 0; i < num_elems; ++i) {
-    host_buf[i] = T(i + rank);
-  }
-  thrust::device_vector<T> device_buf(host_buf);
-  const size_t buf_size_in_bytes = num_elems * sizeof(T);
-  std::vector<T> host_result_buf(num_elems);
-  thrust::device_vector<T> device_result_buf(host_result_buf);
-
-  std::vector<T> host_scratch_buf(num_elems * 8);
-  for (uint32_t i = 0; i < num_elems; ++i) {
-    host_scratch_buf[i] = 1;
-  }
-  thrust::device_vector<T> device_scratch_buf(host_scratch_buf);
-  std::vector<T> host_put_buf(num_elems);
-  thrust::device_vector<T> device_put_buf(host_put_buf);
-
-  mscclpp::UniqueId unique_id;
-  if (rank == 0) unique_id = mscclpp::TcpBootstrap::createUniqueId();
-  MPI_Bcast(&unique_id, sizeof(unique_id), MPI_BYTE, 0, MPI_COMM_WORLD);
-
-  std::vector<int64_t> rank_to_node(nranks);
-  std::vector<int64_t> rank_to_ib(nranks);
-  for (int i = 0; i < nranks; i++) {
-    rank_to_node[i] = i / 8;
-    rank_to_ib[i] = i % 8;
-  }
-
-  cudaStream_t s;
-  CHECK_CUDA_SUCCESS(cudaStreamCreate(&s));
-  CHECK_CUDA_SUCCESS(cudaStreamSynchronize(s));
-  if (nranks == 8) {
-    auto context = std::make_shared<sglang::Msccl1NodeLLcontext>(
-        unique_id,
-        rank,
-        nranks,
-        thrust::raw_pointer_cast(device_scratch_buf.data()),
-        buf_size_in_bytes * 8,
-        rank_to_node,
-        rank_to_ib);
-    printf("rank: %d, Msccl1NodeLLcontext setup.\n", rank);
-    MPI_Barrier(MPI_COMM_WORLD);
-    context->allreduce<T>(
-        s,
-        thrust::raw_pointer_cast(device_buf.data()),
-        thrust::raw_pointer_cast(device_result_buf.data()),
-        device_buf.size());
-  } else if (nranks == 16) {
-    // TODO: this branch is untested since there is something wrong with mpirun in my test machince
-    auto context = std::make_shared<sglang::Msccl2NodeLLcontext>(
-        unique_id,
-        rank,
-        nranks,
-        thrust::raw_pointer_cast(device_scratch_buf.data()),
-        buf_size_in_bytes * 8,
-        thrust::raw_pointer_cast(device_put_buf.data()),
-        buf_size_in_bytes,
-        rank_to_node,
-        rank_to_ib);
-    printf("rank: %d, Msccl2NodeLLcontext setup.\n", rank);
-    MPI_Barrier(MPI_COMM_WORLD);
-    context->allreduce<T>(
-        s,
-        thrust::raw_pointer_cast(device_buf.data()),
-        thrust::raw_pointer_cast(device_result_buf.data()),
-        device_buf.size());
-  }
-
-  // check result correctness
-  thrust::host_vector<T> host_buf_result = device_result_buf;
-  size_t num_results_error_atol_1e_3_rtol_1e_3 = 0;
-  bool nan_detected = false;
-
-  for (uint32_t i = 0; i < num_elems; ++i) {
-    T expected = T(i * nranks + (nranks - 1) * nranks / 2);
-    if (std::isnan(float(host_buf_result[i]))) {
-      nan_detected = true;
-    }
-    if (!isclose(float(host_buf_result[i]), float(expected), 1e-3, 1e-3)) {
-      num_results_error_atol_1e_3_rtol_1e_3++;
-    }
-  }
-  float result_accuracy = 1. - float(num_results_error_atol_1e_3_rtol_1e_3) / float(num_elems);
-
-  printf("rank: %d, nan_detected: %d accuracy: %f\n", rank, nan_detected, result_accuracy);
-
-  CHECK_CUDA_SUCCESS(cudaStreamDestroy(s));
-  MPI_Finalize();
-  return 0;
-}

sgl-kernel/csrc/common_extension.cc

@@ -38,15 +38,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
       "int reg_buffer_sz_bytes) -> ()");
   m.impl("all_reduce", torch::kCUDA, &all_reduce);
 
-  m.def("mscclpp_generate_unique_id", &mscclpp_generate_unique_id);
-  m.def(
-      "mscclpp_init_context(Tensor unique_id, int rank, int world_size, Tensor scratch, Tensor put_buffer, "
-      "int nranks_per_node, int[] rank_to_node, int[] rank_to_ib, int context_selection) -> int");
-  m.impl("mscclpp_init_context", torch::kCUDA, &mscclpp_init_context);
-
-  m.def("mscclpp_allreduce(int context, Tensor inp, Tensor! out, int nthreads, int nblocks) -> ()");
-  m.impl("mscclpp_allreduce", torch::kCUDA, &mscclpp_allreduce);
-
   /*
    * From csrc/attention
    */

sgl-kernel/include/sgl_kernel_ops.h

@@ -89,20 +89,6 @@ std::tuple<std::vector<int64_t>, std::vector<int64_t>> get_graph_buffer_ipc_meta
 void register_buffer(fptr_t _fa, const std::vector<fptr_t>& fake_ipc_ptrs);
 void register_graph_buffers(
     fptr_t _fa, const std::vector<std::vector<int64_t>>& handles, const std::vector<std::vector<int64_t>>& offsets);
-
-// mscclpp
-torch::Tensor mscclpp_generate_unique_id();
-fptr_t mscclpp_init_context(
-    const torch::Tensor& unique_id,
-    const int64_t rank,
-    const int64_t world_size,
-    torch::Tensor& scratch,
-    torch::Tensor& put_buffer,
-    const int64_t nranks_per_node,
-    const std::vector<int64_t>& rank_to_node,
-    const std::vector<int64_t>& rank_to_ib,
-    const int64_t context_selection);
-void mscclpp_allreduce(fptr_t _context, torch::Tensor& inp, torch::Tensor& out, int64_t nthreads, int64_t nblocks);
 #endif
 
 /*

sgl-kernel/python/sgl_kernel/allreduce.py

@@ -92,28 +92,6 @@ if torch.version.hip is not None:
     def qr_max_size() -> int:
         return torch.ops.sgl_kernel.qr_max_size.default()
 
-    # mscclpp
-    def mscclpp_generate_unique_id() -> bytes:
-        raise NotImplementedError()
-
-    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:
-        raise NotImplementedError()
-
-    def mscclpp_allreduce(
-        context: int, inp: torch.Tensor, out: torch.Tensor, nthreads: int, nblocks: int
-    ) -> None:
-        raise NotImplementedError()
-
 else:
 
     def init_custom_ar(
@@ -150,36 +128,3 @@ else:
 
     def meta_size() -> int:
         return torch.ops.sgl_kernel.meta_size.default()
-
-    def mscclpp_generate_unique_id() -> torch.Tensor:
-        return torch.ops.sgl_kernel.mscclpp_generate_unique_id.default()
-
-    def mscclpp_init_context(
-        unique_id: torch.Tensor,
-        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 torch.ops.sgl_kernel.mscclpp_init_context.default(
-            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:
-        torch.ops.sgl_kernel.mscclpp_allreduce.default(
-            context, inp, out, nthreads, nblocks
-        )

sgl-kernel/tests/test_mscclpp.py

@@ -1,146 +0,0 @@
-import multiprocessing as mp
-import os
-import socket
-import unittest
-from enum import IntEnum
-from typing import Any
-
-import sgl_kernel.allreduce as custom_ops
-import torch
-import torch.distributed as dist
-
-
-class MscclContextSelection(IntEnum):
-    MSCCL1SHOT1NODELL = 1
-    MSCCL1SHOT2NODELL = 2
-
-
-def _run_correctness_worker(world_size, rank, distributed_init_port, test_sizes):
-    device = torch.device(f"cuda:{rank % torch.cuda.device_count()}")
-    torch.cuda.set_device(device)
-    distributed_init_method = f"tcp://localhost:{distributed_init_port}"
-    dist.init_process_group(
-        backend="nccl",
-        init_method=distributed_init_method,
-        rank=rank,
-        world_size=world_size,
-    )
-    group = dist.group.WORLD
-    cpu_group = torch.distributed.new_group(list(range(world_size)), backend="gloo")
-    if rank == 0:
-        unique_id = [custom_ops.mscclpp_generate_unique_id()]
-    else:
-        unique_id = [None]
-    dist.broadcast_object_list(
-        unique_id, src=0, device=torch.device("cpu"), group=cpu_group
-    )
-    unique_id = unique_id[0]
-    rank_to_node, rank_to_ib = list(range(world_size)), list(range(world_size))
-    for r in range(world_size):
-        rank_to_node[r] = r // 8
-        rank_to_ib[r] = rank % 8
-    MAX_BYTES = 2**20
-    scratch = torch.empty(
-        MAX_BYTES * 8, dtype=torch.bfloat16, device=torch.cuda.current_device()
-    )
-    put_buffer = torch.empty(
-        MAX_BYTES, dtype=torch.bfloat16, device=torch.cuda.current_device()
-    )
-    print(f"[{rank}] start mscclpp_context init")
-    nranks_per_node = torch.cuda.device_count()
-    selection = int(MscclContextSelection.MSCCL1SHOT1NODELL)
-    mscclpp_context = custom_ops.mscclpp_init_context(
-        unique_id,
-        rank,
-        world_size,
-        scratch,
-        put_buffer,
-        nranks_per_node,
-        rank_to_node,
-        rank_to_ib,
-        selection,
-    )
-    try:
-        test_loop = 10
-        for sz in test_sizes:
-            for dtype in [torch.float32, torch.float16, torch.bfloat16]:
-                if sz * dtype.itemsize > MAX_BYTES:
-                    continue
-                if rank == 0:
-                    print(f"mscclpp allreduce test sz {sz}, dtype {dtype}")
-                for _ in range(test_loop):
-                    inp1 = torch.randint(1, 16, (sz,), dtype=dtype, device=device)
-                    inp1_ref = inp1.clone()
-                    out1 = torch.empty_like(inp1)
-                    custom_ops.mscclpp_allreduce(
-                        mscclpp_context, inp1, out1, nthreads=512, nblocks=21
-                    )
-                    dist.all_reduce(inp1_ref, group=group)
-                    torch.testing.assert_close(out1, inp1_ref)
-    finally:
-        dist.barrier(group=group)
-        dist.destroy_process_group(group=group)
-
-
-def get_open_port() -> int:
-    try:
-        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-            s.bind(("127.0.0.1", 0))
-            return s.getsockname()[1]
-    except OSError:
-        with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
-            s.bind(("::1", 0))
-            return s.getsockname()[1]
-
-
-def multi_process_parallel(
-    world_size: int, test_target: Any, target_args: tuple = ()
-) -> None:
-    mp.set_start_method("spawn", force=True)
-
-    procs = []
-    distributed_init_port = get_open_port()
-    for i in range(world_size):
-        proc_args = (world_size, i, distributed_init_port) + target_args
-        proc = mp.Process(target=test_target, args=proc_args, name=f"Worker-{i}")
-        proc.start()
-        procs.append(proc)
-
-    for i in range(world_size):
-        procs[i].join()
-        assert (
-            procs[i].exitcode == 0
-        ), f"Process {i} failed with exit code {procs[i].exitcode}"
-
-
-class TestMSCCLAllReduce(unittest.TestCase):
-    test_sizes = [
-        512,
-        2560,
-        4096,
-        5120,
-        7680,
-        32768,
-        262144,
-        524288,
-    ]
-    world_sizes = [8]
-
-    def test_correctness(self):
-        for world_size in self.world_sizes:
-            available_gpus = torch.cuda.device_count()
-            if world_size > available_gpus:
-                print(
-                    f"Skipping world_size={world_size}, found {available_gpus} and now ray is not supported here"
-                )
-                continue
-
-            print(f"Running test for world_size={world_size}")
-            multi_process_parallel(
-                world_size, _run_correctness_worker, target_args=(self.test_sizes,)
-            )
-            print(f"custom allreduce tp = {world_size}: OK")
-
-
-if __name__ == "__main__":
-    unittest.main()

test/manual/test_mscclpp.py

@@ -1,196 +0,0 @@
-"""For Now, MSCCL is only supported on TP16 and TP8 case
-
-if [[ $RANK -eq 0  ]]; then
-    ray start --block --head --port=6379 &
-    python3 test_mscclpp.py;
-else
-    ray start --block --address=${MASTER_ADDR}:6379;
-fi
-"""
-
-import os
-import random
-import socket
-import unittest
-from typing import Any
-
-import ray
-import torch
-import torch.distributed as dist
-
-from sglang.srt.distributed import init_distributed_environment
-from sglang.srt.distributed.communication_op import (  # noqa
-    tensor_model_parallel_all_reduce,
-)
-from sglang.srt.distributed.parallel_state import (
-    get_tensor_model_parallel_group,
-    graph_capture,
-    initialize_model_parallel,
-    set_custom_all_reduce,
-    set_mscclpp_all_reduce,
-)
-from sglang.test.test_utils import CustomTestCase
-
-
-def get_open_port() -> int:
-    # try ipv4
-    try:
-        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-            s.bind(("", 0))
-            return s.getsockname()[1]
-    except OSError:
-        # try ipv6
-        with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
-            s.bind(("", 0))
-            return s.getsockname()[1]
-
-
-def multi_process_parallel(
-    world_size: int,
-    master_addr: str,
-    cls: Any,
-    test_target: Any,
-) -> None:
-
-    # Using ray helps debugging the error when it failed
-    # as compared to multiprocessing.
-    # NOTE: We need to set working_dir for distributed tests,
-    # otherwise we may get import errors on ray workers
-
-    ray.init(log_to_driver=True)
-
-    distributed_init_port = get_open_port()
-    refs = []
-    for rank in range(world_size):
-        refs.append(
-            test_target.remote(
-                cls, world_size, master_addr, rank, distributed_init_port
-            )
-        )
-    ray.get(refs)
-
-    ray.shutdown()
-
-
-class TestMSCCLAllReduce(CustomTestCase):
-    @classmethod
-    def setUpClass(cls):
-        random.seed(42)
-        # 1KB to 1MB
-        cls.test_sizes = [512, 4096, 32768, 262144, 524288]
-        cls.world_sizes = [8]
-        TEST_TP16 = int(os.getenv("SGL_MSCCLPP_TEST_TP16", "0"))
-        if TEST_TP16:
-            cls.world_sizes = [16]
-        cls.test_loop = 10
-
-    def test_graph_allreduce(self):
-        TEST_MASTER_ADDR = os.getenv("SGL_MSCCLPP_TEST_MASTER_ADDR", "localhost")
-        for world_size in self.world_sizes:
-            if world_size not in [8, 16]:
-                continue
-            multi_process_parallel(
-                world_size, TEST_MASTER_ADDR, self, self.graph_allreduce
-            )
-
-    def test_eager_allreduce(self):
-        TEST_MASTER_ADDR = os.getenv("SGL_MSCCLPP_TEST_MASTER_ADDR", "localhost")
-        for world_size in self.world_sizes:
-            if world_size not in [8, 16]:
-                continue
-            multi_process_parallel(
-                world_size, TEST_MASTER_ADDR, self, self.eager_allreduce
-            )
-
-    @ray.remote(num_gpus=1, max_calls=1)
-    def graph_allreduce(self, world_size, master_addr, rank, distributed_init_port):
-        del os.environ["CUDA_VISIBLE_DEVICES"]
-        device = torch.device(f"cuda:{rank % torch.cuda.device_count()}")
-        torch.cuda.set_device(device)
-        distributed_init_method = f"tcp://{master_addr}:{distributed_init_port}"
-        set_mscclpp_all_reduce(True)
-        set_custom_all_reduce(False)
-        init_distributed_environment(
-            world_size=world_size,
-            rank=rank,
-            distributed_init_method=distributed_init_method,
-            local_rank=rank % torch.cuda.device_count(),
-        )
-        initialize_model_parallel(tensor_model_parallel_size=world_size)
-        group = get_tensor_model_parallel_group().device_group
-
-        # A small all_reduce for warmup.
-        # this is needed because device communicators might be created lazily
-        # (e.g. NCCL). This will ensure that the communicator is initialized
-        # before any communication happens, so that this group can be used for
-        # graph capture immediately.
-        data = torch.zeros(1)
-        data = data.to(device=device)
-        torch.distributed.all_reduce(data, group=group)
-        torch.cuda.synchronize()
-        del data
-
-        for sz in self.test_sizes:
-            for dtype in [torch.float32, torch.float16, torch.bfloat16]:
-                for _ in range(self.test_loop):
-                    with graph_capture() as graph_capture_context:
-                        # use integers so result matches NCCL exactly
-                        inp1 = torch.randint(
-                            1,
-                            16,
-                            (sz,),
-                            dtype=dtype,
-                            device=torch.cuda.current_device(),
-                        )
-                        inp2 = torch.randint(
-                            1,
-                            16,
-                            (sz,),
-                            dtype=dtype,
-                            device=torch.cuda.current_device(),
-                        )
-                        torch.cuda.synchronize()
-                        graph = torch.cuda.CUDAGraph()
-                        with torch.cuda.graph(
-                            graph, stream=graph_capture_context.stream
-                        ):
-                            out1 = tensor_model_parallel_all_reduce(inp1)
-                            # the input buffer is immediately modified to test
-                            # synchronization
-                            dist.all_reduce(inp1, group=group)
-                            out2 = tensor_model_parallel_all_reduce(inp2)
-                            dist.all_reduce(inp2, group=group)
-                    graph.replay()
-                    torch.testing.assert_close(out1, inp1)
-                    torch.testing.assert_close(out2, inp2)
-
-    @ray.remote(num_gpus=1, max_calls=1)
-    def eager_allreduce(self, world_size, master_addr, rank, distributed_init_port):
-        del os.environ["CUDA_VISIBLE_DEVICES"]
-        device = torch.device(f"cuda:{rank % torch.cuda.device_count()}")
-        torch.cuda.set_device(device)
-        distributed_init_method = f"tcp://{master_addr}:{distributed_init_port}"
-        set_mscclpp_all_reduce(True)
-        set_custom_all_reduce(False)
-        init_distributed_environment(
-            world_size=world_size,
-            rank=rank,
-            distributed_init_method=distributed_init_method,
-            local_rank=rank,
-        )
-        initialize_model_parallel(tensor_model_parallel_size=world_size)
-        group = get_tensor_model_parallel_group().device_group
-
-        for sz in self.test_sizes:
-            for dtype in [torch.float32, torch.float16, torch.bfloat16]:
-                for _ in range(self.test_loop):
-                    inp1 = torch.randint(
-                        1, 16, (sz,), dtype=dtype, device=torch.cuda.current_device()
-                    )
-                    out1 = tensor_model_parallel_all_reduce(inp1)
-                    dist.all_reduce(inp1, group=group)
-                    torch.testing.assert_close(out1, inp1)
-
-
-if __name__ == "__main__":
-    unittest.main()