English 中文

Fused MoE

Fused MoE#

功能概述#

fused_moe 是 MindIE-SD 提供的 MoE 对外入口，用于在 NPU 上完成 MoE 前向推理中的专家选择、Token 分发、专家计算和结果合并。该接口面向开源框架集成场景，调用方传入激活、路由 logits、专家权重和通信配置后，即可通过统一入口完成 routed experts 的前向计算。

MoE 模型中，每个 Token 会根据 router 输出选择少量 experts 参与计算。相比 dense MLP，MoE 可以在扩大模型容量的同时控制单次推理的实际计算量，但也引入了 Token 到 expert 的路由、重排、跨卡通信和结果恢复等额外流程。fused_moe 将这些流程封装在统一接口中，减少框架侧重复适配成本，并便于在不同并行策略下复用同一套 MoE 计算入口。

当前 fused_moe 包含两条路径：融合算子路径和非融合算子路径。融合算子路径当前版本暂不支持，use_fused_op=True 时会回退到非融合算子路径；默认 use_fused_op=False，直接使用非融合算子路径。

接口说明#

from mindiesd import fused_moe

函数签名#

fused_moe(
    hidden_states,
    router_logits,
    num_experts,
    top_k,
    w13_weight,
    w2_weight,
    w13_bias=None,
    w2_bias=None,
    quant_config=None,
    w13_weight_scale=None,
    w2_weight_scale=None,
    tp_group=None,
    ep_group=None,
    dispatcher_type=None,
    inputs_sharded=False,
    k_group=1,
    group_count=1,
    group_select_mode=0,
    routing_method="softmax",
    renormalize=False,
    routed_scaling_factor=1.0,
    custom_routing_function=None,
    reduce_routed_out=True,
    return_dispatcher_type=False,
    use_fused_op=False,
) -> torch.Tensor | tuple[torch.Tensor, str]

参数说明#

参数

类型

必选

默认值

说明

hidden_states

torch.Tensor

是

-

输入激活，形状为 [..., hidden_size]，维度不少于 2。

router_logits

torch.Tensor

是

-

路由 logits，形状为 [..., num_experts]，维度不少于 2，前置维度需与 hidden_states 一致。

num_experts

int

是

-

全局 expert 数量，必须与 router_logits 最后一维一致。使用 EP 时需能被 EP group size 整除。

top_k

int

是

-

每个 Token 选择的 expert 数量，取值范围为 [1, num_experts]。

w13_weight

torch.Tensor

是

-

融合后的 gate/up 投影权重，形状为 [local_experts, hidden_size, 2 * intermediate_size]，必须为 3D Tensor。

w2_weight

torch.Tensor

是

-

down 投影权重，形状为 [local_experts, intermediate_size, hidden_size]，必须与 w13_weight 具有相同的 local_experts。

w13_bias

torch.Tensor / None

否

None

gate/up 投影 bias，形状为 [local_experts, 2 * intermediate_size]，需与 w13_weight 的 expert 和输出维度一致。

w2_bias

torch.Tensor / None

否

None

down 投影 bias，形状为 [local_experts, hidden_size]，需与 w2_weight 的 expert 和输出维度一致。

quant_config

QuantConfig / None

否

None

MindIE-SD 量化配置，用于选择 MoE 前向流程的量化算法。

w13_weight_scale

torch.Tensor / None

否

None

w13_weight 的 quantization scale。

w2_weight_scale

torch.Tensor / None

否

None

w2_weight 的 quantization scale。

tp_group

dist.ProcessGroup / None

否

None

TP 通信组。未启用 EP 且 TP group size 大于 1 时生效。

ep_group

dist.ProcessGroup / None

否

None

EP 通信组。EP group size 大于 1 时优先生效。

dispatcher_type

str / None

否

None

Token 分发策略。可选 "static"、"dynamic"；None 表示根据平台和通信配置自动选择。"dynamic" 仅支持 EP 通信场景。

inputs_sharded

bool

否

False

表示输入是否已经在当前 MoE 通信组内沿 Token 维度 shard。

k_group

int

否

1

分组路由时每个 Token 选择的 expert group 数量，取值范围为 [1, group_count]。

group_count

int

否

1

expert group 总数，num_experts 需要能被 group_count 整除。

group_select_mode

int

否

0

expert group 打分方式。0 表示取组内最大分数，1 表示取组内 top-2 分数之和。使用 1 时每组至少需要 2 个 experts。

routing_method

str

否

