[vLLM] Fused MoE: 라우팅+전문가 연산 융합

들어가며

Mixture of Experts(MoE)는 모델의 총 파라미터 수를 늘리면서도, 각 토큰에 대해 소수의 전문가만 활성화하여 연산량을 억제하는 아키텍처이다. Mixtral 8x7B는 47B 파라미터를 가지지만 토큰당 12B만 활성화하고, DeepSeek-V3/R1은 671B 중 37B만 사용한다. 그러나 MoE의 "라우팅 -> 전문가 선택 -> 분산 연산"이라는 다단계 파이프라인은 GPU에서 비효율적이다. vLLM의 Fused MoE는 이 전체 과정을 하나의 융합 커널로 처리한다.

• 논문:
  • DeepSeek-V3 Technical Report
  • DeepSeekMoE: Towards Ultimate Expert Specialization
• 공식 문서: https://docs.vllm.ai

공식 문서

vLLM 공식 문서: Fused MoE Modular Kernel vLLM 공식 문서: MoE Kernel Features

핵심 구조/코드 분석

FusedMoE 레이어: 진입점

vllm/model_executor/layers/fused_moe/layer.py에 정의된 FusedMoE는 vLLM에서 가장 복잡한 레이어 중 하나이다:

@CustomOp.register("fused_moe")
class FusedMoE(CustomOp):
    """FusedMoE layer for MoE models.
    This layer contains both MergedColumnParallel weights (gate_up_proj /
    w13) and RowParallelLinear weights (down_proj / w2).

    Args:
        num_experts: Number of experts in the model
        top_k: Number of experts selected for each token
        hidden_size: Input hidden state size of the transformer
        intermediate_size: Intermediate size of the experts
    """
    def __init__(
        self,
        num_experts: int,
        top_k: int,
        hidden_size: int,
        intermediate_size: int,
        ...
        enable_eplb: bool = False,
        num_redundant_experts: int = 0,
        n_shared_experts: int | None = None,
    ):
        self.global_num_experts = num_experts + num_redundant_experts
        self.logical_num_experts = num_experts

핵심 설계는 gate_up_proj(w13)과 down_proj(w2) 두 가중치 세트를 하나의 레이어에서 관리하는 것이다. num_redundant_experts는 EPLB(Expert Parallelism Load Balancer)를 위한 여분의 전문가 수이다.

모듈러 아키텍처: 라우터와 러너 분리

FusedMoE의 내부는 라우터(토큰을 전문가에 배정)와 러너(실제 GEMM 수행)로 분리되어 있다:

from vllm.model_executor.layers.fused_moe.router.router_factory import (
    create_fused_moe_router,
)
from vllm.model_executor.layers.fused_moe.runner.moe_runner_factory import (
    create_moe_runner,
)

이 팩토리 패턴 덕분에 라우팅 방식(TopK, Grouped TopK 등)과 실행 백엔드(Triton, CUTLASS, DeepEP 등)를 독립적으로 조합할 수 있다.

Expert Parallelism과 배치 전략

Expert Parallel(EP)에서 전문가를 GPU에 분배하는 determine_expert_map 함수가 핵심이다:

def determine_expert_map(
    ep_size: int,
    ep_rank: int,
    global_num_experts: int,
    expert_placement_strategy: ExpertPlacementStrategy = "linear",
    ...
) -> tuple[int, torch.Tensor | None, torch.Tensor | None]:
    base_experts = global_num_experts // ep_size
    remainder = global_num_experts % ep_size
    local_num_experts = base_experts + 1 if ep_rank < remainder else base_experts

    expert_map = torch.full((global_num_experts,), -1, dtype=torch.int32)
    if expert_placement_strategy == "linear":
        start_idx = ep_rank * base_experts + min(ep_rank, remainder)
        expert_map[start_idx:start_idx + local_num_experts] = torch.arange(
            0, local_num_experts, dtype=torch.int32
        )
    elif expert_placement_strategy == "round_robin":
        local_log_experts = torch.arange(ep_rank, global_num_experts, ep_size)
        expert_map[local_log_experts] = torch.arange(0, local_num_experts)

두 가지 배치 전략이 있다:

• Linear: 연속된 전문가를 같은 GPU에 배치 (기본값)
• Round Robin: 전문가를 순환적으로 분산 배치

