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()