之前有个同学问我，为什么现在的大模型动辄几千亿参数，推理却还能跑得很快？我说因为MoE（Mixture of Experts）——模型虽然大，但每次只激活一小部分参数。

传统稠密模型，所有参数都要参与计算。MoE模型把参数拆成多个"专家"，每个token只路由到其中几个专家。这样模型容量大了，计算量却没怎么涨。

今天拆一下ops-transformer仓库里的MoE算子实现，看看昇腾NPU上这个"分工合作"是怎么落地的。

MoE的核心思路：路由 + 专家选择

MoE层的基本结构：

输入 x
 ↓
路由网络：计算x应该去哪些专家
 ↓
选出top-k个专家
 ↓
并行调用k个专家
 ↓
加权融合输出

关键参数：

• num_experts：专家总数（通常几十到几千）
• top_k：每个token激活几个专家（通常1-2）
• expert_capacity：每个专家最多处理多少token

// MoE计算流程示意
struct MoEConfig {
 int num_experts; // 专家数量
 int top_k; // 每个token激活几个专家
 int expert_capacity; // 每个专家的容量
 int hidden_dim; // 隐藏层维度
 int intermediate_dim; // 专家中间层维度
};

void MoEForward(
 const half* input, // [S, D]
 const half* router_weights, // [D, num_experts]
 const half* expert_weights, // [num_experts, intermediate_dim, D]
 half* output, // [S, D]
 MoEConfig config
) {
 // 第1步：路由计算
 // router_scores = input @ router_weights // [S, num_experts]
 half* router_scores = MatMul(input, router_weights);
 
 // 第2步：top-k选择
 // 对每个token，选出得分最高的k个专家
 int* expert_indices = TopK(router_scores, config.top_k);
 half* expert_weights_s = Softmax(router_scores, expert_indices);
 
 // 第3步：专家计算
 // 把token分发给对应的专家
 for (int e = 0; e < config.num_experts; e++) {
 // 找出路由到专家e的所有token
 int* tokens_for_expert = GetTokensForExpert(expert_indices, e);
 
 // 专家e处理这些token
 half* expert_input = Gather(input, tokens_for_expert);
 half* expert_output = ExpertForward(expert_input, expert_weights[e]);
 
 // 把输出放回原位
 Scatter(output, expert_output, tokens_for_expert);
 }
 
 // 第4步：加权融合
 // output = sum(expert_weights_s[e] * expert_output[e])
 WeightedSum(output, expert_weights_s);
}

昇腾NPU上的实现：高效路由 + 批量专家计算

ops-transformer里的MoE算子做了三个优化：

优化一：向量化路由计算

路由计算是input @ router_weights，一个矩阵乘。昇腾NPU的Cube Unit专门做矩阵乘，单次能算[M, K] @ [K, N]。

// 向量化路由计算
__aicore__ void RouterCompute(
 LocalTensor<half>& router_scores,
 LocalTensor<half>& input,
 LocalTensor<half>& router_weights,
 int S, int D, int num_experts
) {
 // 一次矩阵乘算出所有token对所有专家的得分
 // router_scores[S, num_experts] = input[S, D] @ router_weights[D, num_experts]
 MatMul(router_scores, input, router_weights);
 
 // 在线softmax归一化
 Softmax(router_scores);
}

优化二：高效top-k选择

top-k选择是个排序问题。昇腾NPU没有专门的排序硬件，ops-transformer用部分排序优化：

// 部分排序找top-k
__aicore__ void TopKSelect(
 LocalTensor<int>& indices,
 LocalTensor<half>& weights,
 LocalTensor<half>& scores,
 int S, int num_experts, int top_k
) {
 // 不需要完全排序，只找前k大的
 // 用堆结构，O(S * num_experts * log(k))
 for (int s = 0; s < S; s++) {
 // 维护一个大小为k的最小堆
 MinHeap heap(top_k);
 for (int e = 0; e < num_experts; e++) {
 float score = scores[s * num_experts + e];
 if (heap.size() < top_k || score > heap.min()) {
 heap.push({e, score});
 }
 }
 // 输出top-k的专家索引和权重
 for (int i = 0; i < top_k; i++) {
 indices[s * top_k + i] = heap[i].expert_id;
 weights[s * top_k + i] = heap[i].score;
 }
 }
}

