融合 MoE 核函数特性¶
本文档旨在概述各种 MoE 核函数(包括模块化和非模块化),以便在特定情况下更容易选择合适的核函数集。这包括有关模块化核函数使用的 all2all 后端的信息。
融合 MoE 模块化 All2All 后端¶
存在多个 all2all 通信后端,用于为 FusedMoE 层实现专家并行 (EP)。不同的 FusedMoEPrepareAndFinalizeModular 子类为每个 all2all 后端提供了接口。
下表描述了每个后端的相关特性,即激活格式、支持的量化方案和异步支持。
输出激活格式(标准或批处理)对应于 FusedMoEPrepareAndFinalizeModular 子类的准备步骤的输出,而最终确定步骤需要相同的格式。所有后端 prepare 方法都期望标准格式的激活,并且所有 finalize 方法都返回标准格式的激活。有关格式的更多详细信息,请参见 融合 MoE 模块化核函数 文档。
量化类型和格式列举了每个 FusedMoEPrepareAndFinalizeModular 类支持的量化方案。量化可以在调度之前或之后进行,具体取决于 all2all 后端支持的格式,例如 deepep_high_throughput 仅支持块量化 fp8 格式。任何其他格式都将导致以更高精度进行调度,然后进行量化。每个后端的准备步骤的输出是量化类型。最终确定步骤通常需要与原始激活相同的输入类型,例如,如果原始输入是 bfloat16 且量化方案是带有逐张量缩放的 fp8,则 prepare 将返回 fp8/逐张量缩放激活,而 finalize 将接受 bfloat16 激活。有关 MoE 过程每个步骤中激活的类型和格式的更多详细信息,请参见 融合 MoE 模块化核函数 中的图表。如果未指定量化类型,则核函数在 float16 和/或 bfloat16 上运行。
异步后端支持使用 DBO(双批次重叠)和共享专家重叠(共享专家在组合步骤中计算)。
某些模型要求在 topk==1 时将 topk 权重应用于输入激活而不是输出激活,例如 Llama。对于模块化核函数,此功能由 FusedMoEPrepareAndFinalizeModular 子类支持。对于非模块化核函数,由专家函数处理此标志。
除非另有说明,后端通过 --all2all-backend 命令行参数(或 ParallelConfig 中的 all2all_backend 参数)进行控制。除 flashinfer 外的所有后端仅适用于 EP+DP 或 EP+TP。Flashinfer 可以与 EP 或不带 EP 的 DP 配合使用。
| Backend | 输出激活格式 | 量化类型 | 量化格式 | 异步 | 对输入应用权重 | 子类 |
|---|---|---|---|---|---|---|
| 朴素 | 标准 | 所有1 | G,A,T | 否 | 6 | layer.py |
| deepep_high_throughput | 标准 | fp8 | G(128),A,T2 | 是 | 是 | DeepEPHTPrepareAndFinalize |
| deepep_low_latency | batched | fp8 | G(128),A,T3 | 是 | 是 | DeepEPLLPrepareAndFinalize |
| flashinfer_nvlink_two_sided | 标准 | nvfp4,fp8 | G,A,T | 否 | 否 | FlashInferNVLinkTwoSidedPrepareAndFinalize |
| flashinfer_nvlink_one_sided | 标准 | nvfp4,bf16,mxfp8 | G,A,T | 否 | 否 | FlashInferNVLinkOneSidedPrepareAndFinalize |
图例
- 所有类型:mxfp4, nvfp4, int4, int8, fp8
- A,T 量化发生在调度之后。
- 所有量化都发生在调度之后。
- 由
--moe-backend控制 (flashinfer_cutlass或flashinfer_trtllm) - 这是一个无操作调度器,可与任何模块化专家配对,以生成无需调度或组合即可运行的模块化核函数。这些无法通过环境变量选择。它们通常用于测试或使专家子类适应
fused_expertsAPI。 - 这取决于专家实现。
- G - 分组
- G(N) - 带块大小 N 的分组
- A - 逐激活令牌
- T - 逐张量
模块化核函数由以下 FusedMoEMethodBase 类支持。
ModelOptFp8MoEMethodFp8MoEMethodCompressedTensorsW4A4Nvfp4MoEMethodCompressedTensorsW8A8Fp8MoEMethodGptOssMxfp4MoEMethodUnquantizedFusedMoEMethod
融合专家核函数¶
有许多针对不同量化类型和架构的 MoE 专家核函数实现。大多数遵循基础 Triton fused_experts 函数的通用 API。许多具有模块化核函数适配器,因此它们可以与兼容的 all2all 后端一起使用。此表列出了每个专家核函数及其特定属性。
每个核函数都必须提供一种支持的输入激活格式。某些核函数类型通过不同的入口点支持标准和批处理格式,例如 TritonExperts 和 BatchedTritonExperts。目前,批处理格式的核函数仅在与某些 all2all 后端匹配时才需要,例如 DeepEPLLPrepareAndFinalize。
与后端核函数类似,每个专家核函数仅支持某些量化格式。对于非模块化专家,激活将采用原始类型,并由核函数内部量化。模块化专家将期望激活已处于量化格式。两种类型的专家都将以原始激活类型产生输出。
每个专家核函数支持一个或多个激活函数,例如 silu 或 gelu,它们应用于中间结果。
与后端一样,一些专家支持在输入激活上应用 topk 权重。此表中列出的条目仅适用于非模块化专家。
大多数专家类型都包含一个等效的模块化接口,该接口将是 FusedMoEExpertsModular 的子类。
要与特定的 FusedMoEPrepareAndFinalizeModular 子类一起使用,MoE 核函数必须具有兼容的激活格式、量化类型和量化格式。
| 核函数 | 输入激活格式 | 量化类型 | 量化格式 | 激活函数 | 对输入应用权重 | 模块化 | 来源 |
|---|---|---|---|---|---|---|---|
| triton | 标准 | 所有1 | G,A,T | silu, gelu, swigluoai, silu_no_mul, gelu_no_mul | 是 | 是 | fused_experts,TritonExperts |
| triton (批处理) | batched | 所有1 | G,A,T | silu, gelu | 6 | 是 | BatchedTritonExperts |
| deep gemm | 标准, batched | fp8 | G(128),A,T | silu, gelu | 6 | 是 | DeepGemmExperts,BatchedDeepGemmExperts |
| cutlass_fp4 | 标准, batched | nvfp4 | A,T | silu | 是 | 是 | CutlassExpertsFp4 |
| cutlass_fp8 | 标准, batched | fp8 | A,T | silu, gelu | 是 | 是 | CutlassExpertsFp8,CutlasBatchedExpertsFp8 |
| flashinfer | 标准 | nvfp4, fp8 | T | 5 | 否 | 是 | FlashInferExperts |
| gpt oss triton | 标准 | 不适用 | 不适用 | 5 | 是 | 是 | triton_kernel_fused_experts,OAITritonExperts |
| marlin | 标准, batched | 3 / 不适用 | 3 / 不适用 | silu, swigluoai | 是 | 是 | fused_marlin_moe,MarlinExperts,BatchedMarlinExperts |
| trtllm | 标准 | mxfp4, nvfp4 | G(16),G(32) | 5 | 否 | 是 | TrtLlmMxfp4ExpertsMonolithic,TrtLlmMxfp4ExpertsModular,TrtLlmNvFp4ExpertsMonolithic,TrtLlmNvfp4ExpertsModular |
| hpc | 标准 | fp8 | G(128),T | silu | 是 | 是 | HPCExperts |
| rocm aiter moe | 标准 | mxfp4, fp8 | G(32),G(128),A,T | silu, gelu, swigluoai | 是 | 否 | rocm_aiter_fused_experts,AiterExperts |
| cpu_fused_moe | 标准 | 不适用 | 不适用 | silu | 否 | 否 | CPUFusedMOE |
| 朴素批处理4 | batched | int8, fp8 | G,A,T | silu, gelu | 6 | 是 | NaiveBatchedExperts |
图例
- 所有类型:mxfp4, nvfp4, int4, int8, fp8
- 一个围绕 triton 和 deep gemm 专家的调度器包装器。将根据类型 + 形状 + 量化参数进行选择
- uint4, uint8, fp8, fp4
- 这是一个支持批处理格式的朴素专家实现。主要用于测试。
activation参数被忽略,默认使用 SwiGlu。- 仅由模块化核函数处理或在使用模块化核函数时支持。
模块化核函数“系列”¶
下表显示了旨在协同工作的模块化核函数“系列”。可能存在一些可以工作但尚未经过测试的组合,例如 flashinfer 与其他 fp8 专家。
| backend | FusedMoEPrepareAndFinalizeModular 子类 | FusedMoEExpertsModular 子类 |
|---|---|---|
| deepep_high_throughput | DeepEPHTPrepareAndFinalize | DeepGemmExperts,TritonExperts,TritonOrDeepGemmExperts,CutlassExpertsFp8, MarlinExperts |
| deepep_low_latency | DeepEPLLPrepareAndFinalize | BatchedDeepGemmExperts,BatchedTritonExperts,CutlassBatchedExpertsFp8,BatchedMarlinExperts |
| flashinfer | FlashInferCutlassMoEPrepareAndFinalize | FlashInferExperts |