본문으로 건너뛰기

[vLLM] Shared Experts: 공유 전문가 레이어 (DeepSeek-V2/V3)

들어가며

DeepSeek-V2/V3 모델은 MoE 아키텍처에 **공유 전문가(Shared Expert)**를 도입했다. 일반적인 MoE에서 각 토큰은 top-k 전문가에게만 라우팅되지만, 공유 전문가는 모든 토큰이 항상 거치는 "공통 전문가"다. 이는 모든 토큰이 공유하는 기본 지식을 별도의 전문가로 분리하여, 라우팅 전문가들이 더 특화된 역할을 할 수 있게 한다. vLLM은 이 패턴을 SharedFusedMoE와 통신 오버랩 최적화로 지원한다.

핵심 구조/코드 분석

SharedFusedMoE 클래스

SharedFusedMoE는 FusedMoE를 상속하여 공유 전문가 결과를 함께 반환한다:

class SharedFusedMoE(FusedMoE):
    """
    A FusedMoE operation that also computes the results of shared experts.
    If an all2all communicator is being used the shared expert computation
    can be interleaved with the fused all2all dispatch communication step.
    """

    def forward(
        self,
        hidden_states: torch.Tensor,
        router_logits: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        result = super().forward(
            hidden_states=hidden_states,
            router_logits=router_logits,
        )
        if self.shared_experts is None:
            return None, result
        else:
            return result

핵심은 super().forward()가 이미 공유 전문가를 처리한다는 것이다. FusedMoE 레이어 자체가 shared_experts 속성을 가지고 있어, 라우팅 전문가와 공유 전문가를 한 번에 실행할 수 있다.

FusedMoE의 shared_experts 지원

FusedMoE 레이어는 생성 시 shared_experts를 받아 내부적으로 관리한다:

# layer.py에서
from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
    SharedExperts,
)

SharedExperts 클래스는 공유 전문가의 forward 연산을 담당하며, MoE 커널의 shared_experts_input 파라미터를 통해 연동된다.

통신 오버랩: All2All과 공유 전문가 인터리빙

분산 환경에서 가장 큰 최적화는 all2all 통신과 공유 전문가 연산의 오버랩이다:

시간축 →
[dispatch all2all] [공유 전문가 GEMM] [라우팅 전문가 GEMM] [combine all2all]
   ↑ 통신              ↑ 계산             ↑ 계산              ↑ 통신

all2all dispatch로 토큰을 각 노드의 전문가에게 보내는 동안, 공유 전문가 연산을 동시에 수행한다. 통신 대기 시간을 계산으로 채우는 전략이다.

DeepEP High-Throughput 모드의 Prepare 단계에서 이 오버랩이 발생한다:

# deepep_ht.py에서
# dispatch 시작 (비동기)
# -> 공유 전문가 연산 수행 (dispatch 통신 중)
# -> dispatch 완료 대기
# -> 라우팅 전문가 연산

MoE 커널에서의 shared_experts_input

Modular 커널은 shared_experts_input을 파라미터로 받아 공유 전문가 결과를 최종 출력에 합산한다:

def apply(self, layer, x, topk_weights, topk_ids, shared_experts_input):
    return self.moe_kernel.apply(
        hidden_states=x,
        w1=layer.w13_weight,
        w2=layer.w2_weight,
        topk_weights=topk_weights,
        topk_ids=topk_ids,
        shared_experts_input=shared_experts_input,
    )

shared_experts_input이 None이 아니면, finalize 단계에서 라우팅 전문가 출력에 공유 전문가 출력을 더한다.

왜 이 설계인가

  1. 최소 침습 설계: SharedFusedMoE는 FusedMoE를 상속하되, forward의 반환값만 변경한다. 기존 MoE 인프라를 최대한 재활용하면서 공유 전문가를 지원한다.

  2. 통신 은닉: 전문가 병렬 처리에서 all2all은 주요 병목이다. 공유 전문가 연산을 통신과 오버랩하면 사실상 공유 전문가의 추가 비용을 "무료"로 만들 수 있다.

  3. TODO: 리팩토링 예정: 코드에 # TODO(bnell): Remove this entirely라는 주석이 있다. 현재 SharedFusedMoE는 과도기적 구현이며, 향후 FusedMoE 자체에 공유 전문가 로직이 완전히 통합될 예정이다. 이미 FusedMoE가 shared_experts 속성을 가지고 있어, 별도 서브클래스의 필요성이 줄어들고 있다.

참고 자료

댓글

관련 포스트

vLLM 의 다른글