MoE 专家并行深潜：从 All-to-All 通信瓶颈到 DeepEP 的工程突围

引言：参数不再等价算力

2024 年之后，MoE（Mixture of Experts）几乎成了前沿大模型的默认选择：Mixtral 8x7B、DeepSeek-V3/R1、Qwen3、Llama 4 都在走这条路。背后的工业逻辑很直接——Transformer 里 FFN 占了参数大头，而 FFN 的计算天然可切分。把一层 FFN 拆成 N 个专家，每次前向只激活 Top-K 个，就能让「模型参数量」与「单 token 计算量」解耦。

但这件事一点都不便宜。稀疏性换来的收益，会被一层从天而降的通信开销吃掉很大一块：All-to-All。本文就从这条通信链路入手，讲清楚 MoE 推理在工程上究竟难在哪，以及 DeepEP 这类专用通信库是怎么把它打到近乎隐形的。

一、MoE 一层的真实数据流

一个标准的 MoE 层，前向其实只有四步：

1. Gate 计算：token 隐藏态经过一个线性层得到 N 个专家的打分，取 Top-K（K 通常 1~8）。
2. Dispatch：把每个 token 的隐藏态，按路由结果发送到对应专家所在的 GPU。
3. Expert FFN：每张卡驻留若干专家，对收到的 token 做本地 FFN。
4. Combine：把专家输出发回原 token 所在的 GPU，按 gate 权重加权求和。

第 2、4 步就是两次 All-to-All。注意它不是 All-Reduce：All-Reduce 里每卡的收发总量等于自身数据量，而 All-to-All 里每张卡要向其余所有卡发送互不相同的分片。假设单步本地 token 数为 T，隐藏维度 d，专家并行度 EP=E，那么单次 dispatch 的每卡发送量约为 T·d·(E-1)/E ≈ T·d。

关键在于：这个量随 EP 线性增长，却完全不随 expert batch 摊薄。所以你在单机 8 卡上用 EP=8 很爽，一旦跨节点拉到 EP=64，token 有 63/64 的概率要走出机间 RDMA，时延从 NVLink 的 ~2μs 级别跳到 10μs 以上，带宽也从 900GB/s 掉到 400Gb/s（≈50GB/s）量级。这就是专家并行的第一堵墙。

二、路由不是免费的：负载均衡与容量上限

如果所有 token 都挤到同一个专家，那张卡就成串行瓶颈，其余卡空转，同时该专家还要按最大 token 数 padding。为此工程上有两件必做的事：容量截断与辅助损失。

# 简化版 Top-K 路由：含 capacity 截断与 Switch 风格辅助损失
import torch
import torch.nn.functional as F

class TopKGate(torch.nn.Module):
    def __init__(self, dim, num_experts, top_k=2, capacity_factor=1.25):
        super().__init__()
        self.w = torch.nn.Parameter(torch.randn(dim, num_experts) * 0.02)
        self.top_k = top_k
        self.capacity_factor = capacity_factor

    def forward(self, x):
        # x: [T, d]
        T, E = x.shape[0], self.w.shape[1]
        logits = x.float() @ self.w                     # [T, E]
        probs = torch.softmax(logits, dim=-1)

        topk_w, topk_idx = torch.topk(probs, self.top_k, dim=-1)
        topk_w = topk_w / topk_w.sum(-1, keepdim=True)  # 归一化门控权重

        capacity = int(self.capacity_factor * T * self.top_k / E)
        mask = torch.zeros_like(probs)
        mask.scatter_(1, topk_idx, 1.0)

        # token 在所属专家内的排队序号，超出 capacity 即被丢弃
        cum = mask.cumsum(0) - 1
        keep = (cum < capacity).gather(1, topk_idx).bool()

        # 辅助损失：鼓励 token 在专家间均匀分布
        f = mask.sum(0).float() / (T * self.top_k)      # 实际承载比例
        P = probs.mean(0)                                # 平均路由概率
        aux_loss = E * (f * P).sum()

        return topk_idx, topk_w, keep, aux_loss

这段代码暴露了三个真实的工程取舍：

• capacity_factor 越大，丢弃越少，但 padding 浪费越多。所有专家都要按 capacity 对齐计算。DeepSeek 系列在训练时用受控的 capacity，推理阶段则改为 dropless + 动态 padding，用通信量换精度。
• 辅助损失的系数不能太大。它会和主任务损失抢梯度，模型质量直接往下掉。更优雅的做法是给每个专家加一个可学习的 bias 项做动态偏置（loss-free balancing），而不是硬拉分布。
• 被丢弃的 token 会从残差分支原样穿过。意味着它参与了 Attention 却没经过 FFN 加工——训练与推理的 drop 策略必须一致，否则会出现分布漂移。

三、为什么 EP 比 TP 更吃 Batch

张量并行（TP）里，每卡处理同一批 token 的一段隐藏维度，通信是 All-Reduce，量级约 2·T·d，和 EP 的 dispatch 是同一数量级。真正的差别在并行效率：TP 的切分是确定的，而 MoE 的 token→expert 映射是数据依赖的随机路由。于是：

