# fused_moe.py 分析 本文整理 `src/flag_gems/fused/fused_moe.py` 的实现结构、调用路径、数据类型分支和计算原理，方便从源码角度理解 FlagGems 的 fused MoE 前向。 ## 第一章：MoE 与 fused MoE 的数学原理与计算原理 ### 1.0 MoE 在大模型中的背景 MoE，Mixture of Experts，是大模型中一种常见的稀疏前馈结构。它的核心思想是： - 不再让每个 token 都经过同一组固定的前馈层 - 而是准备多个 expert 子网络 - 再由 router 为每个 token 动态选择少量 expert 参与计算 在 Transformer 大模型中，MoE 通常替代原来的 dense FFN 层。也就是说，在注意力层之后，原本是： ```text Attention -> Dense FFN ``` 引入 MoE 后会变成： ```text Attention -> Router + Experts ``` 这样做的目的，是在尽量不线性增加单 token 计算量的前提下，大幅增加模型总参数量和表示能力。 可以把它理解为： - dense FFN 是“所有 token 共用同一套前馈参数” - MoE 是“不同 token 只激活一小部分 expert 参数” 因此，MoE 在大模型中的价值主要体现在： - 提升模型容量 - 控制单 token 计算成本 - 让不同 token 能使用更有针对性的 expert 子网络 从运行时角度看，MoE 算子负责的就是这一层稀疏前馈计算。它的输入通常是 attention 之后的 hidden states，输出则是经过 router 选择 expert、完成两层 expert MLP 并聚合后的新 hidden states。 更具体地说，这个算子主要做三件事： 1. 根据 router 给出的 `topk_ids` 和 `topk_weights`，确定每个 token 应该送到哪些 expert 2. 对被选中的 expert 执行前馈计算 3. 把多个 expert 的输出按路由权重聚合回 token 维度 因此，从模型结构上看，MoE 算子本质上是在实现： ```text token -> route to experts -> expert MLP -> weighted combine ``` 而 `fused_moe.py` 做的不是重新定义这个算子，而是为这一层 MoE 前向提供一个更高性能的实现，重点优化： - token 到 expert 的重排 - expert GEMM 的执行方式 - 量化路径 - 最终聚合过程 ### 1.1 MoE 的数学原理 ![MoE Principle](./fused_moe_image/moe.png) 结合上图，标准 MoE 层的计算过程可以形式化为“路由选择 -> 分发到专家 -> 专家前馈计算 -> 加权聚合输出”四个阶段。 设第 $t$ 个 token 的隐藏状态表示为： ```math x_t \in \mathbb{R}^{d} ``` 其中 $d$ 为隐藏维度。MoE 层包含 $E$ 个 expert，记第 $e$ 个 expert 对应的前馈函数为 $f_e$。 第一步是路由计算。Router 对输入 $x_t$ 生成对所有 expert 的路由打分。最常见的 router 形式是一个线性映射： ```math s_t = W_g^{\top} x_t + b_g ``` 其中 $W_g \in \mathbb{R}^{d \times E}$、$b_g \in \mathbb{R}^{E}$，因此 router 输出的打分向量满足： ```math s_t = \operatorname{Router}(x_t) \in \mathbb{R}^{E} ``` 通常可进一步写成 softmax 后的路由概率： ```math p_t = \operatorname{softmax}(s_t), \qquad p_t \in \mathbb{R}^{E} ``` 第二步是 Top-k 选择。Router 选择概率最高的 $k$ 个 expert，记其索引集合为 $S_t$： ```math \mathcal{S}_t = \operatorname{TopK}(p_t, k) ``` 对每个被选中的 expert $e \in S_t$，记其对应的路由权重为 $\alpha_{t,e}$。这一步对应图中的 `Top-k Selection`。 第三步是 dispatch。根据集合 $S_t$，token $x_t$ 会被逻辑上分发到这些被选中的 expert；未被选中的 expert 不参与当前 token 的计算。 第四步是 expert 计算。若每个 expert 采用图中所示的两层 MLP 结构 `Linear1 -> Activation(SiLU) -> Linear2`，则第 $e$ 个 expert 的输出可以写为： ```math f_e(x) = W_{2,e}^{\top}\,\mathrm{SiLU}(W_{1,e}^{\top}x + b_{1,e}) + b_{2,e} ``` 其中： - $W_{1,e}$、$W_{2,e}$ 分别为第 $e$ 个 expert 的两层线性变换参数 - $b_{1,e}$、$b_{2,e}$ 为对应偏置 - $\mathrm{SiLU}(\cdot)$ 为图中 expert 内部使用的激活函数 最后一步是聚合输出。对所有被选中的 expert 输出按对应路由权重进行加权求和，得到当前 token 的最终输出： ```math y_t = \sum_{e \in \mathcal{S}_t} \alpha_{t,e}\, f_e(x_t) ``` 因此，上图中的完整计算链条可以概括为： ```math x_t \xrightarrow{\text{Router}} s_t \xrightarrow{\text{Softmax}} p_t \xrightarrow{\text{TopK}} \mathcal{S}_t \xrightarrow{\text{Dispatch}} \{f_e(x_t)\}_{e \in \mathcal{S}_t} \xrightarrow{\text{Weighted Sum}} y_t ``` 也就是说，MoE 的本质是：先由 router 从全部 expert 中选择一小部分参与当前 token 的计算，再将这些被选中 expert 的输出按路由权重加权求和，从而得到最终输出 $y_t$。 ### 1.2 fused_moe 的数学原理 把每个 `(token, topk-slot)` 看作一个 routed item $r$，则： - 对应 token 为 $t(r)$ - 对应 expert 为 $e(r)$ - 对应路由权重为 $a(r)$ 那么 `fused_moe` 实际仍然是在算同一个数学对象。 第一层： ```math h_r = W_{1,e(r)}^{\top}x_{t(r)} + b_{1,e(r)} ``` 激活后： ```math z_r = \phi(h_r) ``` 第二层： ```math o_r = W_{2,e(r)}^{\top}z_r + b_{2,e(r)} ``` 最后对同一个 token 的 routed 输出求和： ```math y_t = \sum_{r: t(r)=t} a(r)\,o_r ``` 代码中有一个开关 `apply_router_weight_on_input`，决定把 router weight 乘在第一层输出之后还是第二层输出之后。由于 `a(r)` 是标量，线性层满足： ```math W^{\top}(a z) = a\,W^{\top}z ``` 因此两种写法数学上等价。 ### 1.3 fused_moe 的计算原理 `fused_moe` 和普通 MoE 在数学上等价，主要差别在工程组织方式。 普通 MoE 更像是按逻辑步骤分开执行： 1. router 选 top-k experts 2. dispatch / gather token 3. 对每个 expert 分别做 `w1` 4. 激活 5. 对每个 expert 分别做 `w2` 6. 乘 router weight 并聚合 而 `fused_moe` 的核心思想是： 1. 先把 routed token 按 expert 重排 2. 把每个 expert 的 token 数量按 `BLOCK_SIZE_M` 补齐 3. 让一个 Triton program 处理“同一 expert 的一块 token × 一块输出列” 4. 用统一 kernel 两次完成 `w1` 和 `w2` 5. 在 kernel 内尽量融合量化、bias、router weight 等附加逻辑 从 routed token 视角看，kernel 做的是： ```math C[r, n] = \sum_{k=1}^{K} A[t(r), k] \cdot B[e(r), n, k] ``` 其中： - $t(r)$ 是 routed item $r$ 对应的原始 token - $e(r)$ 是 routed item $r$ 对应的 expert - $n$ 是输出通道 在实现上，这套组织方式依赖三个关键步骤： - `moe_align_block_size(...)` 把 routed token 按 expert 排列成规则 block - `fused_moe_kernel(...)` 按 block 执行 expert GEMM - `moe_sum(...)` 对每个 token 的 top-k 输出做最终求和 ### 1.4 和普通 MoE 实现的区别 从数学上看，`fused_moe` 和普通 MoE 没有本质差别，仍然是在做： ```math y_t = \sum_{i=1}^{k} \alpha_{t,i}\, f_{e_{t,i}}(x_t) ``` 差别主要在工程组织方式： - 普通 MoE 更像是按逻辑步骤分开执行 - dispatch - expert matmul - activation - expert matmul - combine - fused MoE 则更强调： - 按 expert 重排 token - 让一个 Triton program 处理规则 tile - 把量化、bias、router weight 等附加逻辑尽量揉进 kernel 这样做的收益是： - 减少 kernel launch - 改善权重访问局部性 - 提高大规模 expert 场景下的吞吐 ## 文件定位 `fused_moe.py` 是一个面向 Triton 的 MoE 前向实现文件，核心职责包括： - 选择或生成 MoE kernel 配置 - 对输入激活做量化预处理 - 启动 Triton MoE kernel - 执行两层 expert MLP - 在 top-k expert 输出之间做聚合 它依赖两个关键辅助模块： - `flag_gems.fused.moe_align_block_size` 负责把 routed token 按 expert 重排并按 block 对齐 - `flag_gems.fused.moe_sum` 负责把每个 token 的 top-k expert 输出求和 ## 对外入口 对外可直接调用的入口有两个： - `inplace_fused_experts(...)` - `outplace_fused_experts(...)` 两者都只是薄封装，最终都调用： - `fused_experts_impl(...)` 其中： - `inplace_fused_experts` 把结果直接写回 `hidden_states` - `outplace_fused_experts` 返回新分配的输出张量 ## 文件内主要函数分组 ### 1. 配置选择 - `get_embedded_moe_configs()` 从 `fused_moe_config.yaml` 读取内嵌调优配置 - `_get_device_name()` 获取并规范化设备名 - `get_moe_configs(...)` 按设备名、expert 数、输出维度、dtype、block shape 查询配置 - `try_get_optimal_moe_config(...)` 优先使用内嵌配置，否则回退到启发式配置 - `get_default_config(...)` 默认配置生成逻辑 - `get_moe_wna16_block_config(...)` 针对 WNA16 路径生成 block 配置 - `_ensure_block_size_k_divisible(...)` 修正 `BLOCK_SIZE_K` ### 2. 数据类型与量化辅助 - `_get_config_dtype_str(...)` 把当前 dtype 和量化模式映射成配置查表 key - `_get_config_quant_dtype(...)` 把量化模式映射成激活量化 dtype - `_fp8_quantize(...)` 激活 FP8 量化 - `_int8_quantize(...)` 激活 INT8 量化 - `moe_kernel_quantize_input(...)` 在进入 GEMM 前对输入激活做量化 - `dequant_mxfp4(...)` - `dequant_mxfp6(...)` ### 3. 激活函数 - `MoEActivation` 定义支持的激活枚举 - `apply_moe_activation(...)` 执行激活 - `_silu_and_mul_kernel(...)` Triton 版 SiLU-and-mul ### 4. Triton kernel - `write_zeros_to_output(...)` expert 不在当前 rank 时写零 - `fused_moe_kernel(...)` 通用 MoE kernel，覆盖普通浮点、FP8、INT8 路径 - `fused_moe_kernel_gptq_awq(...)` WNA16 专用 kernel，面向 GPTQ/AWQ 风格的量化权重 ### 5. kernel 启动与分发 - `invoke_fused_moe_triton_kernel(...)` 启动通用 kernel - `invoke_fused_moe_wna16_triton_kernel(...)` 启动 WNA16 kernel - `dispatch_fused_moe_kernel(...)` 根据量化模式分发到具体启动函数 ## 主调用路径 主调用链可以概括为： ```text inplace_fused_experts / outplace_fused_experts -> fused_experts_impl -> _get_config_dtype_str -> _get_config_quant_dtype -> try_get_optimal_moe_config -> chunk loop -> moe_kernel_quantize_input (for w1 input) -> moe_align_block_size 或 naive assignment -> dispatch_fused_moe_kernel (for w1) -> invoke_fused_moe_triton_kernel -> fused_moe_kernel -> apply_moe_activation -> moe_kernel_quantize_input (for w2 input) -> dispatch_fused_moe_kernel (for w2) -> invoke_fused_moe_triton_kernel -> fused_moe_kernel -> moe_sum ``` 如果走 WNA16 特化路径，则中间会改为： ```text dispatch_fused_moe_kernel -> invoke_fused_moe_wna16_triton_kernel -> fused_moe_kernel_gptq_awq ``` ## fused_experts_impl 的执行流程 `fused_experts_impl(...)` 是整个文件的调度核心。 ### 1. 参数校验 它会检查： - `hidden_states`、`w1`、`w2` 的 shape 是否匹配 - `topk_weights` 和 `topk_ids` 形状是否一致 - `hidden_states` 是否连续 - `w1/w2` 最后一维 stride 是否为 1 - 输入 dtype 是否属于： - `torch.float32` - `torch.float16` - `torch.bfloat16` 当前实现还显式限制： - `activation == "silu"` 虽然 `MoEActivation` 定义了多种激活，但当前入口只允许 SiLU 路径。 ### 2. 配置选择 它先生成两个关键变量： - `config_dtype` 用于 kernel 配置查表 - `quant_dtype` 用于激活量化逻辑 之后用 `try_get_optimal_moe_config(...)` 选择当前 chunk 对应的 Triton 配置。 ### 3. 中间缓存分配 会分配三块中间缓存： - `intermediate_cache1` 保存 `w1` 输出，shape 为 `[M, top_k, N]` - `intermediate_cache2` 保存激活后的结果，shape 为 `[M * top_k, activation_out_dim]` - `intermediate_cache3` 保存 `w2` 输出，shape 为 `[M, top_k, K]` 其中 `cache13` 复用了 `cache1` 和 `cache3` 的存储，因为两者生命周期不重叠。 ### 4. 计算 dtype 选择 输入 dtype 与 kernel 中的 `compute_type` 对应关系： - `torch.bfloat16 -> tl.bfloat16` - `torch.float16 -> tl.float16` - `torch.float32 -> tl.float32` ### 5. 特殊量化格式处理 如果启用了 `ocp_mx_scheme`，会先把 MX 权重反量化成普通浮点权重，再走后续通用路径。 对于 `use_int8_w8a16` 和 `use_int4_w4a16`，当前公开入口中也会先把权重反量化成 `hidden_states.dtype`，随后把： - `use_int8_w8a16 = False` - `use_int4_w4a16 = False` 这意味着虽然文件内保留了 WNA16 特化 kernel，但从 `fused_experts_impl(...)` 这条入口实际运行时，常见路径仍然是通用 `fused_moe_kernel(...)`。 ### 6. 分 chunk 执行 实现按 `CHUNK_SIZE = 16 * 1024` 分块处理 token，原因是： - 限制中间缓存大小 - 让 kernel 配置更容易适配当前 chunk 规模 每个 chunk 内部执行： 1. 量化第一层输入 2. 把 routed token 按 expert 重排 3. 跑 `w1` 4. 激活 5. 量化第二层输入 6. 跑 `w2` 7. 对 top-k expert 输出做求和 ## moe_align_block_size 的作用 MoE 的普通数学形式里，每个 token 会被路由到若干 expert。直接按 token 顺序计算时，expert 权重访问会非常离散。 `moe_align_block_size(...)` 的作用是： - 把所有 `(token, topk-slot)` 展平为 routed token - 按 expert 把 routed token 排序 - 对每个 expert 的 token 数量按 `BLOCK_SIZE_M` 向上补齐 它输出三个核心量： - `sorted_token_ids` 重排后的 routed token 下标 - `expert_ids` 每个 M 方向 block 对应的 expert id - `num_tokens_post_padded` 补齐后的 routed token 总数 这样后续 Triton kernel 就可以让一个 program 只处理： - 一块 routed token - 一个固定 expert - 一块输出列 这就是 fused MoE 能把很多零散小计算组织成规则 block GEMM 的关键。 ## fused_moe_kernel 的核心原理 `fused_moe_kernel(...)` 是通用 kernel，也是阅读此文件时最值得重点关注的函数。 ### 1. program 到 tile 的映射 kernel 先根据： - `BLOCK_SIZE_M` - `BLOCK_SIZE_N` - `GROUP_SIZE_M` 把 `pid` 映射到： - `pid_m` - `pid_n` 这里使用 grouped ordering 来改善 L2 reuse。 ### 2. 当前 block 对应哪些 token 和 expert 对每个 `pid_m`： - 如果启用了重排，先从 `sorted_token_ids` 取出当前 block 的 routed token - 从 `expert_ids[pid_m]` 取出当前 block 对应的 expert 如果 `expert_id == -1`，则说明该 expert 不在当前 expert parallel rank，直接调用 `write_zeros_to_output(...)`。 ### 3. 构造 A / B 指针 输入 A 的访问方式为： ```text offs_token // top_k ``` 因为 routed token 是 `(token, slot)` 的展平形式，所以要除以 `top_k` 才能回到原始 token 下标。 权重 B 则按照： - 当前 expert - 当前 K block - 当前 N block 来取对应子块。 ### 4. 主循环 沿 K 维做 block 累加： - 读取 A 子块 - 读取 B 子块 - 走普通浮点或量化分支 - 累加到 `accumulator` `accumulator` 始终先用 `fp32` 保存，再在输出前转成 `compute_type`。 ### 5. 后处理 主循环结束后依次执行： - 可选 dequant - 可选 bias - 可选乘 router weight - cast 到 `compute_type` - 写回 C ## 普通浮点与 FP8 路径 ### 普通浮点 当： - `use_fp8_w8a8 == False` - `use_int8_w8a8 == False` - `A_scale is None` - `B_scale is None` 时，kernel 直接执行普通浮点路径： ```text accumulator += tl.dot(a, b) ``` 这是 `bfloat16`、`float16`、`float32` 共享的主要逻辑。 ### FP8 W8A8 当启用 `use_fp8_w8a8` 时： - Python 侧先调用 `moe_kernel_quantize_input(...)` - `moe_kernel_quantize_input(...)` 会进一步调用 `_fp8_quantize(...)` - kernel 内根据量化粒度决定如何加载 scale 支持的粒度包括： - tensor-wise - per-channel - block-wise block-wise 路径下，会在每个 K block 内读取本 block 对应的： - `a_scale` - `b_scale` 然后执行近似： ```text C_block ~= dot(qA_block, qB_block) * a_scale * b_scale ``` ### bfloat16 支持 文件明确支持： - `torch.bfloat16` 其执行方式不是把整个 kernel 改成 BF16 累加，而是： - 输入 / 输出使用 BF16 - `accumulator` 先用 FP32 - 最后再 cast 回 `tl.bfloat16` ## WNA16 路径 文件内还保留了一条 WNA16 特化路径： - `invoke_fused_moe_wna16_triton_kernel(...)` - `fused_moe_kernel_gptq_awq(...)` 它主要用于： - `int8_w8a16` - `int4_w4a16` 并支持： - group scale - zero point - GPTQ/AWQ 风格权重布局 但需要注意： - 当前 `fused_experts_impl(...)` 中会先把 INT8/INT4 权重反量化成普通浮点 - 因此从当前公开入口来看，这条分支通常不会成为主路径 也就是说，代码结构上支持 WNA16 专用 kernel，但入口行为上更偏向统一走通用 kernel。 ## 调用图 ```text inplace_fused_experts -> fused_experts_impl outplace_fused_experts -> fused_experts_impl fused_experts_impl -> _get_config_dtype_str -> _get_config_quant_dtype -> try_get_optimal_moe_config -> get_moe_configs -> _get_device_name -> get_embedded_moe_configs -> get_default_config -> moe_kernel_quantize_input -> _fp8_quantize / _int8_quantize -> moe_align_block_size -> dispatch_fused_moe_kernel -> invoke_fused_moe_triton_kernel -> fused_moe_kernel -> write_zeros_to_output -> invoke_fused_moe_wna16_triton_kernel -> get_moe_wna16_block_config -> _ensure_block_size_k_divisible -> fused_moe_kernel_gptq_awq -> write_zeros_to_output -> apply_moe_activation -> _silu_and_mul_kernel -> moe_sum ``` ## 当前实现的几个注意点 - 当前入口只允许 `activation == "silu"` - 明确支持输入 dtype： - `float32` - `float16` - `bfloat16` - 通用主路径是 `fused_moe_kernel(...)` - 文件内虽然保留了 WNA16 特化 kernel，但当前公开入口通常会先把 INT8/INT4 权重反量化，实际常走通用 kernel - `moe_align_block_size(...)` 是性能关键，它决定了 routed token 是否能被整理成规则 block ## 小结 `fused_moe.py` 的核心思想可以概括为三点： 1. 把 routed token 按 expert 排序并按 block 对齐 2. 用统一的 Triton kernel 两次完成 expert 的两层 GEMM 3. 在 kernel 内尽量融合量化、bias、router weight 等附加逻辑 从源码阅读角度，推荐优先关注以下函数： - `fused_experts_impl(...)` - `moe_kernel_quantize_input(...)` - `dispatch_fused_moe_kernel(...)` - `fused_moe_kernel(...)` - `moe_align_block_size(...)` 这几处连起来，基本就构成了整个 fused MoE 前向的主干。