[SGLang] NCCL & MSCCL++: 집합 통신 라이브러리 통합

2026년 4월 13일수정: 2026년 4월 13일

들어가며

SGLang의 분산 통신 스택에서 NCCL과 MSCCL++는 Custom AllReduce의 상위 계층, 즉 대규모 텐서와 CUDA Graph 환경에서의 기본 백엔드로 동작한다. pynccl.py는 NVIDIA NCCL을 ctypes로 직접 래핑하고, pymscclpp.py는 Microsoft의 MSCCL++ 라이브러리를 통합한다.

두 라이브러리 모두 torch.distributed를 우회하여 직접 NCCL/MSCCL++ C API를 호출한다. 이를 통해 CUDA Graph 캡처 중에도 통신이 가능하다.

구조도

GroupCoordinator
    │
    ├── pynccl_comm: PyNcclCommunicator
    │       │
    │       ├── NCCLLibrary (ctypes 래퍼)
    │       │     ├── ncclAllReduce
    │       │     ├── ncclAllGather
    │       │     ├── ncclReduceScatter
    │       │     ├── ncclSend / ncclRecv
    │       │     └── ncclBroadcast
    │       │
    │       └── ncclComm_t (커뮤니케이터 핸들)
    │
    └── pymscclpp_comm: PyMscclppCommunicator
            │
            ├── mscclpp_init_context()
            ├── mscclpp_allreduce()
            └── pre_tune_config() (자동 튜닝)

핵심 코드 분석: PyNcclCommunicator

초기화

class PyNcclCommunicator:
    def __init__(self, group, device, library_path=None):
        if not isinstance(group, StatelessProcessGroup):
            assert dist.get_backend(group) != dist.Backend.NCCL, \
                "PyNcclCommunicator should be attached to a non-NCCL group."
            self.rank = dist.get_rank(group)
            self.world_size = dist.get_world_size(group)

        self.nccl = NCCLLibrary(library_path)

        if self.rank == 0:
            self.unique_id = self.nccl.ncclGetUniqueId()
        else:
            self.unique_id = ncclUniqueId()

        # Unique ID를 broadcast로 공유
        tensor = torch.ByteTensor(list(self.unique_id.internal))
        dist.broadcast(tensor, src=ranks[0], group=group)

핵심 설계: Gloo(CPU) 그룹에 부착한다. torch.distributed의 NCCL 백엔드와 별도의 NCCL 커뮤니케이터를 생성함으로써, CUDA Graph 캡처 중에도 독립적으로 동작할 수 있다.

Warmup

초기화 직후 소규모 AllReduce로 워밍업한다.

with torch.cuda.device(device):
    self.comm = self.nccl.ncclCommInitRank(
        self.world_size, self.unique_id, self.rank)
    warmup_stream = torch.cuda.Stream()
    with torch.cuda.stream(warmup_stream):
        data = torch.zeros(1, device=device)
        self.all_reduce(data)
    warmup_stream.synchronize()
    del data

# 기본적으로 비활성 상태 (프로파일링, prefill 단계)
self.disabled = True

생성 후 disabled=True로 설정된다. CUDA Graph 캡처 시 change_state(enable=True) 컨텍스트 매니저로 활성화한다.

AllReduce (in-place)

def all_reduce(self, tensor, op=ReduceOp.SUM):
    if self.disabled:
        return
    assert tensor.device == self.device
    stream = self._resolve_stream()
    self.nccl.ncclAllReduce(
        buffer_type(tensor.data_ptr()),   # sendbuff
        buffer_type(tensor.data_ptr()),   # recvbuff (same = in-place)
        tensor.numel(),
        ncclDataTypeEnum.from_torch(tensor.dtype),
        ncclRedOpTypeEnum.from_torch(op),
        self.comm,
        cudaStream_t(stream.cuda_stream),
    )

sendbuff와 recvbuff가 같은 포인터이므로 in-place 연산이다.

AllReduce (out-of-place)

def outplace_all_reduce(self, in_tensor, out_tensor=None, op=ReduceOp.SUM):
    if out_tensor is None:
        out_tensor = torch.empty_like(in_tensor)
    stream = self._resolve_stream()
    self.nccl.ncclAllReduce(
        buffer_type(in_tensor.data_ptr()),   # sendbuff
        buffer_type(out_tensor.data_ptr()),  # recvbuff - DIFFERENT pointer
        in_tensor.numel(), ...)
    return out_tensor

