DeepInfer 源码拆解（三十五）：MoE 层的 MLP——门控路由、分组 GEMM 与三套专家路径

从第三十四篇接上：注意力闭环之后，剩下的那一半

第三十四篇结束时留下一句伏笔：注意力主线里的「索引器 → 候选块 → topk → 稀疏注意力」已经闭环，下一块是 mlp.py 的 616 行。这一篇就是对它的展开。

先把前面几篇里散落的线索收拢一下。DeepSeek-V4.1-Flash 的 Transformer 块里，注意力只占一半，另一半是 MoE：hidden_size = 5120，n_routed_experts = 384，num_experts_per_tok = 6，n_shared_experts = 1，moe_intermediate_size = 2304。这意味着每处理一个 token，理论上会触达 6 个细粒度专家，每个专家有两个矩阵：gate_up（5120 × 4608）与 down（2304 × 5120）。对一个 batch 来说，路由是逐 token、逐专家独立发生的，任何两个 token 命中的专家组合都不相同。

这就是为什么 MoE 层是整个推理引擎里「形状最不规整」的一层。注意力的形状由 batch 与序列长度决定，是稠密、规则的；MoE 的形状由路由结果决定，是稀疏、离散、不均匀的。前三十四篇里，MoE 专家并行与 EPLB 讲了专家怎么放、冗余副本怎么排、物理专家怎么重排；DP Attention 与 Wide-EP 讲了 EP 与 MegaMoE 的拓扑前提；userbuffer 讲了通信缓冲怎么分配。但它们都没有逐行拆 DeepseekMoE.forward 本身：门控怎么把 hidden 变成路由表、路由结果怎么被 EPLB 映射成物理槽位、三套 GEMM 路径各自执行什么、权重在装载时被重排成什么布局。

这一篇补上这一块。拆的时候会不断回到一个中心问题：同一段代码，为什么要在 fp8×fp8、fp8×fp4、bf16×bf16 三种精度组合下都能走通，又为什么要在标准 EP 路径与 MegaMoE 融合路径之间做选择。 V4.1-Flash 的 quantization_config 写得很直白：quant_method: fp8、weight_block_size: [32, 32]、scale_fmt: ue8m0、expert_dtype: fp4。也就是说激活走 FP8，专家权重走 FP4。把 MoE 层设计成「激活精度 × 权重精度」的矩阵（代码里叫 mma_type），是理解这一层所有分支的第一把钥匙。

一个文件，四种路径：mlp.py 的入口结构

deepinfer/model_executor/models/v4_modules/mlp.py 只有 616 行，但它的职责非常多。文件头注释写得很清楚：

DeepseekMoE adapts the moe.py container into the runtime expert layout and implements EPLB-aware routing. process_moe_weight and _convert_mega_moe_weight repack HF expert weights into the packed FP8/FP4 and mega-MoE layouts; they allocate runtime parameters, not HF-shaped ones.

这句话把职责分成了两半：一半是参数（把 HF 形状的权重重排成运行时布局），另一半是运行时（把一次前向拆成 gate、dispatch、GEMM、combine、shared 五步）。上游的 moe.py 只负责参数容器（FusedExperts、TopKGate、ReplicatedSharedExpert），并且把三套专家计数的语义写在了文件头：

• logical = cfg.moe.n_routed_experts，门控选择的取值范围；
• extra = cfg.moe.n_redundant_experts（由 model_loader._configure_moe_redundancy 在装载时写入），冗余副本数；
• physical = logical + extra，实际分配的参数行数，必须能被 ep_size 整除；
• local = physical / ep_size，本 rank 实际持有的专家个数。

这四个词在后面每一处都会用到：门控按 logical 选，EPLB 把 logical 映射到 physical，GEMM 按 local 的物理专家执行，通信按 physical 的总数在 rank 间分配。

先不急着进代码细节，先把一次前向的骨架画出来。DeepseekMoE.__init__ 里第一个真正的决策是 mma_type：

1
self.use_fp8_act = cfg.get("fp8_mlp", False)
2
self.use_fp4_weight = cfg.get("mxfp4", False)
3
act_type = "fp8" if self.use_fp8_act else "bf16"
4
weight_type = "fp4" if self.use_fp4_weight else "bf16" if moe_module.linear1.weight.dtype == torch.bfloat16 else "fp8"
5
self.mma_type = act_type + "x" + weight_type
6
assert self.mma_type in ("fp8xfp4", "fp8xfp8", "bf16xbf16")

三种组合对应三种硬件行为：

接着是第二个决策，use_mega_moe：

1
self.use_mega_moe = (
2
    (PLATFORM.is_gpu or self.mma_type == "fp8xfp4")
3
    and self.mma_type in ("fp8xfp4", "fp8xfp8", "bf16xbf16")
4
    and PG.is_within_nvl_domain(PG.expert_parallel_group())
5
)

三个条件分别对应三个硬前提：平台是 GPU（或者虽然是 NPU，但精度组合是 fp8×fp4 这种 NPU 上还没有替代实现的情况）、精度组合合法、EP 组落在同一个 NVLink 域内。前两个是「能不能算」，第三个是「能不能用对称内存互访」。三者同时成立才走 MegaMoE。注意第三个条件：NCCL 与 DeepEP 的双边通信可以跨节点，对称内存不行——对称内存要求每个 rank 在相同的虚拟地址上看到其他 rank 的缓冲，跨节点时没有任何机制能保证这一点。

于是整个文件前向的入口其实只有两行判断：

1
def forward(self, hidden_states, hidden_states_bf16, input_metadata):
2
    if self.use_mega_moe:
3
        return self.mega_moe_forward(hidden_states, hidden_states_bf16, input_metadata)
4
    topk_idx, topk_weights = self.gate_forward(hidden_states_bf16, input_metadata)
5
    ...

一个 if 之下，两种完全不同的执行结构。后面几节会分别展开。

标准路径的执行图可以概括成「门控 → dispatch → grouped GEMM → SwiGLU → grouped GEMM → combine」，MegaMoE 路径则把这些全部交给一个融合内核。DeepSeek 的技术报告里给过一张图，把 V4 的 EP 方案与其他工作放在一起对比：Comet 只把 Dispatch 与 Linear-1 重叠、把 Linear-2 与 Combine 重叠，各自独立；V4 则让整条 EP 流水在一个会话里连续进行，前一个专家的通信与后一个专家的计算交叉，而不是按阶段打隔断。DeepInfer 的 MegaMoE 就是这个思路在开源实现里的版本。

（图源：DeepSeek-V4 技术报告 DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence，Figure 5。）

读这张图时抓两件事：横轴是时间，纵轴是各个 rank；深色块是计算，浅色块是通信。Comet 那一行里，计算和通信是「配对出现的正交块」——每段通信只和相邻的一段计算重叠。V4 那一行则是同一条流水上每个专家依次走完 dispatch、GEMM、combine，相邻专家的三个阶段互相错开，整段时间里没有任何一个 rank 完全空闲。图中的 Linear-1 就是 gate_up，Linear-2 就是 down，Activation 就是 SwiGLU——这三者的关系正是下一节 routed_experts_forward 要拆的。

DeepEP 的 README 里也有一张更细的内核级示意图，展示「normal kernel」在低延迟场景之外的高吞吐工作模式：它把 token 按专家分组后连续排布，一次通信把一整个 batch 的所有专家数据打包发出去，而不是每个 token 一封信。这是理解「为什么 MoE 的 GEMM 要写成 contiguous 布局」的视觉背景。

（图源：DeepEP 官方仓库 figures/normal.png。）

门控：从隐藏状态到路由表

门控是 MoE 里最不起眼、却最容易出错的一步。它要做五件事：打分、选 top-k、拿偏置纠正分数、归一化、把 logical 专家号翻译成 physical 槽位。代码把前四件事和第五件事分开了：前四件在 gate_forward 里交给 moe_topk_gate 或 bf16_mega_gate 内核，第五件通过参数 to_physical_map 与 logical_count 一起塞进内核。

先看参数构造。__init__ 里门控被重新封装成 ReplicatedLinear，并且强制输出 FP32：

