본문으로 건너뛰기

[sglang] NVLink를 활용한 sglang의 DSpark 헤드 최적화: 더 빠른 토큰 생성

PR 링크: sgl-project/sglang#39414 상태: Merged | 변경: +1992 / -10

들어가며

최근 sglang 레포지토리의 Pull Request(PR) #36720은 DeepSeek-V4.1 모델의 DSpark draft head에서 발생하는 성능 병목 현상을 해결하는 데 중점을 둡니다. 특히, 각 토큰 생성 단계마다 발생하는 보캡(vocab) 로짓(logits)의 all-gather 연산을 최적화하여 전체적인 추론 속도를 향상시키는 것을 목표로 합니다.

기존 방식에서는 각 GPU(rank)가 129KB 크기의 로짓 샤드(shard)를 NCCL을 통해 all-gather하는 데 약 13-18 마이크로초(µs)가 소요되었습니다. 이는 한 번의 검증 주기(verify cycle)마다 5번씩 발생하여 약 75µs의 오버헤드를 유발했으며, 샘플링된 토큰의 브로드캐스트에도 추가적인 시간이 소요되었습니다. 이 PR은 NVIDIA의 NVLink 기술을 활용한 집단 통신(collectives)을 도입하여 이 all-gather 연산의 속도를 획기적으로 개선하고자 합니다. NVLink의 멀티캐스트(multicast) 기능을 사용하면 동일한 데이터 전송을 약 3.5µs로 단축할 수 있습니다.

본 글에서는 이 PR의 주요 변경 사항을 살펴보고, NVLink 집단 통신이 어떻게 성능 향상을 가져오는지, 그리고 이러한 최적화가 갖는 일반적인 교훈은 무엇인지 코드 diff와 리뷰 내용을 바탕으로 심층적으로 분석해 보겠습니다.

코드 분석

이 PR은 크게 세 가지 주요 영역에서 변경이 이루어졌습니다.

PR의 초기 커밋들은 NVLink 기반의 집단 통신 연산들을 구현하고 통합하는 데 집중했습니다. 이는 nvlink_comm.cuh, ops/communication/nvlink_comm.py 등의 파일을 통해 이루어졌습니다.

주요 변경 사항은 다음과 같습니다:

  • CustomAllReduceV2: NVLink의 멀티캐스트 기능을 활용하는 새로운 통신 방식을 지원합니다. 이는 all_reduce, all_gather, reduce_scatter와 같은 기본적인 집단 통신 연산을 더 효율적으로 수행할 수 있게 합니다.
  • 벤치마크: NVLink 통신 성능을 측정하기 위한 새로운 벤치마크(benchmark/bench_nvlink_comm/bench_nvlink_comm.py)가 추가되었습니다. 이를 통해 다양한 시나리오에서 NVLink와 기존 NCCL 방식의 성능을 비교할 수 있습니다.
