融合MoE模块化内核¶

简介¶

FusedMoEModularKernel 的实现 在这里

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

• 连续/标准/非批处理，以及
• 批处理

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

• 在Contiguous变体中，All2All调度器将激活值作为形状为(M, K)的连续张量返回，同时返回形状为(M, num_topk)的TopK Ids和TopK权重。示例可参考DeepEPHTPrepareAndFinalize。
• 在批处理变体中，All2All调度器将激活值以形状为(num_experts, max_tokens, K)的张量形式返回。这里，订阅同一专家的激活值/令牌会被批量处理在一起。注意，并非张量的所有条目都是有效的。激活张量通常会附带一个大小为num_experts的expert_num_tokens张量，其中expert_num_tokens[i]表示订阅第i个专家的有效令牌数量。可参考PplxPrepareAndFinalize或DeepEPLLPrepareAndFinalize查看示例。

FusedMoE操作通常由多个操作组成，包括连续(Contiguous)和批处理(Batched)两种变体，如下图所示

动机¶

从图中可以看出，存在大量运算操作，且每个操作可能有多种实现方式。将这些操作组合成有效FusedMoE实现的方式集合很快就会变得难以处理。模块化内核框架通过将操作分组为逻辑组件来解决这个问题。这种宽泛的分类使组合变得可控，并避免了代码重复。这也将All2All调度与组合实现从FusedMoE实现中解耦，允许它们独立开发和测试。此外，模块化内核框架为不同组件引入了抽象类，从而为未来实现提供了定义良好的框架结构。

文档的其余部分将重点讨论连续/非批处理的情况。推广到批处理的情况应该是直截了当的。

ModularKernel 组件¶

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

1. TopK权重归约
2. FusedMoEPrepareAndFinalize
3. FusedMoEPermuteExpertsUnpermute

TopKWeightAndReduce¶

TopK权重应用与归约组件发生在Unpermute操作之后、All2All组合之前。请注意FusedMoEPermuteExpertsUnpermute负责Unpermute操作，而FusedMoEPrepareAndFinalize负责All2All组合。在FusedMoEPermuteExpertsUnpermute中执行TopK权重应用与归约是有价值的，但有些实现选择在FusedMoEPrepareAndFinalize中完成。为了实现这种灵活性，我们提供了TopKWeightAndReduce抽象类。

请在此处查看TopKWeightAndReduce的实现 here。

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

• FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl 方法返回 TopKWeightAndReduceNoOp，如果 FusedMoEPermuteExpertsUnpermute 实现自身完成了权重应用和归约操作。
• FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl 方法返回 TopKWeightAndReduceContiguous / TopKWeightAndReduceNaiveBatched / TopKWeightAndReduceDelegate，当 FusedMoEPermuteExpertsUnpermute 实现需要 FusedMoEPrepareAndFinalize::finalize() 来完成权重应用和归约操作时。

FusedMoEPrepareAndFinalize¶

FusedMoEPrepareAndFinalize抽象类公开了prepare和finalize函数。prepare函数负责输入激活量化和All2All分发。finalize函数负责调用All2All合并。此外，finalize函数可能会执行TopK权重应用和归约操作（请参阅TopKWeightAndReduce部分）

FusedMoEPermuteExpertsUnpermute¶

FusedMoEPermuteExpertsUnpermute 类是MoE操作的核心实现所在。该抽象类提供了几个重要函数：

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

apply()¶

apply 方法是实现执行的地方

• 置换
• 与权重W1进行矩阵乘法
• 激活与乘法
• 量化
• 与权重W2进行矩阵乘法
• 取消置换
• 可能应用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对象。

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管理器的目的是设置All2All内核实现。FusedMoEPrepareAndFinalize实现通常从All2All管理器获取内核实现的"句柄"来调用Dispatch和Combine函数。请查看All2All管理器实现这里。

步骤2：添加FusedMoEPrepareAndFinalize类型¶

本节介绍FusedMoEPrepareAndFinalize抽象类所公开的各种函数的作用。

FusedMoEPrepareAndFinalize::prepare(): prepare方法实现了量化(Quantization)和All2All分发(Dispatch)功能。通常这里会调用对应All2All管理器中的分发函数。

FusedMoEPrepareAndFinalize::finalize(): 可能执行TopK权重应用与归约以及All2All合并操作。通常会调用相关All2AllManager中的Combine函数。

FusedMoEPrepareAndFinalize::activation_format(): 如果prepare方法(即All2All分发)的输出是批处理形式，则返回FusedMoEActivationFormat.BatchedExperts。否则返回FusedMoEActivationFormat.Standard。

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

FusedMoEPrepareAndFinalize::max_num_tokens_per_rank(): 这是单次提交至All2All调度的最大令牌数量。

FusedMoEPrepareAndFinalize::num_dispatchers(): 调度单元的总数。该值决定了Dispatch输出的尺寸。Dispatch输出的形状为(num_local_experts, max_num_tokens, K)。其中max_num_tokens = num_dispatchers() * max_num_tokens_per_rank()。

我们建议选择一个与您的All2All实现最匹配的现有FusedMoEPrepareAndFinalize实现作为参考。

如何添加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 类包含2个方法，它们共同负责创建 FusedMoEModularKernel 对象。具体如下：

• select_gemm_impl，以及
• init_prepare_finalize

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和FusedMoEPremuteExpertsUnpermute类型的组合，如果它们兼容，则运行一些正确性测试。如果您正在添加一些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跟踪。示例：python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel --pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts

FusedMoEPrepareAndFinalize 实现¶

下表列出了撰写本文时的FusedMoEPrepareAndFinalize实现方案，

FusedMoEPermuteExpertsUnpermute¶

下表列出了撰写本文时的FusedMoEPermuteExpertsUnpermute实现方案，