1
self.gate = ReplicatedLinear(
2
    moe_module.n_embed,
3
    self.num_routed_experts,
4
    use_fp32_output=True,
5
    weight=moe_module.gate.Wgate.weight,
6
)

为什么是 FP32 输出？因为打分函数的输出要参与 top-k 比较，而 FP8/BF16 的 mantissa 位数不足以在接近的分数之间可靠地区分。V4 的 scoring_func 是 sqrtsoftplus，topk_method 是 noaux_tc，norm_topk_prob = True——这三点在 __init__ 里都被断言：

1
assert cfg.moe.topk_method == "noaux_tc"
2
assert cfg.moe.norm_topk_prob
3
assert self.scoring_func == "sqrtsoftplus"

sqrtsoftplus 的意思是 $\mathrm{Softplus}(x)^{1/2}$，即

其中 $x_i$ 是门控线性层对专家 $i$ 的原始 logit。与单纯的 Softplus 相比，开平方把分布尾部压平：专家之间的分数差异被压缩，top-k 边界附近更容易出现接近的分数，这正是「偏置纠正 + 归一化」这套机制发挥作用的场景。开平方也让梯度在训练时更平缓——但这是训练侧的事情，推理只需知道打分函数的形状。

noaux_tc 是 V3 引入的「无辅助损失 + 专家偏置」的路由：训练时用偏置项的梯度替代传统的辅助负载均衡损失。推理时偏置是一个固定的 FP32 向量 e_score_correction_bias，形状为 [num_routed_experts]。用法分两步：选择时把偏置加到原始分数上，用纠正后的分数做 top-k；加权时仍然用原始分数（代码里叫 affine score）在选中的 $k$ 个专家之间归一化，再乘 routed_scaling_factor。偏置只改变「选谁」，不改变「权重多大」——这是这套机制与辅助损失最关键的区别：如果偏置也进了权重，负载均衡的力度会直接扰动模型输出。gate_kwargs 里同时出现 bias、routed_scaling_factor 与 scoring_func，正是把这三件事都交给门控内核：内核按 scoring_func 打原始分，用 bias 选，用原始分算权重。

V4.1 多了一个东西：图像 token 独立的路由偏置。moe.py 里 TopKGate 同时持有两个偏置：

1
self.e_score_correction_bias = nn.Parameter(torch.zeros(self.n_routed_experts, dtype=torch.float32))
2
self.e_score_correction_bias_vl = nn.Parameter(torch.zeros(self.n_routed_experts, dtype=torch.float32))

gate_forward 根据 input_metadata.images_seq_mask 判断本步是否有图像 token，若有则传 image_bias 与 image_token_mask：

1
images_seq_mask = getattr(input_metadata, "images_seq_mask", None)
2
gate_kwargs = dict(
3
    ...
4
    bias=self.e_score_correction_bias,
5
    image_bias=self.e_score_correction_bias_vl if images_seq_mask is not None else None,
6
    image_token_mask=images_seq_mask,
7
    ...
8
)

为什么图像 token 需要独立偏置？文本 token 与图像 token 在隐藏空间中的分布差异很大：图像 token 来自视觉编码器的回填（视觉编码器 VE 一篇里讲过嵌入怎么被填回序列），它们天然聚集在另一簇区域，如果共用一套偏置，图像 token 会把一部分专家打得更热，把文本 token 挤出去。分离偏置相当于给两类 token 各自留一份负载均衡字典。代价也很明确：路由选择变成了「按 token 类型分支」，内核要额外读一个 mask 和一张偏置表。

除了偏置，gate_kwargs 里还有四个值得单独讲的东西：

1. mask=input_metadata.valid_token_mask。这一步的 batch 里有 padding token（CUDA Graph 的图尺寸对齐造成的），它们不该参与路由统计，也不该污染 MoE 均衡指标。
2. force_random。这是一个与正确性无关、与压测有关的开关。input_metadata.gate_force_random_cp（CP 下的独立开关）与 gate_force_random 决定路由是否用随机值代替真实分数。作用是在 profiling 时把专家的选择分布从真实路由变成均匀随机，从而测出「路由完全均衡」与「路由完全倾斜」两种极端下的 MoE 表现。后面「工程边界」一节会讲一个更强的压测旋钮：DEEPINFER_FORCE_IMBALANCE_RATIO。
3. to_physical_map 与 logical_count。这就是 EPLB 的映射表：门控选出 logical 专家号之后，必须把它转换成 physical 槽位才能取到参数矩阵。eplb/__init__.py 里定义了它的形状：[num_moe_layers, num_logical_experts, num_redundant_experts + 1]，无效位填 -1。第二维为什么是 num_redundant_experts + 1？因为一个 logical 专家至少有一个副本，最多有 1 + extra 个副本；+1 就是主副本的槽位。
4. use_mega_gate = PLATFORM.is_gpu。MegaMoE 路径上，门控自己也被融进了一个融合内核，叫 bf16_mega_gate；它在 GPU 上可用，在 NPU 上退回普通路径。

gate_forward 的最后两行就是真正的分叉：

1
if self.use_mega_gate:
2
    topk_idx, topk_weights = deep_gemm.bf16_mega_gate(hidden_states, self.gate.weight, **gate_kwargs)
3
else:
4
    topk_idx, topk_weights = moe_topk_gate(self.gate(hidden_states), **gate_kwargs)
5
return topk_idx, topk_weights

moe_topk_gate 来自 tile_kernels.moe，输入是已经算好的 FP32 logits；bf16_mega_gate 来自 DeepGEMM，输入是 BF16 hidden 与门控权重，GEMM + top-k 一起在内核里完成。为什么 MegaMoE 连门控都要融合？在标准路径上，门控是一个 [T, 5120] × [5120, 384] 的小 GEMM，输出 [T, 384] 的 FP32 logits，然后 top-k 内核读它。两步之间有一次中间张量落盘（写回 HBM 再读回）。在 decode 阶段，$T$ 很小（几十到几百），这次中间读写的开销与 GEMM 本身相当。融进一个内核之后，logits 留在寄存器或共享内存里，top-k 直接消费。

这里顺便说一句门控输出的 out= 参数。在 mega_moe_forward 里，门控的输出张量不是临时分配的，而是直接写进 sym_buffer 预分配的 topk_idx 与 topk_weights 切片：

1
self.gate_forward(
2
    hidden_states_bf16,
3
    input_metadata,
4
    out=(sym_buffer.topk_idx[:num_tokens], sym_buffer.topk_weights[:num_tokens]),
5
)

为什么？因为 MegaMoE 的融合内核需要从对称内存里读路由结果——每个 rank 都要看到其他 rank 的路由表才能决定往哪送 token。预先放进 sym_buffer，就避免了「先写普通显存、再复制到对称内存」这次额外的拷贝。

EPLB：逻辑专家到物理槽位的那一跳

门控选出来的是 logical expert 编号，范围 [0, 384)。但 GPU 上真正存在的参数行是 physical 槽位，数量是 384 + extra，并且要平均分给 ep_size 个 rank。中间这一跳由 EPLB 的三张表完成：

先看 extra 是怎么定下来的。HF 的 config.json 里并没有 n_redundant_experts 字段——它不是模型结构的一部分，而是部署时根据并行度决定的。model_loader._configure_moe_redundancy 里：

1
def resolve_num_extra_experts(num_routed_experts, min_num_extra_experts, ep_size):
2
    padding = -(num_routed_experts + min_num_extra_experts) % ep_size
3
    return min_num_extra_experts + padding

即「不小于下限、且让 physical 总数能被 ep_size 整除」的最小值。假设 DEEPINFER_OVERRIDE_NUM_EXTRA_EXPERTS = 16、ep_size = 8、logical = 384：384 + 16 = 400，正好被 8 整除，extra = 16。若 ep_size = 32，400 不能被 32 整除（余 16），于是 padding = 16，extra = 32，physical = 416。为什么要向上取整而不是向下？因为冗余副本是负载均衡的「弹性」，少几个不致命，但 physical 不能被整除会让每个 rank 持有的 local 专家数不同，GEMM 的分组布局、通信的 buffer 尺寸、权重的按 rank 物化都会失去统一形状——这比少几个副本严重得多。代码在取整后还会打一条 warning，把实际值写进日志：

