본문으로 건너뛰기

[vLLM] MoE 라우팅 전략: 7종 라우팅 알고리즘 분석

들어가며

Mixture of Experts(MoE) 모델에서 라우팅은 각 토큰이 어떤 전문가에게 할당될지를 결정하는 핵심 메커니즘이다. vLLM은 다양한 MoE 아키텍처(Mixtral, DeepSeek-V2/V3, Llama-4 등)를 지원하기 위해 7종의 라우터 구현체를 제공한다. 이번 글에서는 router/ 디렉토리의 전체 구조를 분석한다.

핵심 구조/코드 분석

BaseRouter: 템플릿 메서드 패턴

모든 라우터는 BaseRouter를 상속하며, select_experts()가 공통 파이프라인을 정의한다:

class BaseRouter(FusedMoERouter):
    def select_experts(self, hidden_states, router_logits):
        # Step 1: EPLB 상태 검증
        self._validate_eplb_state()
        # Step 2: 인덱스 dtype 획득
        indices_type = self._get_indices_type()
        # Step 3: 라우팅 계산 (서브클래스가 구현)
        topk_weights, topk_ids = self._compute_routing(
            hidden_states, router_logits, indices_type)
        # Step 4: EPLB 매핑 (논리적 -> 물리적 전문가)
        topk_ids = self._apply_eplb_mapping(topk_ids)
        # Step 5: dtype 변환
        topk_ids = self._convert_indices_dtype(topk_ids, indices_type)
        return topk_weights, topk_ids

EPLB(Expert Parallelism Load Balancing)는 전문가 복제본 간 부하 분산을 위한 매핑을 수행한다. Triton 커널로 구현되어 논리적 전문가 ID를 물리적 ID로 변환하면서 동시에 부하 통계를 기록한다.

1. FusedTopKRouter: 기본 라우터

가장 기본적인 top-k 라우팅으로, softmax 또는 sigmoid 스코어링을 지원한다:

class FusedTopKRouter(BaseRouter):
    def _compute_routing(self, hidden_states, router_logits, indices_type):
        topk_weights, topk_ids, _ = fused_topk(
            hidden_states=hidden_states,
            gating_output=router_logits,
            topk=self.top_k,
            renormalize=self.renormalize,
            scoring_func=self.scoring_func,
        )
        return topk_weights, topk_ids

ops.topk_softmax는 CUDA C++ 커널로 구현되어, PyTorch의 torch.topk + torch.softmax 조합보다 빠르다.

2. GroupedTopKRouter: DeepSeek-V2/V3 라우터

DeepSeek 모델의 그룹 기반 라우팅을 구현한다. 전문가를 그룹으로 나누고, 먼저 상위 그룹을 선택한 뒤 그 안에서 전문가를 선택하는 2단계 방식이다:

def grouped_topk(hidden_states, gating_output, topk, renormalize,
                 num_expert_group, topk_group, scoring_func,
                 routed_scaling_factor, e_score_correction_bias):
    if scoring_func == "softmax":
        scores = torch.softmax(gating_output, dim=-1)
    elif scoring_func == "sigmoid":
        scores = gating_output.sigmoid()

    if e_score_correction_bias is not None:
        original_scores = scores
        scores = scores + e_score_correction_bias.unsqueeze(0)
        group_scores = scores.view(num_token, num_expert_group, -1) \
            .topk(2, dim=-1)[0].sum(dim=-1)
    # 상위 그룹 선택
    group_idx = torch.topk(group_scores, k=topk_group, dim=-1)[1]
    # 마스킹 후 최종 전문가 선택
    topk_ids = torch.topk(tmp_scores, k=topk, dim=-1)[1]
    # 바이어스 보정 전 원본 스코어로 가중치 복원
    topk_weights = original_scores.gather(1, topk_ids)

e_score_correction_bias는 전문가 선택 시 바이어스를 추가하되, 최종 가중치는 원본 스코어를 사용한다는 점이 중요하다.

3. FusedTopKBiasRouter: 바이어스 보정 라우터

그룹핑 없이 바이어스만 적용하는 간소화 버전이다:

class FusedTopKBiasRouter(BaseRouter):
    def _compute_routing(self, hidden_states, router_logits, indices_type):
        topk_weights, topk_ids = fused_topk_bias(
            hidden_states=hidden_states,
            gating_output=router_logits,
            e_score_correction_bias=self.e_score_correction_bias.data,
            topk=self.top_k,
            renormalize=self.renormalize,
        )

4. CustomRoutingRouter: Llama-4 등

모델별 커스텀 라우팅 함수를 주입할 수 있다:

class CustomRoutingRouter(BaseRouter):
    @property
    def routing_method_type(self):
        from vllm.model_executor.models.llama4 import Llama4MoE
        if self.custom_routing_function == Llama4MoE.custom_routing_function:
            return RoutingMethodType.Llama4
        return RoutingMethodType.Custom

Llama-4의 경우 FlashInfer+TRT-LLM 백엔드가 해당 라우팅을 직접 지원한다.

5. RoutingSimulatorRouter: 테스트용

균일 분포, 정규 분포 등으로 랜덤 라우팅을 시뮬레이션하는 테스트/벤치마크용 라우터다:

class DistributionBasedRouting(RoutingStrategy):
    def route_tokens(self, hidden_states, router_logits, top_k, indices_type):
        topk_ids = self._sample_expert_ids(num_tokens, num_experts, ...)
        topk_weights = self._generate_weights(num_tokens, top_k, ...)
        return topk_weights, topk_ids

라우터 팩토리

create_fused_moe_router가 설정에 따라 적절한 라우터를 생성한다:

def create_fused_moe_router(...) -> FusedMoERouter:
    # 우선순위: Simulator > GroupedTopK > Custom > Bias > FusedTopK
    if routing_strategy != "":
        return RoutingSimulatorRouter(...)
    if use_grouped_topk:
        return GroupedTopKRouter(...)
    if custom_routing_function is not None:
        return CustomRoutingRouter(...)
    if e_score_correction_bias is not None:
        return FusedTopKBiasRouter(...)
    return FusedTopKRouter(...)

왜 이 설계인가

  1. 템플릿 메서드 패턴: 공통 로직(EPLB 매핑, dtype 변환)은 BaseRouter에서 처리하고, 라우팅 알고리즘만 서브클래스에서 구현한다. 새로운 라우팅 방식 추가 시 _compute_routing만 구현하면 된다.

  2. EPLB 통합: 전문가 병렬 처리에서의 부하 분산이 라우터 레벨에서 투명하게 처리된다. Triton 커널로 매핑과 통계 수집을 퓨전하여 오버헤드를 최소화한다.

  3. 배치 불변성: envs.VLLM_BATCH_INVARIANT 설정으로 sorted=True topk를 사용하여, 배치 크기와 무관하게 동일한 결과를 보장할 수 있다.

참고 자료

댓글

관련 포스트

vLLM 의 다른글