跳到内容

融合 MoE 模块化内核

简介

FusedMoEModularKernel 的实现位于 此处。

根据输入激活的格式，FusedMoE 实现主要分为 2 种类型。

  • 连续 / 标准 / 非批量，以及
  • 批量

注意

术语连续、标准和非批量在本文档中可互换使用。

输入激活格式完全取决于所使用的 All2All Dispatch。

  • 在连续变体中，All2All Dispatch 返回一个形状为 (M, K) 的连续张量，以及形状为 (M, num_topk) 的 TopK ID 和 TopK 权重。请参考 DeepEPHTPrepareAndFinalize 作为示例。
  • 在批量变体中，All2All Dispatch 返回一个形状为 (num_experts, max_tokens, K) 的张量。在这里，属于同一专家的激活/Token 被批量处理。请注意，并非该张量的所有条目都有效。激活张量通常伴随一个大小为 num_experts 的 expert_num_tokens 张量，其中 expert_num_tokens[i] 表示属于第 i 个专家的有效 Token 数量。请参考 PplxPrepareAndFinalize 或 DeepEPLLPrepareAndFinalize 作为示例。

FusedMoE 操作通常由多个操作组成，在连续和批量变体中都是如此，如下面的图表所示。

FusedMoE Non-Batched

FusedMoE Batched

注意

在操作方面，批量和非批量情况的主要区别在于 Permute / Unpermute 操作。所有其他操作都保留。

动机

从图表中可以看出，存在大量操作，并且每种操作都可以有多种实现方式。将操作组合起来以创建有效的 FusedMoE 实现的方式集合很快就会变得难以处理。模块化内核框架通过将操作分组到逻辑组件来解决此问题。这种广泛的分类使得组合变得可管理，并防止代码重复。这还将 All2All Dispatch & Combine 实现与 FusedMoE 实现解耦，并允许其独立开发和测试。此外，模块化内核框架为不同组件引入了抽象类，从而为未来的实现提供了明确的骨架。

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

模块化内核组件

FusedMoEModularKernel 将 FusedMoE 操作分为 3 部分：

  1. TopKWeightAndReduce
  2. FusedMoEPrepareAndFinalize
  3. FusedMoEPermuteExpertsUnpermute

TopKWeightAndReduce

TopK 权重应用和归约组件发生在 Unpermute 操作之后和 All2All Combine 之前。请注意，FusedMoEPermuteExpertsUnpermute 负责 Unpermute，FusedMoEPrepareAndFinalize 负责 All2All Combine。在 FusedMoEPermuteExpertsUnpermute 中执行 TopK 权重应用和归约是有益的。但有些实现选择在 FusedMoEPrepareAndFinalize 中执行。为了实现这种灵活性，我们有一个 TopKWeightAndReduce 抽象类。

可以在 此处找到 TopKWeightAndReduce 的实现。

FusedMoEPrepareAndFinalize::finalize() 方法接受一个 TopKWeightAndReduce 参数，该参数在方法内部调用。FusedMoEModularKernel 作为 FusedMoEPermuteExpertsUnpermute 和 FusedMoEPerpareAndFinalize 实现之间的桥梁，以确定 TopK 权重应用和归约发生的位置。

  • 如果 FusedMoEPermuteExpertsUnpermute 实现执行了权重应用和归约，则 FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl 方法将返回 TopKWeightAndReduceNoOp。
  • 如果 FusedMoEPermuteExpertsUnpermute 实现需要 FusedMoEPrepareAndFinalize::finalize() 来执行权重应用和归约，则 FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl 方法将返回 TopKWeightAndReduceContiguous / TopKWeightAndReduceNaiveBatched / TopKWeightAndReduceDelegate。

FusedMoEPrepareAndFinalize

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

FusedMoEPrepareAndFinalize Blocks

FusedMoEPermuteExpertsUnpermute

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

  • 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() 中分配工作空间张量和输出张量，并将它们传递给 FusedMoEPermuteExpertsUnpermute::apply() 方法。然后，工作空间可用作 FusedMoE 实现中的中间缓冲区。

finalize_weight_and_reduce_impl()

有时在 FusedMoEPermuteExpertsUnpermute::apply() 中执行 TopK 权重应用和归约会更有效。在此处找到一个示例 此处。我们有一个 TopKWeightAndReduce 抽象类来促进此类实现。请参阅 TopKWeightAndReduce 部分。FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl() 返回实现希望 FusedMoEPrepareAndFinalize::finalize() 使用的 TopKWeightAndReduce 对象。

FusedMoEPermuteExpertsUnpermute Blocks

FusedMoEModularKernel

FusedMoEModularKernel 由 FusedMoEPrepareAndFinalize 和 FusedMoEPermuteExpertsUnpermute 对象组成。FusedMoEModularKernel 伪代码/草图：

