大模型推理引擎vLLM(27): TP场景moe与EP场景MOE的相同点和不同点、使用modular_kernel.py的差异

TP MoE：比如 8 张卡，每个 rank 都有全部 expert，对这一批 token 都会参与 expert 计算，只是把每个 expert 的权重矩阵按 TP 方式切开（w13 column parallel，w2 row parallel），最后 allreduce 拼结果。
EP MoE：把 expert 分到不同 rank，每个 rank 只存/算某几个 expert 的完整权重；token 通过 all2all 被 dispatch 到对应 rank，算完再 combine 回去。

2 相同点

都有 Router → topk 选 expert
都走 FusedMoE / default_moe_runner 大框架
都可能用到 modular_kernel.py 的 FusedMoEKernel 抽象（Prepare/Finalize + Experts）
最终都要把各 expert 输出按 topk weight 加权合并

3 不同点

维度	普通 TP MoE	EP / DeepEP MoE
专家分布	每 rank 全 expert	每 rank 部分 expert
权重形态	w13 column parallel，w2 row parallel	本地 expert 完整权重
跨 rank 通信	输出 allreduce	token all2all dispatch/combine
`expert_map`	`None`（ep_size=1）	有，映射 global→local expert
Prepare/Finalize	`MoEPrepareAndFinalizeNoDPEPModular`	`DeepEPLLPrepareAndFinalize` 等
activation format	`Standard`	`BatchedExperts`（DeepEP LL）
modular kernel 初始化	quant method 内部自建，或 layer 不包一层	`layer.maybe_init_modular_kernel()` 包成 `FusedMoEModularMethod`

4 怎么使用modular_kernel.py中的类的

4.1 modular_kernel.py文件结构

FusedMoEKernel（总调度）
├── Prepare/Finalize 侧：FusedMoEPrepareAndFinalizeModular
│ prepare(): 量化 +（可选）dispatch
│ finalize(): combine + topk weight reduce
│
└── Expert 计算侧：FusedMoEExpertsModular
apply(): 真正的 GEMM + activation（调 aiter/triton/deepgemm 等）
finalize_weight_and_reduce_impl(): 返回 TopKWeightAndReduce

其实就是这个文件有一个大的类，这个类里面又包含了两个类，其中一个是FusedMoEPrepareAndFinalizeModular用来做前处理和后处理，另一个是FusedMoEExpertsModular用来负责做专家计算。

4.2 EP MOE场景怎么使用的modular_kernel.py文件

4.2.1 什么时候会走这条路

EP 场景要同时满足（见 FusedMoEParallelConfig.use_all2all_kernels）：

use_all2all_kernels = (dp_size > 1) and enable_expert_parallel

典型启动参数：

-dp 8 -tp 1 --enable-expert-parallel --all2all_backend=deepep_low_latency

此时不会让每个 rank 算全部 expert，而是 expert 按 EP rank 切分；跨 rank 通信走 all2all dispatch/combine，不是 TP 的 allreduce。

4.2.2 初始化：layer 层统一组装

EP 路径的关键是：权重 load 完之后，FusedMoE.maybe_init_modular_kernel() 会把原来的 quant_method 替换成 FusedMoEModularMethod：

prepare_communication_buffer_for_model()
  └── FusedMoE.maybe_init_modular_kernel()
        ├── base_quant_method.maybe_make_prepare_finalize()
        │     └── all2all_utils.maybe_make_prepare_finalize()
        │           └── 创建 DeepEPLLPrepareAndFinalize（deepep_low_latency 时）
        └── FusedMoEModularMethod.make(...)
              └── FusedMoEKernel(prepare_finalize, select_gemm_impl(...))

对应代码逻辑：

maybe_make_prepare_finalize() 根据 all2all_backend 创建 Prepare/Finalize，DeepEP LL 时是 DeepEPLLPrepareAndFinalize
FusedMoEModularMethod.make() 里调用 old_quant_method.select_gemm_impl(prepare_finalize, layer) 选 Expert 实现
两者拼成 FusedMoEKernel

4.2.3 运行时一次 forward 在干什么

FusedMoEModularMethod.apply() 最终调到 FusedMoEKernel.apply()，内部是 FusedMoEKernelModularImpl，三步走：

Router（layer 外 / runner 里）
  → topk_weights, topk_ids
  → FusedMoEKernel.apply()
       ① Prepare/Finalize.prepare()
            量化 + DeepEP dispatch（token 发到拥有该 expert 的 rank）
            输出变成 BatchedExperts 格式：(num_local_experts, max_tokens, hidden)
       ② FusedMoEExpertsModular.apply()
            本地 expert 做 GEMM + activation
            例如 BatchedDeepGemmExperts / BatchedTritonExperts
       ③ Prepare/Finalize.finalize()
            DeepEP combine（结果送回 token 原 rank）
            + topk weight reduce

4.2.4 EP 场景用到的类（举例）

组件	DeepEP LL 典型类	文件
Prepare/Finalize	`DeepEPLLPrepareAndFinalize`	`deepep_ll_prepare_finalize.py`
Expert 计算	`BatchedDeepGemmExperts` / `BatchedTritonExperts`	`batched_deep_gemm_moe.py` 等
总调度	`FusedMoEKernel`	`modular_kernel.py`
包装层	`FusedMoEModularMethod`	`fused_moe_modular_method.py`

DeepEPLLPrepareAndFinalize 的 activation_format 是 BatchedExperts，所以 Expert 侧必须配 batched 版本的实现，不能配 Standard 格式的 AiterExperts。

4.3 TP MOE场景怎么使用的modular_kernel.py文件

TP 场景要分几条子路径说，不能一概而论。

4..3.1 路径 A：Oracle 内建 MK（当前主流，Aiter / Triton / FlashInfer / FP8 等）

典型配置：-tp N -dp 1，不开 --enable-expert-parallel；bf16/fp16 走 Aiter 或 Triton。

。

和 EP 的关键区别

维度	TP 路径 A	EP 路径
MK 持有者	`quant_method.moe_kernel`（内部 MK）	同上（已迁移 quant）；或 `FusedMoEModularMethod`（未迁移）
是否换 wrapper	不换 `FusedMoEModularMethod`	未迁移 quant 才换
`maybe_init_modular_kernel`	`supports_internal_mk=True` 时直接 return	同上；未迁移时才执行
Prepare/Finalize	`MoEPrepareAndFinalizeNoDPEPModular`	`DeepEPLLPrepareAndFinalize` 等
activation format	`Standard`	DeepEP LL → `BatchedExperts`

初始化：quant_method 内部组装（不是 layer 包一层）

以 UnquantizedFusedMoEMethod 为例：

FusedMoE.__init__
  └── quant_method = UnquantizedFusedMoEMethod
        └── select_unquantized_moe_backend()  # 选 Aiter / Triton / FlashInfer ...

process_weights_after_loading()
  └── _setup_kernel()
        └── make_unquantized_moe_kernel()          # oracle/unquantized.py
              ├── maybe_make_prepare_finalize(allow_new_interface=True)
              │     └── MoEPrepareAndFinalizeNoDPEPModular   # TP，无 all2all
              └── experts_cls(...)                        # AiterExperts / TritonExperts ...
              └── FusedMoEKernel(prepare_finalize, experts)
        └── self.moe_kernel = kernel   # supports_internal_mk = True