本文基于昇腾CANN和昇腾NPU围绕 MoE Router 技术展开。MoEMixture 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——一个轻量网络全局决策classMoERouter(torch.nn.Module):def__init__(self,hidden_dim4096,num_experts256,top_k8):super().__init__()self.gatetorch.nn.Linear(hidden_dim,num_experts,biasFalse)self.top_ktop_kdefforward(self,x):# x: [batch * seq_len, hidden_dim]logitsself.gate(x)# [B*S, 256]scoresF.softmax(logits,dim-1)topk_weights,topk_indicestorch.topk(scores,self.top_k,dim-1)# 归一化——选中的 8 个 Expert 权重加起来为 1topk_weightstopk_weights/topk_weights.sum(dim-1,keepdimTrue)returntopk_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 通信——跨卡分发 TokenvoidTokenDispatch(int*topk_indices,float*hidden_states,intnum_tokens,intnum_devices){// 统计每张卡要发多少 Tokenintsend_counts[8]{0};for(intt0;tnum_tokens;t)for(intk0;ktop_k;k)send_counts[topk_indices[t*top_kk]/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 推理配方
MoE Router：谁来决定 Token 去哪个 Expert
本文基于昇腾CANN和昇腾NPU围绕 MoE Router 技术展开。MoEMixture 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——一个轻量网络全局决策classMoERouter(torch.nn.Module):def__init__(self,hidden_dim4096,num_experts256,top_k8):super().__init__()self.gatetorch.nn.Linear(hidden_dim,num_experts,biasFalse)self.top_ktop_kdefforward(self,x):# x: [batch * seq_len, hidden_dim]logitsself.gate(x)# [B*S, 256]scoresF.softmax(logits,dim-1)topk_weights,topk_indicestorch.topk(scores,self.top_k,dim-1)# 归一化——选中的 8 个 Expert 权重加起来为 1topk_weightstopk_weights/topk_weights.sum(dim-1,keepdimTrue)returntopk_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 通信——跨卡分发 TokenvoidTokenDispatch(int*topk_indices,float*hidden_states,intnum_tokens,intnum_devices){// 统计每张卡要发多少 Tokenintsend_counts[8]{0};for(intt0;tnum_tokens;t)for(intk0;ktop_k;k)send_counts[topk_indices[t*top_kk]/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 推理配方
