본문으로 건너뛰기

[sglang] [SGLang] MoE Prefill의 혁신: DWDP(Distributed Weight Data Parallelism) 도입 분석

PR 링크: sgl-project/sglang#29778 상태: Merged | 변경: +2147 / -24

들어가며

Mixture of Experts(MoE) 모델은 거대한 파라미터를 효율적으로 활용할 수 있게 해주지만, 분산 환경에서의 추론 성능은 종종 통신 병목에 의해 제한됩니다. 기존의 Expert Parallelism(EP) 방식은 각 토큰을 해당 전문가(Expert)가 있는 GPU로 보내는 All-to-All token dispatch 과정을 거칩니다. 이 과정은 대규모 클러스터에서 상당한 네트워크 동기화 오버헤드와 지연 시간을 발생시킵니다.

최근 SGLang 프로젝트에 추가된 DWDP(Distributed Weight Data Parallelism)는 이 패러다임을 완전히 뒤집습니다. 토큰을 전문가에게 보내는 대신, 전문가의 가중치를 토큰이 있는 GPU로 미리 가져오는(Prefetch) 방식을 택한 것입니다. 이 PR은 NVLink P2P 통신과 CUDA VMM(Virtual Memory Management)을 활용하여 MoE Prefill 단계의 성능을 비약적으로 향상시킨 DWDP 구현을 담고 있습니다.

코드 변경사항 분석

1. 통신 전략의 확장: communicator.py

가장 먼저, DWDP를 SGLang의 공식적인 분산 실행 모드로 통합하는 작업이 이루어졌습니다. LayerCommunicator는 이제 DWDP 활성화 여부를 판단하여 MoE 레이어의 연산 방식을 결정합니다.

[Before]

def _compute_mlp_mode(cls, context: _LayerModeComputationContext):
    if (
        not get_moe_a2a_backend().is_none()
        or should_use_flashinfer_cutlass_moe_fp4_allgather()
    ):
        return ScatterMode.SCATTERED

[After]

def _compute_mlp_mode(cls, context: _LayerModeComputationContext):
    if (
        not get_moe_a2a_backend().is_none()
        or should_use_flashinfer_cutlass_moe_fp4_allgather()
        or enable_dwdp() # DWDP 모드 지원 추가
    ):
        return ScatterMode.SCATTERED

enable_dwdp() 함수는 서버 인자(server_args)에서 dwdp_size가 1보다 큰지 확인하여 활성화 여부를 반환합니다. 이를 통해 DWDP가 활성화되면 토큰을 분산시키지 않고 로컬에서 모든 전문가 연산을 수행할 준비를 합니다.

2. DWDP의 핵심 오케스트레이터: dwdp_manager.py

DwdpManager 클래스는 DWDP의 생명주기를 관리하는 핵심 컴포넌트입니다. 모델의 MoE 레이어를 수집하고, 각 GPU가 소유한 전문가 가중치를 NVLink를 통해 공유할 수 있도록 설정합니다.

[Code Highlight: DWDP Setup]

# dwdp_manager.py 내 setup 메서드 일부
# 1. 전문가 레이아웃 계산
self.layout = DwdpExpertLayout(
    num_routed_experts=num_routed,
    dwdp_size=self.dwdp_size,
    dwdp_rank=self.dwdp_rank,
)

# 2. NVLink 기반 Transport 생성
transport = DWDPTransport.create(
    layer_weight_specs=layer_weight_specs,
    local_params=local_params,
    group=group,
    layout=self.layout,
    device_id=self.device_id,
)

# 3. 가중치 버퍼 및 VMM 주소 바인딩
weight_buffer = WeightBuffer.create(
    layer_weight_specs=layer_weight_specs,
    handles=transport.handle_set,
    # ...
)

여기서 주목할 점은 CUDA VMM의 활용입니다. 각 GPU는 자신의 가중치에 대한 shareable handle을 생성하고 이를 피어(peer) GPU들과 교환합니다. 이후 WeightBuffer는 이 핸들들을 사용하여 모든 전문가 가중치가 마치 하나의 연속된 메모리 공간에 있는 것처럼 보이는 Composite Virtual Address(VA)를 생성합니다. 이를 통해 커널 코드 수정 없이도 원격 가중치에 직접 접근할 수 있습니다.

