HunyuanImage 3.0 的 MoE 阅读
背景
HunyuanImage 3.0 是目前开源生图模型中少有的使用了MoE diffusion 结构的模型，我觉得这是未来的主流结构，以它为例子学习了一下 MoE diffusion.
概述
HunyuanImage 3.0 MoE 的核心思路，是在 Transformer Block 中用路由器、多个 routed experts，以及可选的 shared MLP 替换原本的 Dense FFN。
代码框架同时保留了 Dense MLP 和 MoE 两种实现。构造 decoder layer 时，如果 num_experts 是大于 1 的整数，或者它是一个列表且 max(config.num_experts) > 1，并且当前 layer_idx 没有被 moe_layer_num_skipped 跳过，则使用 HunyuanMoE；否则使用普通的 HunyuanMLP：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
if (
(
isinstance(config.num_experts, int)
and config.num_experts > 1
)
or (
isinstance(config.num_experts, list)
and max(config.num_experts) > 1
)
) and layer_idx >= config.moe_layer_num_skipped:
self.mlp = HunyuanMoE(
config,
layer_idx=layer_idx,
)
else:
self.mlp = HunyuanMLP(
config,
layer_idx=layer_idx,
is_shared_mlp=False,
is_moe=False,
)
Dense Transformer Block
一个普通的 Pre-Norm Transformer Block 可以简化为：
\[h' = h + \mathrm{Attention}\bigl(\mathrm{Norm}(h)\bigr)\] \[h'' = h' + \mathrm{FFN}\bigl(\mathrm{Norm}(h')\bigr)\]对应代码为：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# h' = h + Attention(Norm(h))
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
hidden_states, self_attn_weights, present_key_value = self.self_attn(
hidden_states=hidden_states,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_value=past_key_value,
custom_pos_emb=custom_pos_emb,
)
hidden_states = residual + hidden_states
# h'' = h' + FFN(Norm(h'))
residual = hidden_states
hidden_states = self.post_attention_layernorm(hidden_states)
hidden_states = self.mlp(hidden_states)
hidden_states = residual + hidden_states
Transformer 模型中的大量参数通常集中在 FFN 子层中。
在 Dense 模型里，每个 token 都经过同一组 FFN 参数：
1
2
3
4
5
6
token
│
└── Dense FFN
├── gate/up projection
├── activation
└── down projection
所有 token 激活相同的参数，因此计算量和模型参数量通常会同步增长。
Sparse MoE Block
Hunyuan MoE 使用一个路由器、多个 routed experts，以及可选的 shared MLP 替换 Dense FFN：
1
2
3
4
5
HunyuanImage3DecoderLayer
└── HunyuanMoE
├── HunyuanTopKGate
├── shared_mlp: HunyuanMLP（仅当 use_mixed_mlp_moe=True）
└── experts: ModuleList[HunyuanMLP × num_experts]
公开的 HunyuanImage 3.0 配置使用 64 个 routed experts，num_experts = 64。
整体数据流可以表示为：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
hidden_states [B, S, H]
│
├── Router Linear: H → E
│ │
│ └── Softmax + Top-K
│ 得到专家 ID 和路由权重
│
├── Routed Experts
│ │
│ ├── 按 expert ID 对 token 分组
│ ├── 执行选中的专家 FFN
│ └── 按路由权重加权求和
│
└── Shared MLP（仅当 use_mixed_mlp_moe=True）
│
└── 对所有 token 执行
HunyuanTopKGate easy_topk 推理路径
下面展示当前推理使用的 easy_topk 调用路径，不展开 default 对应的 topkgating 路径：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
HunyuanTopKGate(nn.Module)
├── __init__(config, layer_idx)
│ ├── moe_topk = K # 当前层的标量；从 config 中直接读取或按 layer_idx 取值
│ └── wg: Linear(H → E, bias=False, float32) # weight shape: [E, H]
│
└── forward(hidden_states, topk_impl)
├── hidden_states # shape: [B, S, H]
├── hidden_states.reshape(-1, H) # shape: [B × S, H]
├── if wg.weight.dtype == float32
│ └── hidden_states.float() # shape: [B × S, H], dtype: float32
├── logits = wg(hidden_states) # shape: [B × S, E]
└── topk_impl == "easy"
└── easy_topk(logits, moe_topk) # input: [B × S, E], K
├── gates = softmax(logits, dim=1) # shape: [B × S, E]
├── topk：选择概率最高的 K 个专家
│ ├── topk_weight_1 # shape: [B × S, K]
│ └── expert_index # shape: [B × S, K]
├── weight_sums = topk_weight_1.sum(
│ dim=1, keepdim=True) # shape: [B × S, 1]
├── weight_sums = torch.clamp(
│ weight_sums, min=1e-8) # shape: [B × S, 1]
├── topk_weight = topk_weight_1
│ / weight_sums # shape: [B × S, K]
└── return topk_weight, expert_index # shapes: [B × S, K], [B × S, K]
Router Linear 将 hidden state 从 hidden_size 维投影到 num_experts 维，得到每个 token 对所有专家的 router logits。当前实现显式设置了 bias=False，因此该线性层未使用 bias：
对应源码中的 HunyuanTopKGate.wg：
1
2
3
4
5
6
self.wg = nn.Linear(
config.hidden_size,
num_experts,
bias=False,
dtype=torch.float32,
)
然后由 softmax 转为概率分布：
1
gates = F.softmax(logits, dim=1)
logits 和 gates 的 shape 都是：
1
[B × S, E]
其中第一维表示展平后的 token，第二维表示专家。dim=1 是专家维，因此 gates[token_idx, :] 表示一个 token 对全部 num_experts 个专家的概率分布，并且这些概率之和为 1。
Routed Experts 推理路径
HunyuanMoE.forward 实际上实现了两种专家执行路径：
flashinfer：将所有 Expert 权重堆叠为带专家维度的 Tensor，然后调用flashinfer.fused_moe.cutlass_fused_moe，由 fused MoE Kernel 统一完成选中专家的计算和输出组合，避免在 Python 中逐个遍历 Expert。- 非
flashinfer：使用 eager 路径。先将每个 token 复制 $K$ 次，再根据topk_idx为每个 Expert 生成 mask，通过 Python 循环分别执行 Expert MLP，最后使用topk_weights对 $K$ 个专家输出加权求和。
flashinfer 路径 | 非 flashinfer 路径（eager） |
|---|---|
Expert 0 weights ─┐
Expert 1 weights ─┤
... ├── stack
Expert 63 weights ┘
│
├── moe_weight
└── moe_weight_2
│
hidden_states + topk_index + topk_weight
│
▼
cutlass_fused_moe
│
├── 根据 topk_index 调度
│ 选中的 token–expert 计算
├── 不在 Python 中遍历 Expert
├── 不计算未选中的 token–expert
└── fused 输出组合
│
▼
combined_output [B, S, H]
|
hidden_states + topk_idx + topk_weights
│
├── 每个 token 复制 K 次
└── 根据 topk_idx 生成 mask
│
▼
for i in range(num_experts)
│
├── 取出发送给 Expert i 的 token
├── experts[i](selected_inputs)
└── 将结果写回 expert_outputs
│
▼
所有 Expert 输出收集完成
│
├── reshape [B × S, K, H]
├── 乘以 topk_weights
└── 沿 K 维求和
│
▼
combined_output [B, S, H]
|
下面继续非 flashinfer 的 eager 路径。
每个 routed expert 都是一个独立的 HunyuanMLP：
1
2
3
4
5
6
7
8
9
10
11
self.experts = nn.ModuleList(
[
HunyuanMLP(
config,
layer_idx=layer_idx,
is_shared_mlp=False,
is_moe=True,
)
for _ in range(self.num_experts)
]
)
ModuleList 中的每个 HunyuanMLP 都拥有独立参数。公开配置中 num_experts = 64，因此每个 MoE 层会创建 64 个 routed experts。
Gate 返回：
1
2
3
4
topk_weights, topk_idx = self.gate(
hidden_states,
topk_impl="easy",
)
为了按照专家下标分发 token，首先将 topk_idx 展平：
1
flat_topk_idx = topk_idx.view(-1)
得到：
1
flat_topk_idx.shape = [B × S × K]
每个 token 会被发送给 $K$ 个专家，因此输入也需要沿 token 维复制 $K$ 次：
1
2
3
4
5
6
7
8
9
hidden_states_flat = input_hidden_states.view(
-1,
hidden_size,
)
hidden_states_repeated = hidden_states_flat.repeat_interleave(
self.moe_topk,
dim=0,
)
hidden_states_repeated 与 flat_topk_idx 的第一维一一对应：
1
hidden_states_repeated[j] → flat_topk_idx[j]
第 j 个 token 副本需要发送给下标为 flat_topk_idx[j] 的专家。
接下来遍历所有专家，每个专家推理激活它的 token：
1
2
3
4
5
6
7
8
9
10
11
12
13
expert_outputs = torch.zeros_like(
hidden_states_repeated,
)
for i in range(self.num_experts):
expert_mask = flat_topk_idx == i
selected_inputs = hidden_states_repeated[
expert_mask
]
expert_output = self.experts[i](
selected_inputs
)
expert_outputs[expert_mask] = expert_output
会找出所有被路由到 Expert i 的 token 副本。selected_inputs 的 token 数量由当前专家实际接收到的 token 副本数决定，因此不同专家的 selected_inputs.shape[0] 可能不同。
这里的遍历不同 Expert 的计算之间没有数据依赖
所有专家执行完成后，expert_outputs 的 shape 仍然是：
1
[B × S × K, H]
并且每个位置保存对应 token 副本经过目标专家后的输出。
最后将每个原始 token 对应的 $K$ 个专家输出重新组合：
1
2
3
4
5
6
7
8
combined_output = (
expert_outputs.view(
bsz * seq_len,
self.moe_topk,
hidden_size,
)
* topk_weights.unsqueeze(-1)
).sum(dim=1)
其中：
1
2
expert_outputs.view(...).shape = [B × S, K, H]
topk_weights.unsqueeze(-1).shape = [B × S, K, 1]
路由权重会广播到 hidden size 维，然后在 $K$ 个专家维度上加权求和，得到：
1
combined_output.shape = [B × S, H]
flashinfer 对上述推理过程的优化：
eager 路径将 Expert 保存在 ModuleList 中，并通过 Python 循环依次调用：
1
2
3
expert_output = self.experts[i](
selected_inputs
)
flashinfer 路径不会逐个调用 self.experts[i]。第一次进入 fused MoE 路径时，会先将所有 Expert 的权重堆叠为带专家维度的 Tensor：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
if not self._weights_initialized:
self._initialize_weights_on_device(
hidden_states.device
)
def _initialize_weights_on_device():
...
self.moe_weight = torch.stack( # [num_experts, 1.5 * H, H]
expert_weights_gate_up
).contiguous()
self.moe_weight_2 = torch.stack( # [num_experts, H, 0.75 * H]
expert_weights_down
).contiguous()
...
公开配置中 hidden_size = 4096，moe_intermediate_size = 3072 = 0.75H。当前 Expert 使用 SiLU/SwiGLU，融合的 gate/up projection 输出维度为 $2 \times 0.75H = 1.5H$，因此堆叠后的权重 shape 为：
1
2
moe_weight.shape = [E, 1.5 × H, H]
moe_weight_2.shape = [E, H, 0.75 × H]
第一维 $E$ 是专家维。因此，这里不是将 64 个 Expert 拼成一个二维大矩阵，而是构造两个带专家维度的权重 Tensor：
1
2
3
4
5
6
7
8
9
10
11
moe_weight
├── Expert 0 gate/up weight
├── Expert 1 gate/up weight
├── ...
└── Expert E-1 gate/up weight
moe_weight_2
├── Expert 0 down weight
├── Expert 1 down weight
├── ...
└── Expert E-1 down weight
最后将展平后的 token、Top-K 专家下标、Top-K 路由权重，以及堆叠后的 Expert 权重统一传给 fused MoE Kernel：
1
2
3
4
5
6
7
8
9
10
flashinfer.fused_moe.cutlass_fused_moe(
reshaped_input.contiguous(),
topk_index.to(torch.int).contiguous(),
topk_weight.to(torch.float).contiguous(),
self.moe_weight,
self.moe_weight_2,
torch.bfloat16,
output=combined_output,
quant_scales=None,
)
fused MoE Kernel 根据 topk_index[token_idx, k] 找到每个 token 需要使用的 Expert 权重，并使用 topk_weight[token_idx, k] 对选中的 $K$ 个 Expert 输出进行加权组合。
FlashInfer 会将 $B \times S \times K$ 个 token–expert 计算任务按照 Expert 分组，再通过 CUTLASS Grouped GEMM 统一提交给 GPU，由不同 CUDA Block 并行处理。这样可以避免在 Python 中通过 for 循环逐个调用 Expert，同时只计算 Top-K 选中的 Expert，而不是计算全部 Expert 后再使用 Mask 过滤。
更多实现细节可以参考 FlashInfer cutlass_fused_moe 官方文档。
Shared MLP
当 use_mixed_mlp_moe=True 时，HunyuanMoE 还会创建一个 shared_mlp。它不经过 Router，也不进行 Top-K 选择，而是让所有 token 都经过同一个 HunyuanMLP：
1
2
3
┌── Shared MLP ───────────┐
hidden_states ───────┤ ├── 相加 → output
└── Router → Top-K Experts┘
Shared MLP 和 Routed Experts 接收相同的 hidden_states，二者是并行分支，不存在前后依赖关系。代码中先计算 Shared MLP，再计算 Routed Experts，最后将两条分支的输出相加。