--- /dev/null
+++ b/benchmark/bench_nvlink_comm/bench_nvlink_comm.py
@@ -0,0 +1,256 @@
+from __future__ import annotations
+... (중략) ...
+@marker.parametrize("residual", [False, True])
+@marker.parametrize("op", OPS)
+@marker.parametrize("tokens", [2**n for n in range(15)])
+@marker.benchmark("provider", PROVIDERS)
+def benchmark(op: str, tokens: int, residual: bool, provider: str):
+    cpu_group, gpu_group, pynccl_coord, device = _init_groups()
+    world_size = dist.get_world_size(cpu_group)
+    nvlink, v2 = _init_comms()
+    ...
+    if provider.startswith("nvlink-push"):
+        # Only the gather spans the plane; all-reduce and the scatter each put a
+        # sender's whole contribution into a single slot.
+        cap = nvlink.max_push_size * (world_size if op == "all_gather" else 1)
+        if out_bytes > cap:
+            marker.skip(f"{out_bytes} B exceeds the {cap} B push plane")
+    ...
+    return marker.do_bench(
+        fn,
+        ...
+    )
+```

위 코드는 NVLink 기반의 `all_gather` 연산이 기존 NCCL 방식보다 훨씬 빠른 속도를 보이는 것을 보여줍니다. 특히 `nvlink-push` 프로바이더는 NVLink의 푸시(push) 기능을 활용하여 통신 대역폭을 극대화합니다.

### 2. `VocabGather` 인터페이스

`srt/distributed/device_communicators/vocab_gather.py` 파일에는 `VocabGather`라는 새로운 인터페이스가 도입되었습니다. 이 인터페이스는 TP(Tensor Parallelism) 환경에서 보캡 로우(row)를 효율적으로 gather하는 역할을 수행합니다.

*   **다양한 전송 방식 지원**: `VocabGather`는 초기화 시점에 최적의 전송 방식을 자동으로 선택합니다. NVLink가 지원되는 환경에서는 `NVLinkVocabGather`를 사용하고, 그렇지 않은 경우에는 기존의 `NcclVocabGather`를 사용합니다. 이를 통해 호출자는 복잡한 분기 로직 없이 일관된 인터페이스를 사용할 수 있습니다.
*   **NVLink 최적화**: `NVLinkVocabGather`는 NVLink의 푸시(push) 및 풀(pull) 커널을 활용합니다. 데이터 크기가 푸시 슬롯(push slot)에 맞으면 푸시 커널을 사용하고, 그렇지 않으면 풀 커널을 사용하여 데이터를 효율적으로 전송합니다. 또한, NCCL 기반의 전송도 지원하여 유연성을 높였습니다.

```python
# srt/distributed/device_communicators/vocab_gather.py (개념적 예시)
class VocabGather:
    def __init__(self, group_coordinator):
        if group_coordinator.has_nvlink_multicast:
            self.gather = NVLinkVocabGather(group_coordinator)
        else:
            self.gather = NcclVocabGather(group_coordinator)

    def __call__(self, data):
        return self.gather(data)

이 인터페이스는 보캡 로우를 [rows, local] -> [rows, world_size * local] 형태로 효율적으로 재구성하며, NVLink의 멀티캐스트 기능을 활용하여 통신 오버헤드를 크게 줄입니다.

3. DSpark 통합

DSparkV4MarkovHead 클래스에서 configure_tp_shard 메서드를 통해 VocabGather를 빌드하고, _apply_step_logits_sharded 메서드에서 이를 호출하여 실제 로직에 통합합니다. SGLANG_DSPARK_NVLINK_VOCAB_GATHER 환경 변수(기본값 활성화)를 통해 NVLink 기반 보캡 개더의 사용 여부를 제어할 수 있습니다.

--- a/python/sglang/srt/distributed/device_communicators/vocab_gather.py
+++ b/python/sglang/srt/distributed/device_communicators/vocab_gather.py
@@ -10,7 +10,7 @@
     NVLINK_COMM_PULL_KERNEL_MAX_SIZE = 128 << 20  # 128 MiB
 
 
-def make_vocab_gather(group_coordinator: GroupCoordinator):
+def make_vocab_gather(group_coordinator: GroupCoordinator) -> VocabGather:
     if group_coordinator.world_size == 1:
         return LocalVocabGather()
 
@@ -22,7 +22,7 @@
         return NcclVocabGather(group_coordinator)
 
 
-class VocabGather:
+class VocabGather(torch.nn.Module):
     def __init__(self, group_coordinator: GroupCoordinator):
         self.gather = make_vocab_gather(group_coordinator)
 
@@ -30,7 +30,7 @@
         return self.gather(data)
 
 
-class LocalVocabGather:
+class LocalVocabGather(VocabGather):
     def __init__(self):
         pass
 
@@ -40,7 +40,7 @@
         return data
 
 
-class NVLinkVocabGather:
+class NVLinkVocabGather(VocabGather):
     def __init__(self, group_coordinator: GroupCoordinator):
         self.group_coordinator = group_coordinator
         self.device = group_coordinator.device
@@ -107,7 +107,7 @@
         return stacked_by_rank
 
 
-class NcclVocabGather:
+class NcclVocabGather(VocabGather):
     def __init__(self, group_coordinator: GroupCoordinator):
         self.group_coordinator = group_coordinator
         self.device = group_coordinator.device