"softmax"

router logits 的打分方式，可选 "softmax" 或 "sigmoid"。

renormalize

bool

否

False

是否对 softmax 路由选中的 top-k routing weights 重新归一化。sigmoid 路由按 NPU gating top-k 算子语义输出归一化后的权重。

routed_scaling_factor

float

否

1.0

路由权重缩放系数，在专家选择阶段生效。

custom_routing_function

callable / None

否

None

自定义路由函数，调用形式为 custom_routing_function(hidden_states, gating_output, topk, renormalize)，需返回 (topk_weights, topk_ids)。

reduce_routed_out

bool

否

True

static dispatcher 在 inputs_sharded=False 时是否对 routed output 做 all_reduce。

return_dispatcher_type

bool

否

False

是否在输出中返回实际解析到的 dispatcher 类型，返回值为 "static" 或 "dynamic"。

use_fused_op

bool

否

False

是否优先启用融合算子路径。当前版本暂不支持该路径，设置为 True 时会回退到非融合算子路径；默认 False 直接使用非融合算子路径。

返回值#

torch.Tensor：MoE 前向计算结果，形状与 hidden_states 一致。 若 return_dispatcher_type=True，返回 (output, dispatcher_type)。

融合算子路径（预留）#

当前版本暂不支持融合算子路径。若设置 use_fused_op=True，接口会回退到非融合算子路径完成计算。

非融合算子路径（当前默认）#

该路径将 MoE 前向拆成多个阶段完成，支持非量化和量化 MoE 推理，覆盖单卡、TP 和 EP 场景，并提供 static 和 dynamic 两种 Token 分发方式。static dispatcher 适用于单卡、TP 以及部分 EP 场景；dynamic dispatcher 面向 EP 场景下的 all-to-all Token 交换。

执行流程#

该路径将 MoE 前向过程拆分为以下阶段：

  1. prepare：整理输入激活和 router logits，并根据输入 layout 准备后续计算所需的数据。

  2. select_experts：根据 router 输出，为每个 Token 选择对应的 top-k experts，并生成 routing weights。

  3. dispatch：根据专家选择结果，将 Token 按 expert 路由并重排，生成专家计算所需的输入。

  4. mlp：执行专家侧 grouped MLP 计算，完成 routed experts 的前馈计算。

  5. combine：将专家输出按原 Token 顺序合并，恢复 routed MoE 的输出结果。

  6. finalize：完成输出恢复和通信后的收尾处理。

Dispatcher 策略#

支持两类 Token 分发策略：

  • static dispatcher：使用静态 Token 分发路径，适用于单卡、TP 场景，以及部分 EP 场景。该路径通过 NPU MoE routing 算子完成 Token 排序、expert token 统计和结果恢复。

  • dynamic dispatcher：使用动态 Token 分发路径，适用于 EP 场景。该路径会根据 Token 到 expert 的分布执行 all-to-all 通信，并在专家计算前后完成 Token 顺序恢复。

当 dispatcher_type=None 时，接口根据通信模式和 NPU 型号自动选择：EP 场景下 Atlas 800I A3 超节点服务器 / Ascend 950PR / Ascend 950DT 使用 dynamic dispatcher，Atlas 800I A2 推理服务器使用 static dispatcher；非 EP 场景（单卡/TP）始终使用 static dispatcher。也可以通过 dispatcher_type="static" 或 dispatcher_type="dynamic" 显式指定。

通信配置#

通过 tp_group 和 ep_group 支持分布式 MoE 推理。当前接口一次仅启用一种 MoE 通信策略：

  • 当传入有效 ep_group 时，优先使用 EP 通信策略。

  • 当未启用 EP 且传入有效 tp_group 时，使用 TP 通信策略。

  • 当二者均未启用时，按单卡路径执行。

调用方需要保证输入激活和 router logits 的 layout 与 inputs_sharded 配置一致。设置 inputs_sharded=True 时，输入应已在当前 MoE 通信组内沿 Token 维度完成 shard，否则可能产生重复 Token 计算。

量化配置#

未传入 quant_config，或 quant_config.quant_algo 为 None / NO_QUANT 时，按非量化方式执行 MoE 前向流程。当前支持以下量化配置：

  • QuantConfig(quant_algo=QuantAlgorithm.W8A8_DYNAMIC)：W8A8 dynamic quantization，支持 Atlas 800I A2 推理服务器 / Atlas 800I A3 超节点服务器。

  • QuantConfig(quant_algo=QuantAlgorithm.W8A8_MXFP8)：W8A8 MXFP8 quantization，支持Ascend 950PR / Ascend 950DT。

