[flashinfer] FlashInfer Unified MoE API: NVFP4 백엔드 통합 및 자동 튜닝 최적화

2026년 6월 10일수정: 2026년 6월 10일

PR 링크: flashinfer-ai/flashinfer#3093 상태: Merged | 변경: +4441 / -13

들어가며

최신 대규모 언어 모델(LLM)에서 Mixture-of-Experts(MoE) 구조는 필수적이지만, 다양한 하드웨어 백엔드(CuteDSL, TRTLLM 등)마다 최적화된 커널 구현이 달라 파편화 문제가 심각했습니다. 특히 NVFP4와 같은 저정밀도 연산 환경에서는 백엔드별로 최적의 성능을 내는 파라미터가 달라 수동 튜닝이 매우 어렵습니다. 이번 PR은 MoELayer를 중심으로 여러 백엔드를 통합하고, 런타임에 GPU 시간을 측정하여 가장 빠른 백엔드를 동적으로 선택하는 Unified MoE API를 도입하여 이 문제를 해결합니다.

코드 분석

1. Unified MoE API 설계 (`flashinfer/fused_moe/api.py`)

기존의 파편화된 데이터 구조를 MoEConfig, MoEActivationPack, MoEWeightPack으로 추상화했습니다. 이를 통해 백엔드 간의 인터페이스를 통일했습니다.

# Before: 파편화된 3축 dtype/granularity/variant 분리
# After: 단일 QuantVariant enum으로 통합
class MoEConfig(FrozenDataclass):
    routing: RoutingConfig
    quant: QuantConfig
    experts: ExpertConfig
    backend: BackendOptions

2. 백엔드 어댑터 (`flashinfer/fused_moe/runners.py`)

CuteDslNvfp4Runner와 TrtllmFp4RoutedRunner는 TunableRunner를 상속받아, 각 백엔드의 고유한 텐서 리스트를 MoEActivationPack으로부터 변환(pack/unpack)하는 역할을 수행합니다.

# Before: 각 러너가 독립적인 tuning_config를 가짐
# After: 인스턴스별로 tuning_config를 생성하여 유연성 확보
class TrtllmFp4RoutedRunner(TunableRunner):
    def pack_inputs(self, act, weights):
        # 백엔드 네이티브 텐서 리스트로 변환
        return self._translate_to_native(act, weights)

3. MoELayer 디스패처 (`flashinfer/fused_moe/layer.py`)

MoELayer는 첫 호출 시 각 백엔드의 성능을 벤치마킹하여 '승자'를 결정하고, 이후 호출에서는 캐싱된 결과를 사용하여 오버헤드를 최소화합니다.

# MoELayer 내부의 동적 선택 로직
class MoELayer:
    def __call__(self, act_pack, weight_pack):
        if self.winner_backend is None:
            self._select_winner(act_pack, weight_pack)
        return self.runners[self.winner_backend].forward(act_pack, weight_pack)

왜 이게 좋은가

성능 최적화: 고정된 백엔드 사용 대신, 런타임에 bench_gpu_time을 통해 토큰 수와 하드웨어 아키텍처(SM100 등)에 최적화된 백엔드를 선택합니다. 예를 들어, TRTLLM은 소규모 토큰에서, CuteDSL은 대규모 토큰에서 더 나은 성능을 보이는 경향을 자동으로 반영합니다.
유지보수성: prepare_weights 헬퍼를 통해 가중치 전처리 로직을 중앙화하여, 테스트와 벤치마크 코드의 중복을 제거했습니다.
안정성: refcheck 기능을 통해 모든 백엔드가 동일한 bf16 레퍼런스 결과와 일치하는지 검증하여, 저정밀도 연산에서의 정확도 저하를 방지합니다.

리뷰 피드백 반영

CUDA Graph 호환성: pack_inputs 내부에 디바이스-호스트 동기화가 필요한 로직을 추가하자는 제안이 있었으나, 그래프 캡처 시 성능 저하와 에러를 방지하기 위해 이를 거부하고 별도의 검증 로직으로 분리했습니다.
중복 제거: tllm_enums.py와 API 간의 중복을 제거하고, MoETensors 대신 Pack 기반의 명확한 데이터 구조로 리팩토링하여 코드 가독성을 높였습니다.