이 변경은 DSpark 모델의 핵심 연산에 NVLink 기반의 고속 통신을 적용하여, 특히 대규모 언어 모델에서 빈번하게 발생하는 보캡 로짓 통신을 최적화합니다.

왜 이게 좋은가?

성능 향상

PR에서 제공된 성능 측정 결과는 NVLink 도입의 효과를 명확하게 보여줍니다.

4x B200, TP4 환경에서 단일 토큰 생성 단계의 로짓 샤드([rows, 32320] fp32) 전송 성능 비교:

rows path NVLink NCCL all_gather(dim=-1)
1 push 3.6 µs 16.5 µs
6 push 10.0 µs 27.4 µs
32 pull 41.9 µs 70.9 µs
128 pull 142.4 µs 154.8 µs

위 표에서 볼 수 있듯이, NVLink는 특히 작은 크기(rows=1, 6)의 데이터 전송에서 NCCL 대비 4배 이상의 성능 향상을 보입니다. 큰 크기에서도 상당한 속도 개선이 이루어졌습니다. 이는 각 토큰 생성 단계에서 발생하는 수많은 all-gather 연산의 누적 시간을 크게 단축시켜, 전체 추론 처리량(throughput)을 향상시킵니다.

리뷰어의 추적 결과에 따르면, nvlink_push_kernel은 기존 NCCL AllGather 대비 약 37.7µs 대 72.8µs로, 약 35µs의 시간 단축 효과를 보였습니다. 이는 전체 검증 주기(verify cycle) 시간을 6404µs에서 6453µs로 줄이는 데 기여했습니다. 비록 전체 시간에서 차지하는 비율이 크지 않아 보일 수 있지만, 초당 수십, 수백 개의 토큰을 생성하는 LLM 환경에서는 이러한 미세한 시간 단축이 모여 상당한 성능 향상으로 이어집니다.

일반적인 교훈

  1. 하드웨어 특화 최적화의 중요성: NVLink와 같은 고속 인터커넥트는 GPU 간 통신에서 상당한 성능 이점을 제공합니다. 모델의 핵심 연산이 통신에 병목이 있다면, 하드웨어의 특성을 최대한 활용하는 최적화가 필수적입니다.
  2. 유연한 통신 추상화: VocabGather와 같이 다양한 통신 백엔드(NVLink, NCCL)를 지원하고 자동으로 최적의 방식을 선택하는 추상화 계층은 코드의 재사용성과 유지보수성을 높입니다. 개발자는 하드웨어별 복잡성을 신경 쓰지 않고 고수준의 로직에 집중할 수 있습니다.
  3. 정량적 성능 측정 및 검증: 새로운 최적화 기법을 도입할 때는 반드시 벤치마크와 실제 사용 사례(end-to-end test)를 통해 성능 향상을 정량적으로 입증해야 합니다. 리뷰어의 피드백은 이러한 검증 과정을 더욱 강화합니다.
  4. 점진적 통합: PR의 첫 6개 커밋이 NVLink 집단 통신의 기반을 마련하고, 이후 커밋에서 이를 VocabGather 및 DSpark 헤드에 통합하는 방식은 복잡한 변경 사항을 관리하고 검증하는 데 효과적인 전략입니다.

결론

이 PR은 sglang의 DSpark 헤드에서 NVLink 집단 통신을 활용하여 보캡 로짓의 all-gather 연산을 최적화함으로써, DeepSeek-V4.1 모델의 토큰 생성 속도를 크게 향상시켰습니다. VocabGather라는 유연한 인터페이스와 NVLink의 고성능 통신 기능을 결합하여, LLM 추론 성능을 한 단계 끌어올리는 중요한 개선을 이루었습니다. 이는 향후 고성능 분산 딥러닝 시스템 설계에 있어 하드웨어 특화 최적화와 효율적인 통신 추상화의 중요성을 다시 한번 강조하는 사례라 할 수 있습니다.

참고 자료

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.

댓글

관련 포스트

PR Analysis 의 다른글