Fused MoE 커널 기능

Fused MoE 커널 기능 (Fused MoE Kernel Features)

이 문서는 vLLM에서 제공하는 다양한 MoE 커널(모듈러/비모듈러)을 한눈에 정리한 것입니다. 어떤 상황에서 어떤 커널 조합을 골라야 할지 판단할 때 이 문서가 출발점이 되어 줍니다. 특히 모듈러 커널이 사용하는 all2all(all-to-all) 백엔드에 대한 정보도 함께 다룹니다.

출처: 문서

본문

Fused MoE 모듈러 All2All 백엔드 (Fused MoE Modular All2All backends)

MoERunner 레이어에서 전문가 병렬화(expert parallelism, EP)를 구현하기 위해 여러 all2all 통신 백엔드가 사용됩니다. 서로 다른 FusedMoEPrepareAndFinalizeModular 하위 클래스들이 각 all2all 백엔드의 인터페이스를 제공합니다.

아래 표는 각 백엔드의 관련 기능, 즉 활성화 형식(activation format), 지원하는 양자화 방식, 비동기(async) 지원 여부를 정리한 것입니다.

출력 활성화 형식(standard 또는 batched)은 FusedMoEPrepareAndFinalizeModular 하위 클래스의 prepare 단계 출력을 의미하며, finalize 단계는 같은 형식을 요구합니다. 모든 백엔드의 prepare 메서드는 standard 형식의 활성화를 입력으로 받고, 모든 finalize 메서드는 standard 형식의 활성화를 반환합니다. 형식에 대한 자세한 내용은 Fused MoE Modular Kernel 문서를 참고하세요.

양자화 타입과 형식은 각 FusedMoEPrepareAndFinalizeModular 클래스가 지원하는 양자화 방식을 나타냅니다. 양자화는 all2all 백엔드가 지원하는 형식에 따라 dispatch 전이나 후에 일어날 수 있습니다. 예를 들어 deepep_high_throughput은 block-quantized fp8 형식만 지원하므로, 다른 형식이면 더 높은 정밀도로 dispatch한 뒤 나중에 양자화합니다. 각 백엔드의 prepare 단계 출력이 바로 양자화된 타입입니다. finalize 단계는 일반적으로 원래 활성화와 같은 입력 타입을 요구합니다. 예를 들어 원래 입력이 bfloat16이고 양자화 방식이 per-tensor 스케일의 fp8이라면, prepare는 fp8/per-tensor 스케일 활성화를 반환하고 finalize는 bfloat16 활성화를 받습니다. MoE 처리의 각 단계에서 활성화의 타입·형식에 대한 상세는 Fused MoE Modular Kernel의 다이어그램을 참고하세요. 양자화 타입을 지정하지 않으면 커널은 float16 및/또는 bfloat16으로 동작합니다.

비동기 백엔드는 DBO(Dual Batch Overlap)와 shared expert overlap(combine 단계에서 shared expert를 함께 계산하는 방식)을 지원합니다.

특정 모델(Llama 등)은 topk==1일 때 topk 가중치를 출력 활성화가 아니라 입력 활성화에 적용해야 합니다. 모듈러 커널에서는 FusedMoEPrepareAndFinalizeModular 하위 클래스가 이 기능을 지원하며, 비모듈러 커널에서는 experts 함수가 이 플래그를 처리해야 합니다.

별도로 명시하지 않는 한, 백엔드는 --all2all-backend 커맨드라인 인자(또는 ParallelConfig의 all2all_backend 파라미터)로 제어합니다. flashinfer를 제외한 모든 백엔드는 EP+DP 또는 EP+TP에서만 동작합니다. Flashinfer는 EP 없이 EP 또는 DP로 동작할 수 있습니다.

td {
  padding: 0.5rem !important;
  white-space: nowrap;
}

th {
  padding: 0.5rem !important;
  min-width: 0 !important;
}
Backend Output act. format Quant. types Quant. format Async Apply Weight On Input Subclass
naive standard all1 G,A,T N 6 MoERunner
deepep_high_throughput standard fp8 G(128),A,T2 Y Y DeepEPHTPrepareAndFinalize
deepep_low_latency batched fp8 G(128),A,T3 Y Y DeepEPLLPrepareAndFinalize
flashinfer_nvlink_two_sided standard nvfp4,fp8 G,A,T N N FlashInferNVLinkTwoSidedPrepareAndFinalize
flashinfer_nvlink_one_sided standard nvfp4,bf16,mxfp8,fp8 G,A,T N N FlashInferNVLinkOneSidedPrepareAndFinalize

표 해석 (Table key)

  • All types: mxfp4, nvfp4, int4, int8, fp8
  • A,T: 양자화가 dispatch 후에 일어납니다.
  • All: 모든 양자화가 dispatch 후에 일어납니다.
  • --moe-backend(flashinfer_cutlass 또는 flashinfer_trtllm)로 제어합니다.
  • 이는 dispatch나 combine 없이 실행되는 모듈러 커널을 만들기 위해 어떤 모듈러 experts와도 짝지을 수 있는 no-op 디스패처입니다. 환경변수로는 선택할 수 없으며, 주로 테스트 용도나 expert 하위 클래스를 fused_experts API에 적응시키는 데 사용합니다.
  • experts 구현에 따라 달라집니다.
    • G - Grouped
  • G(N) - 블록 크기 N의 Grouped
  • A - Per activation token (활성화 토큰별)
  • T - Per tensor (텐서별)

모듈러 커널은 다음 FusedMoEMethodBase 클래스들이 지원합니다.

  • ModelOptFp8MoEMethod
  • Fp8MoEMethod
  • CompressedTensorsW4A4Nvfp4MoEMethod
  • CompressedTensorsW8A8Fp8MoEMethod
  • GptOssMxfp4MoEMethod
  • UnquantizedFusedMoEMethod

