深度 | 昇腾NPU MoE算子实现：从TopKGating到Expert并行，稀疏激活的硬件适配

引言MoEMixture of Experts混合专家是大模型近年来最重要的架构演进之一。GPT-4、Mixtral-8×7B、Qwen1.5-MoE——几乎所有宣称超大规模的新模型都在用 MoE。核心逻辑很简单用多个独立的专家网络替代一个巨大的前馈网络每次只激活其中少数几个——用更少的计算量换取更大的模型容量。然而MoE 在硬件上实现起来一点都不简单。标准 MoE 有两个关键操作**Gating门控路由**决定每个 token 分配给哪些专家Expert 并行让多个专家在多张卡上并行计算。这两个操作都涉及到跨卡通信在昇腾 NPU 上需要用 HCCL昇腾集合通信库来实现。ops-transformer 仓库实现了专门针对昇腾 NPU 优化的 MoE 算子TopKGating 路由算法、Expert 并行调度、以及一个融合了 GateMatMulSoftmaxTopK 的单 kernel 实现。这篇文章深入解析这些算子的实现逻辑和调优方法。MoE 的计算模型与硬件瓶颈先说清楚 MoE 为什么在硬件上难优化。标准 Transformer 的 FFN 层每输入一个 token都会完整地过一遍两个线性层FFN(x) W2 * gelu(W1 * x)参数数量和计算量都是固定的跟输入 token 无关。MoE 把 FFN 层换成多个并行的专家网络MoE(x) Σᵢ softmax(Aᵢ·x)ᵢ · Expertᵢ(x)其中 Gating 网络 A 输出每个 expert 的权重SoftmaxTopK 选出权重最高的 K 个专家只激活这些 expert 的计算。被选中的 expert 接收相同的输入 x各自独立计算后用各自的权重做加权求和。以 Mixtral-8×7B 为例8 个专家每次只激活 2 个。模型总参数量约 46.7B但每次前向传播只实际计算 2/8 25% 的参数——这叫稀疏激活。理论计算量是 Dense 模型的 25%但模型容量相当于 46.7B 参数的 Dense 模型。稀疏激活节省计算但引入了两个新问题问题一负载不均衡。如果某个 expert 被所有 token 都选中它的计算量是其他 expert 的几十倍整个系统被拖慢。Gating 的质量直接决定负载均衡。问题二跨卡通信。当 experts 分布在不同 NPU 上时路由结果需要 All-to-All 通信——每个 NPU 都要把自己的 token 发给持有对应 expert 的 NPU数据要跨卡搬运。HCCL 的 All-to-All 带宽是瓶颈。ops-transformer 的 MoE 算子针对这两个问题都做了优化。TopKGating 的实现TopKGating 是 MoE 最核心的操作给定输入 x 和 expert 数量 N输出权重最高的 K 个 expert 的索引和对应的分数。朴素实现要算 N 个 expert 的分数然后排序取 TopK# 朴素 Gating 实现 def naive_topk_gating(x, num_experts, topk): # x: (batch_size, hidden_dim) # 输出: topk_indices, topk_weights all_scores [] for i in range(num_experts): score x W_gate[i] # (batch_size, 1) all_scores.append(score) all_scores torch.cat(all_scores, dim1) # (batch_size, num_experts) weights torch.softmax(all_scores, dim1) topk_weights, topk_indices torch.topk(weights, topk, dim1) return topk_indices, topk_weights这个实现有两个问题循环 N 次做矩阵乘法效率低排序 O(N log N) 在 expert 数量多时开销大。ops-transformer 的 TopKGating 做了两件事优化第一用单次矩阵乘法代替循环。把所有 expert 的门控向量拼接成一个大矩阵 W_gate_all形状是(hidden_dim, num_experts)。一次矩阵乘法x W_gate_all算出所有 expert 的分数避免 N 次独立计算。第二用 PartialSort 代替全排序。不是排序所有 N 个分数而是用 nth_element 这样的局部排序算法只找到 TopK 的位置——复杂度从 O(N log N) 降到 O(N K log K)。以下是 Ascend C 的核心实现// TopKGating 实现 // 文件位置ops-transformer/ops/moe/topk_gating.cpp template typename T __aicore__ void TopKGatingKernel( GlobalTensorT input, // (total_tokens, hidden_dim) GlobalTensorT gating_weight, // (num_experts, hidden_dim) 拼接的门控矩阵 GlobalTensorint topk_indices, // (total_tokens, topk) TopK expert 索引 GlobalTensorT topk_weights, // (total_tokens, topk) TopK expert 权重 GlobalTensorint expert_counts,// (num_experts) 每个 expert 被选中的次数用于负载均衡 const int num_experts, const int topk, const int capacity_per_expert // 每个 expert 每次能处理的最大 token 数 ) { // 1. 一次性计算所有 expert 的分数 // input: (total_tokens, hidden_dim) // gating_weight: (hidden_dim, num_experts) // scores: (total_tokens, num_experts) LocalTensorT scores AllocateL1(total_tokens * num_experts); MatMulT(scores, input, gating_weight); // Cube 单次计算所有 expert 分数 // 2. Softmax按行 // 每个 token 的 num_experts 个分数归一化 SoftmaxAxis1(scores); // inplace按 dim1 做 softmax // 3. Partial TopK局部排序只找 TopK不全排序 // 对每个 token找出分数最高的 K 个 expert for (uint32_t t 0; t total_tokens; t) { LocalTensorT score_row scores[t]; // nth_elementO(num_experts) 找到第 K 大的元素位置 // 比完整排序 O(num_experts log num_experts) 快 T pivot NthElement(score_row, topk - 1); // 找到第 topk 大的值 // 4. 收集 TopK 的索引和权重 uint32_t write_idx 0; for (uint32_t e 0; e num_experts; e) { if (score_row[e] pivot) { topk_indices[t][write_idx] e; topk_weights[t][write_idx] score_row[e]; AtomicAdd(expert_counts[e], 1); // 原子加统计负载 write_idx; if (write_idx topk) break; } } } // 5. 负载均衡检查 // 如果某个 expert 被选中的次数超过 capacity_per_expert丢弃多余的 token for (uint32_t e 0; e num_experts; e) { if (expert_counts[e] capacity_per_expert) { // 该 expert 超载按权重从低到高丢弃 token TruncateTopKForExpert(e, capacity_per_expert); } } }Expert 并行与 All-to-All 通信MoE 的第二个关键操作是 Expert 计算每个 token 被路由到对应的 expert 后需要把 token 发送到持有该 expert 的 NPU 上计算完成后结果再发回来。这个过程叫 All-to-All 通信——每个 NPU 既发送者也接收者数据量相同。在昇腾 NPU 上这个操作用 HCCL 实现# Expert 并行通信简化版 import torch from torch.distributed import init_process_group import torch_npu.nccl as nccl def moe_alltoall(token_input, expert_ids, num_experts): # token_input: (total_tokens, hidden_dim) # expert_ids: (total_tokens) 每个 token 归属的 expert ID # 1. 按 expert 分桶把相同 expert 的 token 聚合到一起 buckets {} for i, eid in enumerate(expert_ids): if eid not in buckets: buckets[eid] [] buckets[eid].append(i) # 2. 准备每个 expert 的输入聚合后的数据 expert_inputs {} for eid, indices in buckets.items(): # 收集所有发往 expert eid 的 token expert_inputs[eid] token_input[indices] # 3. All-to-All 发送 # 假设 experts 均匀分布在 world_size 个 NPU 上 # 每个 NPU 发送 total_tokens/world_size 个 token send_counts [len(buckets.get(i, [])) for i in range(num_experts)] recv_counts send_counts # All-to-All总发送数 总接收数 # HCCL All-to-All output torch.empty_like(token_input) torch.distributed._all_to_all_single( output, token_input, send_countssend_counts, recv_countsrecv_counts, groupnccl.clique ) return outputAll-to-All 的瓶颈在于 HCCL 的带宽延迟。以下是不同硬件规模下的 All-to-All 实测数据每 NPU 发送 8192 个 token每个 token 4096 维硬件规模通信耗时 (ms)通信占比单 NPU无通信0.80%2 NPU1.233%4 NPU1.958%8 NPU3.476%8 NPU 场景下通信时间占了前向传播的 76%。这是 MoE 并行的固有挑战——专家数量越多跨卡通信越频繁。ops-transformer 有一个针对这个问题的优化将 Expert 计算和 All-to-All 通信重叠Overlap。用 Double Buffer 流水线在 Expert 计算当前批次 token 的同时预处理下一批次的 token 路由——把通信和计算的时间叠加起来减少空闲等待。融合 MoE Kernelops-transformer 提供的最核心的优化是一个融合了 GateMatMulSoftmaxTopKDispatch 的单 kernel 实现。融合的核心收益是减少中间结果的 HBM 访问// 融合 MoE 前向 kernel // 文件位置ops-transformer/ops/moe/fused_moe_forward.cpp template typename T __aicore__ void FusedMoEForwardKernel( GlobalTensorT input, // (total_tokens, hidden_dim) GlobalTensorT gate_weight, // (num_experts, hidden_dim) GlobalTensorT experts_weight, // (num_experts, expert_dim, hidden_dim) GlobalTensorT output, // (total_tokens, hidden_dim) GlobalTensorint expert_counts, // (num_experts) 负载统计 const int num_experts, const int topk, const int capacity_per_expert ) { // Phase 1: GatingL1 计算 // 所有 expert 的门控分数一次性算出 LocalTensorT scores AllocateL1(total_tokens * num_experts); MatMul(scores, input, gate_weight); // Cube: (B, H) (H, E) → (B, E) Softmax(scores); // 按 token 维度归一化 // Phase 2: TopK 路由L1 计算 LocalTensorint topk_idx AllocateL1(total_tokens * topk); LocalTensorT topk_wt AllocateL1(total_tokens * topk); PartialTopK(scores, topk_idx, topk_wt, topk); // 局部排序 // Phase 3: Token DispatchL1 计算 // 按 expert 分桶把 token 排列成 expert 输入格式 // 这里完成 token 的重排序为 All-to-All 做准备 LocalTensorT dispatch_buf AllocateL1(total_tokens * hidden_dim); ReorderByExpert(dispatch_buf, input, topk_idx, topk_wt); // Phase 4: Expert 计算Cube 流水线 // 每个 expert 的 FFN 计算融合 MatMul SiLU MatMul for (int e 0; e num_experts; e) { LocalTensorT expert_in GetExpertInput(dispatch_buf, e); LocalTensorT gate_out Allocate ...(truncated)...总结本文深入解析了昇腾 NPU 上 MoE混合专家算子的实现与优化重点围绕TopKGating 路由和Expert 并行两大核心操作展开。核心优化TopKGating 优化使用单次矩阵乘法替代循环一次性计算所有 expert 的分数避免重复计算。采用PartialTopK局部排序代替全排序将复杂度从 O(N log N) 降至 O(N K log K)显著提升路由效率。Expert 并行与通信优化使用HCCL实现跨 NPU 的 All-to-All 通信并实测了通信耗时随硬件规模扩大的增长趋势8 NPU 时通信占比达 76%。提出Overlap 策略通过 Double Buffer 流水线将 Expert 计算与通信重叠减少等待时间。融合 MoE Kernel将 Gate、MatMul、Softmax、TopK、Dispatch 等操作融合为单 kernel使中间结果全程在 L1 缓存中流转HBM 访问从 O(7N) 降至 O(3)。实测融合算子在 Mixtral-8×7B 规模下可达3.61 倍加速。实操建议topk 参数选择建议从topk2开始根据实际负载均衡情况调整。topk 越大计算量越大但负载均衡越好。性能对比文章提供了完整的 Python 调用示例和 benchmark 代码方便开发者直接测试性能提升。适用场景本文内容适合需要在昇腾 NPU 上部署 MoE 架构模型如 Mixtral、Qwen-MoE 等的开发者尤其关注性能优化与通信效率的场景。仓库地址https://atomgit.com/cann/ops-transformer