support 1 shot allreduce in 1-node and 2-node using mscclpp

[zyksir]

Motivation

For now we can only use 1-shot all reduce in 1 node. However, there are times that we need TP16 and 2-node allreduce. Recently mscclpp has some examples about 1 shot all reduce in 2-node, which can improve the latency of 2-node allreduce for small messages.

Modifications

• update on 5.21: mscclpp only provide python based demo for all-reduce, which need to compile kernel in runtime. Here I rewrite the python API using cpp so that we can use pybind11 to call mscclpp. For now the code can compile and run correctly. "Autotune for best nBlocks and nThreads" is left unfinished, which means the performance might not be the best.
• update on 5.22: mscclpp allreduce pass all the unit tests. The original mscclpp demo has bugs in cuda graph for device memory handles. The result of benchmark on A100 is listed as follows. "Autotune for best nBlocks and nThreads" need to be added for better performance on large message sizes. In addition, there is something wrong with my test machine and I failed to run any code with ray on 2-node. So Both test_mscclpp.py files is not tested on TP16 case. The accuracy is tested by benchmark_mscclpp.py is launched by torchrun.

kernel benchmark

This is tested on A100 with 8 GPUs and 4NIC per node.

msg_size	torch eager time	msccl eager time	msccl graph time	pynccl graph time
2.0 KiB	61.44	33.5872	22.8557	38.2976
4.0 KiB	44.6464	33.9968	22.0979	42.3834
8.0 KiB	44.6464	31.5392	22.487	43.9296
16.0 KiB	57.9584	31.3344	23.4598	44.1344
32.0 KiB	55.5008	32.4608	24.0947	45.8035
64.0 KiB	55.6032	33.28	26.4704	53.3606
128.0 KiB	65.4336	39.936	32.6144	61.0406
256.0 KiB	81.6128	49.7664	43.9603	81.0394
512.0 KiB	93.3888	71.3728	66.2221	95.744
1.0 MiB	126.157	110.797	110.469	125.471

This is tested on H100 with 8GPUs and 2NIC per node. We can see that for 1MB message, MSCCLPP is slower than NCCL. This is because the current implementation of MSCCLPP use LLPacket, which will waste half of the bandwidth. The test machine only has two NIC for 8 GPUs, which means the bandwidth will be easier to become bottleneck.

msg_size	torch eager time	msccl eager time	msccl graph time	pynccl graph time
2.0 KiB	82.04	38.86	21.38	35.09
4.0 KiB	73.45	37.56	22.10	36.59
8.0 KiB	65.30	43.71	26.77	37.98
16.0 KiB	62.65	37.97	23.35	45.06
32.0 KiB	80.16	45.32	24.69	40.19
64.0 KiB	70.45	38.46	28.43	40.25
128.0 KiB	70.25	38.56	33.61	45.55
256.0 KiB	67.01	46.38	45.70	59.32
512.0 KiB	85.04	69.51	68.76	74.79
1.0 MiB	91.49	129.99	114.86	76.35

end-to-end accuray test

This is tested on A100 and QWen-72B with TP16
Baseline

MSCCLPP

This is tested on H100 on DeepSeek-V3 with TP16

mmlu result

Parallel	Baseline	Mscclpp
64	17.08s/acc87.1%	32.87s/acc87.4%
32	22.23s/acc87.1%	21.51s/acc87.3%
16	39.03s/acc87.4%	38.29s/acc87.3%
8	76.29s/acc87.1%	67.05s/acc87.4%
4	118.86s/acc87.2%	119.66s/acc87.1%
2	134.81s/acc87.4%	129.57s/acc87.4%
1	243.15s/acc87.3%	238.17s/acc87.3%

gsm8k result

Parallel	Baseline	Mscclpp
64	30.19s/acc94.0%	26.41s/acc94.5%
32	42.63s/acc94.5%	38.29s/acc93.0%
16	61.18s/acc94.5%	54.19s/acc93.5%
8	88.85s/acc95.0%	79.52s/acc93.5%
4	140.89s/acc95.5%	128.08s/acc94.5%
2	247.56s/acc96.0%	222.89s/acc95.0%
1	448.45s/acc94.0%	407.65s/acc93.0%

end-to-end performance test

This is tested on H100 and DeepSeek-V3 with TP16. We can see that the improvement is larger on low concurrency, where the message size is smaller. This is aligned with the benchmark result mentioned above.

concurrency	Baseline	Mscclpp
4	22.77	20.47
16	29.30	26.39
32	32.48	31.81

existing issues on mscclpp

There are some issues that need to be solved in the future

1. the graph launch time of mscclpp is longer than nccl.
2. the mscclpp is unstable on large message sizes