DEEPINFER_OVERRIDE_NUM_EXTRA_EXPERTS=16 已向上取整为 32，以使 ... 能被 ep_size=32 整除

extra 的放置是 EPLB 的在线职责。eplb/policy.py 的 rebalance_experts_global 做两件事：先给热点 logical 专家分配冗余副本，再把所有 physical 槽位均衡地摊到各 rank。

第一件事是一个贪心循环。初始时，前 384 个 physical 槽位一一对应 logical 专家，logical_count 全为 1；剩下的 extra 个槽位是空的。每一轮，它计算「如果把下一个副本给专家 j，所有历史 step 中每副本负载的最大值会变成多少」，然后选使这个最大值最小的 j：

1
score = tokens_per_expert_all_diff / logical_count
2
score_added = tokens_per_expert / (logical_count + 1)
3
values, indices = score.max(-1, keepdim=True)
4
values = values.expand_as(score).contiguous()
5
score.scatter_(-1, indices, score_added.gather(-1, indices))
6
values.scatter_(-1, indices, score.max(-1, keepdim=True).values)
7
redundancy_indices = values.sum(0).argmin(-1)

这段写法值得解释。score 是「每个 step、每个 logical 专家的每副本平均负载」；score_added 是加了一个副本之后的新均值。values 先取每一行（每个 step）的最大值——因为瓶颈由最热的专家决定；scatter 把最大值那个位置替换成 score_added 后再取一次 max，就得到「把副本给 argmax 专家之后，这一行的最大值」。最后按 step 求和、取 argmin，选出让总瓶颈最小的专家。tokens_per_expert_all_diff 里加了一个 arange * 1e-4 的微小扰动，只为让分数互不相等、让 argmax 的结果确定——不然并列时取哪个专家会随实现而变。

分配完副本之后，第二次贪心把所有 physical 槽位分给各 rank。这一步的目标写得很清楚：

也就是让「历史各 step 中，最忙 rank 的负载」之和最小。它不是按当前这一步的瞬时负载贪心，而是按一段历史窗口的负载向量做多目标权衡——因为专家放置一旦改变，权重要重排、缓存要失效，不能每步都动。DEEPINFER_EXPERT_COUNT_WINDOW_SIZE 默认 100，DEEPINFER_MOE_LB_STAT_INTERVAL 默认 128，这两个参数决定统计窗口与重排触发周期。

两张映射表在门控内核里的用法，可以用一个例子讲清楚。假设 logical 专家 17 有两个副本，logical_to_physical_map[layer, 17, 0] = 93、[layer, 17, 1] = 401，logical_count[layer, 17] = 2。门控选出 17 之后，内核要在 93 与 401 之间挑一个：优先挑本 rank 持有的那个——这样这个 token 不需要跨 rank 发送，dispatch 的通信量直接少一份。如果两个副本都在别的 rank 上，就只能跨秩发送，选哪一个由内核的固定规则决定，关键是所有 rank 对同一 token 得出相同的结论，否则 combine 的归约会错位。

ep_rank 参数存在的意义正在于此：映射表只描述「哪个逻辑专家有哪几个物理槽位」，不描述「这些槽位在谁的卡上」；槽位号本身是按 rank 均分的（rank $r$ 持有 [r * n_local, (r+1) * n_local)），所以内核拿 ep_rank 与 num_local_physical_experts 一除，就知道每个候选槽位是不是本地的。num_physical_topk 仍然是 6，但选出来的是 physical 槽位，后面 dispatch、GEMM、combine 全程都用这个编号。logical_count 则告诉内核「这个逻辑专家最多有几个候选」：副本数为 1 时只需要读映射表的第 0 列，副本数大于 1 时才去读后续列。

expert_count 是另一张表：[num_moe_layers, num_local_physical_experts]，每个 step 清零，dispatch 时由 DeepEP 的 cumulative_local_expert_recv_stats 参数累加。它既是 EPLB 重排的输入，也是「MoE 均衡指标」的来源——第二十三篇里的 [stats] 日志中，MoE 负载比就是拿这张表算的。

标准路径：dispatch、两组 GEMM、combine

标准路径的代码集中在 routed_experts_forward 与 forward 里。先看 dispatch：

1
event_overlap = self.ep_buffer.dispatch(
2
    x=hidden_states,
3
    num_sms=envs.DEEPINFER_USERBUFFER_COMM_SMS,
4
    topk_idx=topk_idx,
5
    topk_weights=topk_weights,
6
    cumulative_local_expert_recv_stats=self._eplb_layer().expert_count,
7
    num_experts=self.num_physical_experts,
8
    num_max_tokens_per_rank=self.max_num_prefilling_tokens if input_metadata.is_prompt else self.max_num_decoding_tokens,
9
    expert_alignment=deep_gemm.get_mk_alignment_for_contiguous_layout(),
10
    previous_event=dispatch_pre_evt,
11
    async_with_compute_stream=True,
12
    allocate_on_comm_stream=dispatch_pre_evt is not None,
13
    do_handle_copy=False,
14
    do_cpu_sync=input_metadata.is_prompt and envs.DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE != "compact",
15
    do_expand=True,
16
    use_tma_aligned_col_major_sf=True,
17
    defer_epilogue=True,
18
)

逐参数解释，因为每一个都对应一个设计决策：

• num_sms 给通信划出 SM 配额。MoE 的 dispatch/combine 与 GEMM 争同一批流处理器；给通信留多少 SM，是「通信延迟」与「计算吞吐」之间的直接取舍。环境变量 DEEPINFER_USERBUFFER_COMM_SMS 默认由 userbuffer 模块决定；第二十八篇已经讲过这里的 buffer 分配。
• num_max_tokens_per_rank 决定通信缓冲的上界。prefill 与 decode 用的是两个不同的值：max_num_prefilling_tokens 远大于 max_num_decoding_tokens。DeepEP 需要这个上界来预先切分缓冲、避免运行时 realloc。
• expert_alignment 是分组 GEMM 的 M 维对齐。DeepEP 在 dispatch 时按这个值给每个专家的 token 段补齐，使后续 grouped GEMM 的每个专家段都从对齐边界开始。对齐值来自 deep_gemm.get_mk_alignment_for_contiguous_layout()，在 __init__ 里已经统一设置过：

1
expert_alignment = deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout()
2
deep_gemm.set_mk_alignment_for_contiguous_layout(expert_alignment)
3
tile_kernels.set_token_alignment(expert_alignment)

三行代码把同一个对齐值写进三个库：DeepGEMM 的 GEMM 分组、TileKernels 的 SwiGLU token 对齐、DeepEP 的 dispatch 补齐。任何一处不一致，SwiGLU 内核就会把 padding 行当成真实 token 去算，或者 GEMM 读到错位的段边界。

• do_expand=True。这是最容易被忽略的参数。路由之后，一个 token 要送给 6 个专家，每个专家看到的是「这个 token 的 hidden」这一行。do_expand 让 DeepEP 在 dispatch 时把每个 token 复制 top-k 份，按目标专家分组排布。于是 GEMM 的 A 矩阵行数从 $T$ 变成 $6T$。不做 expand 的话，GEMM 就得自己按 top-k 索引 gather，反而更慢。
• do_cpu_sync 只在 prefill 且 YoCo 不是 compact 模式时为 True。decode 阶段每步都要同步会直接杀掉流水；prefill 阶段 token 多、同步一次摊薄得多，而且调度需要精确的接收计数来切 chunk。
• defer_epilogue=True 与 use_tma_aligned_col_major_sf=True 都是把「收尾」推给下游内核：dispatch 只搬数据，不在通信流上做量化/布局转换；scale factor 的排布直接按 TMA 对齐的列主序写，供 GEMM 的 TMA 加载器直接使用。

返回值是一个 EventOverlap。event_overlap.current_stream_wait() 让计算流等待通信完成，同时拿到四个东西：

1
moe_in_hidden, _dispatched_topk_idx, dispatched_topk_weights, handle = event_overlap.current_stream_wait()

moe_in_hidden 的形状是 [num_expanded_tokens, hidden]，num_expanded_tokens 可以为零（一个专家都没收到 token 时也允许）。handle 携带 psum_num_recv_tokens_per_expert——每个本地物理专家收到的 token 数，再做前缀和。这个张量就是接下来两个 GEMM 与 SwiGLU 的「分段目录」。

