본문으로 건너뛰기

[vLLM] MoE Oracle & Prepare/Finalize: 백엔드 선택과 분산 데이터 교환

들어가며

vLLM의 MoE 시스템은 크게 세 부분으로 나뉜다: Oracle(어떤 백엔드를 쓸지 결정), Prepare/Finalize(토큰을 전문가에게 전달하고 결과를 수집), Experts(실제 GEMM 연산). 이번 글에서는 Oracle 시스템의 백엔드 선택 로직과, Prepare/Finalize의 분산 데이터 교환 패턴을 분석한다.

핵심 구조/코드 분석

Oracle: 백엔드 자동 선택

oracle/ 디렉토리에는 양자화 타입별 백엔드 선택 로직이 있다:

oracle/
├── __init__.py
├── fp8.py          # FP8 MoE 백엔드 선택
├── mxfp4.py        # MXFP4 MoE 백엔드 선택
├── mxfp8.py        # MXFP8 MoE 백엔드 선택
├── nvfp4.py        # NVFP4 MoE 백엔드 선택
└── unquantized.py  # 비양자화 MoE 백엔드 선택

각 Oracle은 모델 설정(전문가 수, hidden_size, 분산 설정 등)을 입력받아 최적의 커널 백엔드를 반환한다. 예를 들어 MXFP8의 경우:

from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import (
    select_mxfp8_moe_backend,
)
self.fp8_backend, self.experts_cls = select_mxfp8_moe_backend(config=self.moe)

반환값은 백엔드 열거형과 해당 전문가 클래스의 튜플이다. 이를 통해 같은 양자화 방식이라도 하드웨어, 텐서 병렬 설정, 전문가 병렬 설정에 따라 다른 커널을 선택할 수 있다.

Prepare/Finalize: 분산 교환 패턴

prepare_finalize/ 디렉토리는 MoE의 토큰 분배/수집 전략을 구현한다:

prepare_finalize/
├── no_dp_ep.py                    # 단일 노드, DP/EP 없음
├── deepep_ht.py                   # DeepEP High-Throughput
├── deepep_ll.py                   # DeepEP Low-Latency
├── flashinfer_nvlink_one_sided.py # FlashInfer NVLink 단방향
├── flashinfer_nvlink_two_sided.py # FlashInfer NVLink 양방향
└── naive_dp_ep.py                 # 기본 DP+EP

No DP/EP: 가장 단순한 경우

분산 처리가 없을 때의 Prepare/Finalize는 직관적이다:

class MoEPrepareAndFinalizeNoDPEPModular(mk.FusedMoEPrepareAndFinalizeModular):
    def prepare(self, a1, topk_weights, topk_ids, num_experts,
                expert_map, apply_router_weight_on_input, quant_config):
        if apply_router_weight_on_input:
            assert topk == 1
            a1 = a1 * topk_weights.to(a1.dtype)
        a1q, a1q_scale = _quantize_input(a1, quant_config)
        return a1q, a1q_scale, None, None, None

    def finalize(self, output, fused_expert_output,
                 topk_weights, topk_ids, apply_router_weight_on_input,
                 weight_and_reduce_impl):
        weight_and_reduce_impl.apply(
            output=output,
            fused_expert_output=fused_expert_output,
            topk_weights=topk_weights,
            topk_ids=topk_ids,
        )

prepare에서 입력을 양자화하고, finalize에서 전문가 출력을 가중 합산한다.

DeepEP High-Throughput: 통신 오버랩

DeepEP는 전문가 병렬 처리를 위한 고성능 all-to-all 통신 라이브러리다:

class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
    def __init__(self, buffer: deep_ep.Buffer, num_dispatchers: int,
                 dp_size: int, rank_expert_offset: int):
        self.buffer = buffer
        self.dp_size = dp_size
        self.rank_expert_offset = rank_expert_offset

DeepEP HT(High-Throughput) 모드에서는 토큰을 전문가에게 분배하는 dispatch와 결과를 수집하는 combine을 비동기로 수행하여, 계산과 통신을 오버랩한다. hidden_size가 512바이트 단위로 정렬되어야 한다는 제약이 있다:

@staticmethod
def maybe_roundup_layer_hidden_size(hidden_size: int, dtype: torch.dtype):
    hidden_size_bytes = hidden_size * dtype.itemsize
    xfer_atom_size = 512  # 32 * 16 (size(int4))
    if hidden_size_bytes % xfer_atom_size == 0:
        return hidden_size
    hidden_size_bytes = round_up(hidden_size_bytes, xfer_atom_size)
    return hidden_size_bytes // dtype.itemsize

Modular vs Monolithic 커널

Prepare/Finalize는 두 가지 모드를 지원한다:

  • Modular: Prepare -> Experts -> Finalize를 개별 단계로 실행. 통신 오버랩에 유리.
  • Monolithic: 전체를 하나의 커널로 퓨전. FlashInfer+TRT-LLM 같은 통합 백엔드에서 사용.
def make_moe_prepare_and_finalize_no_dp_ep(use_monolithic: bool):
    return (MoEPrepareAndFinalizeNoDPEPMonolithic()
            if use_monolithic
            else MoEPrepareAndFinalizeNoDPEPModular())

왜 이 설계인가

  1. 관심사 분리: Oracle(백엔드 선택), Prepare/Finalize(데이터 이동), Experts(연산)를 명확히 분리하여, 각 구성요소를 독립적으로 교체/최적화할 수 있다.

  2. 통신-계산 오버랩: DeepEP 통합으로 전문가 병렬 처리 시 통신 오버헤드를 숨길 수 있다. Prepare에서 dispatch를 시작하고, 다른 연산을 수행한 뒤 Finalize에서 결과를 수집하는 파이프라인이 가능하다.

  3. 하드웨어 적응: NVLink 유무, GPU 세대, 분산 설정에 따라 최적의 Prepare/Finalize 전략이 달라진다. Oracle이 이를 자동으로 선택하므로 사용자는 신경 쓸 필요가 없다.

참고 자료

댓글

관련 포스트

vLLM 의 다른글