深度 | 昇腾NPU MoE算子实现：从TopKGating到Expert并行，稀疏激活的硬件适配
引言
MoE（Mixture of Experts，混合专家）是大模型近年来最重要的架构演进之一。GPT-4、Mixtral-8×7B、Qwen1.5-MoE——几乎所有宣称"超大规模"的新模型都在用 MoE。核心逻辑很简单：用多个独立的"专家"网络替代一个巨大的前馈网络，每次只激活其中少数几个——用更少的计算量换取更大的模型容量。
然而，MoE 在硬件上实现起来一点都不简单。标准 MoE 有两个关键操作：**Gating（门控路由）**决定每个 token 分配给哪些专家，Expert 并行让多个专家在多张卡上并行计算。这两个操作都涉及到跨卡通信，在昇腾 NPU 上需要用 HCCL（昇腾集合通信库）来实现。
ops-transformer 仓库实现了专门针对昇腾 NPU 优化的 MoE 算子：TopKGating 路由算法、Expert 并行调度、以及一个融合了 Gate+MatMul+Softmax+TopK 的单 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 的权重，Softmax+TopK 选出权重最高的 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, dim=1) # (batch_size, num_experts) weights = torch.softmax(all_scores, dim=1) topk_weights, topk_indices = torch.topk(weights, topk, dim=1) 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( GlobalTensor<T> input, // (total_tokens, hidden_dim) GlobalTensor<T> gating_weight, // (num_experts, hidden_dim) 拼接的门控矩阵 GlobalTensor<int> topk_indices, // (total_tokens, topk) TopK expert 索引 GlobalTensor<T> topk_weights, // (total_tokens, topk) TopK expert 权重 GlobalTensor<int> 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) LocalTensor<T> scores = Allocate<L1>(total_tokens * num_experts); MatMul<T>(scores, input, gating_weight); // Cube 单次计算所有 expert 分数 // 2. Softmax（按行） // 每个 token 的 num_experts 个分数归一化 SoftmaxAxis1(scores); // inplace，按 dim=1 做 softmax // 3. Partial TopK（局部排序，只找 TopK，不全排序） // 对每个 token，找出分数最高的 K 个 expert for (uint32_t t = 0; t < total_tokens; ++t) { LocalTensor<T> score_row = scores[t]; // nth_element：O(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_counts=send_counts, recv_counts=recv_counts, group=nccl.clique ) return outputAll-to-All 的瓶颈在于 HCCL 的带宽延迟。以下是不同硬件规模下的 All-to-All 实测数据（每 NPU 发送 8192 个 token，每个 token 4096 维）：
| 硬件规模 | 通信耗时 (ms) | 通信占比 |
|---|---|---|
| 单 NPU（无通信） | 0.8 | 0% |
| 2 NPU | 1.2 | 33% |
| 4 NPU | 1.9 | 58% |
| 8 NPU | 3.4 | 76% |
8 NPU 场景下，通信时间占了前向传播的 76%。这是 MoE 并行的固有挑战——专家数量越多，跨卡通信越频繁。
ops-transformer 有一个针对这个问题的优化：将 Expert 计算和 All-to-All 通信重叠（Overlap）。用 Double Buffer 流水线，在 Expert 计算当前批次 token 的同时，预处理下一批次的 token 路由——把通信和计算的时间叠加起来，减少空闲等待。
融合 MoE Kernel
ops-transformer 提供的最核心的优化是一个融合了 Gate+MatMul+Softmax+TopK+Dispatch 的单 kernel 实现。融合的核心收益是减少中间结果的 HBM 访问：
// 融合 MoE 前向 kernel // 文件位置：ops-transformer/ops/moe/fused_moe_forward.cpp template <typename T> __aicore__ void FusedMoEForwardKernel( GlobalTensor<T> input, // (total_tokens, hidden_dim) GlobalTensor<T> gate_weight, // (num_experts, hidden_dim) GlobalTensor<T> experts_weight, // (num_experts, expert_dim, hidden_dim) GlobalTensor<T> output, // (total_tokens, hidden_dim) GlobalTensor<int> expert_counts, // (num_experts) 负载统计 const int num_experts, const int topk, const int capacity_per_expert ) { // ===== Phase 1: Gating（L1 计算） ===== // 所有 expert 的门控分数一次性算出 LocalTensor<T> scores = Allocate<L1>(total_tokens * num_experts); MatMul(scores, input, gate_weight); // Cube: (B, H) @ (H, E) → (B, E) Softmax(scores); // 按 token 维度归一化 // ===== Phase 2: TopK 路由（L1 计算） ===== LocalTensor<int> topk_idx = Allocate<L1>(total_tokens * topk); LocalTensor<T> topk_wt = Allocate<L1>(total_tokens * topk); PartialTopK(scores, topk_idx, topk_wt, topk); // 局部排序 // ===== Phase 3: Token Dispatch（L1 计算） ===== // 按 expert 分桶，把 token 排列成 expert 输入格式 // 这里完成 token 的重排序，为 All-to-All 做准备 LocalTensor<T> dispatch_buf = Allocate<L1>(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) { LocalTensor<T> expert_in = GetExpertInput(dispatch_buf, e); LocalTensor<T> 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 参数选择：建议从
topk=2开始，根据实际负载均衡情况调整。topk 越大，计算量越大，但负载均衡越好。 - 性能对比：文章提供了完整的 Python 调用示例和 benchmark 代码，方便开发者直接测试性能提升。
适用场景
本文内容适合需要在昇腾 NPU 上部署 MoE 架构模型（如 Mixtral、Qwen-MoE 等）的开发者，尤其关注性能优化与通信效率的场景。
仓库地址：https://atomgit.com/cann/ops-transformer