跳到內容

融合 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 以獲取所有可用模組化專家的列表。