第一组 GEMM 由 _grouped_gemm 发出：

1
gate_up_proj_out = torch.empty(
2
    (num_expanded_tokens, self.gate_up_proj_weight.shape[1]),
3
    dtype=torch.bfloat16, device="cuda",
4
)
5
_grouped_gemm(
6
    hidden_states,
7
    self._weight_and_sf(self.gate_up_proj_weight),
8
    gate_up_proj_out,
9
    psum_num_recv_tokens_per_expert,
10
    expected_m,
11
    self.use_fp8_act,
12
    self.use_fp4_weight,
13
)

_grouped_gemm 内部按精度分两条路：

1
if fp8_a:
2
    deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous(
3
        a, b, d, psum_num_recv_tokens_per_expert,
4
        recipe_a=(1, get_fp8_casting_granularity()),
5
        recipe_b=(1, 32 if fp4_b else get_fp8_casting_granularity()),
6
        use_psum_layout=True,
7
        expected_m_for_psum_layout=expected_m,
8
    )
9
else:
10
    deep_gemm.m_grouped_bf16_gemm_nt_contiguous(...)

m_grouped_*_contiguous 的语义是「只沿 M 轴分组，N 与 K 固定」。这正是 MoE 的形状：每个专家的 gate_up 权重都是 $K \times N$ 的同一形状，只有收到的 token 数 $M_e$ 不同。token 按专家连续排布，psum_num_recv_tokens_per_expert 给出每段的结束位置，内核按段切换权重指针。recipe_a=(1, 32) 表示激活按每 32 个元素一个 scale 做 FP8 分块量化；recipe_b=(1,32) 对应权重 32×32 块，与 config.json 里的 weight_block_size: [32, 32] 一致；fp4_b 时权重块是 32，但元素是 e2m1。

use_psum_layout=True 值得单独说。传统的 grouped GEMM 需要一张 m_indices（每个 token 属于哪个专家）来分组；psum 布局则只给「累计 token 数」，内核用二分或线性扫描把 token 映射到段。少读一张索引表，少一次 HBM 往返。代价是 token 必须在段内连续且按对齐补齐——这正是 do_expand 与 expert_alignment 保证的。

GEMM 之后有一段看起来奇怪的代码：

1
if isinstance(hidden_states, tuple):
2
    hidden_states[0].data = hidden_states[1].data = torch.Tensor()
3
else:
4
    hidden_states.data = torch.Tensor()

它把输入张量的 storage 指针换成一个空 Tensor，提前释放 moe_in_hidden 占用的显存。为什么不直接 del？因为 hidden_states 是函数参数，调用方还持有引用；把 .data 指向空张量相当于「就地掏空」，让这块显存立刻可被 caching allocator 复用。MoE 层的峰值显存往往就出现在「输入还在、gate_up 输出已分配」的那一刻，提前释放输入能把峰值压低一个 [6T, 5120] 的量级。这个技巧在 decode 时尤其重要：batch 小，但每层的中间张量仍然按 6 倍膨胀。

接下来是 SwiGLU 与第二次量化，由 TileKernels 的一个融合内核完成：

1
down_in_hidden = swiglu_forward_and_per_token_cast(
2
    gate_up_proj_out,
3
    **_get_swiglu_kwargs(self.use_fp8_act),
4
    psum_num_tokens_per_expert=psum_num_recv_tokens_per_expert,
5
    topk_weights=topk_weights,
6
    clamp_value=self.swiglu_clamp_value,
7
)

这一步同时干了五件事：把 [6T, 2*2304] 切成 gate/up 两半、做 silu(gate) * up、按 top-k 权重缩放、对结果做每 token 的 FP8 量化、按专家段写入连续布局。后面「SwiGLU 与 clamp」一节会展开。

第二组 GEMM 把 [6T, 2304] 压回 [6T, 5120]，然后进入 combine。注意 hidden_states 此时是 bf16，即使 down 的输入是 FP8、权重是 FP4，输出仍然回到 bf16——因为 combine 要把 6 份结果相加，FP8 累加会丢精度。

combine 的调用在 forward 的后半段，比 dispatch 多了一个 bias：

1
event_overlap = self.combine(
2
    moe_out_hidden,
3
    shared_out,
4
    handle,
5
    combine_pre_evt,
6
)

shared_out 作为 bias 传给 DeepEP 的 combine。这不是「先 combine 再相加」，而是在 combine 的归约末尾直接把共享专家的输出加进去，省掉一次 [T, 5120] 的读写。共享专家为什么能在这里算完？看它前面的三行：

1
combine_pre_evt = self.ep_buffer.capture()
2
with limit_num_sms(envs.DEEPINFER_USERBUFFER_COMM_SMS, enabled=PLATFORM.is_gpu):
3
    shared_out = self.shared_forward(hidden_states)

capture() 在当前流上打一个事件，limit_num_sms 把计算流的 SM 配额压到和通信一样多，然后执行共享专家的 MLP。这段时间里，combine 还没发出（它在等 combine_pre_evt），但通信流已经被 dispatch 的句柄占着；等 shared_forward 一结束，combine 立刻带着「通信 + 共享结果」一起发出。效果是：共享专家的计算被塞进「routed 专家的 combine 尚未开始」的空隙里，而 combine 的归约又把它的结果免费加进去。

limit_num_sms 的存在是为了避免共享专家把 SM 全占满，导致通信流的 kernel 无法启动。comm 与 compute 共享同一组 SM，配额是硬约束；如果不限，shared_forward 的 GEMM 会把 SM 抢光，combine 只能排队，重叠就失效了。它的实现只有四行：把 deep_gemm 与 tile_kernels 各自的 num_sms 调低 num_sms 个，退出时再恢复。这四行之所以管用，是因为后面所有 GEMM 与激活内核的 launch 配置都读这两个全局值来决定网格大小——不是真的把 SM 隔离出去，而是「让内核自报少用一些 SM」。

最后 current_stream_wait() 返回 moe_combined_hidden，形状 [T, 5120]。至此一次 MoE 前向结束。

标准路径的时间线可以画成一条窄带：gate → dispatch(通信流) → 等待 → GEMM1 → SwiGLU+量化 → GEMM2 → [combine(通信流) ∥ shared]。整条链里只有 combine ∥ shared 一处重叠，dispatch 与 GEMM 之间是硬等待。为什么不能把 dispatch 与上一层 GEMM 重叠？因为路由依赖上一层的 hidden，必须等上一层算完。这就是标准路径的下限，也是 MegaMoE 要动的地方。

权重重排：process_moe_weight 在准备什么

MoE 层的权重不是「原样加载、按需切片」。HF 检查点里每个专家都是独立的两三个矩阵，而运行时需要的是「专家维打头、每个专家内部满足 GEMM 布局」的打包张量。process_moe_weight 就是这条转换线。

先看参数容器的形状约定。linear.py 里 ParallelRoutedLinear 的注释写得很明确：

Expert-major grouped GEMM metadata: (n_local_experts * h_in, h_out). HF materialization allocates packed FP4 directly. Dummy initialization uses this 2D representation before converting to the runtime layout.

这里有一个容易把形状搞反的地方：ParallelRoutedLinear(h_in, h_out, n_local_experts) 调用的是 nn.Linear(h_out, h_in * n_local_experts)，于是 self.weight 的 PyTorch 形状是 [h_in * E, h_out]——输入维在前、输出维在后，专家沿第一维首尾相接。这正是注释里写的 (n_local_experts * h_in, h_out)。FusedExperts 的两层实例化印证了它：

1
self.linear1 = ParallelRoutedLinear(cfg.hidden_size, self.hidden_dim * 2, self.n_local_routed_experts, use_fp8_gemm=cfg.fp8_mlp)
2
self.linear2 = ParallelRoutedLinear(self.hidden_dim, cfg.hidden_size, self.n_local_routed_experts, use_fp8_gemm=cfg.fp8_mlp)

FusedExperts 把 n_routed_experts（physical 总数）除以 ep_size 得到 n_local_routed_experts，再传给两层——所以 linear1 的 weight 是 [n_local × 5120, 4608]，linear2 是 [n_local × 2304, 5120]。process_moe_weight 的第一段就是对着这个布局拆（n_experts 取的是 self.num_local_physical_experts）：