class FusedMoEModularKernel:
    def __init__(self,
                 prepare_finalize: FusedMoEPrepareAndFinalize,
                 fused_experts: FusedMoEPermuteExpertsUnpermute):

        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

如何操作

如何添加 FusedMoEPrepareAndFinalize 类型

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

  • PplxPrepareAndFinalize 类型由 Pplx All2All 内核支持；
  • DeepEPHTPrepareAndFinalize 类型由 DeepEP 高吞吐量 All2All 内核支持；
  • DeepEPLLPrepareAndFinalize 类型由 DeepEP 低延迟 All2All 内核支持。

步骤 1：添加 All2All 管理器

All2All Manager 的目的是设置 All2All 内核实现。FusedMoEPrepareAndFinalize 实现通常从 All2All Manager 获取内核实现“句柄”来调用 Dispatch 和 Combine 函数。请查看 All2All Manager 的实现 此处。

步骤 2：添加 FusedMoEPrepareAndFinalize 类型

本节介绍了 FusedMoEPrepareAndFinalize 抽象类公开的各种函数的重要性。

FusedMoEPrepareAndFinalize::prepare()：prepare 方法实现量化和 All2All Dispatch。通常会调用相关 All2All Manager 的 Dispatch 函数。

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

FusedMoEPrepareAndFinalize::prepare_no_receive()：prepare_no_receive 方法实现量化和 All2All Dispatch。它不会等待 dispatch 操作的结果，而是返回一个可用于等待最终结果的 thunk。通常会调用相关 All2All Manager 的 Dispatch 函数。

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

FusedMoEPrepareAndFinalize::activation_format()：如果 prepare 方法的输出（即 All2All dispatch）是 Batch 型，则返回 FusedMoEActivationFormat.BatchedExperts。否则返回 FusedMoEActivationFormat.Standard。

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

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

FusedMoEPrepareAndFinalize::num_dispatchers()：总分派单元数。此值决定了 Dispatch 输出的大小。Dispatch 输出的形状为 (num_local_experts, max_num_tokens, K)。其中 max_num_tokens = num_dispatchers() * max_num_tokens_per_rank()。

我们建议选择一个已经存在的 FusedMoEPrepareAndFinalize 实现，该实现与您的 All2All 实现非常相似，并将其作为参考。

如何添加 FusedMoEPermuteExpertsUnpermute 类型

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

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

FusedMoEPermuteExpertsUnpermute::supports_chunking()：如果实现支持分块，则返回 True。通常，输入 FusedMoEActivationFormat.Standard 的实现支持分块，而 FusedMoEActivationFormat.BatchedExperts 不支持。

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

FusedMoEPermuteExpertsUnpermute::workspace_shapes() / FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl / FusedMoEPermuteExpertsUnpermute::apply：参考上面的 FusedMoEPermuteExpertsUnpermute 部分。

FusedMoEModularKernel 初始化

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

  • maybe_make_prepare_finalize,
  • select_gemm_impl，和
  • init_prepare_finalize

maybe_make_prepare_finalize

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

  • ModelOptNvFp4FusedMoE

select_gemm_impl

select_gemm_impl 方法在基类中是未定义的。派生类负责实现一个构造有效/适当的 FusedMoEPermuteExpertsUnpermute 对象的方法。请参考以下实现：

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

init_prepare_finalize

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

请查看 init_prepare_finalize。重要提示：FusedMoEMethodBase 派生类在其 apply 方法中使用 FusedMoEMethodBase::fused_experts 对象。当设置允许构建有效的 FusedMoEModularKernel 对象时，我们会将其覆盖 FusedMoEMethodBase::fused_experts。这基本上使派生类不关心使用哪个 FusedMoE 实现。

如何进行单元测试

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

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

  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(), Config::is_fe_supports_chunking() 方法。

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

如何检查 FusedMoEPrepareAndFinalize & FusedMoEPermuteExpertsUnpermute 兼容性

单元测试文件 test_modular_kernel_combinations.py 也可以作为独立脚本执行。例如：python3 -m tests.kernels.moe.test_modular_kernel_combinations --pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts 作为副作用，此脚本可用于测试 FusedMoEPrepareAndFinalize & FusedMoEPermuteExpertsUnpermute 的兼容性。当使用不兼容的类型调用时，脚本将出错。

如何进行性能剖析

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

FusedMoEPrepareAndFinalize 实现

请参阅 Fused MoE Kernel features 以获取所有可用模块化 prepare 和 finalize 子类的列表。

FusedMoEPermuteExpertsUnpermute

请参阅 Fused MoE Kernel features 以获取所有可用模块化专家的列表。