路由选择#

默认路由使用 NPU gating top-k 算子完成专家选择。routing_method="softmax" 时对 router logits 做 softmax，routing_method="sigmoid" 时做 sigmoid。softmax 路由可通过 renormalize=True 对选中的 top-k weights 做重新归一化；sigmoid 路由按 NPU gating top-k 算子语义输出归一化后的权重。

未启用分组路由时保持默认 k_group=1、group_count=1 即可。启用分组路由时，会先按 group_select_mode 选择 expert group，再在选中的 group 中选择每个 Token 的 top-k experts。

分组路由需要满足以下约束：

  • num_experts 能被 group_count 整除。

  • k_group 取值范围为 [1, group_count]。

  • top_k 不超过被选中 expert group 内的 expert 总数。

  • group_select_mode=1 时，每个 expert group 至少包含 2 个 experts。

使用示例#

用例说明#

实际接入框架时，hidden_states、router_logits 和专家权重通常来自模型 forward 和权重加载流程。 分布式示例默认 tp_group 或 ep_group 已完成初始化。

单卡 MoE#

单卡场景不需要传入通信组，使用默认策略即可。

import torch
from mindiesd import fused_moe

num_tokens = 8
hidden_size = 4096
intermediate_size = 14336
num_experts = 8
top_k = 2
dtype = torch.bfloat16
device = "npu"

hidden_states = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype)
router_logits = torch.randn(num_tokens, num_experts, device=device, dtype=dtype)
w13_weight = torch.randn(num_experts, hidden_size, 2 * intermediate_size, device=device, dtype=dtype)
w2_weight = torch.randn(num_experts, intermediate_size, hidden_size, device=device, dtype=dtype)
w13_bias = torch.randn(num_experts, 2 * intermediate_size, device=device, dtype=dtype)
w2_bias = torch.randn(num_experts, hidden_size, device=device, dtype=dtype)

out = fused_moe(
    hidden_states=hidden_states,
    router_logits=router_logits,
    num_experts=num_experts,
    top_k=top_k,
    w13_weight=w13_weight,
    w2_weight=w2_weight,
    w13_bias=w13_bias,
    w2_bias=w2_bias,
    renormalize=True,
)

分组路由 MoE#

当模型需要先选择 expert group，再从选中 group 中选择 top-k experts 时，可以配置分组路由参数。下面示例使用 sigmoid 路由，并通过 routed_scaling_factor 对路由权重做缩放。

import torch
from mindiesd import fused_moe

num_tokens = 8
hidden_size = 4096
intermediate_size = 14336
num_experts = 16
top_k = 2
dtype = torch.bfloat16
device = "npu"

hidden_states = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype)
router_logits = torch.randn(num_tokens, num_experts, device=device, dtype=dtype)
w13_weight = torch.randn(num_experts, hidden_size, 2 * intermediate_size, device=device, dtype=dtype)
w2_weight = torch.randn(num_experts, intermediate_size, hidden_size, device=device, dtype=dtype)

out = fused_moe(
    hidden_states=hidden_states,
    router_logits=router_logits,
    num_experts=num_experts,
    top_k=top_k,
    w13_weight=w13_weight,
    w2_weight=w2_weight,
    k_group=1,
    group_count=4,
    group_select_mode=1,
    routing_method="sigmoid",
    routed_scaling_factor=0.5,
)

W8A8 dynamic quant MoE（Atlas 800I A2 推理服务器 / Atlas 800I A3 超节点服务器）#

W8A8 dynamic quant 路径要求 w13_weight 和 w2_weight 为 torch.int8，并传入对应的 quantization scale。MindIE-SD 会在 MLP 计算前检查权重格式；若权重不是 NPU NZ 格式，会自动转换为 NZ 后再调用 INT8 grouped MLP 算子。

import torch
from mindiesd import fused_moe
from mindiesd.quantization.config import QuantConfig
from mindiesd.quantization.mode import QuantAlgorithm

num_tokens = 8
hidden_size = 4096
intermediate_size = 14336
num_experts = 8
top_k = 2
dtype = torch.bfloat16
device = "npu"