1
deepinfer_w = src_w.view(n_experts, -1, src_w.shape[1]).mT.contiguous()
2
src_w.untyped_storage().resize_(0)

以 linear1 为例：view(n_local, -1, 4608) 把首尾相接的本地专家拆成 [n_local, 5120, 4608]；mT 只转置最后两维，得到 [n_local, 4608, 5120]；contiguous() 让每个专家的矩阵在显存里连续。为什么目标形状是 [E, N, K] 而不是 [E, K, N]？因为 DeepGEMM 的 m_grouped_*_nt_contiguous 采用 NT 布局：A 是 [M, K] 行主序，B 以 [N, K] 存储，计算等价于 $A \cdot B^\top$。每个专家的 B 必须是 [N, K] 且独立连续，grouped GEMM 才能靠「切换 B 的基址」在专家之间跳转，而不必为每个专家复制一份描述符。

第二段处理 scale factor。FP8 权重的 scale 在检查点里是按块存的，形状 [E * (K/32), N/32]——对 linear1 是 [n_local * 160, 144]；同样先 view 再 mT，得到 [n_local, 144, 160]。接着按精度分两条路：

1
if use_fp4:
2
    assert ((deepinfer_w.view(torch.int32) & 0x03030303) == 0).all(), "Only support FP4-QAT checkpoint"
3
    _deepinfer_w, _deepinfer_sf = tile_kernels.quant.per_token_cast(
4
        (deepinfer_w.flatten(0, 1), deepinfer_sf.flatten(0, 1)),
5
        "e2m1", 32,
6
        x_block_size=(fp8_block_size, fp8_block_size),
7
        round_sf=True,
8
        use_tma_aligned_col_major_sf=True,
9
        use_packed_ue8m0=True,
10
    )
11
else:
12
    deepinfer_sf = deep_gemm.transform_sf_into_required_layout(
13
        deepinfer_sf, n, k,
14
        recipe=(1, fp8_block_size, fp8_block_size),
15
        num_groups=n_experts,
16
        is_sfa=False,
17
        disable_ue8m0_cast=False,
18
    )

FP4 那条路的断言值得停下来看。(w.view(int32) & 0x03030303) == 0 检查每个 FP4 值的最低两位是否为零。FP8-E4M3 的尾数有 3 位，而 FP4-E2M1 的尾数只有 1 位；一个由 QAT 训练出来的模型，权重本来就落在 FP4 可表示的格点上，低两位恒为 0。如果不是 QAT 检查点，直接截断会把权重悄悄改掉，模型质量下降却没有任何报错——所以这里选择「直接断言失败」。这也是为什么代码里写 Only support FP4-QAT checkpoint 而不是做一次四舍五入：错的是输入，不该由推理引擎替训练侧兜底。

per_token_cast 的 "e2m1"、块宽 32、x_block_size=(32,32) 表示沿 K 每 32 个元素一组、沿 N 每 32 行一块，算一个 FP32 scale，再把 scale 打包成 UE8M0（无符号 8 位指数）。use_tma_aligned_col_major_sf=True 让 scale 张量按 TMA 对齐的列主序写，GEMM 的 TMA 加载器可以直接搬。FP8 那条路则调用 transform_sf_into_required_layout，is_sfa=False 表示这是 B 侧（权重）的 scale，num_groups=n_experts 让转换按专家分组进行，disable_ue8m0_cast=False 允许把 scale 转成 UE8M0。

两个分支都做完后，源权重的 storage 被 resize_(0) 释放。这一步不是可选的优化，而是加载 384 个专家的必需品：V4.1-Flash 每个专家的 gate_up 与 down 合起来是 4608×5120 + 5120×2304 个元素，按原检查点的 FP8 字节（或包装后的 FP4 字节，取较大一侧）估算单个专家也有数十 MB 量级；384 个专家合起来就是十几 GB，如果不边转边放，加载阶段要同时持有源和目标两份，峰值翻倍。

FP4 路径还有一个标记：

1
if use_fp4:
2
    deepinfer_w.mxfp4 = True

mxfp4 是挂在 Parameter 上的一个属性，用来告诉后续的 kernel 选择逻辑「这个权重是 MXFP4 布局」；_weight_and_sf 返回 (w, w.sf) 还是 w，也由 use_fp8_act 决定——激活是 FP8 时，权重必须带 scale 一起传；激活是 BF16 时，权重本身就是 BF16，没有 scale。

MegaMoE 的权重还要再转一次。post_load_weights 里：

1
self.gate_up_proj_weight, self.down_proj_weight = _convert_mega_moe_weight(
2
    self.gate_up_proj_weight, self.down_proj_weight
3
)

deep_gemm.transform_weights_for_mega_moe 把 L1/L2 两块权重转成融合内核要求的交错布局。为什么不能复用标准路径的布局？因为 MegaMoE 的 GEMM 不是「一个专家一个 B 矩阵」的连续分组，而是把 gate/up 的交错、专家间对称内存寻址、共享专家的权重全部编进一个内核的寻址算术里。布局变了，权重就必须重排；_convert_mega_moe_weight 在转换后 clone() 一份、再把源 storage 释放，理由与前面相同。

这里有一处容易误读：post_load_weights 在 __init__ 里被调用（仅 use_mega_moe 时），但 __init__ 阶段权重还是 dummy。真正的权重加载路径在 model_loader 里，HF 加载器按 physical 槽位直接写进已分配的打包布局；post_load_weights 在加载完成后还会被调用一次，完成 MegaMoE 的二次转换。源码里 if self.use_fp4_weight: 分支选择 reconstruct_param 而不是 process_moe_weight，正是因为 HF 加载已经直接填好了 FP4 打包布局，这里只需要重建 Parameter 壳、保留 .sf 引用。

共享专家：为什么是 ReplicatedLinear

V4.1-Flash 只有 1 个共享专家（n_shared_experts = 1），但它的实现方式很值得看。moe.py 里：

1
hidden_dim = cfg.moe.moe_intermediate_size
2
hidden_dim1 = hidden_dim * 2 * n_shared_experts
3
hidden_dim2 = hidden_dim * n_shared_experts
4
self.linear1 = nn.Linear(cfg.hidden_size, hidden_dim1, bias=False)
5
self.linear2 = nn.Linear(hidden_dim2, cfg.hidden_size, bias=False)

共享专家不是 n_shared_experts 个独立模块，而是沿中间维拼接成一个宽矩阵。1 个专家时就是普通的 SwiGLU MLP；多个时，linear1 的输出是 [T, 2*dim*n]，SwiGLU 之后是 [T, dim*n]，linear2 再压回 hidden。这样做的好处是共享专家永远只有一次 GEMM 调用，不需要按专家循环，也不需要 grouped GEMM 的分段元数据。

运行时把它包成 ReplicatedLinear：

1
shared_gate_up_proj = _replicated_linear(shared_expert.linear1.weight)
2
shared_down_proj = _replicated_linear(shared_expert.linear2.weight)

ReplicatedLinear 是「完整权重、每 rank 一份」的线性层，不做张量并行切分。为什么共享专家要复制而不是切分？

1. 它不参与 all-to-all。 路由专家必须按 EP 分布，因为 384 个专家放不进一张卡；共享专家只有一个，参数量是 3 * hidden * dim ≈ 3 × 5120 × 2304 ≈ 35M（FP8 下约 35 MB），每卡一份完全放得下。复制它，就省掉了「dispatch 到某个 rank、算完再 combine 回来」的一整圈通信。
2. 每个 rank 都需要完整结果。 combine 的 bias 要加在本地那份 routed 输出上；如果共享专家是切分的，每个 rank 只有一部分中间维，必须先 all-reduce 才能得到完整的 [T, 5120]，这又是一次集合通信。
3. 它可以和 combine 重叠。 因为每卡算的都是同一份完整结果，shared_forward 可以在通信流跑 combine 的同时独立执行；切分版本则必须等 reduce，重叠不起来。