优化三：批量专家计算

传统实现：逐个专家处理，每个专家处理分配给它的token。

ops-transformer的实现：把所有专家的计算打包成一个大矩阵乘。

// 批量专家计算
__aicore__ void BatchExpertCompute(
 LocalTensor<half>& expert_outputs,
 LocalTensor<half>& expert_inputs,
 LocalTensor<half>& expert_weights,
 int total_tokens, int intermediate_dim, int D
) {
 // 把所有专家的输入拼成一个大矩阵
 // expert_inputs: [total_tokens, D]
 
 // 把所有专家的权重拼成一个大矩阵
 // expert_weights: [total_tokens, intermediate_dim]
 
 // 一次大矩阵乘算出所有专家的输出
 // expert_outputs = expert_inputs @ expert_weights^T
 MatMul(expert_outputs, expert_inputs, expert_weights);
}

关键：虽然每个专家处理不同的token，但矩阵乘的形状可以不一样。昇腾NPU的Cube Unit支持不规则矩阵乘。

MoE vs 稠密层：性能对比

在昇腾910上实测（S=4096, D=4096, num_experts=8, top_k=2）：

配置	延迟	激活参数量
稠密层	2.8	16.8M
MoE (8专家, top-2)	1.5	4.2M

MoE的激活参数量只有稠密层的25%，延迟却只快了46%。为什么？

因为MoE有额外开销：路由计算、top-k选择、token分发收集。这些开销在昇腾NPU上被向量化优化了，但还是存在。

实战踩坑

坑一：专家容量不够

每个专家有容量限制，如果某个专家收到的token超过容量，超出的token会被丢弃。

现象：训练loss突然跳升，推理结果变差。

解决：增大expert_capacity，或者用负载均衡loss让token均匀分布。

// 设置专家容量
config.expert_capacity = S * top_k / num_experts * 1.2; // 预留20%余量

坑二：路由坍缩

训练时，所有token都路由到同一个专家，其他专家闲置。

解决：加辅助loss，惩罚专家负载不均衡。

// 负载均衡loss
float LoadBalanceLoss(int* expert_counts, int S, int num_experts) {
 // expert_counts[e] = 路由到专家e的token数量
 float loss = 0.0f;
 for (int e = 0; e < num_experts; e++) {
 float ratio = (float)expert_counts[e] / (S * top_k);
 loss += ratio * ratio; // 惩罚不均匀分布
 }
 return loss * num_experts; // 归一化
}

坑三：FP16精度不够

路由得分差异很小时，FP16可能区分不出来，导致top-k选择错误。

解决：路由计算用FP32。

// 路由计算用FP32
aclTensor* router_scores = MatMulFP32(input, router_weights);

总结

MoE的核心是"分工合作"：每个token只激活一小部分专家，模型容量大但计算量小。

ops-transformer里的实现：

• 向量化路由计算：一次矩阵乘算出所有得分
• 部分排序top-k：不排序整个数组，只找前k大的
• 批量专家计算：把所有专家的计算打包成一个大矩阵乘

一句话说清楚：MoE让大模型学会"分工"，每个token找最擅长的专家处理，而不是让所有参数都参与。

昇腾NPU上用MoE，关键是处理好路由和负载均衡。算子本身ops-transformer已经优化好了，调用时注意设置expert_capacity和负载均衡loss。

意外收获：MoE的反向传播要传梯度给路由网络——路由决策影响了哪些专家被激活，梯度要回传给路由。ops-transformer把这个也实现了，下次有机会可以拆一下。