hidden_states = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype)
router_logits = torch.randn(num_tokens, num_experts, device=device, dtype=dtype)
w13_weight = torch.randint(
    -128,
    127,
    (num_experts, hidden_size, 2 * intermediate_size),
    device=device,
    dtype=torch.int8,
)
w2_weight = torch.randint(
    -128,
    127,
    (num_experts, intermediate_size, hidden_size),
    device=device,
    dtype=torch.int8,
)
w13_weight_scale = torch.rand(num_experts, 2 * intermediate_size, device=device, dtype=dtype)
w2_weight_scale = torch.rand(num_experts, hidden_size, device=device, dtype=dtype)

out = fused_moe(
    hidden_states=hidden_states,
    router_logits=router_logits,
    num_experts=num_experts,
    top_k=top_k,
    w13_weight=w13_weight,
    w2_weight=w2_weight,
    quant_config=QuantConfig(quant_algo=QuantAlgorithm.W8A8_DYNAMIC),
    w13_weight_scale=w13_weight_scale,
    w2_weight_scale=w2_weight_scale,
    dispatcher_type="static",
    inputs_sharded=False,
)

MXFP8 dynamic quant MoE（Ascend 950PR / Ascend 950DT）#

MXFP8 路径使用 QuantAlgorithm.W8A8_MXFP8 量化配置，w13_weight 和 w2_weight 使用 torch.float8_e4m3fn，并传入对应的 quantization scale。

import torch
import torch_npu
from mindiesd import fused_moe
from mindiesd.quantization.config import QuantConfig
from mindiesd.quantization.mode import QuantAlgorithm

num_tokens = 8
hidden_size = 4096
intermediate_size = 14336
num_experts = 8
top_k = 2
dtype = torch.bfloat16
device = "npu"


def prepare_mxfp8_moe_weight(weight):
    num_experts, _, n_size = weight.shape
    quant_weight, weight_scale = torch_npu.npu_dynamic_mx_quant(
        weight.transpose(1, 2),
        dst_type=torch.float8_e4m3fn,
    )
    weight_scale = weight_scale.reshape(num_experts, n_size, -1, 2)
    return quant_weight.transpose(1, 2), weight_scale.transpose(1, 2)


hidden_states = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype)
router_logits = torch.randn(num_tokens, num_experts, device=device, dtype=dtype)
w13_weight_fp = torch.randn(
    num_experts,
    hidden_size,
    2 * intermediate_size,
    device=device,
    dtype=dtype,
)
w2_weight_fp = torch.randn(
    num_experts,
    intermediate_size,
    hidden_size,
    device=device,
    dtype=dtype,
)
w13_weight, w13_weight_scale = prepare_mxfp8_moe_weight(w13_weight_fp)
w2_weight, w2_weight_scale = prepare_mxfp8_moe_weight(w2_weight_fp)

out = fused_moe(
    hidden_states=hidden_states,
    router_logits=router_logits,
    num_experts=num_experts,
    top_k=top_k,
    w13_weight=w13_weight,
    w2_weight=w2_weight,
    quant_config=QuantConfig(quant_algo=QuantAlgorithm.W8A8_MXFP8),
    w13_weight_scale=w13_weight_scale,
    w2_weight_scale=w2_weight_scale,
    dispatcher_type="static",
    inputs_sharded=False,
)

TP static MoE#

TP 场景传入 tp_group。当 inputs_sharded=False 时，表示每个 rank 输入在当前 MoE 通信组内未沿 Token 维度切分。

import torch
from mindiesd import fused_moe

num_tokens = 8
hidden_size = 4096
intermediate_size = 14336
num_experts = 8
top_k = 2
dtype = torch.bfloat16
device = "npu"

hidden_states = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype)
router_logits = torch.randn(num_tokens, num_experts, device=device, dtype=dtype)
w13_weight = torch.randn(num_experts, hidden_size, 2 * intermediate_size, device=device, dtype=dtype)
w2_weight = torch.randn(num_experts, intermediate_size, hidden_size, device=device, dtype=dtype)

out = fused_moe(
    hidden_states=hidden_states,
    router_logits=router_logits,
    num_experts=num_experts,
    top_k=top_k,
    w13_weight=w13_weight,
    w2_weight=w2_weight,
    tp_group=tp_group,
    dispatcher_type="static",
    inputs_sharded=False,
    reduce_routed_out=True,
)

EP static MoE#

EP 场景也可以显式指定 static dispatcher。此时 w13_weight 和 w2_weight 表示当前 rank 持有的本地 expert 权重，ep_group 用于完成 MoE 通信。

import torch
from mindiesd import fused_moe