代价是每个 rank 都重复算同样的 35M 参数。在 EP64 的部署里，这意味着 64 份重复计算；但相对 384 个路由专家分摊到 64 个 rank（每 rank 6 个专家）的计算量，共享专家只占约 1/6 的 MLP 计算，而且它没有通信。这笔账在 decode 阶段尤其划算：decode 的瓶颈是通信延迟与 kernel 启动，不是这几百微秒的 GEMM。

shared_forward 的实现只有五行：

1
hidden_states = self.shared_gate_up_proj(hidden_states)
2
down_in_hidden = swiglu_forward_and_per_token_cast(
3
    hidden_states, **_get_swiglu_kwargs(self.use_fp8_act),
4
    clamp_value=self.swiglu_clamp_value,
5
)
6
return self.shared_down_proj(down_in_hidden, out_acc=out_acc)

注意它没有传 topk_weights——共享专家的输出不按路由权重缩放，它以 1.0 的系数直接加到 routed 结果上。out_acc 是「累加输出」：MegaMoE 路径里，shared_down_proj(..., out_acc=out) 把共享结果直接累加到已经算好的 routed 输出上，省掉一次单独的张量分配与加法。标准路径则由 combine 的 bias 完成同样的融合。

SwiGLU 与 clamp：一个被写进每一路 GEMM 的常数

SwiGLU 的定义不复杂：

其中 $\mathrm{SiLU}(z) = z \cdot \sigma(z) = \dfrac{z}{1 + e^{-z}}$，$\odot$ 是逐元素乘。gate_up_proj 一次算出 $[xW_{\text{gate}}, xW_{\text{up}}]$ 两半，SwiGLU 把前半过 SiLU、与后半相乘，输出维度减半。之后再过 down_proj 回到 hidden。

V4.1 在 SwiGLU 上加了一个 swiglu_limit = 10.0。技术报告里写得很具体：

我们在 DeepSeek-V4-Flash 与 DeepSeek-V4-Pro 的训练中，把 SwiGLU 的线性分量 clamp 到 $[-10, 10]$，同时把 gate 分量的上界压到 10。

对应到 TileKernels 的 swiglu_forward 内核，clamp_value 的语义正是这两条：

1
val_l = T.min(val_l, clamp_value)                       # gate：只压上界
2
val_r = T.max(T.min(val_r, clamp_value), -clamp_value)  # linear：双向 clamp
3
val = val_l / (1 + T.exp(-val_l)) * val_r

为什么 gate 只压上界、linear 要双向压？SiLU 在负半轴趋近 0，不会产生大幅值；真正会爆的是正半轴的 $z \cdot \sigma(z)$，所以 gate 只需要上界。而 linear 分量直接参与乘法，正负都可能出现离群值，必须双向夹住。这个区别不是实现随意，而是数学形状决定的。

这个 clamp 最初是为了训练稳定（报告明确说它「消除离群值、显著帮助稳定训练」），但它被写进了模型函数：推理时必须原样执行，否则输出分布会偏移。到了推理侧，clamp 还多了一层作用：它限制了进入 FP8 量化的动态范围。swiglu_forward_and_per_token_cast 在同一个内核里先算 SwiGLU、再按 token 做 FP8 量化；per-token scale 由该行最大绝对值决定，如果某个元素是 1e4 量级的离群值，整行的 scale 会被它绑架，其余元素只剩两三个有效位。clamp 到 10 之后，行内动态范围可控，8 位尾数的分辨率才真正用在有效数据上。

内核里还有一个 count_clamp 分支，用来统计被 clamp 的元素个数（clamped_count 是长度 4 的 int64：gate 上界、linear 上界、linear 下界、总元素数）。这不是推理必需，而是给训练/调试观察「clamp 触发频率」用的；生产推理传 None，走无统计的快路径。

topk_weights 也在这一步乘进去：

1
val = val_l / (1 + T.exp(-val_l)) * val_r * topk_weights_var

为什么在 SwiGLU 之后、down GEMM 之前乘，而不是 combine 之后乘？因为每个展开后的 token 行对应「原 token × 某个专家」，combine 要把 top-k 份相加。先乘权重，combine 就是纯加法；后乘则要先还原「哪几行属于同一个 token」，再做加权和，多一次索引与一次乘加。把缩放塞进 SwiGLU 的 epilogue，等于零成本。

routed_scaling_factor = 1.5 没有单独传进 SwiGLU，而是已经在门控里乘进了 topk_weights。这一点从 gate_kwargs 里的 routed_scaling_factor=self.routed_scaling_factor 可以看出：门控内核负责「归一化 top-k 概率 × 缩放因子」，SwiGLU 只负责用最终的权重缩放。两处都不重复乘，否则结果会大 1.5 倍。

还有一处细节：use_fast_certificate。当输入是 BF16、输出是 E4M3 且 scale 取整时，内核用 __expf 与 __fdividef 两个近似指令替代精确的 exp/除法，并在同一个 kernel 里做「证书」检验：取近似结果尾数的低 19 位，检查它离 E4M3 的舍入边界是否足够远；只要有一个值落在边界附近，这个 warp 就整体回退到精确路径重算。这样既拿到了近似指令的吞吐，又保证量化结果与精确路径逐位一致。这个设计出现在激活与量化的融合处，是因为这里的输出直接决定后续 GEMM 的输入——一行的舍入错一位，误差会沿两层 GEMM 放大。

MegaMoE：五步融成一个内核

MegaMoE 路径的代码比标准路径短得多，因为大部分工作被推给了 deep_gemm.fp8_fp4_mega_moe（或 deep_gemm.bf16_mega_moe）。mega_moe_forward 只做四件事：门控、把输入放进对称内存、预分配输出、发内核。

1
def mega_moe_forward(self, hidden_states, hidden_states_bf16, input_metadata):
2
    sym_buffer = self.sym_buffer
3
    num_tokens = hidden_states_bf16.shape[0]
4
    self.gate_forward(
5
        hidden_states_bf16, input_metadata,
6
        out=(sym_buffer.topk_idx[:num_tokens], sym_buffer.topk_weights[:num_tokens]),
7
    )
8
    if self.use_fp8_act:
9
        assert hidden_states[0].data_ptr() == sym_buffer.x.data_ptr()
10
        assert hidden_states[1].data_ptr() == sym_buffer.x_sf.data_ptr()
11
    else:
12
        sym_buffer.x[:num_tokens].copy_(hidden_states_bf16)
13
    out = torch.empty((num_tokens, self.hidden_size), dtype=torch.bfloat16, device="cuda")
14

15
    self.mega_moe_func(
16
        out,
17
        self._weight_and_sf(self.gate_up_proj_weight),
18
        self._weight_and_sf(self.down_proj_weight),
19
        sym_buffer,
20
        **self.shared_weights_kwargs,
21
        cumulative_local_expert_recv_stats=self._eplb_layer().expert_count,
22
        activation=self.activation_type,
23
        activation_clamp=self.swiglu_clamp_value,
24
        fast_math=True,
25
    )
26
    if not self.fuse_shared:
27
        self.shared_forward(hidden_states, out_acc=out)
28
    return out

assert hidden_states[0].data_ptr() == sym_buffer.x.data_ptr() 是这段代码里最值得解释的一行。FP8 激活路径上，输入不是临时张量，而是直接落在对称内存 sym_buffer.x 里。上一层的输出（通常是融合的 norm+cast 内核）被要求写到对称内存；如果这条链路成立，mega_moe_forward 里就没有任何拷贝。怎么做到的？fused_norm_cast_kwargs 提供了这个钩子：

1
def fused_norm_cast_kwargs(self, num_tokens):
2
    if not (self.use_mega_moe and self.use_fp8_act):
3
        return {}
4
    sym_buffer = self.sym_buffer
5
    kwargs = {
6
        "out": (sym_buffer.x[:num_tokens], sym_buffer.x_sf[:num_tokens]),
7
        "fp8_cast_params": {"fmt": "e4m3", "num_per_channels": 32, "round_sf": True,
8
                            "use_packed_ue8m0": True, "use_tma_aligned_col_major_sf": False},
9
        "extra_sf_layout": "mega_moe",
10
    }
11
    if self.fuse_shared and PLATFORM.is_gpu:
12
        kwargs["extra_sf_out"] = sym_buffer.shared_l1_acts_sf[:num_tokens]
