融合 MoE 模块化内核

简介

FusedMoEModularKernel 的实现位于 此处

基于输入激活格式，FusedMoE 的实现大致分为 2 种类型。

• 连续型（Contiguous / Standard / Non-Batched），以及
• 分批型（Batched）

注意

“连续型”、“标准型”和“非分批型”这几个术语在本文档中可互换使用。

输入激活格式完全取决于所使用的 All2All Dispatch（全对全调度）。

• 在连续型变体中，All2All Dispatch 返回形状为 (M, K) 的连续张量，以及形状为 (M, num_topk) 的 TopK Ids 和 TopK 权重。示例可参考 DeepEPHTPrepareAndFinalize。
• 在分批型变体中，All2All Dispatch 返回形状为 (num_experts, max_tokens, K) 的张量。在此格式中，订阅同一专家的激活/Token 被分批在一起。请注意，张量并非所有条目都有效。激活张量通常伴随一个大小为 num_experts 的 expert_num_tokens 张量，其中 expert_num_tokens[i] 表示订阅第 i 个专家的有效 Token 数量。示例可参考 DeepEPLLPrepareAndFinalize。

无论是在连续型还是分批型变体中，FusedMoE 操作通常由多个步骤组成，如下面的图示所示。

注意

在操作层面，分批与非分批情况的主要区别在于 Permute（置换）/ Unpermute（反置换）操作，其他操作保持不变。

动机

从图中可以看出，这里涉及大量操作，且每种操作可能有多种实现方式。将这些操作组合起来构成有效 FusedMoE 实现的方法组合很快就会变得难以处理。模块化内核框架通过将操作归类为逻辑组件解决了这个问题。这种广泛的分类使组合变得可控，并防止了代码重复。这也实现了 All2All Dispatch 和 Combine 实现与 FusedMoE 实现的解耦，允许进行独立开发和测试。此外，模块化内核框架为不同组件引入了抽象类，从而为未来的实现提供了定义良好的骨架。

本文档的其余部分将重点介绍连续型/非分批情况。推导到分批情况应该是直接的。

ModularKernel 组件

FusedMoEModularKernel 将 FusedMoE 操作分为 3 个部分，

1. TopKWeightAndReduce
2. FusedMoEPrepareAndFinalizeModular
3. FusedMoEExpertsModular

TopKWeightAndReduce

TopK 权重应用和规约组件在 Unpermute 操作之后、All2All Combine 之前执行。请注意，FusedMoEExpertsModular 负责 Unpermute，而 FusedMoEPrepareAndFinalizeModular 负责 All2All Combine。在 FusedMoEExpertsModular 中进行 TopK 权重应用和规约是有价值的，但一些实现选择在 FusedMoEPrepareAndFinalizeModular 中进行。为了实现这种灵活性，我们提供了 TopKWeightAndReduce 抽象类。

TopKWeightAndReduce 的实现请查看 此处。

FusedMoEPrepareAndFinalizeModular::finalize() 方法接受一个 TopKWeightAndReduce 参数，该参数在方法内部被调用。FusedMoEModularKernel 充当 FusedMoEExpertsModular 和 FusedMoEPrepareAndFinalize 实现之间的桥梁，以确定 TopK 权重应用和规约发生的具体位置。

• 如果 FusedMoEExpertsModular 实现自行处理权重应用和规约，FusedMoEExpertsModular::finalize_weight_and_reduce_impl 方法会返回 TopKWeightAndReduceNoOp。
• 如果 FusedMoEExpertsModular 实现需要 FusedMoEPrepareAndFinalizeModular::finalize() 来执行权重应用和规约，则 FusedMoEExpertsModular::finalize_weight_and_reduce_impl 方法会返回 TopKWeightAndReduceContiguous / TopKWeightAndReduceNaiveBatched / TopKWeightAndReduceDelegate。

FusedMoEPrepareAndFinalizeModular

抽象类 FusedMoEPrepareAndFinalizeModular 公开了 prepare、prepare_no_receive 和 finalize 函数。prepare 函数负责输入激活的量化和 All2All Dispatch。如果实现，prepare_no_receive 与 prepare 类似，但它不会等待从其他工作节点接收结果。相反，它返回一个“接收器”回调，必须调用该回调以等待工作节点的最终结果。并非所有 FusedMoEPrepareAndFinalizeModular 类都必须支持此方法，但如果可用，它可以用于将工作与初始全对全通信交错进行，例如将共享专家与融合专家交错。finalize 函数负责调用 All2All Combine。此外，finalize 函数可能会也可能不会执行 TopK 权重应用和规约（请参阅 TopKWeightAndReduce 部分）。

FusedMoEExpertsModular

FusedMoEExpertsModular 类是 MoE 操作的核心所在。该抽象类公开了几个重要函数：

• apply()
• workspace_shapes()
• finalize_weight_and_reduce_impl()

apply()

apply 方法是实现执行以下操作的地方：

• Permute（置换）
• 与权重 W1 进行矩阵乘法
• 激活 + 乘法（Act + Mul）
• 量化
• 与权重 W2 进行矩阵乘法
• Unpermute（反置换）
• 可选的 TopK 权重应用 + 规约