Adaptation

This part records some issues while I develop this feature.

C++ MSCCLPP API

The mscclpp only provide a naive demo, where the input is fixed and each time the kernel is compiled runtime. I rewrite the code in sgl-kernel to support different input and cuda graph. The correctness of this code is tested by the unit-test and benchmark, but for unknown reason, the performance of cpp code is slower than the demo. I haved create an issue in mscclpp repo.

_GLIBCXX_USE_CXX11_ABI issue

bug: while load pybind11 extension, "undefined symbol: _ZN7mscclpp9CudaErrorC1ERKSsi". To find the root cause, command c++filt <symbol> is used to find the original definition, and command nm -D <path-to-so> | grep CudaError is used to find symbol defined in so file. By doing so we can see find the root cause: libtorch.so is compiled with -D_GLIBCXX_USE_CXX11_ABI=0 and libmscclpp.so is compiled with -D_GLIBCXX_USE_CXX11_ABI=1, which causes conflict.

ps: using python3 -c "import torch; print(int(torch._C._GLIBCXX_USE_CXX11_ABI))" to select the value of _GLIBCXX_USE_CXX11_ABI is more robust since in the future this value might differ.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[jon-chuang]

Just wondering if they can support standalone one shot reduce scatter, all gather and a2a as well? This just involves choosing whether to send entire buffer or slice thereof and choosing either to sum over the buffers or concat, from which one gets a 2 x 2 matrix. E.g.

Allreduce - full, sum
Reduce scatter - slice, sum
A2A - slice, cat
Allgather - full, cat

I wonder how the a2a would compare with deepseek DeepEP and bytedance M2N (not open sourced) libs

[zyksir]

Base on the provided mscclpp API, mscclpp can be used to implemented all these kernels. As long as all reduce is supported in sglang, the only thing left is to change the kernel implementation.

However, the contributors of mscclpp does not provide demo code, extra work will be needed to debug the cuda implementation.

In addition, the DeepEP api is not an alltoall since it fuses some permutation kernel. I think we need to carefully test the performance between mscclpp and nvshmem.

[zyksir]

DeepEP low latency kernel uses IBGDA to reduce the latency. For now I don't think that mscclpp will use IBGDA. I will check this as well

[Edenzzzz]

MSCCLPP can support other kernels as implemented in Nanoflow. However, I think its key advantage is customizability instead of performance? In my tests for large messages, it performs ~20% better when we manually control the block size for compute-comm overlap

[zyksir]

I agree that customizability is the key advantage. In the future we can implement allreduce/layernorm fusion based on this.

[Alcanderian]

I found that we may have to support two algo to cover large batch size decode.
https://github.com/microsoft/mscclpp/blob/main/python/mscclpp_benchmark/allreduce_bench.py#L182

        if memory.nbytes < 2**22:
            mscclpp_algos = [MscclppAllReduce5(mscclpp_group, memory, memory_out, N_GPUS_PER_NODE, proxy_service)]
        else:
            mscclpp_algos = [MscclppAllReduce4(mscclpp_group, memory, N_GPUS_PER_NODE, proxy_service)]

[zyksir]

I found that we may have to support two algo to cover large batch size decode. https://github.com/microsoft/mscclpp/blob/main/python/mscclpp_benchmark/allreduce_bench.py#L182

        if memory.nbytes < 2**22:
            mscclpp_algos = [MscclppAllReduce5(mscclpp_group, memory, memory_out, N_GPUS_PER_NODE, proxy_service)]
        else:
            mscclpp_algos = [MscclppAllReduce4(mscclpp_group, memory, N_GPUS_PER_NODE, proxy_service)]

I found that even in the benchmark script provided by mscclpp, the performance on large message is slower than NCCL. So for now I am working on improvement on small message sizes.

[Edenzzzz]

Btw, by two-shot do you mean intra-node all-reduce + inter-node all-reduce? I think NCCL doesn't do hierarchical all-reduce yet?

[zyksir]

Btw, by two-shot do you mean intra-node all-reduce + inter-node all-reduce? I think NCCL doesn't do hierarchical all-reduce yet?

The 2-node allreduce is actually intra-node reduce-scatter + inter-node send-recv + intra-node broadcast. I am not sure how NCCL does allreduce. but I think this is a different way.

[Edenzzzz]

I think NCCL just uses either ring or tree, don't know if that counts as one-shot
NVIDIA/nccl#530

[Fridge003]

Please fix the compilation error in https://github.com/sgl-project/sglang/actions/runs/15343913091/job/43175858908?pr=6277

[zyksir]

@Fridge003 Thank you! I have fixed this compile error