目录

abstract:

0 引言

1 文件整体介绍

2 Monolithic 和 Modular的区别

2.1 Modular：FusedMoEKernelModularImpl

2.2 Monolithic：FusedMoEKernelMonolithicImpl

2.3 个人理解以及疑问

3 topk_ids, topk_weights的由来以及内容

4 class TopKWeightAndReduce(ABC):

5 PrepareResultType与ReceiverType

6 class FusedMoEPrepareAndFinalize(ABC):

6.1 def post_init_setup(self, fused_experts: "FusedMoEPermuteExpertsUnpermute"):

6.2 def activation_format(self) -> FusedMoEActivationFormat:

6.3 def topk_indices_dtype(self) -> torch.dtype | None:

6.4 def max_num_tokens_per_rank(self) -> int | None:

7 class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize):

7.1     def prepare(

7.2 def prepare_async(

7.3     def finalize(

7.3.1 output: torch.Tensor

7.3.2 fused_expert_output: torch.Tensor

7.3.3 topk_weights: torch.Tensor

7.3.4. topk_ids: torch.Tensor

7.3.5 apply_router_weight_on_input: bool

7.3.6. weight_and_reduce_impl: TopKWeightAndReduce

8 class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalize):

9 class FusedMoEExperts(ABC):

9.1     def is_supported_config(

10 class FusedMoEExpertsModular(FusedMoEExperts):

10.1 def moe_problem_size(

10.2  def workspace_shapes(

10.3    def apply(

11 class FusedMoEKernelModularImpl:

11.1 def _allocate_buffers(

11.2 def _maybe_apply_shared_experts(

11.3 def _prepare(

11.4 _fused_experts(

11.5 def _finalize(

11.6 def apply(

12 class FusedMoEKernel:

参考文献


abstract:

modular_kernel.py

modular  英/ ˈmɒdjələ(r) /  组合式的；模块化的；模数的；有标准组件的。

monolithic 英/ ˌmɒnəˈlɪθɪk /  整体的；巨石的，庞大的；完全统一的, 单块集成电路，单片电路。

router_logits是由线性层矩阵乘法得到，然后再做softmax和topk的操作，

topk_ids [M, topk]：每个 token 选哪些专家（整数 ID）
topk_weights [M, topk]：topk专家权重（浮点，用来加权求和）

tuple 英/ tjʊpəl; ˈtʌpəl / 元组不可变

 @property函数调用不用加括号。

indices   英/ ˈɪndɪsiːz /  目录（index的复数）

expert_map,全局专家到本地专家的映射

expected_m 每个专家要处理的期望token数量，主要是给 GEMM 库（DeepGEMM 这类 masked grouped GEMM）当「规模提示」，用来做 kernel 选型/调度优化，不是用来决定算哪些 token。

token数M K ，权重E N K的转置

Standard、BatchedExperts

调用普通实例函数自动传实例；调用@classmethod 自动传类；调用@staticmethod函数什么都不传；

  • 纯 DBO（无共享专家）：发起 combine（异步，后台跑）---立刻 register hook + dbo_yield() → 让给 batch B---receiver() 等 combine 完成
  • DBO + 共享专家 overlap：发起 combine（异步，后台跑）---算共享专家（前台，和 combine 并行） 多出来的这一步---register hook + dbo_yield()  ，才让给 batch B---receiver() 等 combine 完成

本来 combine 一发起就让给 B；现在有共享专家时，先不让，在 combine 还在后台跑的时候做共享专家，共享专家算完再让出给 B。

0 引言

之前基于vllm015读过这段代码，但是那时候是刚开始第一遍读，有些理解的不太对，另一个就是新版本代码正好有了一些改变，这次再详细的过一遍。

1 文件整体介绍

#
# This file defines a set of base classes used to make MoE kernels more modular.
# The goal is to be able to utilize different communication mechanisms with
# any fused MoE kernel without needing to have combinatoric implementations.
#
# The fused moe kernels are broken down into the following components:
#
# [Router] → [Quantize-Dispatch] → [Permute-Experts-Unpermute] → [Combine]
#
# Each component will be independent of (but may inform) the others except for
# [Quantize-Dispatch] and `[Combine] (see below). The components can then be
# mixed and matched with so that DP+EP can be supported easily for multiple
# MoE kernel implementations.
#
# The following main classes are defined:
# * FusedMoEPrepareAndFinalizeModular - an abstract base class for preparation of MoE
#   inputs (e.g. quantization, distribution) and finalization of Moe outputs.
#   The prepare method must take care of any needed quantization and the
#   finalize method, informed by the FusedMoEExpertsModular method,
#   may apply weights and/or do the final reduction of the output.
# * FusedMoEExpertsModular - an abstract base class for the main fused
#   MoE operation, i.e matmul + act_mul + optionally quant + matmul.
#   Some FusedMoEExpertsModular implementations may choose to do
#   the weight application and/or reduction. The class communicates this
#   to [Finalize] via a TopKWeightAndReduce object.
# * FusedMoEModularKernel - an interface class that combines a
#   FusedMoEPrepareAndFinalizeModular and a FusedMoEExpertsModular to
#   provide the standard fused MoE kernel interface.
# * TopKWeightAndReduce - A TopKWeightAndReduce implementation chosen
#   by the FusedMoEExpertsModular implementation that is passed
#   on to [Finalize].
#
# [Quantize-Prepare] and [Finalize] functionality are bundled into a single
# class `FusedMoEPrepareAndFinalizeModular` since they could use collective
# communication mechanisms that need to be consistent.
#

个人理解：

其实这里就是说，这个modular_kernel.py是为了让moe这块代码模块化，这样后面要是比如替换了什么通信库，或者替换了专家计算的kernel，那么这块代码不用修改，也就是模块化，减少耦合，然后让不同的通信库和计算库可以任何组装，然后这个文件定义了几个基类，

然后这个文件主要是下面的这三个类

  • FusedMoEPrepareAndFinalizeModular：这个类就是负责前处理和后处理的，其实就是负责dispatch和combine的；
  • FusedMoEExpertsModular：这个是负责那两次专家计算以及激活的；
  • FusedMoEModularKernel：这个是对外的，也就是用户只需要用这个类就可以了，然后这个类里面其实就包含了上面说的两个类。

但是其实看了后面代码会发现，这个代码注释是错误的，老的，其实真正对外开放的类是FusedMoEKernel,然后这个类里面是分了两种实例，

  • FusedMoEKernelModularImpl
  • FusedMoEKernelMonolithicImpl
       # Initialize the implementation (monolithic or modular).
        self.impl: FusedMoEKernelModularImpl | FusedMoEKernelMonolithicImpl
        if isinstance(
            prepare_finalize, FusedMoEPrepareAndFinalizeModular
        ) and isinstance(fused_experts, FusedMoEExpertsModular):
            self.impl = FusedMoEKernelModularImpl(
                prepare_finalize,
                fused_experts,
                shared_experts,
                inplace,
                N,
                K,
            )

        elif isinstance(
            prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic
        ) and isinstance(fused_experts, FusedMoEExpertsMonolithic):
            assert not inplace
            self.impl = FusedMoEKernelMonolithicImpl(
                prepare_finalize,
                fused_experts,
            )

2 Monolithic 和 Modular的区别

Modular Monolithic

路由输入

已经算好的 topk_ids + topk_weights

原始的 router_logits

Router 谁做

框架里先 select_experts，再进 kernel

往往 fuse 在 kernel 里面（或整条链路按 logits 走）

典型场景

Triton、DeepEP、大多数通用 MoE

FlashInfer TRTLLM 等 router+experts 一体 的 kernel

对外接口

FusedMoEKernel.apply(..., topk_ids, topk_weights)

FusedMoEKernel.apply_monolithic(..., router_logits)

2.1 Modular：FusedMoEKernelModularImpl

上层 Router 先算 topk
    ↓
prepare(topk_ids, topk_weights)     # dispatch / 量化
    ↓
experts.apply(topk_ids, topk_weights)   # w13 + 激活 + w2
    ↓
finalize(..., TopKWeightAndReduce)  # combine + 加权 + reduce

2.2 Monolithic：FusedMoEKernelMonolithicImpl

整条链路按「router logits 还没拆成 topk」来设计：

        # TODO(rob): add inplace support.
        a1q, a1q_scale, router_logits = self.prepare_finalize.prepare(
            hidden_states,
            router_logits=router_logits,
            quant_config=self.fused_experts.quant_config,
            defer_input_quant=self.fused_experts.expects_unquantized_inputs,
        )

        fused_out = self.fused_experts.apply(
            hidden_states=a1q,
            w1=w1,
            w2=w2,
            router_logits=router_logits,
            activation=activation,
            global_num_experts=global_num_experts,
            expert_map=expert_map,
            apply_router_weight_on_input=apply_router_weight_on_input,
            a1q_scale=a1q_scale,
            # grouped topk + fused topk bias parameters
            num_expert_group=num_expert_group,
            e_score_correction_bias=e_score_correction_bias,
            routed_scaling_factor=routed_scaling_factor,
            topk_group=topk_group,
        )

        output = self.prepare_finalize.finalize(fused_out)

        return output

和 Modular 对比：

步骤 Modular Monolithic

prepare 入参

topk_ids, topk_weights

router_logits

experts 入参

topk_ids, topk_weights

router_logits

finalize 入参

output, fused_out, topk, TopKWeightAndReduce

只要 fused_out

注释里写得很直白：给「router 和 experts 已经 fuse 在一起」的 kernel 用，例如 FLASHINFER_TRTLLM。

这类 kernel 内部自己完成 routing（topk、分组 topk 等），框架不必先在外面 select_experts 再传 topk_ids。

2.3 个人理解以及疑问

Router 线性层
    ↓
router_logits  [M, num_experts]   ← 原始打分

Modular:
    select_experts(router_logits)
        → softmax/topk 等
        → topk_ids, topk_weights
    → prepare → experts → finalize

Monolithic:
    跳过 select_experts
    → 把 router_logits 和 hidden_states 一起交给 fused kernel
    → kernel 内部自己做 routing + 专家计算

router_logits是由线性层矩阵乘法得到，然后再做softmax和topk的操作，

3 topk_ids, topk_weights的由来以及内容

vllm021/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py文件中


def fused_topk(
    hidden_states: torch.Tensor,
    gating_output: torch.Tensor,
    topk: int,
    renormalize: bool,
    indices_type: torch.dtype | None = None,
    scoring_func: str = "softmax",
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    assert hidden_states.size(0) == gating_output.size(0), "Number of tokens mismatch"

    M, _ = hidden_states.size()

    topk_weights = torch.empty(
        M, topk, dtype=torch.float32, device=hidden_states.device
    )
    topk_ids = torch.empty(
        M,
        topk,
        dtype=torch.int32 if indices_type is None else indices_type,
        device=hidden_states.device,
    )
    token_expert_indices = torch.empty(
        M, topk, dtype=torch.int32, device=hidden_states.device
    )

    if scoring_func == "softmax":
        topk_func = dispatch_topk_softmax_func(
            use_rocm_aiter=rocm_aiter_ops.is_fused_moe_enabled()
        )
        topk_weights, topk_ids = topk_func(
            topk_weights, topk_ids, token_expert_indices, gating_output, renormalize
        )

        return topk_weights, topk_ids, token_expert_indices
    elif scoring_func == "sigmoid":
        topk_func = dispatch_topk_sigmoid_func(
            use_rocm_aiter=rocm_aiter_ops.is_fused_moe_enabled()
        )
        topk_weights, topk_ids = topk_func(
            topk_weights, topk_ids, token_expert_indices, gating_output, renormalize
        )

        return topk_weights, topk_ids, token_expert_indices
    else:
        raise ValueError(f"Unsupported scoring function: {scoring_func}")

router_logits [M, num_experts]
        ↓  softmax/sigmoid + 选 top-k
topk_ids      [M, topk]     选中了哪几个专家
topk_weights  [M, topk]     这几个专家各多大权重

4 class TopKWeightAndReduce(ABC):

class TopKWeightAndReduce(ABC):
    """
    An abstract base class for weight application and reduction implementations.
    """

    @abstractmethod
    def apply(
        self,
        output: torch.Tensor | None,
        fused_expert_output: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        apply_router_weight_on_input: bool,
    ) -> torch.Tensor:
        """
        Apply topk_weights to the fused_experts_outputs and/or reduce.
        If an output tensor is not passed, it will be created in the
        function.
        """
        raise NotImplementedError

这是个基类，然后用来计算topk加权求和的，

MoE 里每个 token 会走 topk 个专家，专家输出形状一般是

fused_expert_output: [M, topk, K]

最后要变成：

output: [M, K]

需要：乘 topk_weights + 在 topk 维上求和。

但不同实现差别很大：

  • 有的 experts 里已经乘权重、已经 reduce 了
  • 有的 只算未加权结果，要等 finalize/combine 再做
  • Batched 布局和 Standard 布局 reduce 方式也不同

所以抽象出 TopKWeightAndReduce，由 Experts 实现选一个策略，交给 Finalize 用。

参数 含义

fused_expert_output

专家输出，通常 [M, topk, K]，可能还没加权/归约

topk_weights / topk_ids

路由权重和专家 ID

apply_router_weight_on_input

权重是否已在输入上乘过

output

写入的最终 [M, K]，可原地写

5 PrepareResultType与ReceiverType

#
# PrepareResultType is a tuple of:
# - quantized + dispatched a.
# - quantized + dispatched a1_scales.
# - Optional ExpertTokensMetadata containing gpu/cpu tensors
#   as big as the number of local experts with the information about the
#   number of tokens assigned to each local expert.
# - Optional dispatched expert topk IDs
# - Optional dispatched expert topk weight
#
# See `prepare` method below.
#
PrepareResultType = tuple[
    torch.Tensor,
    torch.Tensor | None,
    ExpertTokensMetadata | None,
    torch.Tensor | None,
    torch.Tensor | None,
]

ReceiverType = Callable[[], PrepareResultType]

这里ReceiverType = Callable[[], PrepareResultType]的意思就是定义了一个可调用类似，他的参数是空的，返回值类型是PrepareResultType。然后这个PrepareResultType是个元组，包含下面五个。

位置 你的理解 补充一句

0

激活

dispatch（+ 可能量化）后的 a1，给 experts 算 GEMM 用

1

激活的量化参数

a1_scales，没量化时为 None

2

每个专家分到多少 token

ExpertTokensMetadata，里面是长度 E 的 expert_num_tokens，不是 topk 表本身，就是说每个专家分到的token数量。

3

每个 token 选哪些专家

dispatch 之后 的 topk_ids（可能和 dispatch 前对齐方式不同），有的路径为 None

4

topk 专家权重

dispatch 之后 的 topk_weights，同样可能为 None

6 class FusedMoEPrepareAndFinalize(ABC):

这个是基类，然后下面的两个类，都会继承这个类。

  • class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize):
  • class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalize):

6.1 def post_init_setup(self, fused_experts: "FusedMoEPermuteExpertsUnpermute"):

    def post_init_setup(self, fused_experts: "FusedMoEPermuteExpertsUnpermute"):
        """
        Initialize FusedMoEPrepareAndFinalize settings that depend on
        FusedMoEPermuteExpertsUnpermute experts object.
        The FusedMoEPrepareAndFinalize implementations that have such
        dependencies may choose to override this function.
        """
        return

这个函数是后初始化函数，为什么有个post，是因为这里的FusedMoEPermuteExpertsUnpermute类和FusedMoEPrepareAndFinalize类是分别单独创建的，所以这里相当于等FusedMoEPermuteExpertsUnpermute创建完之后，才能知道他又哪些功能，然后再根据他去对FusedMoEPrepareAndFinalize类进行相应的配置。

6.2 def activation_format(self) -> FusedMoEActivationFormat:


    @property
    @abstractmethod
    def activation_format(self) -> FusedMoEActivationFormat:
        """
        A property indicating the output format of the activations for the
        'prepare' method.
        """
        raise NotImplementedError

这个函数就是获取dispatch之后激活的张量布局，MoE 的 prepare() 会把 token 量化并 dispatch 到各 EP rank。dispatch 之后，activation 有两种常见排布：

枚举值 形状 直观含义

Standard

[M, K]

所有 token 打平成一维，按 token 顺序排

BatchedExperts

[E, M, K]

按 local expert 分批，每个 expert 一行 buffer

6.3 def topk_indices_dtype(self) -> torch.dtype | None:


    @abstractmethod
    def topk_indices_dtype(self) -> torch.dtype | None:
        """
        The PrepareFinalize All2All implementations generally constrain the
        dtype of the topk_ids they support. This function returns the
        required topk indices dtype so it can be respected.
        Return None if there are no such restrictions.
        """
        raise NotImplementedError

这个函数其实就是返回topk_ids的类型。

6.4 def max_num_tokens_per_rank(self) -> int | None:

    @abstractmethod
    def max_num_tokens_per_rank(self) -> int | None:
        """
        Some PrepareFinalize All2All implementations are batched. Meaning,
        they can process only as set of tokens at a time. This
        function returns the batch size i.e the maximum number of tokens
        the implementation can process at a time.
        Return None if there are no such restrictions.
        """
        raise NotImplementedError

这个就是每个rank上处理的token上限，这个是低延迟的要用，

7 class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize):

7.1     def prepare(

    @abstractmethod
    def prepare(
        self,
        a1: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        num_experts: int,
        expert_map: torch.Tensor | None,
        apply_router_weight_on_input: bool,
        quant_config: FusedMoEQuantConfig,
    ) -> PrepareResultType:
        """
        Perform any quantization (and/or) dispatching needed for this kernel.
        - a1: The (unquantized) input to the MoE layer.
        - topk_ids: The topk ids.
        - topk_weights: The topk weights.
        - num_experts: The total number of experts in the global expert space.
        - expert_map: A tensor mapping expert indices from the global expert
          space to the local expert space of the expert parallel shard.
        - apply_router_weight_on_input: When True, apply the weights to the
          activations, before quantization + dispatching.
        - quant_config: Quantization info provided by the fused experts.

        Returns a tuple of:
        - quantized + dispatched a.
        - Optional quantized + dispatched a1_scales.
        - Optional ExpertTokensMetadata containing gpu/cpu tensors
          as big as the number of local experts with the information about the
          number of tokens assigned to each local expert.
        - Optional dispatched expert topk IDs
        - Optional dispatched expert topk weight
        """
        raise NotImplementedError
参数 类型 含义

a1

Tensor

MoE 层输入激活，一般是 [M, hidden]（M 个 token，hidden 维）。注释写 unquantized，是否先量化由实现和 quant_config 决定。

topk_weights

Tensor

[M, topk]，每个 token 对选中专家的权重（router 之后）。

topk_ids

Tensor

[M, topk]，每个 token 选中的专家 ID（全局专家空间）。dispatch 靠它决定 token 发往哪。

num_experts

int

全局专家总数（所有 rank 上的专家数，不是本 rank 本地专家数）。

expert_map

Tensor | None

EP 时用：全局专家 ID → 本 rank 本地专家 ID 的映射；不在本 rank 的专家会映到 -1 等。无 EP 时常为 None。

apply_router_weight_on_input

bool

是否在 dispatch 之前 把 topk_weights 乘到 a1 上。一般为 False；为 True 时常要求 topk=1。

quant_config

FusedMoEQuantConfig

量化配置（FP8/INT8、scale、是否 per-token 等），由 Experts 侧提供，Prepare 按此决定何时、如何量化。

这里面唯一一个注意的是， - expert_map: A tensor mapping expert indices from the global expert space to the local expert space of the expert parallel shard.这个是全局专家到本地专家的映射，什么意思，这是个一维的，然后它的长度就是全局专家个数，某个位置用0 1 2 3表示全局专家在当前rank的编号，某个专家-1表示这个专家不在当前rank，所以这个expert_map是每个rank都有一份。举例子，[0,  -1, 1, 2, -1, 3, 4 ,5,]

序号 返回值 含义

0

Tensor

量化 + dispatch 后的激活。本 rank 上要算的专家输入；形状可能是 [M', K]（Standard）或 [E, max_M, K]（BatchedExperts）。

1

Tensor | None

激活的 scale（a1_scales）。量化时有；未量化或未在此步量化时为 None。

2

ExpertTokensMetadata | None

每个本地专家分到多少 token（长度 E 的计数），不是「具体哪几个 token」。Batched / DeepEP 路径常见；简单路径可为 None。

3

Tensor | None

dispatch 后的 topk_ids。EP 后 token 顺序可能变，需要和激活对齐的 topk；有的实现为 None，Experts 仍用外面的 topk。

4

Tensor | None

dispatch 后的 topk_weights，与第 3 项同理。

7.2 def prepare_async(

    def prepare_async(
        self,
        a1: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        num_experts: int,
        expert_map: torch.Tensor | None,
        apply_router_weight_on_input: bool,
        quant_config: FusedMoEQuantConfig,
    ) -> tuple[Callable, ReceiverType] | ReceiverType:
        """
        Perform any quantization (and/or) dispatching needed for this kernel
        but do not wait for results from other workers.
        - a1: The (unquantized) input to the MoE layer.
        - a1_scale: Optional scales for a1
        - a2_scale: Optional scales for the second MoE gemm.  Required to make
          sure the quantization is consistent for both gemms.
        - topk_ids: The topk ids.
        - topk_weights: The topk weights.
        - num_experts: The total number of experts in the global expert space.
        - expert_map: A tensor mapping expert indices from the global expert
          space to the local expert space of the expert parallel shard.
        - apply_router_weight_on_input: When True, apply the weights to the
          activations, before quantization + dispatching.

        Returns a callback or a hook callback pair that when invoked waits for
        results from other workers and has the same return signature as
        `prepare`, if a hook is returned this is more lightweight check that
        the recv is complete without doing extra work (used by DBO, will be
        refactored in the very near future)

        e.g.

        ret = obj.prepare_async(...)

        if isinstance(ret, tuple):
            hook, receiver = ret
            hook()

        if hook is not None:
        a, a_scales, expert_meta, topk_ids, topk_weights = receiver()

        is equivalent to:

        a, a_scales, expert_meta, topk_ids, topk_weights = obj.prepare(...)
        """
        raise NotImplementedError

这个就是异步的prepare函数，其他的没啥。

7.3     def finalize(

    @abstractmethod
    def finalize(
        self,
        output: torch.Tensor,
        fused_expert_output: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        apply_router_weight_on_input: bool,
        weight_and_reduce_impl: TopKWeightAndReduce,
    ) -> None:
        """
        Perform any combine plus apply weights and perform a reduction on the
        fused experts output.
        - output: The output tensor, written in place.  Must be (M, K) shape.
        - fused_expert_output: The unweighted, unreduced output of the fused
          experts, it will have (M, topk, K) shape.
        - topk_weights: The weights to be applied to the fused_experts_output.
        - topk_ids: The topk_ids.
        - apply_router_weight_on_input: When False, apply the weights to
          fused_expert_output.
        - weight_and_reduce_impl: An optional TopKWeightAndReduce
          implementation.
        """
        raise NotImplementedError
  1. All2All Combine：把各 rank 上算完的 expert 输出，按 token 路由 送回原 rank
  2. 乘 topk_weights：每个 token 的 topk 路输出按路由权重加权
  3. Reduce：把 topk 路合成一路，[M, topk, K] → [M, K]

7.3.1 output: torch.Tensor

MoE 层的最终输出 buffer，原地写入。

属性 说明

形状

[M, K]，M = token 数，K = hidden_dim

谁分配

FusedMoEKernelModularImpl.apply() 里 empty_like(hidden_states) 或 inplace 复用 hidden_states

谁写入

finalize() 或其调用的 combine / weight_and_reduce

DeepEP LL 里直接传给 combine kernel：

self.buffer.low_latency_combine(..., out=output)

无 EP 时由 TopKWeightAndReduceContiguous 写入：

ops.moe_sum(fused_expert_output, output) # 结果进 output

7.3.2 fused_expert_output: torch.Tensor

Experts 算完、尚未 combine / 加权 / reduce 的原始输出。

属性 说明

文档形状

[M, topk, K]

含义

每个 token 的 topk 路 expert 输出，还没乘权重、没把 topk 维加起来

来源

_fused_experts() / FusedMoEExpertsModular.apply() 的返回值

数学上大致是：

fused_expert_output[m, i, :] = 第 m 个 token 的第 i 个选中 expert 的输出向量

DeepEP combine 会拿它和 topk_ids、topk_weights 一起做 combine + 加权 + 聚合，直接写入 output。

7.3.3 topk_weights: torch.Tensor

路由权重，softmax 后 topk 得到的浮点权重。

属性 说明

形状

[M, topk]

含义

每个 token 对每个选中 expert 的权重（和为 1 或经 scaling）

用途

finalize 里乘到 fused_expert_output 上（若还没乘过）

TopKWeightAndReduceContiguous 里：

fused_expert_output.mul_(topk_weights.view(m, -1, 1)) # 广播到 K 维

ops.moe_sum(fused_expert_output, output) # topk 维求和

DeepEP combine 也会把 topk_weights 传进 low_latency_combine。

7.3.4. topk_ids: torch.Tensor

每个 token 选了哪些 expert（全局或物理 id）。

属性 说明

形状

[M, topk]

含义

与 fused_expert_output 的 topk 维一一对应

用途

combine 时知道每路输出该送回哪个 token、对应哪条 expert 槽位

DeepEP LL 会先映射 id：

combine_topk_ids = self._map_global_to_physical_ids(topk_ids)

combine kernel 用 topk_ids 做 unpermute / 路由还原，把分散在各 rank 的结果拼回 [M, topk, K] 再聚合。

7.3.5 apply_router_weight_on_input: bool

路由权重是否已经乘在输入 activation 上（一般仅 topk=1 时支持）。

值 含义 finalize 行为

False（常见）

权重还没乘到输入

在 finalize 里对 fused_expert_output 乘 topk_weights

True

prepare 里已对 a1 乘过权重

finalize 不应再乘，避免乘两次

DeepEP LL 的处理：

if apply_router_weight_on_input:
    combine_topk_weights = torch.ones_like(topk_weights)  # 当作全 1
else:
    combine_topk_weights = topk_weights

TopKWeightAndReduceContiguous 同样：

if not apply_router_weight_on_input:
    fused_expert_output.mul_(topk_weights.view(m, -1, 1))

7.3.6. weight_and_reduce_impl: TopKWeightAndReduce

对于使用deepep的场景，这个weight_and_reduce_impl函数是用不到的，是在deepep的combine接口内部做的reduce

8 class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalize):

这个类先不看了，目前先用的是前面的modular那个类。

9 class FusedMoEExperts(ABC):

这个也是基类，然后下面的两个类都会继承这个类

  • class FusedMoEExpertsModular(FusedMoEExperts):
  • class FusedMoEExpertsMonolithic(FusedMoEExperts):

9.1     def is_supported_config(

这里面有个这个函数，这里的意思就是，当选择某个expert kernel的时候，只有满足了这个函数的各种条件之后，才会选择这个expert kernel类。

10 class FusedMoEExpertsModular(FusedMoEExperts):

这个类继承了刚才的FusedMoEExperts类，然后

10.1 def moe_problem_size(

    def moe_problem_size(
        self,
        a1: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        topk_ids: torch.Tensor,
    ) -> tuple[int, int, int, int, int]:
        """
        Extract the MoE problem size from the given tensor arguments:
        - a: The hidden states, input to the MoE layer.
        - w1: The first set of expert weights.
        - w2: The second set of expert weights.
        - topk_ids: The topk ids.

        Note: extracting the problem shape from the weight and activation
        tensors is not obvious.  It needs to be done this way specifically
        due to subtle issues with particular kernels, e.g. the int4 kernels
        divide the trailing dimension by two, so it's not "correct" to
        extract N or K from the trailing dimension of w1 or w2.  Similarly,
        some kernels transpose the weights, so this needs to be kept in mind.

        Note: This implementation covers most cases. However, if experts
        require a specialized implementation, like MarlinExperts, they are free
        to override this function.
        """
        if w1.dim() == 6:
            assert w2.dim() == 6
            E = w1.size(0)
            # [E, K/64, N/16, 4, 16, 16] packed layout
            N = w1.size(2) * w1.size(4)
            K = w1.size(1) * w1.size(3) * w1.size(5)
        else:
            assert len(w1.shape) == 3 and len(w2.shape) == 3
            E, N, _ = w1.shape
            K = a1.size(-1)

        if a1.dim() == 2:
            # Make sure we are using the correct a1 (pre-permute).
            assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}"
            M = a1.size(0)
        else:
            assert a1.dim() == 3
            assert a1.size(0) == E, f"{a1.size(0)} == {E}"
            M = a1.size(1)  # This is max_num_tokens

        assert topk_ids.dim() == 2
        topk = topk_ids.size(1)

        return E, M, N, K, topk

这个类其实就是计算出来了E M N K topk, 

10.2  def workspace_shapes(


    @abstractmethod
    def workspace_shapes(
        self,
        M: int,
        N: int,
        K: int,
        topk: int,
        global_num_experts: int,
        local_num_experts: int,
        expert_tokens_meta: ExpertTokensMetadata | None,
        activation: MoEActivation,
    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
        """
        Compute the shapes for the temporary and final outputs of the two gemms
        and activation in the fused expert function.  Since the gemms are
        independent, the workspace for the first gemm can be shared with the
        workspace for the last gemm.

        Inputs:
        - M: number of tokens.
        - N: Row (or column) dimension of expert weights.
        - K: hidden dimension
        - topk: The number of top-k experts to select.
        - global_num_experts: global number of experts.
        - local_num_experts: local number of experts due to DP/EP.
        - expert_tokens_meta: number of tokens per expert metadata for batched
                              format.

        Returns a tuple of:
        - workspace13 shape tuple: must be large enough to hold the
          result of either expert gemm.
        - workspace2 shape tuple: must be large enough to hold the
          result of the activation function.
        - output shape tuple: must be exact size of the final gemm output.
        - Note: workspace shapes can be 0 if the workspace is not needed.
          But in order for activation chunking to work, the first dimension
          of each tuple must be the number of tokens when the shape is
          not 0.
        """
        raise NotImplementedError

这里其实就是根据之前计算的E N K 等一些信息，计算两次gemm计算需要的空间大小，

10.3    def apply(

    @abstractmethod
    def apply(
        self,
        output: torch.Tensor,
        hidden_states: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        activation: MoEActivation,
        global_num_experts: int,
        expert_map: torch.Tensor | None,
        a1q_scale: torch.Tensor | None,
        a2_scale: torch.Tensor | None,
        workspace13: torch.Tensor,
        workspace2: torch.Tensor,
        expert_tokens_meta: ExpertTokensMetadata | None,
        apply_router_weight_on_input: bool,
    ) -> None:
        """
        This function computes the intermediate result of a Mixture of Experts
        (MoE) layer using two sets of weights, w1 and w2.

        Parameters:
        - output: (torch.Tensor): The unweighted, unreduced output tensor.
        - hidden_states: (torch.Tensor): The (quantized) input tensor to the MoE
          layer.
        - w1 (torch.Tensor): The first set of expert weights.
        - w2 (torch.Tensor): The second set of expert weights.
        - topk_weights: A map of row to expert weights. Some implementations
          choose to do weight application.
        - topk_ids (torch.Tensor): A map of row to expert id.
        - activation (str): The activation function to apply after the first
          MoE layer.
        - global_num_experts (int): The total number of experts in the global
          expert space.
        - expert_map (Optional[torch.Tensor]):  A tensor mapping expert indices
          from the global expert space to the local expert space of the expert
          parallel shard.
        - a1q_scale (Optional[torch.Tensor]): Optional quantized scale to be
          used for a1.  Result of quantization from prepare/finalize and not
          from the FusedMoEQuantConfig.
        - workspace13 (torch.Tensor): A scratch tensor used for gemm outputs
          must be large enough to hold output of either MoE gemm.
        - workspace2 (torch.Tensor): A scratch tensor used for the activation
          function.
        - expert_tokens_meta (Optional[ExpertTokensMetadata]) - An optional
          ExpertTokensMetadata object containing gpu/cpu tensors
          as big as the number of local experts with the information about the
          number of tokens assigned to each local expert.
        - apply_router_weight_on_input: True if router weights are already
          applied on the input. This is relevant if the implementation
          chooses to do weight application.
        """
        raise NotImplementedError

这个函数就是负责做那两次矩阵计算和激活的，也就是gemm1---激活----gemm2。

参数 类型 典型形状 / 格式 谁提供 作用

output

Tensor

HT: (M, K)
LL: (E, max_tokens, K)

_allocate_buffers

w2 GEMM 输出 buffer；未加权、未跨 rank reduce；结果 in-place 写入

hidden_states

Tensor

HT: (M, K)
LL: (E, max_tokens×dispatchers, K)

Prepare（a1q）

Expert 输入；可能已量化、已 dispatch 到本 rank

w1

Tensor

(E_local, N, K)

MoE 层权重

gate+up 合并权重（w13）；第一次 GEMM

w2

Tensor

(E_local, K, N_act)

MoE 层权重

down 权重；第二次 GEMM；N_act 为激活后宽度

topk_weights

Tensor

(M, topk)

Router + Prepare

每个 token 对各 expert 的 router 权重；部分实现在 apply 内用，LL 多在 finalize 用

topk_ids

Tensor

(M, topk)

Router + Prepare

每个 token 选中的 全局 expert id；HT 用于 permute/unpermute

activation

MoEActivation

枚举值

模型配置

w13 后激活类型（SiLU 等）；影响中间维是否 N//2

global_num_experts

int

标量

上层传入

全局 expert 总数；用于负载估计、路由映射等

expert_map

Tensor | None

长度 = 全局 E

EP 配置

global id → local id（无则为 -1）；HT permute 需要

a1q_scale

Tensor | None

随量化方案

Prepare 动态量化

第一次 GEMM 输入的 scale；与 hidden_states 配对

a2_scale

Tensor | None

随量化方案

quant_config.a2_scale

第二次 GEMM 输入的静态 scale；多数路径由激活后动态 quant 替代

workspace13

Tensor

见 workspace_shapes

_allocate_buffers

w13 / permute 等 scratch；需能容纳较大的一次 GEMM 输出

workspace2

Tensor

见 workspace_shapes

_allocate_buffers

激活阶段 scratch；宽度通常为 N 或 N//2

expert_tokens_meta

ExpertTokensMetadata | None

每 expert 一个计数

Prepare（DeepEP LL）

各本地 expert 实际 token 数；LL masked GEMM 必需

apply_router_weight_on_input

bool

标量

模型 / 配置

True 表示权重已在 Prepare 乘到输入；Experts 不应再乘

11 class FusedMoEKernelModularImpl:

这个就是具体的实施类，这里面会用到这两个类分别作前后处理和专家计算

  •         prepare_finalize: FusedMoEPrepareAndFinalizeModular,
  •         fused_experts: FusedMoEExpertsModular,

11.1 def _allocate_buffers(

    def _allocate_buffers(
        self,
        out_dtype: torch.dtype,
        device: torch.device,
        M_chunk: int,
        M_full: int,
        N: int,
        K: int,
        top_k: int,
        global_num_experts: int,
        local_num_experts: int,
        expert_tokens_meta: ExpertTokensMetadata | None,
        activation: MoEActivation,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Allocate temporary and output buffers for the fused experts op.
        Inputs:
        - out_dtype: output type of workspace and output tensors.
        - device: the device of the workspace and output tensors.
        See `workspace_shapes` for a description of the remainder of arguments.
        Returns a tuple of (workspace13, workspace2, output) tensors.
        """
        assert M_full > 0 and M_chunk > 0

        workspace_dtype = self.fused_experts.workspace_dtype(out_dtype)

        # Get intermediate workspace shapes based off the chunked M size.
        workspace13_shape, workspace2_shape, _ = self.fused_experts.workspace_shapes(
            M_chunk,
            N,
            K,
            top_k,
            global_num_experts,
            local_num_experts,
            expert_tokens_meta,
            activation,
        )

        # Get final output shape based on the full M size.
        _, _, fused_out_shape = self.fused_experts.workspace_shapes(
            M_full,
            N,
            K,
            top_k,
            global_num_experts,
            local_num_experts,
            expert_tokens_meta,
            activation,
        )

        # We can reuse the memory between cache1 and cache3 because by the
        # time we need cache3, we're done with cache1.
        # Reuse workspace13 for the output since there is only one chunk.
        max_shape_size = max(prod(workspace13_shape), prod(fused_out_shape))
        common_workspace, workspace2 = current_workspace_manager().get_simultaneous(
            ((max_shape_size,), workspace_dtype),
            (workspace2_shape, workspace_dtype),
        )
        workspace13 = _resize_cache(common_workspace, workspace13_shape)
        fused_out = _resize_cache(common_workspace, fused_out_shape)

        return workspace13, workspace2, fused_out

这个函数在 _fused_experts() 里、调用 fused_experts.apply() 之前执行，负责为 MoE 两次 GEMM + 激活阶段申请三块 buffer，并尽量复用同一块物理内存。

参数 含义

out_dtype

输出 / workspace 的数据类型，一般等于 hidden_states.dtype（调用时传的是 in_dtype）

device

张量所在设备，来自 a1q.device

M_chunk

中间 workspace 按多大 token 数算形状；当前实现 = M_full

M_full

最终输出 fused_out 按多大 token 数算形状；dispatch 后本 rank 实际要算的 token 数

N

专家中间维（w1 的 N，gate+up 合并后的宽度）

K

hidden dim（输入/输出特征维）

top_k

每个 token 选几个专家

global_num_experts

全局专家总数

local_num_experts

本 rank 本地专家数（通常 w1.shape[0]）

expert_tokens_meta

BatchedExperts 布局下每个专家分到的 token 数；Standard 路径常为 None

activation

激活类型（SiLU 等），影响中间 buffer 宽度（是否 N//2）

返回值（三个张量）：

返回值 用途

workspace13

第 1 次 GEMM 和第 3 次 GEMM 的共享 scratch（名字里的 1 和 3）

workspace2

激活阶段的 scratch

fused_out

专家计算的最终输出 buffer，作为 apply(output=...) 写入目标

另外，

        # We can reuse the memory between cache1 and cache3 because by the
        # time we need cache3, we're done with cache1.
        # Reuse workspace13 for the output since there is only one chunk.
        max_shape_size = max(prod(workspace13_shape), prod(fused_out_shape))
        common_workspace, workspace2 = current_workspace_manager().get_simultaneous(
            ((max_shape_size,), workspace_dtype),
            (workspace2_shape, workspace_dtype),
        )
        workspace13 = _resize_cache(common_workspace, workspace13_shape)
        fused_out = _resize_cache(common_workspace, fused_out_shape)

这里的意思，从workspace13和fused_out_shape中选出来一个最大的，按照这个最大的去申请内存，然后后面将复用做workspace13和fused_out

11.2 def _maybe_apply_shared_experts(

    def _maybe_apply_shared_experts(
        self,
        shared_experts_input: torch.Tensor | None,
        x_and_scale_quanted: tuple[torch.Tensor, torch.Tensor] | None = None,
    ):
        if self.shared_experts is not None:
            assert shared_experts_input is not None
            self.shared_experts.apply(
                shared_experts_input,
                SharedExpertsOrder.MK_INTERNAL_OVERLAPPED,
                x_and_scale_quanted=x_and_scale_quanted,
            )

这里是如果有共享专家的话，那么这里共享专家会和combin进行overlap。

11.3 def _prepare(

    def _prepare(
        self,
        hidden_states: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        global_num_experts: int,
        expert_map: torch.Tensor | None,
        apply_router_weight_on_input: bool,
        prequanted_a1_scale: torch.Tensor | None = None,
        prequanted_a1: bool = False,
    ) -> tuple[
        torch.Tensor,
        torch.Tensor | None,
        ExpertTokensMetadata | None,
        torch.Tensor,
        torch.Tensor,
    ]:
        """
        The _prepare method is a wrapper around self.prepare_finalize.prepare
        that handles DBO and async.
        """
        use_all2all_kernels = (
            self.moe_parallel_config is not None
            and self.moe_parallel_config.use_all2all_kernels
        )
        if (
            prequanted_a1
            and not use_all2all_kernels
        ):
            return hidden_states, prequanted_a1_scale, None, topk_ids, topk_weights

        if (
            self.fused_experts.activation_format()
            == FusedMoEActivationFormat.BatchedExperts
        ):
            num_tokens = hidden_states.shape[0] if hidden_states is not None else None
            num_dispatchers = getattr(self.fused_experts, "num_dispatchers", None)
            topk = topk_ids.shape[1] if topk_ids is not None else None

            if (
                num_tokens is not None
                and num_dispatchers is not None
                and topk is not None
                and global_num_experts is not None
                and global_num_experts > 0
            ):
                expected_m = (
                    num_tokens * num_dispatchers * topk + global_num_experts
                ) // global_num_experts

                self.fused_experts.set_expected_m(expected_m)
            else:
                logger.warning_once(
                    "Skip set_expected_m because required values are invalid: "
                    "num_tokens=%s, num_dispatchers=%s, topk=%s, "
                    "global_num_experts=%s",
                    num_tokens,
                    num_dispatchers,
                    topk,
                    global_num_experts,
                )

        if not self.prepare_finalize.supports_async():
            # We shouldn't be running an a2a kernel that doesn't
            # support async prepare/finalize
            # TODO(lucas): enable in follow-up
            assert not dbo_enabled()

            (
                a1q,
                a1q_scale,
                expert_tokens_meta,
                _expert_topk_ids,
                _expert_topk_weights,
            ) = self.prepare_finalize.prepare(
                hidden_states,
                topk_weights,
                topk_ids,
                global_num_experts,
                expert_map,
                apply_router_weight_on_input,
                self.fused_experts.quant_config,
                defer_input_quant=self.fused_experts.expects_unquantized_inputs,
            )
        else:
            # Overlap shared expert compute with all2all dispatch.
            dbo_maybe_run_recv_hook()
            prepare_ret = self.prepare_finalize.prepare_async(
                hidden_states,
                topk_weights,
                topk_ids,
                global_num_experts,
                expert_map,
                apply_router_weight_on_input,
                self.fused_experts.quant_config,
                defer_input_quant=self.fused_experts.expects_unquantized_inputs,
            )

            # TODO(lucas): refactor this in the alternative schedules followup
            # currently unpack if we have hook + receiver pair or just
            # receiver (see finalize_async docstring)
            hook, receiver = (
                prepare_ret if isinstance(prepare_ret, tuple) else (None, prepare_ret)
            )

            if hook is not None:
                if dbo_enabled():
                    # If DBO is being used, register the hook with the ubatch
                    # context and call it in dbo_maybe_run_recv_hook instead of
                    #  passing it to the receiver.
                    dbo_register_recv_hook(hook)
                    dbo_yield()
                else:
                    hook()

            (
                a1q,
                a1q_scale,
                expert_tokens_meta,
                _expert_topk_ids,
                _expert_topk_weights,
            ) = receiver()

        # Maybe prepare gathered topk_ids and topk_weights from other EP ranks.
        topk_ids = topk_ids if _expert_topk_ids is None else _expert_topk_ids
        topk_weights = (
            topk_weights if _expert_topk_weights is None else _expert_topk_weights
        )

        return a1q, a1q_scale, expert_tokens_meta, topk_ids, topk_weights

_prepare 是 FusedMoEKernelModularImpl 里的内部包装函数，不直接对外暴露，我看他其实就是封装了一层，本质上其实就还是调用的self.prepare_finalize.prepare(，只不过增加了一下其他的处理工作，比如DBO相关的代码，从注释这里也能看出来是这样的。

        The _prepare method is a wrapper around self.prepare_finalize.prepare

        that handles DBO and async.

这里面

                expected_m = (
                    num_tokens * num_dispatchers * topk + global_num_experts
                ) // global_num_experts

就是当前rank上有num_tokens，那么乘以num_dispatchers就是说估算下这次所有rank一共有多少个token，然后乘以topk，也就是说我们这个每个专家可能要负责多个token，这三个值乘积其实就是计算出来所有专家一共要做多少个token的计算，然后除以global_num_experts就是估算出来每个专家大约要处理多少个token,

expected_m 用来给 DeepGEMM 等 masked grouped GEMM 做 kernel 选型/调度提示，不决定实际算哪些 token。

11.4 _fused_experts(

    def _fused_experts(
        self,
        in_dtype: torch.dtype,
        a1q: torch.Tensor,
        a1q_scale: torch.Tensor | None,
        w1: torch.Tensor,
        w2: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        activation: MoEActivation,
        global_num_experts: int,
        local_num_experts: int,
        expert_map: torch.Tensor | None,
        apply_router_weight_on_input: bool,
        expert_tokens_meta: ExpertTokensMetadata | None,
        output_alias: torch.Tensor | None = None,
    ) -> torch.Tensor:
        _, M_full, N, K, top_k = self.fused_experts.moe_problem_size(
            a1q, w1, w2, topk_ids
        )
        if self.N > 0:
            N = self.N
            K = self.K

        # This happens when none of the tokens from the all2all reach this
        # EP rank. Also, note that this is only relevant for CUDAGraph
        # incompatible all2all kernels like the DeepEP high-throughput
        # kernels. CUDAGraph compatible all2all kernels like the DeepEP
        # low-latency kernels are always batched and can never run into
        # the tensor.numel() == 0 case.
        if M_full == 0:
            return torch.empty_like(a1q, dtype=in_dtype)

        workspace13, workspace2, fused_out = self._allocate_buffers(
            in_dtype,
            a1q.device,
            M_full,
            M_full,
            N,
            K,
            top_k,
            global_num_experts,
            local_num_experts,
            expert_tokens_meta,
            activation,
        )

        # If caller's output buffer already matches fused_out shape/dtype, alias
        # to skip the redundant copy in TopKWeightAndReduceNoOP.apply downstream.
        # This eliminates ~94% of __amd_rocclr_copyBuffer events (Copy 2 of the
        # double-copy MoE write-back path).
        if current_platform.is_rocm():
            from vllm._aiter_ops import rocm_aiter_ops

            if (
                rocm_aiter_ops.is_fused_moe_enabled()
                and output_alias is not None
                and output_alias.shape == fused_out.shape
                and output_alias.dtype == fused_out.dtype
                and output_alias.device == fused_out.device
                and output_alias.is_contiguous()
            ):
                fused_out = output_alias

        self.fused_experts.apply(
            output=fused_out,
            hidden_states=a1q,
            w1=w1,
            w2=w2,
            topk_weights=topk_weights,
            topk_ids=topk_ids,
            activation=activation,
            global_num_experts=global_num_experts,
            expert_map=expert_map,
            a1q_scale=a1q_scale,
            a2_scale=self.fused_experts.a2_scale,
            workspace13=workspace13,
            workspace2=workspace2,
            expert_tokens_meta=expert_tokens_meta,
            apply_router_weight_on_input=apply_router_weight_on_input,
        )

        return fused_out

这个函数其实就是做专家计算的，也是算对fused_experts.apply(的封装调用，

        # If caller's output buffer already matches fused_out shape/dtype, alias
        # to skip the redundant copy in TopKWeightAndReduceNoOP.apply downstream.
        # This eliminates ~94% of __amd_rocclr_copyBuffer events (Copy 2 of the
        # double-copy MoE write-back path).
        if current_platform.is_rocm():
            from vllm._aiter_ops import rocm_aiter_ops

            if (
                rocm_aiter_ops.is_fused_moe_enabled()
                and output_alias is not None
                and output_alias.shape == fused_out.shape
                and output_alias.dtype == fused_out.dtype
                and output_alias.device == fused_out.device
                and output_alias.is_contiguous()
            ):
                fused_out = output_alias

11.5 def _finalize(

    def _finalize(
        self,
        output: torch.Tensor,
        fused_out: torch.Tensor,
        hidden_states: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        apply_router_weight_on_input: bool,
        shared_experts_input: torch.Tensor | None,
        x_and_scale_quanted: tuple[torch.Tensor, torch.Tensor] | None = None,
    ) -> torch.Tensor:
        """
        The _finalize method is a wrapper around self.prepare_finalize.finalize
        that handles DBO, async and shared expert overlap.

        Args:
            shared_experts_input: Optional separate input for shared experts.
                When latent MoE is used, hidden_states is the latent-projected
                tensor (smaller dimension) used by routed experts, while
                shared_experts_input is the original hidden_states (full
                dimension) needed by the shared expert MLP.
        """
        if not self.prepare_finalize.supports_async():
            assert not dbo_enabled()

            self.prepare_finalize.finalize(
                output,
                fused_out,
                topk_weights,
                topk_ids,
                apply_router_weight_on_input,
                self.fused_experts.finalize_weight_and_reduce_impl(),
            )
        else:
            finalize_ret = self.prepare_finalize.finalize_async(
                output,
                fused_out,
                topk_weights,
                topk_ids,
                apply_router_weight_on_input,
                self.fused_experts.finalize_weight_and_reduce_impl(),
            )
            self._maybe_apply_shared_experts(shared_experts_input, x_and_scale_quanted)

            # TODO(lucas): refactor this in the alternative schedules followup
            # currently unpack if we have hook + receiver pair or just
            # receiver (see finalize_async docstring)
            hook, receiver = (
                finalize_ret
                if isinstance(finalize_ret, tuple)
                else (None, finalize_ret)
            )

            if hook is not None:
                if dbo_enabled():
                    # If DBO is being used, register the hook with the ubatch
                    # context and call it in dbo_maybe_run_recv_hook instead of
                    #  passing it to the receiver.
                    dbo_register_recv_hook(hook)
                    dbo_yield()
                else:
                    hook()

            receiver()

        return output

从这段注释就能看出来这个函数的作用

        The _finalize method is a wrapper around self.prepare_finalize.finalize

        that handles DBO, async and shared expert overlap.

这个函数也是个封装层，主要就是调用了finalize函数，还有就是处理了DBO以及共享专家的overlap。

这里除了DBO，其实里面还加上了共享专家计算的overlao，

  • 纯 DBO（无共享专家）：发起 combine（异步，后台跑）---立刻 register hook + dbo_yield() → 让给 batch B---receiver() 等 combine 完成
  • DBO + 共享专家 overlap：发起 combine（异步，后台跑）---算共享专家（前台，和 combine 并行） 多出来的这一步---register hook + dbo_yield()  ，才让给 batch B---receiver() 等 combine 完成

本来 combine 一发起就让给 B；现在有共享专家时，先不让，在 combine 还在后台跑的时候做共享专家，共享专家算完再让出给 B。

11.6 def apply(

    def apply(
        self,
        hidden_states: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        topk_ids: torch.Tensor,
        topk_weights: torch.Tensor,
        activation: MoEActivation = MoEActivation.SILU,
        global_num_experts: int = -1,
        expert_map: torch.Tensor | None = None,
        apply_router_weight_on_input: bool = False,
        shared_experts_input: torch.Tensor | None = None,
        quanted_hidden_states: torch.Tensor | None = None,
        scale: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """
        This function computes a Mixture of Experts (MoE) layer using two sets
        of weights, w1 and w2, and top-k gating mechanism.

        Parameters:
        - hidden_states: (torch.Tensor): The input tensor to the MoE layer.
        - w1 (torch.Tensor): The first set of expert weights.
        - w2 (torch.Tensor): The second set of expert weights.
        - topk_weights (torch.Tensor): The topk weights applied at the end of the layer.
        - topk_ids (torch.Tensor): A map of row to expert id.
        - activation (MoEActivation): The activation function to apply after the first
          MoE layer.
        - global_num_experts (int): The total number of experts in the global
          expert space.
        - expert_map (Optional[torch.Tensor]):  A tensor mapping expert indices
          from the global expert space to the local expert space of the expert
          parallel shard.
        - apply_router_weight_on_input (bool): When true, the topk weights are
          applied directly on the inputs. This is only applicable when topk is
          1.
        - shared_experts_input (Optional[torch.Tensor]): Optional separate
          input for shared experts. For latent MoE, this is the original
          hidden_states before latent projection.

        Returns:
        - torch.Tensor: The output tensor after applying the MoE layer.
        """
        if self.inplace:
            assert self.shared_experts is None
            assert not disable_inplace()
            output = hidden_states
        else:
            output = torch.empty_like(hidden_states)

        local_num_experts = w1.shape[0]
        if global_num_experts == -1:
            global_num_experts = local_num_experts

        a1q, a1q_scale, expert_tokens_meta, topk_ids, topk_weights = self._prepare(
            quanted_hidden_states
            if quanted_hidden_states is not None
            else hidden_states,
            topk_weights,
            topk_ids,
            global_num_experts,
            expert_map,
            apply_router_weight_on_input,
            prequanted_a1_scale=scale,
            prequanted_a1=quanted_hidden_states is not None,
        )

        fused_out = self._fused_experts(
            in_dtype=hidden_states.dtype,
            a1q=a1q,
            a1q_scale=a1q_scale,
            w1=w1,
            w2=w2,
            topk_weights=topk_weights,
            topk_ids=topk_ids,
            activation=activation,
            global_num_experts=global_num_experts,
            local_num_experts=local_num_experts,
            expert_map=expert_map,
            apply_router_weight_on_input=apply_router_weight_on_input,
            expert_tokens_meta=expert_tokens_meta,
            output_alias=output,
        )

        return self._finalize(
            output,
            fused_out,
            hidden_states,
            topk_weights,
            topk_ids,
            apply_router_weight_on_input,
            shared_experts_input=shared_experts_input,
            x_and_scale_quanted=(
                (quanted_hidden_states, scale)
                if quanted_hidden_states is not None and scale is not None
                else None
            ),
        )

这个没啥，这个其实就是调用了前面的那三个函数，就是prepare, fused_experts, finalize，

12 class FusedMoEKernel:


@final
class FusedMoEKernel:
    def __init__(
        self,
        prepare_finalize: FusedMoEPrepareAndFinalize,
        fused_experts: FusedMoEExperts,
        shared_experts: SharedExperts | None = None,
        inplace: bool = False,
        N: int = -1,
        K: int = -1,
    ):
        super().__init__()

        # Initialize the implementation (monolithic or modular).
        self.impl: FusedMoEKernelModularImpl | FusedMoEKernelMonolithicImpl
        if isinstance(
            prepare_finalize, FusedMoEPrepareAndFinalizeModular
        ) and isinstance(fused_experts, FusedMoEExpertsModular):
            self.impl = FusedMoEKernelModularImpl(
                prepare_finalize,
                fused_experts,
                shared_experts,
                inplace,
                N,
                K,
            )

        elif isinstance(
            prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic
        ) and isinstance(fused_experts, FusedMoEExpertsMonolithic):
            assert not inplace
            self.impl = FusedMoEKernelMonolithicImpl(
                prepare_finalize,
                fused_experts,
            )

        else:
            raise ValueError(
                "prepare_finalize and fused_experts must both be either monolithic "
                f"or non-monolithic but got {prepare_finalize.__class__.__name__} "
                f"and {fused_experts.__class__.__name__}"
            )

        self._post_init_setup()

    @property
    def owns_shared_experts(self) -> bool:
        if isinstance(self.impl, FusedMoEKernelModularImpl):
            return self.impl.shared_experts is not None
        else:
            return False

    @property
    def is_monolithic(self) -> bool:
        return isinstance(self.impl, FusedMoEKernelMonolithicImpl)

    @property
    def prepare_finalize(self) -> FusedMoEPrepareAndFinalize:
        return self.impl.prepare_finalize

    @property
    def fused_experts(self) -> FusedMoEExperts:
        return self.impl.fused_experts

    def supports_lora(self) -> bool:
        return self.fused_experts.supports_lora()

    def _post_init_setup(self):
        """
        Resolve any leftover setup dependencies between self.prepare_finalize
        and self.fused_experts here.
        """
        self.prepare_finalize.post_init_setup(self.impl.fused_experts)
        assert (
            self.prepare_finalize.activation_format
            == self.fused_experts.activation_format()
        )

    def supports_expert_map(self) -> bool:
        """
        A flag indicating whether or not this class supports expert maps.
        """
        return self.fused_experts.supports_expert_map()

    def output_is_reduced(self) -> bool:
        """
        Indicates whether or not the output of fused MoE kernel
        is reduced across all ranks.
        """
        return self.prepare_finalize.output_is_reduced()

    def apply_monolithic(
        self,
        hidden_states: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        router_logits: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
        activation: MoEActivation,
        global_num_experts: int,
        expert_map: torch.Tensor | None,
        apply_router_weight_on_input: bool,
        # grouped topk + fused topk bias parameters
        num_expert_group: int | None = None,
        e_score_correction_bias: torch.Tensor | None = None,
        routed_scaling_factor: float | None = None,
        topk_group: int | None = None,
    ) -> torch.Tensor:
        assert isinstance(self.impl, FusedMoEKernelMonolithicImpl)
        return self.impl.apply(
            hidden_states=hidden_states,
            w1=w1,
            w2=w2,
            router_logits=router_logits,
            activation=activation,
            global_num_experts=global_num_experts,
            expert_map=expert_map,
            apply_router_weight_on_input=apply_router_weight_on_input,
            num_expert_group=num_expert_group,
            e_score_correction_bias=e_score_correction_bias,
            routed_scaling_factor=routed_scaling_factor,
            topk_group=topk_group,
        )

    def apply(
        self,
        hidden_states: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        activation: MoEActivation,
        global_num_experts: int,
        expert_map: torch.Tensor | None,
        apply_router_weight_on_input: bool,
        shared_experts_input: torch.Tensor | None = None,
        quanted_hidden_states: torch.Tensor | None = None,
        scale: torch.Tensor | None = None,
    ) -> torch.Tensor:
        assert isinstance(self.impl, FusedMoEKernelModularImpl)
        return self.impl.apply(
            hidden_states=hidden_states,
            w1=w1,
            w2=w2,
            topk_weights=topk_weights,
            topk_ids=topk_ids,
            activation=activation,
            global_num_experts=global_num_experts,
            expert_map=expert_map,
            apply_router_weight_on_input=apply_router_weight_on_input,
            shared_experts_input=shared_experts_input,
            quanted_hidden_states=quanted_hidden_states,
            scale=scale,
        )

这里其实就是根据init时候传入的不同实例创建不同的implLei ,

prepare_finalize  +  fused_experts
        ↓
都是 Modular  →  FusedMoEKernelModularImpl
都是 Monolithic →  FusedMoEKernelMonolithicImpl
类型不匹配    →  报错

参考文献

大模型推理引擎vLLM(17): vllm015中model_executor/layers/fused_moe/modular_kernel.py代码阅读笔记_vllm ascend fsdp moe fused-CSDN博客

Logo
https://edu.csdn.net/learn/39067/627173?utm_source=2019755004

汇聚全球AI编程工具，助力开发者即刻编程。

更多推荐

cover

AI审计手记 #06：沙箱逃逸漏洞与“不修复“的防御策略转移

avatar AI编程社区
cover

Codex 和 Claude Code 怎么订阅？套餐、API计费与付款避坑指南

avatar AI编程社区
cover

手绘草图秒变网页：GPT-Image + Claude 前端开发实战与选型指南

avatar AI编程社区