Fused Experts 커널 (Fused Experts Kernels)

다양한 양자화 타입과 아키텍처에 대해 여러 MoE experts 커널 구현이 존재합니다. 대부분 기본 Triton fused_experts 함수의 일반 API를 따릅니다. 상당수가 모듈러 커널 어댑터를 갖추고 있어 호환되는 all2all 백엔드와 함께 쓸 수 있습니다. 아래 표는 각 experts 커널과 그 특성을 정리한 것입니다.

각 커널은 지원되는 입력 활성화 형식 중 하나를 제공받아야 합니다. 일부 커널 종류는 서로 다른 진입점을 통해 standard·batched 두 형식을 모두 지원합니다(예: TritonExperts와 BatchedTritonExperts). batched 형식 커널은 현재 특정 all2all 백엔드(예: DeepEPLLPrepareAndFinalize)와 짝을 이루는 경우에만 필요합니다.

백엔드 커널과 마찬가지로, 각 experts 커널도 특정 양자화 형식만 지원합니다. 비모듈러 experts의 경우 활성화는 원래 타입으로 제공되고 커널 내부에서 양자화됩니다. 모듈러 experts는 활성화가 이미 양자화된 형식으로 제공되기를 기대합니다. 두 종류의 experts 모두 원래 활성화 타입으로 출력을 냅니다.

각 experts 커널은 중간 결과에 적용되는 하나 이상의 활성화 함수(예: silu, gelu)를 지원합니다.

백엔드와 마찬가지로 일부 experts는 topk 가중치를 입력 활성화에 적용하는 것을 지원합니다. 이 표의 해당 컬럼 항목은 비모듈러 experts에만 적용됩니다.

대부분의 experts 종류에는 FusedMoEExpertsModular의 하위 클래스인 모듈러 인터페이스가 함께 제공됩니다.

특정 FusedMoEPrepareAndFinalizeModular 하위 클래스와 함께 사용하려면 MoE 커널이 호환되는 활성화 형식·양자화 타입·양자화 형식을 가져야 합니다.

Kernel Input act. format Quant. types Quant. format Activation function Apply Weight On Input Modular Source
triton standard all1 G,A,T silu, gelu,swigluoai,silu_no_mul,gelu_no_mul Y Y fused_experts,TritonExperts
triton (batched) batched all1 G,A,T silu, gelu 6 Y BatchedTritonExperts
deep gemm standard,batched fp8 G(128),A,T silu, gelu 6 Y DeepGemmExperts,BatchedDeepGemmExperts
cutlass_fp4 standard,batched nvfp4 A,T silu Y Y CutlassExpertsFp4
cutlass_fp8 standard,batched fp8 A,T silu, gelu Y Y CutlassExpertsFp8,CutlasBatchedExpertsFp8
flashinfer standard nvfp4,fp8 T 5 N Y FlashInferExperts
gpt oss triton standard N/A N/A 5 Y Y triton_kernel_fused_experts,OAITritonExperts
marlin standard,batched 3 / N/A 3 / N/A silu,swigluoai Y Y fused_marlin_moe,MarlinExperts,BatchedMarlinExperts
trtllm standard mxfp4,nvfp4 G(16),G(32) 5 N Y TrtLlmMxfp4ExpertsMonolithic,TrtLlmMxfp4ExpertsModular,TrtLlmNvFp4ExpertsMonolithic,TrtLlmNvfp4ExpertsModular
hpc standard fp8 G(128),T silu Y Y HPCExperts
rocm aiter moe standard mxfp4,fp8 G(32),G(128),A,T silu, gelu,swigluoai Y N rocm_aiter_fused_experts,AiterExperts
cpu_moe standard N/A N/A silu, gelu,gelu_tanh,swigluoai Y N X86CPUUnquantizedExperts,ArmCPUUnquantizedExperts,CPUUnquantizedExperts
naive batched4 batched int8,fp8 G,A,T silu, gelu 6 Y NaiveBatchedExperts

표 해석 (Table key)

  • All types: mxfp4, nvfp4, int4, int8, fp8
  • triton과 deep gemm experts를 감싸는 디스패처 래퍼입니다. 타입·형상·양자화 파라미터에 따라 선택합니다.
  • uint4, uint8, fp8, fp4
  • batched 형식을 지원하는 naive experts 구현입니다. 주로 테스트에 사용합니다.
  • activation 파라미터는 무시되고 기본적으로 SwiGlu가 사용됩니다.
  • 모듈러 커널과 함께 쓸 때만 처리·지원됩니다.

모듈러 커널 "패밀리" (Modular Kernel "families")

아래 표는 함께 동작하도록 설계된 모듈러 커널 "패밀리"를 보여줍니다. 동작할 수도 있지만 아직 테스트되지 않은 조합도 있습니다(예: 다른 fp8 experts와 함께 쓰는 flashinfer).

backend FusedMoEPrepareAndFinalizeModular subclasses FusedMoEExpertsModular subclasses
deepep_high_throughput DeepEPHTPrepareAndFinalize DeepGemmExperts,TritonExperts,TritonOrDeepGemmExperts,CutlassExpertsFp8, MarlinExperts
deepep_low_latency DeepEPLLPrepareAndFinalize BatchedDeepGemmExperts,BatchedTritonExperts,CutlassBatchedExpertsFp8,BatchedMarlinExperts
flashinfer FlashInferCutlassMoEPrepareAndFinalize FlashInferExperts

더 알아보기 (Learn more)

  • Fused MoE Modular Kernel 문서 — 모듈러 커널의 활성화 타입·형식별 단계 다이어그램
  • ParallelConfig의 all2all_backend / --moe-backend 옵션