3. 더블 버퍼링을 통한 오버랩 최적화

DWDP의 진정한 강력함은 연산과 통신의 오버랩에서 나옵니다. DwdpWeightManager는 현재 레이어($l$)의 MoE 연산을 수행하는 동안 다음 레이어($l+1$)의 가중치를 미리 가져오는 Double-buffered prefetch를 수행합니다.

[Code Highlight: Prefetch Logic]

def prefetch_first_layers(self) -> None:
    if self._weight_manager is not None:
        self._weight_manager.prefetch_first_layers()

def record_compute_and_prefetch_next(self, layer_idx: int) -> None:
    if self._weight_manager is not None:
        self._weight_manager.record_compute_and_prefetch_next(layer_idx)

이 구조 덕분에 GPU는 원격 가중치가 도착할 때까지 기다리지 않고(Zero sync), NVLink의 높은 대역폭을 활용하여 백그라운드에서 데이터를 채웁니다.

4. 리소스 정리: parallel_state.py

분산 환경에서 리소스 누수는 치명적입니다. PR은 모델 파라미터 파괴 시 DWDP 매니저도 함께 정리되도록 보장합니다.

[Before]

def destroy_model_parallel():
    global _TP
    if _TP:
        _TP.destroy()

[After]

def destroy_model_parallel():
    dwdp_mgr = get_global_dwdp_manager()
    if dwdp_mgr is not None:
        dwdp_mgr.cleanup() # VMM 핸들 및 메모리 해제
        set_global_dwdp_manager(None)

    global _TP
    # ... 기존 로직

왜 이게 좋은 최적화인가?

1. 압도적인 성능 향상

벤치마크 결과에 따르면, 4개의 B200 GPU에서 gpt-oss-120b 모델을 테스트했을 때 DWDP는 기존 DEP(EP + All-Gather) 방식 대비 최대 1.92배(ISL=32K)의 성능 향상을 보여주었습니다. 특히 시퀀스 길이가 길어질수록(ISL이 클수록) All-to-All 통신 비용이 급증하는데, DWDP는 이를 가중치 프리페칭으로 대체하여 병목을 효과적으로 제거했습니다.

2. Zero Synchronization

기존 EP 방식은 모든 GPU가 토큰을 주고받기 위해 동기화 지점(Barrier)을 가져야 했습니다. DWDP는 각 GPU가 독립적으로 자신의 토큰을 처리하며 필요한 가중치만 NVLink로 당겨오기 때문에, 랭크 간 동기화가 거의 필요 없습니다. 이는 SCHEDULER_SKIP_ALL_GATHER와 같은 최적화를 가능하게 합니다.

최신 GPU 아키텍처(H100, B200)는 매우 높은 NVLink 대역폭을 제공합니다. DWDP는 이 하드웨어 특성을 극대화하여, 네트워크 카드(NIC)를 통한 이더넷/인피니밴드 통신보다 훨씬 빠른 P2P 메모리 복사를 수행합니다.

리뷰어 피드백 및 논의

리뷰 과정에서 silencejade는 DWDP 구현이 CUDA 전용 API(VMM 등)를 사용하기 때문에 비 CUDA 장치(예: ROCm, XPU)에서 실행이 깨질 수 있음을 지적했습니다.

이에 대해 OrangeRedeng은 해당 모듈을 DWDP가 활성화된 경우에만 지연 임포트(Lazy Import)하거나 별도의 PR(#29778)을 통해 호환성을 확보하는 방향으로 논의를 진행했습니다. 이는 고성능 최적화 기법을 도입할 때 프레임워크의 범용성을 유지하는 것이 얼마나 중요한지를 보여주는 좋은 사례입니다.

결론

이번 DWDP 도입은 MoE 모델 서빙의 패러다임을 'Data-to-Compute'에서 'Weight-to-Compute'로 전환한 중요한 이정표입니다. 특히 Prefill 단계에서 발생하는 대규모 통신 지연을 하드웨어 친화적인 방식으로 해결했다는 점에서 시니어 엔지니어들에게 시사하는 바가 큽니다. 대규모 MoE 모델을 운영하는 팀이라면 SGLang의 DWDP 설정을 적극적으로 검토해 보시기 바랍니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글