# 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 前向的主干。