CUDA Graph에서는 out-of-place가 필요하다. Graph 캡처 시 입력 텐서의 주소가 고정되어야 하므로, 결과를 별도 텐서에 기록한다.

Send/Recv (P2P 통신)

Pipeline Parallelism에서 사용되는 점대점 통신이다.

def send(self, tensor, dst):
    stream = self._resolve_stream()
    self.nccl.ncclSend(
        buffer_type(tensor.data_ptr()), tensor.numel(),
        ncclDataTypeEnum.from_torch(tensor.dtype),
        dst, self.comm, cudaStream_t(stream.cuda_stream))

def recv(self, tensor, src):
    stream = self._resolve_stream()
    self.nccl.ncclRecv(
        buffer_type(tensor.data_ptr()), tensor.numel(),
        ncclDataTypeEnum.from_torch(tensor.dtype),
        src, self.comm, cudaStream_t(stream.cuda_stream))

상태 전환

@contextmanager
def change_state(self, enable=None):
    if enable is None:
        enable = self.available
    old_disable = self.disabled
    self.disabled = not enable
    try:
        yield
    finally:
        self.disabled = old_disable

핵심 코드 분석: PyMscclppCommunicator

지원 범위

class PyMscclppCommunicator:
    _SUPPORTED_WORLD_SIZES = [8, 16]
    _MAX_BYTES = mscclpp_convert_to_bytes(os.getenv("SGLANG_MSCCLPP_MAX_BYTES", "1MB"))
    _SUPPORTED_DTYPE = [torch.float, torch.float16, torch.bfloat16]

MSCCL++는 8-GPU 또는 16-GPU 구성에서만 동작한다. 최대 크기 기본값은 1MB로, NCCL보다 빠른 영역이 매우 작다.

자동 튜닝

MSCCL++의 핵심 차별점은 nthreads x nblocks 조합을 자동 튜닝하는 것이다.

def pre_tune_config(self, dtype=torch.bfloat16):
    nthreads_to_try = [256, 512, 1024]
    nblocks_to_try = [21, 42, 84]
    for msg_size in self.msg_size_for_finetune:
        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

메시지 크기별로 9가지 조합(3 x 3)을 벤치마크하여 최적 구성을 찾는다. 이 결과는 rank 0에서 broadcast되어 모든 랭크가 동일한 설정을 사용한다.

AllReduce 실행

def all_reduce(self, tensor, op=ReduceOp.SUM):
    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

bisect로 가장 가까운 튜닝된 크기를 찾아 해당 설정으로 실행한다.

비교: NCCL vs MSCCL++

특성	PyNcclCommunicator	PyMscclppCommunicator
지원 world_size	제한 없음	8, 16만
최적 크기 범위	대규모	1MB 이하
자동 튜닝	없음	nthreads x nblocks 벤치마크
CUDA Graph	지원 (change_state)	지원 (change_state)
Eager 모드	기본 비활성	기본 비활성

설계 근거

왜 torch.distributed를 우회하는가? torch.distributed는 CUDA Graph 캡처 중 사용할 수 없다. NCCL C API를 직접 호출함으로써 Graph 내에서도 통신 커널을 녹화할 수 있다.

왜 Gloo 그룹에 부착하는가? 초기화 시 unique ID를 broadcast해야 하는데, 이미 NCCL 그룹이 있으면 충돌한다. Gloo는 CPU 기반이므로 독립적으로 조율 가능하다.

참고

NCCL 소스코드: python/sglang/srt/distributed/device_communicators/pynccl.py
MSCCL++ 소스코드: python/sglang/srt/distributed/device_communicators/pymscclpp.py
NVIDIA NCCL 공식 문서
Microsoft MSCCL++ GitHub

SGLang 의 다른글

이전글 [SGLang] Custom All-Reduce: NCCL 너머의 최적화된 집합 통신
현재글 : [SGLang] NCCL & MSCCL++: 집합 통신 라이브러리 통합
다음글 [SGLang] Data Parallel Controller: 다중 인스턴스 조율