本文基于昇腾CANN和昇腾NPU，围绕 MoE Router 技术展开。

MoE（Mixture of Experts）把一个大 FFN 拆成几十到几百个"小专家"，每个 Token 只激活其中 Top-K 个。Router 就是这个分发系统——它看到每个 Token 后决定：你去 Expert 15，你去 Expert 88。CANN 在昇腾NPU上把 Router 的 Gate 计算和 Token 分发绑在一起，用 HCCL 的 AlltoAll 完成跨卡通信。

Router 就是一个小型网络

Router 本身有多简单？一个 Linear(hidden_dim, num_experts) 加 Softmax。参数量 = 4096 × 256 ≈ 1M——不到整个 MoE 层 2B 参数的 0.05%。但它要做的事是全局性的：为每一个 Token 从 256 个 Expert 里选出 Top-8。

# MoE Router——一个轻量网络，全局决策

class MoERouter(torch.nn.Module):
    def __init__(self, hidden_dim=4096, num_experts=256, top_k=8):
        super().__init__()
        self.gate = torch.nn.Linear(hidden_dim, num_experts, bias=False)
        self.top_k = top_k

    def forward(self, x):
        # x: [batch * seq_len, hidden_dim]
        logits = self.gate(x)       # [B*S, 256]
        scores = F.softmax(logits, dim=-1)
        topk_weights, topk_indices = torch.topk(scores, self.top_k, dim=-1)
        # 归一化——选中的 8 个 Expert 权重加起来为 1
        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
        return topk_weights, topk_indices

为什么是 Top-8 而不是 Top-1？因为单个 Expert 通常只在某个特定子领域很强——Expert 15 专长代码、Expert 88 专长数学。一个问题可能同时涉及代码和数学，需要两个 Expert 共同回答。Top-K 让每个 Token 可以用多个 Expert 的组合表示。

Router 的训练是通过辅助 Load Balancing Loss 来约束的。没有这个约束，Router 会倾向于把大部分 Token 发给同一个 Expert——模型的容量利用率极低。Load Balancing Loss 惩罚 Token 分配的不均匀——让 256 个 Expert 各收到约 1/256 的 Token。

Token 分发的通信瓶颈

Router 算完后，Token 要被实际发送到对应的 Expert 所在的卡上。假设 8 卡，每卡存 32 个 Expert。一个 Token 选中的 Expert 可能分布在 8 张不同卡上——Router 的输出是一个 Device 级别的"分发计划"。

// MoE 的 AlltoAll 通信——跨卡分发 Token

void TokenDispatch(int* topk_indices, float* hidden_states,
                   int num_tokens, int num_devices) {
    // 统计每张卡要发多少 Token
    int send_counts[8] = {0};
    for (int t = 0; t < num_tokens; t++)
        for (int k = 0; k < top_k; k++)
            send_counts[topk_indices[t * top_k + k] / experts_per_device]++;

    // HCCL AlltoAllV——不等长收发，因为每卡负载不同
    HcclAlltoAllV(hidden_states, send_counts, send_displs,
                  recv_buffer, recv_counts, recv_displs,
                  HCCL_FLOAT, num_devices, hccl_comm);
}

分发用的是 HCCL 的 AlltoAllV——不等长的 All-to-All。每卡收到的 Token 数量不同——卡 1 可能收到 200 个 Token，卡 2 只收到 150 个。AlltoAllV 允许每对 (src, dst) 之间发送不同长度的数据。4096 Token 的 Batch × Top-8 Expert = 32768 次 Expert 访问，每次发送 8KB 向量，总共约 256MB 的 AlltoAll 数据量。8 卡 800GB/s HCCS 互联下，通信时间约 0.3ms。

昇腾NPU的通信计算重叠

CANN 的优化策略是：不等 AlltoAll 完全完成就开始计算。收到一部分 Token 后，Expert 前向就开始了——边收边算。HCCL 的 AlltoAllV 是异步的——CPU 提交后立刻返回，DMA Engine 在后台搬数据。还有一个策略是"容量限制"——每个 Expert 最多处理 C 个 Token，超出的 Token 不计入该 Expert。C = (total_tokens / num_experts) × 1.25。超出的 Token 被丢弃——走 Residual 连接跳过 MoE 计算。这个丢弃对精度影响很小——超出容量限制的 Token 本来就是 Expert 的边缘负载。

参考仓库

MoE Router 等 Transformer 算子

HCCL 集合通信库

DeepSeek-R1 推理配方