diff --git a/benchmark/kernels/all_reduce/README.md b/benchmark/kernels/all_reduce/README.md new file mode 100644 index 000000000..80d0cd2e5 --- /dev/null +++ b/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= +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. diff --git a/benchmark/kernels/all_reduce/benchmark_mscclpp.py b/benchmark/kernels/all_reduce/benchmark_mscclpp.py index eebbd00ce..5b72d2f46 100644 --- a/benchmark/kernels/all_reduce/benchmark_mscclpp.py +++ b/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() diff --git a/docker/Dockerfile b/docker/Dockerfile index 3efc0dc90..0e5369d2d 100644 --- a/docker/Dockerfile +++ b/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 \ diff --git a/python/sglang/bench_one_batch.py b/python/sglang/bench_one_batch.py index da26065ec..57f3eb3e3 100644 --- a/python/sglang/bench_one_batch.py +++ b/python/sglang/bench_one_batch.py @@ -65,7 +65,10 @@ import torch import torch.distributed as dist from sglang.srt.configs.model_config import ModelConfig -from sglang.srt.distributed.parallel_state import destroy_distributed_environment +from sglang.srt.distributed.parallel_state import ( + destroy_distributed_environment, + destroy_model_parallel, +) from sglang.srt.entrypoints.engine import _set_envs_and_config from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.layers.moe import initialize_moe_config @@ -936,6 +939,7 @@ def latency_test( fout.write(json.dumps(result) + "\n") if server_args.tp_size > 1: + destroy_model_parallel() destroy_distributed_environment() @@ -957,12 +961,17 @@ def main(server_args, bench_args): port_args = PortArgs.init_new(server_args) + # Calculate local ranks for multi-node setup + nranks_per_node = server_args.tp_size // server_args.nnodes + local_rank_start = server_args.node_rank * nranks_per_node + local_rank_end = local_rank_start + nranks_per_node + if server_args.tp_size == 1: work_func(server_args, port_args, bench_args, 0, 0) else: workers = [] - for tp_rank in range(server_args.tp_size): - with maybe_reindex_device_id(tp_rank) as gpu_id: + for tp_rank in range(local_rank_start, local_rank_end): + with maybe_reindex_device_id(tp_rank - local_rank_start) as gpu_id: proc = multiprocessing.Process( target=work_func, args=( diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py index 09b981f9b..3459165b7 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py +++ b/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) diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_ops.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_ops.py index ea1c1a3ed..915bc5d97 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_ops.py +++ b/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 diff --git a/python/sglang/srt/distributed/device_communicators/pymscclpp.py b/python/sglang/srt/distributed/device_communicators/pymscclpp.py index e45093c78..ae62c4d90 100644 --- a/python/sglang/srt/distributed/device_communicators/pymscclpp.py +++ b/python/sglang/srt/distributed/device_communicators/pymscclpp.py @@ -1,108 +1,247 @@ -import bisect +import importlib import logging -import math -import os from contextlib import contextmanager -from enum import IntEnum from typing import Optional, Union import torch import torch.distributed as dist from torch.distributed import ProcessGroup, ReduceOp -import sglang.srt.distributed.device_communicators.custom_all_reduce_ops as ops -from sglang.srt.utils import is_hip +from sglang.srt.compilation.piecewise_context_manager import ( + get_pcg_capture_stream, + is_in_pcg_torch_compile, + is_in_piecewise_cuda_graph, +) +from sglang.srt.server_args import get_global_server_args logger = logging.getLogger(__name__) -_is_hip = is_hip() - - -class MscclContextSelection(IntEnum): - MSCCL1SHOT1NODELL = 1 - MSCCL1SHOT2NODELL = 2 - - -def mscclpp_is_weak_contiguous(inp: torch.Tensor): - return inp.is_contiguous() or ( - inp.storage().nbytes() - inp.storage_offset() * inp.element_size() - == inp.numel() * inp.element_size() - ) - - -def mscclpp_convert_to_bytes(size_str): - """ - Converts a human-readable size string (e.g., "1MB", "2.5kb", "3 GB") - into the equivalent number of bytes using binary units. - - Args: - size_str (str): A string representing size with unit (KB, MB, GB). - - Returns: - int: Number of bytes. - """ - size_str = size_str.strip().lower() - - if not size_str: - raise ValueError("Empty input string") - - # Extract numeric part and unit - for i in range(len(size_str)): - if not size_str[i].isdigit() and size_str[i] != ".": - break - num_str = size_str[:i] - unit = size_str[i:].strip() - - try: - num = float(num_str) - except ValueError: - raise ValueError(f"Invalid numeric value in '{size_str}'") - - # Conversion factors - if unit == "b": - return int(num) - elif unit == "kb": - return int(num * 1024) - elif unit == "mb": - return int(num * 1024 * 1024) - elif unit == "gb": - return int(num * 1024 * 1024 * 1024) - else: - raise ValueError(f"Unsupported unit: {unit}, support B, KB, MB, GB only") - - -def mscclpp_bench_time(func, test_niter: int = 10, warmup_niter: int = 2): - # warmup - for _ in range(warmup_niter): - func() - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - torch.cuda.synchronize() - dist.barrier() - start_event.record() - for _ in range(test_niter): - func() - end_event.record() - end_event.synchronize() - func_cost_us = start_event.elapsed_time(end_event) / test_niter * 1000 - return func_cost_us - class PyMscclppCommunicator: - _SUPPORTED_WORLD_SIZES = [8, 16] - _MAX_BYTES = mscclpp_convert_to_bytes(os.getenv("SGLANG_MSCCLPP_MAX_BYTES", "1MB")) + _SUPPORTED_WORLD_SIZES = [8, 16, 32] _SUPPORTED_DTYPE = [torch.float, torch.float16, torch.bfloat16] - # max_bytes: max supported mscclpp allreduce size - # in A100 mscclpp is faster than nccl only under condition of msg size smaller than1MB + def _is_symm_mem_enabled(self) -> bool: + try: + return get_global_server_args().enable_symm_mem + except ValueError: + return False + + def _is_weak_contiguous(self, inp: torch.Tensor): + return inp.is_contiguous() or ( + inp.storage().nbytes() - inp.storage_offset() * inp.element_size() + == inp.numel() * inp.element_size() + ) + + def _get_tuned_config(self, size): + if size <= 512: + target_size = 512 + elif size > 256 * 1024 * 1024: + target_size = 256 * 1024 * 1024 + else: + target_size = 1 << (size - 1).bit_length() + return self.best_configs.get(target_size) + + def _create_dsl_algorithms(self): + dsl_algos_config = [] + n_nodes = self.world_size // self.nranks_per_node + if n_nodes == 2 or n_nodes == 4: + for tbg in [1, 2, 4, 8]: + for num_threads_per_block in [256, 512, 768, 1024]: + spec = self.mscclpp.language.AlgoSpec( + name=f"allreduce_{n_nodes}node_{tbg}TBG_{num_threads_per_block}TPB", + collective=self.mscclpp.language.collectives.AllReduce( + self.world_size, 1, True + ), + nranks_per_node=self.nranks_per_node, + world_size=self.world_size, + in_place=True, + instances=1, + protocol="LL", + auto_sync=False, + num_threads_per_block=num_threads_per_block, + reuse_resources=True, + use_double_scratch_buffer=True, + min_message_size=tbg * (1 << 10), + max_message_size=8 << 20, + tags={"default": 1}, + ) + algo = self.mscclpp.compile( + self.def_algo.allreduce_multi_nodes, + spec, + self.rank, + thread_block_group_size=tbg, + ) + dsl_algos_config.append((algo, [0], [0])) + return dsl_algos_config + + def _create_native_algorithms(self): + navitve_algorithms_config = [] + dlpack = self.mscclpp.RawGpuBuffer(1 << 27).to_dlpack( + data_type=str(torch.float16) + ) + self.scratch_buffer = torch.utils.dlpack.from_dlpack(dlpack) + self.flag_buffer = torch.ones(128, dtype=torch.uint32, device="cuda") + algos = self.mscclpp_ext.AlgorithmCollectionBuilder().build_default_algorithms( + scratch_buffer=self.scratch_buffer.data_ptr(), + scratch_buffer_size=self.scratch_buffer.nbytes, + rank=self.rank, + ) + + for algo in algos: + if algo.name == "default_allreduce_nvls_packet": + algo.set_message_size_range(0, 512 << 10) + navitve_algorithms_config.append( + (algo, [4, 8, 12, 16], [256, 512, 768, 1024]) + ) + if algo.name == "default_allreduce_packet": + algo.set_message_size_range(0, 2 << 20) + navitve_algorithms_config.append( + (algo, [14, 21, 28, 42, 56], [256, 512, 768, 1024]) + ) + if algo.name == "default_allreduce_rsag_zero_copy": + algo.set_message_size_range(512 << 10, 4 << 30) + navitve_algorithms_config.append( + (algo, [32, 48, 64, 128], [256, 512, 768, 1024]) + ) + if ( + self.symm_mem_enabled + and algo.name == "default_allreduce_nvls_zero_copy" + ): + algo.set_message_size_range(512 << 10, 4 << 30) + navitve_algorithms_config.append( + (algo, [4, 8, 12, 16, 32], [256, 512, 768, 1024]) + ) + + return navitve_algorithms_config + + def _create_algorithms(self): + if self.world_size == 8: + self.algos_config = self._create_native_algorithms() + self._tune(5, 10, 20, self.algos_config) + elif self.world_size == 16 or self.world_size == 32: + self.dsl_algos_config = self._create_dsl_algorithms() + self._tune(5, 10, 20, self.dsl_algos_config) + + def _get_time( + self, + algo, + tune_tensor, + size, + nb, + nt, + n_warmup, + n_graph_launches, + n_ops_per_graph, + ): + # Check if the algorithm can run with the given configuration + if self._run_algo(algo, tune_tensor, size, nb, nt, True) != 0: + return float("inf") + + # Warmup iterations to stabilize performance + for _ in range(n_warmup): + self._run_algo(algo, tune_tensor, size, nb, nt, True) + + # Warmup on capture stream + capture_stream = torch.cuda.Stream() + capture_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(capture_stream): + self._run_algo(algo, tune_tensor, size, nb, nt, True) + capture_stream.synchronize() + + # Capture the algorithm execution in a CUDA graph + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, stream=capture_stream): + for _ in range(n_ops_per_graph): + self._run_algo(algo, tune_tensor, size, nb, nt, True) + + # Measure the execution time of the captured graph + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + start_event.record(capture_stream) + with torch.cuda.stream(capture_stream): + for _ in range(n_graph_launches): + g.replay() + end_event.record(capture_stream) + end_event.synchronize() + elapsed = start_event.elapsed_time(end_event) + + # Synchronize timing results across all ranks to ensure consistent algorithm selection + # replicate n times such due to algo limitations + time_tensor = torch.full( + (self.world_size,), elapsed, dtype=torch.float64, device="cuda" + ).to(dtype=torch.float32) + torch.cuda.current_stream().wait_stream(capture_stream) + if self.rank == 0: + avg_time = time_tensor[self.rank].item() / self.world_size + tensor = torch.tensor([avg_time]) + else: + tensor = torch.empty(1) + dist.broadcast(tensor, src=0, group=self.group) + avg_time = tensor.item() + + return avg_time + + def _tune(self, n_warmup, n_graph_launches, n_ops_per_graph, algos_config): + sizes = [1 << i for i in range(9, 24)] + dlpack = self.mscclpp.RawGpuBuffer(1 << 27).to_dlpack( + data_type=str(torch.float16) + ) + tune_tensor = torch.utils.dlpack.from_dlpack(dlpack) + + for size in sizes: + best_time = float("inf") + best_config = None + for i in range(len(algos_config)): + algo, candidates_nblocks, candidates_nthreads = algos_config[i] + if ( + size >= algo.message_size_range[0] + and size <= algo.message_size_range[1] + ): + for nb in candidates_nblocks: + for nt in candidates_nthreads: + avg_time = self._get_time( + algo, + tune_tensor, + size, + nb, + nt, + n_warmup, + n_graph_launches, + n_ops_per_graph, + ) + if avg_time < best_time: + best_time = avg_time + best_config = (algo, nb, nt) + if best_config: + self.best_configs[size] = best_config + + torch.cuda.synchronize() + for algo, _, _ in algos_config: + algo.reset() + + def _run_algo(self, algo, tensor, size, nblocks, nthreads, sym_mem_enabled=False): + return algo.execute( + comm=self.comm.communicator, + executor=self.executor, + input_buffer=tensor.data_ptr(), + output_buffer=tensor.data_ptr(), + input_size=size, + output_size=size, + dtype=self.dtype_to_mscclpp_dtype(tensor.dtype), + op=self.mscclpp.ReduceOp.SUM, + stream=torch.cuda.current_stream().cuda_stream, + nblocks=nblocks, + nthreads_per_block=nthreads, + symmetric_memory=sym_mem_enabled, + ) + def __init__( self, group: ProcessGroup, device: Union[int, str, torch.device], - max_bytes=_MAX_BYTES, ) -> None: - """ - Args: + """Args: group: the process group to work on. If None, it will use the default process group. device: the device to bind the CustomAllreduce to. If None, @@ -114,11 +253,16 @@ class PyMscclppCommunicator: self._IS_CAPTURING = False self.disabled = True - if not ops.IS_MSCCLPP_AR_AVAILABLE: - # disable because of missing mscclpp library - # e.g. in a non-cuda environment + try: + self.mscclpp = importlib.import_module("mscclpp") + self.mscclpp_ext = importlib.import_module("mscclpp.ext") + self.def_algo = importlib.import_module("mscclpp.default_algos") + except ImportError: + self.available = False + self.mscclpp = None return + self.available = True self.group = group assert ( @@ -161,137 +305,83 @@ class PyMscclppCommunicator: assert isinstance(device, torch.device) self.device = device - self.max_bytes = max_bytes self.rank = rank self.world_size = world_size - - if dist.get_rank(group) == 0: - unique_id = [ops.mscclpp_generate_unique_id()] - else: - unique_id = [None] - dist.broadcast_object_list(unique_id, src=self.ranks[0], group=self.group) - self.unique_id = unique_id[0] - self.rank_to_node, self.rank_to_ib = list(range(world_size)), list( - range(world_size) + self.comm = self.mscclpp.CommGroup( + torch_group=self.group, rank=rank, size=world_size ) - for r in range(world_size): - self.rank_to_node[r] = r // 8 - self.rank_to_ib[r] = self.rank % 8 + self.executor = self.mscclpp.Executor(self.comm.communicator) + self.symm_mem_enabled = self._is_symm_mem_enabled() + self.best_configs = {} + self._create_algorithms() - self._context = None - self.context_selection = None - self.msg_size_for_finetune = [ - 2**i for i in range(10, math.floor(math.log2(self.max_bytes)) + 1) - ] - self.msg_size2best_config = {} - if world_size == 8: - self.context_selection = MscclContextSelection.MSCCL1SHOT1NODELL - elif world_size == 16: - self.context_selection = MscclContextSelection.MSCCL1SHOT2NODELL - if not _is_hip: - self.scratch = torch.empty( - self.max_bytes * 8, - dtype=torch.uint8, - device=self.device, - ) - self.put_buffer = torch.empty( - self.max_bytes * 8 // self.nranks_per_node, - dtype=torch.uint8, - device=self.device, - ) - self._context = ops.mscclpp_init_context( - self.unique_id, - self.rank, - self.world_size, - self.scratch, - self.put_buffer, - self.nranks_per_node, - self.rank_to_node, - self.rank_to_ib, - int(self.context_selection), - ) - else: - raise NotImplementedError("HIP Mscclpp is not supported yet.") - - self.msg_size2best_config = {} - self.pre_tune_config() - if dist.get_rank(group) == 0: - msg_size2best_config = [self.msg_size2best_config] - else: - msg_size2best_config = [None] - dist.broadcast_object_list( - msg_size2best_config, src=self.ranks[0], group=self.group - ) - self.msg_size2best_config = msg_size2best_config[0] - - # PyMscclpp is enabled only in cuda graph - self.disabled = True - - def pre_tune_config(self, dtype=torch.bfloat16) -> bool: - logger.debug(f"start to pre-tune configs for rank {self.rank}") - nthreads_to_try = [256, 512, 1024] - nblocks_to_try = [21, 42, 84] - inp_randn = torch.ones( - self.msg_size_for_finetune[-1] // dtype.itemsize, dtype=dtype, device="cuda" - ) - oup_randn = torch.empty_like(inp_randn) - for msg_size in self.msg_size_for_finetune: - mock_inp, mock_outp = ( - inp_randn[: msg_size // dtype.itemsize], - oup_randn[: msg_size // dtype.itemsize], - ) - best_config, best_time = None, None - for nthreads in nthreads_to_try: - for nblocks in nblocks_to_try: - cur_cost = mscclpp_bench_time( - lambda: ops.mscclpp_allreduce( - self._context, mock_inp, mock_outp, nthreads, nblocks - ) - ) - if best_time is None or cur_cost < best_time: - best_config = (nthreads, nblocks) - best_time = cur_cost - self.msg_size2best_config[msg_size] = best_config - if self.rank == 0: - logger.debug( - f"for msg_size {msg_size}, best_config: {best_config}, best_time: {best_time}us" - ) + def destroy(self): + self.algos_config = None + self.best_configs = None + self.executor = None + self.scratch_buffer = None + self.flag_buffer = None + self.comm = None def should_mscclpp_allreduce( self, inp: torch.Tensor, op: ReduceOp = ReduceOp.SUM ) -> bool: - if self.disabled or self._context is None: + if ( + self.disabled + or self.world_size not in PyMscclppCommunicator._SUPPORTED_WORLD_SIZES + ): return False if inp.dtype not in PyMscclppCommunicator._SUPPORTED_DTYPE: return False - if not mscclpp_is_weak_contiguous(inp): + if not self._is_weak_contiguous(inp): return False - # only support sum op - if op != ReduceOp.SUM: + if op is not ReduceOp.SUM: return False - if inp.numel() * inp.element_size() > self.max_bytes: + if self._get_tuned_config(inp.numel() * inp.element_size()) is None: + return False + # mscclpp must not be used during any piecewise CUDA graph phase + # (compile, capture, or replay) as it changes the allreduce dispatch + # path and triggers recompilation. + if ( + is_in_piecewise_cuda_graph() + or is_in_pcg_torch_compile() + or get_pcg_capture_stream() is not None + ): return False return True - def all_reduce(self, tensor: torch.Tensor, op: ReduceOp = ReduceOp.SUM): - if self._IS_CAPTURING: - if torch.cuda.is_current_stream_capturing(): - self.graph_input_set.add((tensor.dtype, tensor.numel())) - msg_size = tensor.numel() * tensor.itemsize - index = bisect.bisect_left(self.msg_size_for_finetune, msg_size) - msg_size_finetune = self.msg_size_for_finetune[index] - nthreads, nblocks = self.msg_size2best_config[msg_size_finetune] - result = torch.empty_like(tensor) - ops.mscclpp_allreduce(self._context, tensor, result, nthreads, nblocks) - return result + def dtype_to_mscclpp_dtype(self, dtype: torch.dtype): + if dtype == torch.float16: + return self.mscclpp.DataType.float16 + elif dtype == torch.float32: + return self.mscclpp.DataType.float32 + elif dtype == torch.int32: + return self.mscclpp.DataType.int32 + elif dtype == torch.bfloat16: + return self.mscclpp.DataType.bfloat16 + else: + raise ValueError(f"Unknown data type: {dtype}") + + def all_reduce( + self, + tensor: torch.Tensor, + op: ReduceOp = ReduceOp.SUM, + stream: torch.cuda.Stream = None, + ): + assert op == torch.distributed.ReduceOp.SUM + nbytes = tensor.numel() * tensor.element_size() + algo, nblocks, nthreads = self._get_tuned_config(nbytes) + self._run_algo(algo, tensor, nbytes, nblocks, nthreads, self.symm_mem_enabled) + return tensor @contextmanager def change_state( self, enable: Optional[bool] = None, ): - if enable is None: + if enable is None or self.available is False: # guess a default value when not specified + # DO: Decided if raise an exception here or not enable = self.available old_disable = self.disabled diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 8489e9f3f..a4cc81ee2 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -605,7 +605,15 @@ class GroupCoordinator: if self.npu_communicator is not None and not self.npu_communicator.disabled: return self.npu_communicator.all_reduce(input_) - if self.pynccl_comm is not None and self.is_symmetric_memory_enabled(): + should_use_pymscclpp_allreduce = ( + self.pymscclpp_comm is not None + and self.pymscclpp_comm.should_mscclpp_allreduce(input_) + ) + if ( + self.pynccl_comm is not None + and self.is_symmetric_memory_enabled() + and not should_use_pymscclpp_allreduce + ): self.debug_check_symmetric_mempool(self, {"input": input_}, "all_reduce") with self.pynccl_comm.change_state(enable=True): self.pynccl_comm.all_reduce(input_) @@ -615,6 +623,7 @@ class GroupCoordinator: if ( self.ca_comm is not None and not self.ca_comm.disabled + and not should_use_pymscclpp_allreduce and self.ca_comm.should_custom_ar(input_) ): outplace_all_reduce_method = "ca" @@ -624,11 +633,7 @@ class GroupCoordinator: and self.qr_comm.should_quick_allreduce(input_) ): outplace_all_reduce_method = "qr" - elif ( - self.pymscclpp_comm is not None - and not self.pymscclpp_comm.disabled - and self.pymscclpp_comm.should_mscclpp_allreduce(input_) - ): + elif self.pymscclpp_comm is not None and should_use_pymscclpp_allreduce: outplace_all_reduce_method = "pymscclpp" elif ( self.torch_symm_mem_comm is not None @@ -1473,6 +1478,8 @@ class GroupCoordinator: self.cpu_group = None if self.pynccl_comm is not None: self.pynccl_comm = None + if self.pymscclpp_comm is not None: + self.pymscclpp_comm.destroy() if self.ca_comm is not None: self.ca_comm = None if self.mq_broadcaster is not None: diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index 1653d540d..1f376d535 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/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 -set(MSCCLPP_USE_CUDA ON) -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) +target_link_libraries(common_ops_sm90_build PRIVATE ${TORCH_LIBRARIES} c10 cuda cublas cublasLt) +target_link_libraries(common_ops_sm100_build PRIVATE ${TORCH_LIBRARIES} c10 cuda cublas cublasLt) # sparse flash attention target_compile_definitions(common_ops_sm90_build PRIVATE diff --git a/sgl-kernel/csrc/allreduce/mscclpp_allreduce.cu b/sgl-kernel/csrc/allreduce/mscclpp_allreduce.cu deleted file mode 100644 index d9cda22f0..000000000 --- a/sgl-kernel/csrc/allreduce/mscclpp_allreduce.cu +++ /dev/null @@ -1,140 +0,0 @@ -#include -#include -#include -#include - -#include "mscclpp_allreduce.cuh" - -enum MscclContextSelection { - MSCCL1NODELL = 1, - MSCCL2NODELL = 2, -}; - -class MscclContext { - public: - MscclContextSelection selection_; - std::shared_ptr msccl_1nodeLL_context; - std::shared_ptr msccl_2nodeLL_context; - MscclContext(MscclContextSelection selection) : selection_(selection) {} - template - 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(stream, input, output, input_numel, threads, block_limit); - } else if (selection_ == MSCCL2NODELL) { - msccl_2nodeLL_context->allreduce(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(unique_id.size())}, options); - std::memcpy(tensor.data_ptr(), 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(), 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& rank_to_node, - const std::vector& rank_to_ib, - const int64_t context_selection) { - MscclContext* context_ptr = new MscclContext(static_cast(context_selection)); - mscclpp::UniqueId uid = _tensor2unique_id(unique_id); - if (context_selection == MSCCL1NODELL) { - void* scratch_ptr = reinterpret_cast(scratch.data_ptr()); - const size_t scratch_bytes = scratch.numel() * scratch.element_size(); - context_ptr->msccl_1nodeLL_context = std::make_shared( - 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(scratch.data_ptr()); - const size_t scratch_bytes = scratch.numel() * scratch.element_size(); - void* put_buffer_ptr = reinterpret_cast(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( - 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(_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( - stream, - reinterpret_cast(inp.data_ptr()), - reinterpret_cast(out.data_ptr()), - inp.numel(), - nthreads, - nblocks); - break; - } - case at::ScalarType::Half: { - context->allreduce( - stream, - reinterpret_cast(inp.data_ptr()), - reinterpret_cast(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"); - } -} diff --git a/sgl-kernel/csrc/allreduce/mscclpp_allreduce.cuh b/sgl-kernel/csrc/allreduce/mscclpp_allreduce.cuh deleted file mode 100644 index ba0bc33fd..000000000 --- a/sgl-kernel/csrc/allreduce/mscclpp_allreduce.cuh +++ /dev/null @@ -1,779 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT license. -#pragma once -#ifdef USE_ROCM -#include -#else -#include -#include -#endif - -#include -#include -#include -#include -#include -#include -#include - -// 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 -__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 -__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 -__forceinline__ __device__ int4 add_vectors_helper(int4 a, int4 b) { - int4 ret; - ret.w = bit_cast(add_elements(bit_cast(a.w), bit_cast(b.w))); - ret.x = bit_cast(add_elements(bit_cast(a.x), bit_cast(b.x))); - ret.y = bit_cast(add_elements(bit_cast(a.y), bit_cast(b.y))); - ret.z = bit_cast(add_elements(bit_cast(a.z), bit_cast(b.z))); - return ret; -} - -template -__forceinline__ __device__ int4 add_vectors(int4 a, int4 b) { - return add_vectors_helper(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 -__forceinline__ __device__ uint2 add_vectors_helper(uint2 a, uint2 b) { - uint2 ret; - ret.x = bit_cast(add_elements(bit_cast(a.x), bit_cast(b.x))); - ret.y = bit_cast(add_elements(bit_cast(a.y), bit_cast(b.y))); - return ret; -} - -template -__forceinline__ __device__ uint2 add_vectors(uint2 a, uint2 b) { - return add_vectors_helper(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 -__forceinline__ __device__ int add_vectors_helper(int a, int b) { - return bit_cast(add_elements(bit_cast(a), bit_cast(b))); -} - -template -__forceinline__ __device__ int add_vectors(int a, int b) { - return add_vectors_helper(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 -__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(val, data); - } - data = add_vectors(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 -__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(val, data); - } - data = add_vectors(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(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 comm_; - const size_t rank_; - const size_t world_size_; - const std::vector rank_to_node_; - const std::vector rank_to_ib_; - MscclCommGroup( - mscclpp::UniqueId unique_id, - const size_t rank, - const size_t world_size, - const std::vector& rank_to_node, - const std::vector& 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(rank, world_size); - bootstrap->initialize(unique_id); - comm_ = std::make_shared(bootstrap); - } - template - 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>& same_node_connections, - std::unordered_map>& cross_node_connections) { - same_node_connections.clear(); - cross_node_connections.clear(); - std::unordered_map>> 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>& connections, - std::unordered_map>& semaphores, - std::unordered_map& registered_memories, - std::unordered_map& channels) { - channels.clear(); - make_semaphores(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 proxyService, - void* tensor_ptr, - const size_t tensor_bytes, - void* scratch_ptr, - const size_t scratch_bytes, - const std::unordered_map>& connections, - std::unordered_map>& semaphores, - std::unordered_map& registered_memories, - std::unordered_map& channels) { - channels.clear(); - make_semaphores(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 semaphore_ids; - std::unordered_map 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 - void make_semaphores( - const std::unordered_map>& connections, - std::unordered_map>& semaphores) { - semaphores.clear(); - for (const auto& [peer, conn] : connections) { - semaphores[peer] = std::make_shared(*comm_, conn); - } - comm_->setup(); - } - - void register_tensor_with_connections( - void* tensor_ptr, - size_t tensor_bytes, - const std::unordered_map>& connections, - std::unordered_map& 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> 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& old_memory_channels, - std::unordered_map& registered_sm_memories, - std::unordered_map>& memory_semaphores, - std::unordered_map& memory_channels, - mscclpp::GpuBuffer& 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 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 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( - device_memory_handle.data(), - memory_channel_handlers.data(), - memory_channel_handlers.size(), - stream, - cudaMemcpyHostToDevice); - } -}; - -class Msccl1NodeLLcontext { - private: - std::shared_ptr comm_group_ = nullptr; - void* scratch_; - const size_t scratch_bytes_; - std::unordered_map> same_node_connections_; - std::unordered_map> cross_node_connections_; - - std::unordered_map registered_sm_memories_; - std::unordered_map> memory_semaphores_; - std::unordered_map memory_channels_; - mscclpp::GpuBuffer d_memHandles_; - std::unordered_map> input_ptr2memory_channels_; - std::unordered_map> 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& rank_to_node, - const std::vector& 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(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 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 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( - d_memHandles_.data(), memory_channel_handlers.data(), memory_channel_handlers.size(), cudaMemcpyHostToDevice); - } - - ~Msccl1NodeLLcontext() { - CHECK_CUDA_SUCCESS(cudaStreamDestroy(h2d_stream)); - } - - template - 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 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(input); - if (input_ptr2d_memHandles_.find(input_void_ptr) == input_ptr2d_memHandles_.end()) { - std::unordered_map memory_channels; - mscclpp::GpuBuffer 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<<>>( - 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 comm_group_ = nullptr; - void* scratch_; - const size_t scratch_bytes_; - void* put_buffer_; - const size_t put_buffer_bytes_; - std::unordered_map> same_node_connections_; - std::unordered_map> cross_node_connections_; - - std::unordered_map registered_sm_memories_; - std::unordered_map registered_port_memories_; - - std::unordered_map> memory_semaphores_; - std::unordered_map> port_semaphores_; - - std::unordered_map memory_channels_; - std::unordered_map port_channels_; - - mscclpp::GpuBuffer d_memHandles_; - mscclpp::GpuBuffer d_portHandles_; - - std::shared_ptr proxyService; - cudaStream_t h2d_stream; - const size_t nranks_per_node_; - - std::unordered_map> input_ptr2memory_channels_; - std::unordered_map> 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& rank_to_node, - const std::vector& 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(unique_id, rank, world_size, rank_to_node, rank_to_ib); - proxyService = std::make_shared(); - 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 memory_channels_list; - std::vector 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 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( - d_memHandles_.data(), memory_channel_handlers.data(), memory_channel_handlers.size(), cudaMemcpyHostToDevice); - - std::vector 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( - d_portHandles_.data(), port_channel_handlers.data(), port_channel_handlers.size(), cudaMemcpyHostToDevice); - } - - ~Msccl2NodeLLcontext() { - CHECK_CUDA_SUCCESS(cudaStreamDestroy(h2d_stream)); - if (proxyService) { - proxyService->stopProxy(); - } - } - - template - 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 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(input); - if (input_ptr2d_memHandles_.find(input_void_ptr) == input_ptr2d_memHandles_.end()) { - std::unordered_map memory_channels; - mscclpp::GpuBuffer 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<<>>( - 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 diff --git a/sgl-kernel/csrc/allreduce/test_mscclpp_allreduce.cu b/sgl-kernel/csrc/allreduce/test_mscclpp_allreduce.cu deleted file mode 100644 index 4ca0c5739..000000000 --- a/sgl-kernel/csrc/allreduce/test_mscclpp_allreduce.cu +++ /dev/null @@ -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 -#include -#include -#include - -#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 - -#include "mscclpp_allreduce.cuh" - -template -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 host_buf(num_elems); - for (uint32_t i = 0; i < num_elems; ++i) { - host_buf[i] = T(i + rank); - } - thrust::device_vector device_buf(host_buf); - const size_t buf_size_in_bytes = num_elems * sizeof(T); - std::vector host_result_buf(num_elems); - thrust::device_vector device_result_buf(host_result_buf); - - std::vector host_scratch_buf(num_elems * 8); - for (uint32_t i = 0; i < num_elems; ++i) { - host_scratch_buf[i] = 1; - } - thrust::device_vector device_scratch_buf(host_scratch_buf); - std::vector host_put_buf(num_elems); - thrust::device_vector 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 rank_to_node(nranks); - std::vector 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( - 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( - 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( - 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( - 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 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; -} diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index b50687abc..8aff53976 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/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 */ diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index b16252007..61c39b6b3 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -89,20 +89,6 @@ std::tuple, std::vector> get_graph_buffer_ipc_meta void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs); void register_graph_buffers( fptr_t _fa, const std::vector>& handles, const std::vector>& 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& rank_to_node, - const std::vector& 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 /* diff --git a/sgl-kernel/python/sgl_kernel/allreduce.py b/sgl-kernel/python/sgl_kernel/allreduce.py index b02e8166d..17e9fe6c8 100644 --- a/sgl-kernel/python/sgl_kernel/allreduce.py +++ b/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 - ) diff --git a/sgl-kernel/tests/test_mscclpp.py b/sgl-kernel/tests/test_mscclpp.py deleted file mode 100644 index 0a4332bd3..000000000 --- a/sgl-kernel/tests/test_mscclpp.py +++ /dev/null @@ -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() diff --git a/test/manual/test_mscclpp.py b/test/manual/test_mscclpp.py deleted file mode 100644 index c30dbe00d..000000000 --- a/test/manual/test_mscclpp.py +++ /dev/null @@ -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()