13
        kwargs["extra_sf_block_m"] = deep_gemm.get_block_m_for_mega_moe(...)
14
    return kwargs

hc.py 里的 _fused_cast_kwargs 会读这个方法，让 mHC 的 norm 内核把输出直接写进 sym_buffer.x，并把 scale 按 extra_sf_layout="mega_moe" 排布。第十篇讲过 mHC 主干：每个子层前都要对它做混合与归一化，这里正好复用同一个内核，把「归一化 → FP8 量化 → 写对称内存」压成一步。如果不成立（比如输入已经是 FP8 但来自别的来源），断言会直接失败——这比静默拷贝一份更安全，因为它保证 MegaMoE 永远读的是对称内存。

fuse_shared 是 self.has_shared 的别名：有共享专家时，post_load_weights 把共享权重的 kwargs 也打包进 shared_weights_kwargs，融合内核连共享专家一起算，out 里已经是「routed + shared」的完整结果；没有共享专家时，内核算完 out 后再调一次 shared_forward(..., out_acc=out)。V4.1-Flash 有 1 个共享专家，所以走前者。

融合内核要同时干这些事：门控结果读取（从 sym_buffer.topk_idx）、dispatch（对称内存直写）、第一组 GEMM、SwiGLU 与 clamp、激活量化、第二组 GEMM、combine、共享专家（可选）。它不需要一个协调进程去「串」这些步骤，因为每一步的依赖都编码在同一个 kernel 的 grid 与 block 划分里：每个 block 负责某个 rank 的某个专家段，按自己的 slice 把数据从远端对称内存拉过来、算完写回去。

融合的前提是三件事同时成立，缺一不可：

1. 对称内存可寻址。 每个 rank 用 SymmBuffer 分配，得到一组在所有 rank 上地址相同的指针；内核里直接按 base + rank * stride 访问其他 rank 的缓冲。跨节点时这个假设不成立，所以 use_mega_moe 的第三个条件是 is_within_nvl_domain。
2. 布局统一。 权重在装载时被转成融合布局；激活在上一层 norm 时就写进对称内存并按 mega_moe 的 scale 布局排好；门控输出写进 sym_buffer 的固定切片。三道准备缺一道，内核里就得多一次转换或拷贝。
3. 接收计数在线可得。 标准路径靠 dispatch 返回的 handle 拿 psum_num_recv_tokens_per_expert；融合内核在内部自己按路由表统计每个专家段的长度，所以不需要外部句柄。

cumulative_local_expert_recv_stats=self._eplb_layer().expert_count 仍然要传，因为 EPLB 与 [stats] 日志都要读专家 token 数；融合内核把统计累加写进这张表，与标准路径的语义一致。

前向时序：一次 decode 与一次 prefill 的差别

把两条路径放在同一张时间轴上，差别就清楚了。

decode。 假设 batch 256，每步每 token 出 1 个（或投机出 1+nextn 个）。展开后 num_expanded_tokens ≈ 1536。标准路径上，这条链是 gate(小 GEMM) → dispatch(all-to-all) → GEMM1 → SwiGLU → GEMM2 → combine(all-to-all)。其中两次 all-to-all 各是几十到上百微秒，GEMM 在 M=1536 时只需十几微秒，SwiGLU 是带宽受限的几微秒。整步时间由通信主导，而且 dispatch 与 GEMM 之间不可重叠。moe-ep-overlap 那张图里的 V4 方案要解决的正是这一段：让一个专家的 dispatch 与另一个专家的 GEMM 交叉。

MegaMoE 把「通信 + 计算」塞进一个 kernel 后，GPU 内部可以用 block 级流水重叠：一部分 warp 在等远端数据，另一部分在算本地专家段。这就是为什么 decode 阶段收益最明显——它本来就在等通信。

prefill。 假设 8192 个 token，展开后约 49152 行。GEMM 的 $M$ 达到万级，计算占主导；dispatch 的 num_max_tokens_per_rank 取 max_num_prefilling_tokens，通信量随 token 数增长但相对计算摊薄。此时两条路径的差距主要不来自「重叠」而来自「中间张量的读写」：标准路径要把 [6T, 5120]、[6T, 2304] 多次写回 HBM，融合内核则尽量让中间结果留在片上或对称内存。do_cpu_sync 在 prefill 默认开启（除非 YoCo 是 compact 模式），就是要让调度拿到精确的接收计数、配合 chunked prefill 切分。

max_num_tokens 也随之不同：

1
self.max_num_tokens = self.max_num_decoding_tokens if self.engine_type == EngineType.COMPLETION else self.max_num_prefilling_tokens

max_num_decoding_tokens = max_num_seqs * (1 + DEEPINFER_USE_NEXTN_DECODE)——decode 引擎按「最多序列数 × 每序列最多出几个 token」定上界，默认 1 + 5 = 6（nextn 默认 5）；prefill 引擎按 max_num_batched_tokens 定。这两种上界直接决定 userbuffer.allocate_ep 与 allocate_sym_buffer 的缓冲大小：decode 只需要 [256 × 6, 5120] 量级的缓冲，prefill 则需要 [8192, 5120] 量级，差一个数量级以上。这也是为什么部署时 P/D 分离的机器要分别配置这两个上界，而不是共用一个「够大」的值——多分配的那部分显存会挤占 KV 块。

工程边界：对齐、空专家与载荷不均

这一节收集几处容易踩、又不容易从日志里看出来的边界。

其一，对齐三方必须一致。 expert_alignment 同时写进 DeepGEMM 与 TileKernels，dispatch 的补齐也用它。SwiGLU 内核里有一段扫描逻辑，专门跨过对齐填充行：

1
while token_id >= boundary and expert_id < num_experts:
2
    expert_id += 1
3
    token_id += align(boundary, alignment) - boundary
4
    if expert_id < num_experts:
5
        boundary = psum_shared[expert_id]

如果一个 rank 的 expert_alignment 与 GEMM 的不同，这段逻辑会跳过错误的行数，结果是把某个专家的末尾几行算成下一个专家的开头——不会崩，只会悄悄算错。代码在 __init__ 里一次性把三处设成同一个值，正是为了避免这种不一致；任何在运行时单独改 set_mk_alignment_for_contiguous_layout 的操作都属于危险动作。

其二，空专家与零 token。 num_expanded_tokens can be 0 是代码注释里明确写下的。某些专家在整个 batch 中可能一个 token 都没收到，SwiGLU 内核在 num_expanded_tokens == 0 时直接跳过 launch；Combine 时这些段长度为 0，不参与归约。EPLB 的均衡指标里，这些专家的 expert_count 为 0，是重排时「把副本挪走」的依据。

其三，强制不均衡压测。 gate_forward 开头有一段：

1
if envs.DEEPINFER_FORCE_IMBALANCE_RATIO > 0 and self.ep_size > 1:
2
    assert envs.DEEPINFER_FORCE_IMBALANCE_RATIO >= 1.0
3
    factor = self._map_unbalanced_ratio_to_factor(envs.DEEPINFER_FORCE_IMBALANCE_RATIO)
4
    return self._get_selections_by_factor(num_tokens, factor, out=out)

它把真实路由替换成「偏向第 0 个 EP rank 的随机选择」，偏斜程度由 factor 控制。_map_unbalanced_ratio_to_factor 的推导值得看：

其中 $r$ 是期望的「本 rank 负载 / 其他 rank 平均负载」倍数，$S$ 是 ep_size，$k$ 是 top-k，$P$ 是物理专家总数，$n_{\text{local}} = P/S$。思路是均场近似：设本 rank 专家分数服从 $U(0, f)$、其余服从 $U(0,1)$，top-k 的分数阈值 $t$ 满足「本 rank 选中数 / 总选中数 = $x$」与「其余 rank 选中数 / 总选中数 = $1-x$」两个约束，解出 $t$ 与 $f$。当 $x k \ge n_{\text{local}}$（本地专家不够承受这么多选择）时返回 100.0 饱和——这是工程上的「尽力倾斜」，不是数学上的解。

这个旋钮的价值在于：MoE 的生产风险往往不是平均负载，而是热点专家把某个 rank 打爆。用随机路由强制构造倾斜工况，可以在上线前测出 EPLB 的重排是否跟得上、combine 是否在热点 rank 上排队。