Round Robin은 전문가 그룹이 여러 개이고 EPLB가 비활성화된 경우에만 사용할 수 있다.

양자화 통합

FusedMoE는 FusedMoeWeightScaleSupported 열거형으로 지원하는 양자화 스케일 유형을 정의한다:

class FusedMoeWeightScaleSupported(Enum):
    TENSOR = "tensor"
    CHANNEL = "channel"
    GROUP = "group"
    BLOCK = "block"

GPTQ, AWQ, FP8 등의 양자화 기법이 각각 get_quant_method에서 FusedMoE를 감지하고 적절한 MoE 양자화 메서드를 반환한다. 이 설계 덕분에 양자화와 MoE가 직교적으로 조합된다.

왜 이 설계인가

1. 커널 융합의 필요성: MoE를 순진하게 구현하면 라우팅, 토큰 셔플링, 전문가별 GEMM, 결과 합산이 각각 별도의 커널 런치를 필요로 한다. 이 과정에서 커널 런치 오버헤드와 메모리 왕복이 성능을 심각하게 저하시킨다. 융합 커널은 이를 단일 패스로 처리한다.

2. Shared Experts: DeepSeek 계열 모델은 라우팅된 전문가 외에 모든 토큰이 통과하는 "공유 전문가(shared expert)"를 사용한다. n_shared_experts와 SharedExperts 클래스가 이를 지원하며, ROCm AITER에서는 공유 전문가까지 융합한다.

3. CustomOp 등록: @CustomOp.register("fused_moe")는 torch.compile과의 호환성을 위한 것이다. vLLM의 컴파일 파이프라인이 이 레이어를 인식하고 최적화할 수 있게 한다.

4. 유연한 병렬화: TP(텐서 병렬), EP(전문가 병렬), DP(데이터 병렬), PCP(파이프라인 청크 병렬)를 모두 지원한다. FusedMoEParallelConfig.make()가 이 복잡한 조합을 하나의 설정 객체로 캡슐화한다.

논문 핵심 내용

Mixtral of Experts (2401.04088)

Mixtral of Experts 논문은 Sparse Mixture of Experts 언어 모델의 실용성을 입증했다.

핵심 아이디어: Mixtral 8x7B는 각 레이어에 8개의 전문가 FFN 블록을 두고, 라우터가 토큰마다 2개의 전문가를 선택하는 구조다. 총 47B 파라미터이지만 추론 시 토큰당 13B만 활성화되어, 70B Dense 모델과 동등한 성능을 훨씬 적은 연산으로 달성한다.

메트릭	Mixtral 8x7B	비교
총 파라미터	47B	-
활성 파라미터/토큰	13B	Llama 2 70B의 ~19%
컨텍스트 길이	32K 토큰	-
vs Llama 2 70B	동등 또는 우위	모든 벤치마크
vs GPT-3.5	동등 또는 우위	대부분 벤치마크
수학/코드/다국어	크게 우위	vs Llama 2 70B

Mixtral은 특히 수학, 코드 생성, 다국어 벤치마크에서 Llama 2 70B를 크게 앞섰다. Instruction-tuned 버전인 Mixtral 8x7B-Instruct는 GPT-3.5 Turbo, Claude-2.1, Gemini Pro를 인간 평가에서 능가했다. Apache 2.0 라이선스로 공개되어 vLLM의 Fused MoE 커널 개발에 가장 중요한 타겟 모델이 됐다.

DeepSeekMoE (2412.19437)

DeepSeek-V3는 이 MoE 개념을 극한까지 확장하여 256개 전문가 + 공유 전문가 구조를 채택했다. 671B 파라미터 중 37B만 활성화하는 이 초대규모 MoE 모델이 vLLM의 Fused MoE 레이어가 처리해야 하는 가장 도전적인 워크로드다.

마무리

vLLM의 Fused MoE는 단순한 커널 융합을 넘어, 라우터-러너 분리 아키텍처, 4종 병렬화 지원, 양자화 직교 조합이라는 정교한 설계를 갖추고 있다. DeepSeek-R1의 256 전문가를 효율적으로 서빙할 수 있는 것은 이 레이어의 모듈러 설계 덕분이다.