workspace_shapes()

核心 FusedMoE 实现会执行一系列操作。分别为这些操作创建输出内存效率很低。因此，要求实现声明 2 个工作空间形状、工作空间数据类型以及 FusedMoE 输出形状作为 workspace_shapes() 方法的输出。这些信息用于在 FusedMoEModularKernel::forward() 中分配工作空间张量和输出张量，并传递给 FusedMoEExpertsModular::apply() 方法。随后工作空间可用作 FusedMoE 实现中的中间缓冲区。

finalize_weight_and_reduce_impl()

在 FusedMoEExpertsModular::apply() 内部执行 TopK 权重应用和规约有时效率更高。点击 此处查看示例。我们有一个 TopKWeightAndReduce 抽象类来促进此类实现。请参考 TopKWeightAndReduce 部分。FusedMoEExpertsModular::finalize_weight_and_reduce_impl() 返回实现希望 FusedMoEPrepareAndFinalizeModular::finalize() 使用的 TopKWeightAndReduce 对象。

FusedMoEModularKernel

FusedMoEModularKernel 由 FusedMoEPrepareAndFinalizeModular 和 FusedMoEExpertsModular 对象组成。FusedMoEModularKernel 的伪代码/草图如下：

class FusedMoEModularKernel:
    def __init__(self,
                 prepare_finalize: FusedMoEPrepareAndFinalizeModular,
                 fused_experts: FusedMoEExpertsModular):

        self.prepare_finalize = prepare_finalize
        self.fused_experts = fused_experts

    def forward(self, DP_A):

        Aq, A_scale, _, _, _ = self.prepare_finalize.prepare(DP_A, ...)

        workspace13_shape, workspace2_shape, _, _ = self.fused_experts.workspace_shapes(...)

        # allocate workspaces
        workspace_13 = torch.empty(workspace13_shape, ...)
        workspace_2 = torch.empty(workspace2_shape, ...)

        # execute fused_experts
        fe_out = self.fused_experts.apply(Aq, A_scale, workspace13, workspace2, ...)

        # war_impl is an object of type TopKWeightAndReduceNoOp if the fused_experts implementations
        # performs the TopK Weight Application and Reduction.
        war_impl = self.fused_experts.finalize_weight_and_reduce_impl()

        output = self.prepare_finalize.finalize(fe_out, war_impl,...)

        return output

操作指南

如何添加 FusedMoEPrepareAndFinalizeModular 类型

通常，FusedMoEPrepareAndFinalizeModular 类型由 All2All Dispatch & Combine 实现/内核支持。例如：

• DeepEPHTPrepareAndFinalize 类型由 DeepEP 高吞吐量 All2All 内核支持，以及
• DeepEPLLPrepareAndFinalize 类型由 DeepEP 低延迟 All2All 内核支持。

第一步：添加 All2All 管理器

All2All 管理器的目的是设置 All2All 内核实现。FusedMoEPrepareAndFinalizeModular 实现通常从 All2All 管理器获取内核实现“句柄”以调用 Dispatch 和 Combine 函数。请在此处查看 All2All 管理器实现： 此处。

第二步：添加 FusedMoEPrepareAndFinalizeModular 类型

本节描述了 FusedMoEPrepareAndFinalizeModular 抽象类公开的各种函数的重要性。

FusedMoEPrepareAndFinalizeModular::prepare(): prepare 方法实现量化和 All2All Dispatch。通常会调用相关 All2All 管理器中的 Dispatch 函数。

FusedMoEPrepareAndFinalizeModular::has_prepare_no_receive(): 指示此子类是否实现了 prepare_no_receive。默认为 False。

FusedMoEPrepareAndFinalizeModular::prepare_no_receive(): prepare_no_receive 方法实现量化和 All2All Dispatch。它不等待调度操作的结果，而是返回一个可以被调用以等待最终结果的 thunk。通常会调用相关 All2All 管理器中的 Dispatch 函数。

FusedMoEPrepareAndFinalizeModular::finalize(): 可能执行 TopK 权重应用和规约以及 All2All Combine。通常会调用相关 All2AllManager 中的 Combine 函数。

FusedMoEPrepareAndFinalizeModular::activation_format(): 如果 prepare 方法（即 All2All 调度）的输出是分批的，则返回 FusedMoEActivationFormat.BatchedExperts。否则返回 FusedMoEActivationFormat.Standard。

FusedMoEPrepareAndFinalizeModular::topk_indices_dtype(): TopK ID 的数据类型。一些 All2All 内核对 TopK ID 的数据类型有严格要求。此要求被传递给 FusedMoe::select_experts 函数，以便可以遵守该要求。如果没有严格要求，则返回 None。

FusedMoEPrepareAndFinalizeModular::max_num_tokens_per_rank(): 这是一次提交给 All2All Dispatch 的最大 Token 数量。

