引言

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 output

All-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 并行 两大核心操作展开。

核心优化

1. TopKGating 优化

• 使用单次矩阵乘法替代循环，一次性计算所有 expert 的分数，避免重复计算。
• 采用 PartialTopK（局部排序） 代替全排序，将复杂度从 O(N log N) 降至 O(N + K log K)，显著提升路由效率。

1. Expert 并行与通信优化

• 使用 HCCL 实现跨 NPU 的 All-to-All 通信，并实测了通信耗时随硬件规模扩大的增长趋势（8 NPU 时通信占比达 76%）。
• 提出 Overlap 策略，通过 Double Buffer 流水线将 Expert 计算与通信重叠，减少等待时间。

1. 融合 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