其四，force_random 与 valid_token_mask 的语义不同。 force_random 是压测用的「路由变随机」；valid_token_mask 是正确性用的「padding token 不参与」。前者改变结果，后者保持结果不变。两者都在 gate_kwargs 里，但含义相反，读代码时要分清。

其五，CP 下的门控。 prefill 且 CP > 1 时，gate_force_random_cp 替代 gate_force_random：

1
if input_metadata.is_prompt and PG.context_parallel_group().size() > 1:
2
    gate_force_random = input_metadata.gate_force_random_cp
3
else:
4
    gate_force_random = input_metadata.gate_force_random

原因在第二十五篇讲过：CP 下每个 token 的注意力只在某个 rank 上算，但 MoE 要对全部 token 计算，于是门控结果必须在 CP 组内广播，且所有 rank 必须得到同一份路由。如果各 rank 独立用随机路由，广播前就已经分叉了。gate_force_random_cp 由 CP 的 rank 0 生成后再广播，保证一致性。

其六，显存峰值的两条线。 标准路径上，峰值出现在「moe_in_hidden 未释放、gate_up_proj_out 已分配」的时刻，约为 6T × 5120 + 6T × 4608 个元素；MegaMoE 则把中间结果放在 sym_buffer 里，峰值由 SymmBuffer 的 num_max_tokens_per_rank 决定。部署时要按引擎类型分别估算——用 prefill 的上界去跑 decode 引擎，会在启动 profile 阶段把显存算错，最后报「KV 块数为 0」。

其七，expected_m 只在 MegaMoE 路径非空。 标准路径里 routed_experts_forward 被调用时 expected_m=None，因为专家段长度是运行时才知道的；MegaMoE 的 launch 需要给 GEMM 一个「预期 M」来做 block 划分与 TMA 预取，deep_gemm.get_block_m_for_mega_moe(...) 会结合 ep_size、物理专家数、每 rank token 上界、top-k 与实际 token 数给出一个估计值。估大了浪费，估小了回退到更保守的划分——它影响性能，不影响正确性。

其八，EPLB 映射表的缓存一致性。 logical_to_physical_map 与 logical_count 是 GPU 张量，门控内核每步都要读。重排发生时，EPLB 在 CPU 上算好新的放置，再通过拷贝更新这两张表并 shuffle 权重。拷贝与「下一次门控读表」之间没有任何同步原语，依赖的是调度器的重排时机：重排只发生在引擎处理完一步、下一 step 还没排的间隙里。如果外部有别的线程在门控执行时改了表，读端可能看到「映射已更新、权重还没搬完」的中间状态。这和 C++ 调度器里 SchedulerOutputs 借用缓冲的生命周期约束是同一类问题：省掉拷贝和锁的代价，是把时序假设写进了调用约定。

其九，FP4 与对齐的隐藏前提。 部署文档里的约束不是一句「用 FP4 就快」：mma_type 只接受三种组合，FP4 权重只接受 QAT 检查点，M 侧分段对齐、K 侧 32 元素量化块、TMA 对齐的列主序 scale，三者都由 DeepGEMM 的接口约定固定。换模型或换精度时，最先失效的往往不是算子本身，而是这些「写在参数里、却没有断言」的形状假设；排查时应从 config.json 的 moe_intermediate_size、hidden_size 与 weight_block_size 逐项对起。

与其他框架的对照

把 DeepInfer 的设计放到主流框架里，能看到一些有意思的分歧。

最本质的差别在「谁替 MoE 做布局」这件事上。通用框架的 MoE 层要兼容几百种模型，因此把分组 GEMM、通信、权重量化都做成可插拔组件，每个组件之间用显式的中间张量衔接；DeepInfer 因为只服务自家模型，把「激活写进对称内存」「scale 按 mega_moe 排布」「门控输出写进同一块 buffer」这些准备步骤前移到上一层 kernel 里，让 MoE 层本身只剩下「发一个内核」。

另一个差别是 EPLB 的深度。多数框架的冗余专家是静态的、按训练统计预先分配；DeepInfer 把它做成了在线系统：expert_count 每步累加、窗口内统计、周期触发重排、权重 shuffle 与映射表热更新。第三篇、第四篇已经讲过它的运行时细节，这里只强调一点：所有这些开销都落在 MoE 层的「非关键路径」上，门控读映射表是常数时间，重排发生在调度空闲时刻。

小结

DeepInfer 的 mlp.py 用 616 行写完了一个 MoE 层的全部职责，线索可以收成三条。

第一条是精度矩阵。 mma_type = 激活 × 权重 决定了 GEMM 内核、scale 布局、权重打包格式与 MegaMoE 是否可用。V4.1-Flash 走的是 FP8 激活 × FP4 权重，对应 fp8_fp4_mega_moe；断言 Only support FP4-QAT checkpoint 把「模型本来就是 FP4」这个前提写进了代码。

第二条是路由的三张表。 门控选出 logical 专家，EPLB 用 logical_to_physical_map 与 logical_count 把它翻译成 physical 槽位，physical_to_logical_map 反过来服务重排与装载。偏置纠正（e_score_correction_bias 与图像专用的 _vl 版本）决定选谁，映射表决定送到哪张卡的哪个槽位。

第三条是两条执行路径。 标准路径是 dispatch → grouped GEMM → SwiGLU+量化 → grouped GEMM → combine ∥ shared，每一步都是显式的内核与中间张量；MegaMoE 把整条链压进一个融合内核，前提是对称内存、统一布局与 NVLink 域。选择由 use_mega_moe 的三条件决定，不是性能调优的结果，而是硬件拓扑的硬约束。

再往下一层还有两个「顺手带走」的设计：共享专家用 ReplicatedLinear 复制到每个 rank，让它不参与 all-to-all、还能与 combine 重叠；SwiGLU 的 clamp 与 top-k 缩放都塞进激活内核的 epilogue，让 down GEMM 的输入直接就是量化好的、按专家对齐的连续布局。这些不是孤立优化，而是同一条「减少中间张量与通信」的思路在不同位置的落地。

至此，「注意力 → MLP」这一整条 Transformer 块的主干已经在源码层面补齐。下一篇转向 Python 与 C++ 的边界：deepinfer._C.engine 如何把 Scheduler、Sequence、SchedulerOutputs 暴露给 Python，零拷贝视图的生命周期约束是怎么写进接口的，对象池如何避免每步分配，以及 GIL 在哪些调用点上被释放。这条边界不参与计算，却决定了上面这些内核能不能被按时发出去。

参考资料

1. DeepInfer 官方仓库（DeepSeek V4.1-Flash 推理框架）
2. mlp.py：MoE 运行时、门控、分组 GEMM 与 MegaMoE 路径
3. moe.py：FusedExperts、TopKGate 与 ReplicatedSharedExpert 参数容器
4. eplb/policy.py：冗余副本分配与逐层贪心放置
5. model_loader.py：物理专家冗余数与按 EP rank 物化
6. DeepGEMM：分组 GEMM、对称内存与 MegaMoE 接口
7. DeepEP：MoE dispatch/combine 通信库
8. TileKernels：swiglu_forward_and_per_token_cast 激活与量化融合内核
9. DeepSeek-V4 技术报告（arXiv:2606.19348）：EP 方案、SwiGLU clamp 与 FP4 量化
10. DeepSeek-V3 技术报告（arXiv:2412.19437）：无辅助损失负载均衡与 MoE 架构
11. DeepSeek-V4.1-Flash-1017 模型配置（config.json）
12. DeepEP 官方图：normal 工作模式的连续布局通信
13. 站内：DeepInfer 源码完全拆解（四）：MoE 专家并行与 EPLB
14. 站内：DeepInfer 源码拆解（十三）：DP Attention 与 Wide-EP
15. 站内：DeepInfer 源码拆解（二十八）：userbuffer，四类通信缓冲与对称内存
16. 站内：DeepInfer 源码拆解（三十二）：norm、RoPE、量化与写 KV 合成一个 kernel
17. 站内：DeepGEMM 完全拆解：连续布局、掩码布局与 FP8 分组 GEMM