• EP 的通信 shape 每个 step 都不同，NCCL 的显存 buffer 必须按最坏情况预留；
• 单卡 batch 越小，泊松分布带来的专家间负载越不均衡，padding 占比越高；
• 因此小 batch 低延迟场景（在线对话）不适合大 EP，而大 batch 吞吐场景（离线批处理、长思考链解码）才划算。

一句话：EP 是用延迟换吞吐。而 PD 分离恰好把这两种场景切到了不同的物理设备上——Prefill 用大 EP 冲吞吐，Decode 用小 TP 保延迟。这也是当前主流推理栈的默认编排方式。

四、通信量的量化模型

设 EP=E，隐藏维 d，本地 token 数 T，dtype 字节数 b（FP8 为 1，BF16 为 2），每次前向两次 All-to-All，则单卡通信字节数近似：

Comm ≈ 2 · T · d · b · K · (E - 1) / E

比值 Comm/Comp 决定了一切。你能动的旋钮只有四个：降 K（激活专家数）、降 E（少跨节点）、降 b（FP8 dispatch）、加大 T。前两个分别由训练配置和显存容量锁死，所以工程上真正能抠的只有后两条——这正是 DeepEP 的发力点。

五、DeepEP：把通信塞进计算的阴影里

DeepEP 解决的核心问题有两个：跨节点低延迟，以及通信与计算的重叠。三处设计值得单独拎出来：

1. 双域 buffer 分离。 把 NVLink（机内、高带宽）和 RDMA（机间、相对低带宽）的收发路径拆成两条流水线。dispatch 时先经 NVLink 汇聚到本机同号 GPU，再由该 GPU 统一做 RDMA 转发，避免「一卡对多目标」导致的 RDMA QP 爆炸与队头阻塞。配合 Node-Limited Routing——每个 token 最多只被路由到 M 个节点（M 通常取 4）——用路由约束换通信可控。

2. 纯 RDMA 的 low-latency kernel。 传统 RDMA 路径要经 GPU kernel → CPU 代理线程 → NIC。DeepEP 在 decode 路径上让 GPU warp 直接向 NIC 的 doorbell 写 WQE，绕开 CPU 与额外显存拷贝，把单次 dispatch 压到亚毫秒级。

3. 通信 hook 与计算重叠。 上层框架（SGLang / vLLM）用一对 hook 把它包起来：

# SGLang 风格伪代码：dispatch 与专家计算的重叠
def moe_forward(x, gate, experts, ep_group):
    topk_idx, topk_w, _ = gate(x)

    handle = ep_group.dispatch(x, topk_idx, async_op=True)   # 立即返回 handle
    # —— 与通信重叠的部分：稠密分支、共享专家、输出投影 ——
    shared_out = experts.shared(x)

    recv_x, recv_idx = handle.wait_and_get()                 # 等待本次 dispatch 落地
    expert_out = experts.local_ffn(recv_x, recv_idx)         # 本地专家计算

    out_handle = ep_group.combine(expert_out, async_op=True)
    y = residual_norm(x + shared_out)
    return y + out_handle.wait()

这里的关键不是 async_op=True 这么简单，而是你要往通信窗口里塞足够多的独立计算。如果 MoE 层前后没有稠密分支、没有共享专家可算，这个 overlap 就是空的。这也是为什么 DeepSeek-V3 这类模型要设计「共享专家 + 路由专家」的混合结构——客观上它给通信留出了可重叠的计算预算。

# DeepEP 典型调用：机内 NVLink + 机间 RDMA 双域
buffer = deep_ep.Buffer(group,
                        num_nvl_bytes=int(2e9),
                        num_rdma_bytes=int(1e9))
buffer.set_dispatch_layout(topk_idx, num_experts)

recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens, handle, event = \
    buffer.dispatch(hidden_states, topk_idx, topk_weights,
                    num_experts, previous_event=event)

注意 num_recv_tokens 是动态的，buffer 必须按最大容量分配。这是 EP 显存开销里最不起眼、却最容易 OOM 的一块。经验值：

per_rank_rdma_buffer ≈ max_tokens_per_batch × hidden × top_k × sizeof(dtype) × 1.2

六、生产落地 checklist

• 并行组合：节点内 TP+EP（NVLink 域）、节点间纯 EP（RDMA 域）是主流；Prefill 大 EP 冲吞吐，Decode 大 TP/PP 压延迟，二者通过 PD 分离部署到不同实例池。
• 精度：dispatch/combine 用 FP8、专家权重用 W8A8 已基本是标配；但 gate 的 softmax 必须在 BF16/FP32 上完成，量化 gate 会显著掉点。
• 观测：必须单独埋点计时三部分——gate、dispatch wait、expert compute。若 dispatch wait 占比超过 30%，说明你已经摸到 EP 规模的上限，继续加卡只会更慢。
• 回退：常备 E=1（纯 TP/PP）的配置。流量低谷、batch 很小时，大 EP 的 padding 与同步开销会让单 token 成本反高于稠密部署。

七、结语

MoE 的本质是一次取舍：用通信换算力，用路由的不确定性换参数规模。它不是一个「免费获得更大模型」的开关，而是一整套关于路由、缓冲、精度与部署拓扑的工程。真正决定 MoE 推理性价比的，往往不是模型团队把 K 设成多少，而是 infra 团队能否把那两次 All-to-All 藏进计算的阴影里。