num_tokens = 8
hidden_size = 4096
intermediate_size = 14336
num_experts = 8
top_k = 2
dtype = torch.bfloat16
device = "npu"
ep_world_size = torch.distributed.get_world_size(ep_group)
local_experts = num_experts // ep_world_size

hidden_states = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype)
router_logits = torch.randn(num_tokens, num_experts, device=device, dtype=dtype)
w13_weight = torch.randn(local_experts, hidden_size, 2 * intermediate_size, device=device, dtype=dtype)
w2_weight = torch.randn(local_experts, intermediate_size, hidden_size, device=device, dtype=dtype)

out = fused_moe(
    hidden_states=hidden_states,
    router_logits=router_logits,
    num_experts=num_experts,
    top_k=top_k,
    w13_weight=w13_weight,
    w2_weight=w2_weight,
    ep_group=ep_group,
    dispatcher_type="static",
    inputs_sharded=False,
    reduce_routed_out=True,
    renormalize=True,
)

EP dynamic MoE#

EP dynamic 路径需要传入 ep_group。当每个 rank 输入在当前 MoE 通信组内 shard 后的本地 Token 时， 设置 inputs_sharded=True；此时 local_hidden_states 和 local_router_logits 表示当前 rank 的本地 Token 输入，w13_weight 和 w2_weight 仍表示当前 rank 的 本地 expert 权重。输入应已在该组内沿 Token 维度完成 shard，否则可能产生重复 Token 计算。

import torch
from mindiesd import fused_moe

num_tokens = 8
hidden_size = 4096
intermediate_size = 14336
num_experts = 8
top_k = 2
dtype = torch.bfloat16
device = "npu"
ep_world_size = torch.distributed.get_world_size(ep_group)
local_experts = num_experts // ep_world_size
local_num_tokens = num_tokens // ep_world_size
local_hidden_states = torch.randn(local_num_tokens, hidden_size, device=device, dtype=dtype)
local_router_logits = torch.randn(local_num_tokens, num_experts, device=device, dtype=dtype)
w13_weight = torch.randn(local_experts, hidden_size, 2 * intermediate_size, device=device, dtype=dtype)
w2_weight = torch.randn(local_experts, intermediate_size, hidden_size, device=device, dtype=dtype)

out = fused_moe(
    hidden_states=local_hidden_states,
    router_logits=local_router_logits,
    num_experts=num_experts,
    top_k=top_k,
    w13_weight=w13_weight,
    w2_weight=w2_weight,
    ep_group=ep_group,
    dispatcher_type="dynamic",
    inputs_sharded=True,
    renormalize=True,
)

自定义 routing function#

当模型已有自定义路由逻辑时，可通过 custom_routing_function 接入。该函数需要返回 (topk_weights, topk_ids)。

import torch
from mindiesd import fused_moe

num_tokens = 8
hidden_size = 4096
intermediate_size = 14336
num_experts = 8
top_k = 2
dtype = torch.bfloat16
device = "npu"

hidden_states = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype)
router_logits = torch.randn(num_tokens, num_experts, device=device, dtype=dtype)
w13_weight = torch.randn(num_experts, hidden_size, 2 * intermediate_size, device=device, dtype=dtype)
w2_weight = torch.randn(num_experts, intermediate_size, hidden_size, device=device, dtype=dtype)

def custom_routing_function(hidden_states, gating_output, topk, renormalize):
    topk_result = gating_output.softmax(dim=-1).topk(topk, dim=-1)
    topk_weights = topk_result.values
    topk_ids = topk_result.indices
    if renormalize:
        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
    return topk_weights, topk_ids


out = fused_moe(
    hidden_states=hidden_states,
    router_logits=router_logits,
    num_experts=num_experts,
    top_k=top_k,
    w13_weight=w13_weight,
    w2_weight=w2_weight,
    custom_routing_function=custom_routing_function,
    renormalize=True,
)

注意事项#

  • 当前接口仅支持前向推理，不支持反向梯度计算。

  • 当前支持非量化 MoE、W8A8 dynamic quant MoE 和 W8A8 MXFP8 quant MoE 路径。

  • 非量化权重需要使用 grouped matmul 所需布局；量化权重和 scale 需要与 quant_config 匹配。

  • router_logits 的最后一维必须等于 num_experts。

  • dynamic dispatcher 需要在 EP 通信场景下使用。

  • 使用分布式通信时，调用方需要提前完成通信组初始化，并保证各 rank 上的输入 layout 一致。

目录