FusedMoEPrepareAndFinalizeModular::num_dispatchers(): 调度单元的总数。此值决定了 Dispatch 输出的大小。Dispatch 输出的形状为 (num_local_experts, max_num_tokens, K)。这里 max_num_tokens = num_dispatchers() * max_num_tokens_per_rank()。

我们建议选择一个与您的 All2All 实现紧密匹配的现有 FusedMoEPrepareAndFinalizeModular 实现，并将其用作参考。

如何添加 FusedMoEExpertsModular 类型

FusedMoEExpertsModular 执行 FusedMoE 操作的核心。抽象类公开的各种函数及其重要性如下：

FusedMoEExpertsModular::activation_formats(): 返回支持的输入和输出激活格式。即连续/分批格式。

FusedMoEExpertsModular::supports_expert_map(): 如果实现支持专家映射，则返回 True。

FusedMoEExpertsModular::workspace_shapes() / FusedMoEExpertsModular::finalize_weight_and_reduce_impl / FusedMoEExpertsModular::apply: 请参阅上方的 FusedMoEExpertsModular 部分。

FusedMoEModularKernel 初始化

FusedMoEMethodBase 类有 3 个方法，共同负责创建 FusedMoEModularKernel 对象。它们是：

• maybe_make_prepare_finalize,
• select_gemm_impl, 以及
• init_prepare_finalize

maybe_make_prepare_finalize

maybe_make_prepare_finalize 方法负责在适当的时候，根据当前的 all2all 后端（例如启用 EP + DP 时）构造 FusedMoEPrepareAndFinalizeModular 的实例。基类方法目前为 EP+DP 情况构建所有 FusedMoEPrepareAndFinalizeModular 对象。派生类可以覆盖此方法以针对不同场景构建 prepare/finalize 对象，例如 ModelOptNvFp4FusedMoE 可以为 EP+TP 情况构建 FlashInferCutlassMoEPrepareAndFinalize。请参考以下实现：

• ModelOptNvFp4FusedMoE

select_gemm_impl

select_gemm_impl 方法在基类中未定义。派生类有责任实现一种构建有效/适当 FusedMoEExpertsModular 对象的方法。请参考以下实现：

• UnquantizedFusedMoEMethod
• CompressedTensorsW8A8Fp8MoEMethod
• CompressedTensorsW8A8Fp8MoECutlassMethod
• Fp8MoEMethod
• ModelOptNvFp4FusedMoE 派生类。

init_prepare_finalize

根据输入和环境变量设置，init_prepare_finalize 方法创建适当的 FusedMoEPrepareAndFinalizeModular 对象。随后该方法查询 select_gemm_impl 以获取适当的 FusedMoEExpertsModular 对象，并构建 FusedMoEModularKernel 对象。

请查看 init_prepare_finalize。重要：FusedMoEMethodBase 派生类在其 apply 方法中使用 FusedMoEMethodBase::fused_experts 对象。当设置允许构建有效的 FusedMoEModularKernel 对象时，我们会用它来覆盖 FusedMoEMethodBase::fused_experts。这本质上使派生类对所使用的 FusedMoE 实现保持不可知。

如何进行单元测试

我们在 test_modular_kernel_combinations.py 中有 FusedMoEModularKernel 单元测试。

该单元测试遍历 FusedMoEPrepareAndFinalizeModular 和 FusedMoEPremuteExpertsUnpermute 类型的所有组合，如果它们兼容，则运行一些正确性测试。如果您正在添加一些 FusedMoEPrepareAndFinalizeModular / FusedMoEExpertsModular 实现：

1. 分别将实现类型添加到 mk_objects.py 中的 MK_ALL_PREPARE_FINALIZE_TYPES 和 MK_FUSED_EXPERT_TYPES。
2. 更新 /tests/kernels/moe/modular_kernel_tools/common.py 中的 Config::is_batched_prepare_finalize()、Config::is_batched_fused_experts()、Config::is_standard_fused_experts()、Config::is_fe_16bit_supported()、Config::is_fe_fp8_supported()、Config::is_fe_block_fp8_supported() 方法。

这样做会将新实现添加到测试套件中。

如何检查 FusedMoEPrepareAndFinalizeModular & FusedMoEExpertsModular 兼容性

单元测试文件 test_modular_kernel_combinations.py 也可以作为独立脚本执行。示例：python3 -m tests.kernels.moe.test_modular_kernel_combinations --pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts。作为副作用，该脚本可用于测试 FusedMoEPrepareAndFinalizeModular & FusedMoEExpertsModular 的兼容性。当以不兼容的类型调用时，脚本会报错。

如何进行性能分析

请查看 profile_modular_kernel.py。该脚本可用于为任何兼容的 FusedMoEPrepareAndFinalizeModular 和 FusedMoEExpertsModular 类型生成单个 FusedMoEModularKernel::forward() 调用的 Torch 轨迹。示例：python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel --pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts

FusedMoEPrepareAndFinalizeModular 实现

查看 融合 MoE 内核特性以获取所有可用的模块化 prepare 和 finalize 子类的列表。

FusedMoEExpertsModular

查看 融合 MoE 内核特性以获取所有可用的模块化专家内核的列表。