[sglang] DeepSeek V3/R1 추론 최적화: DeepEP 공유 전문가(Shared Expert) 융합 기술 분석

PR 링크: sgl-project/sglang#20089 상태: Merged | 변경: +None / -None

들어가며

DeepSeek V3 및 R1과 같은 대규모 MoE(Mixture-of-Experts) 모델은 효율적인 추론을 위해 전문가 병렬화(Expert Parallelism, EP)를 사용합니다. 기존 SGLang 구현에서는 '공유 전문가(Shared Expert)'가 모든 랭크(Rank)에서 독립적으로 실행되어, DeepEP의 디스패치 및 연산 파이프라인과 분리되어 있었습니다. 이는 불필요한 연산 오버헤드를 발생시키고 파이프라인의 병목을 유발합니다. 본 PR은 공유 전문가를 마치 9번째 전문가처럼 취급하여 DeepEP의 디스패치 경로로 통합(Fuse)함으로써, 전체 추론 효율을 극대화하는 최적화를 다룹니다.

코드 분석

1. python/sglang/srt/layers/moe/fused_moe_triton/layer.py: 전문가 ID 매핑 최적화

DeepEP 환경에서는 각 랭크가 고유한 공유 전문가 슬롯을 가져야 합니다. 이를 위해 기존의 고정된 전문가 ID 매핑 로직을 확장하여, 공유 전문가가 각 랭크의 로컬 슬롯에 할당되도록 변경했습니다.

# Before
start_idx = self.moe_ep_rank * num_local_routed_experts
end_idx = (self.moe_ep_rank + 1) * num_local_routed_experts

# After
start_idx = self.moe_ep_rank * self._num_local_routed
end_idx = start_idx + self._num_local_routed
if start_idx <= expert_id < end_idx:
    return expert_id - start_idx
elif self._has_fused_shared and expert_id >= self._num_global_routed:
    return expert_id - self._num_global_routed + self._num_local_routed

2. python/sglang/srt/layers/moe/topk.py: TopK 출력 재구성

DeepEP 디스패치를 위해 TopK 선택 단계에서 공유 전문가를 포함하도록 출력을 패딩하고, ID를 재매핑하는 로직이 추가되었습니다.

# 공유 전문가 ID를 각 랭크의 고유 슬롯으로 재매핑
topk_ids[:, -num_fused_shared_experts:] = (
    ep_rank * num_local_experts
    + num_local_routed
    + torch.arange(num_fused_shared_experts, device=topk_ids.device)
)

왜 이게 좋은가

이번 최적화의 핵심은 **'연산 파이프라인의 단일화'**입니다. 기존에는 공유 전문가를 위해 별도의 Forward Pass를 수행해야 했으나, 이제는 DeepEP의 dispatch -> compute -> combine 파이프라인 안에서 공유 전문가가 처리됩니다.

1. 데이터 이동 최소화: 공유 전문가 연산이 로컬 랭크에서 MoE 연산과 함께 수행되므로, 별도의 동기화나 데이터 전송이 줄어듭니다.
2. 로드 밸런싱 준비: 이 PR은 1/2 단계로, 향후 'Waterfill' 알고리즘을 통해 공유 전문가를 가장 부하가 적은 랭크에 동적으로 할당할 수 있는 기반을 마련했습니다.

일반적인 교훈으로, 분산 시스템에서 특정 컴포넌트가 독립적으로 실행될 때 발생하는 오버헤드는 이를 메인 파이프라인에 '융합(Fuse)'함으로써 커널 실행 횟수를 줄이고 하드웨어 활용도를 높일 수 있음을 보여줍니다.

리뷰어 피드백 반영

리뷰 과정에서 ch-wan은 TBO/SBO와 같은 다른 최적화 기법과의 충돌 가능성을 지적했습니다. 이에 따라 is_deepep_class_backend()와 같은 유틸리티 함수를 통해 융합 로직이 안전하게 적용되도록 제어 로직을 강화했습니다. 또한, 불필요한 주석을 제거하고 모델 설정값(n_shared_experts)을 직접 참조하도록 수정하여 코드의 유지보수성을 높였습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.Tensor.expand.html
• https://docs.sglang.ai/developer_guide/contribution_guide.html

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