随着 DeepSeek-V3 和 Mixtral 8x7B 的爆火，MoE（Mixture of Experts，混合专家） 架构成为了大模型领域绝对的主流。它解决了“模型越大，推理越慢”的悖论，实现了“万亿参数模型，仅需百亿参数激活”的奇迹。

然而，MoE 的工程实现难度极高。它将原本整齐划一的矩阵运算，打散成了碎片的、动态的稀疏计算。

这一篇，我们将深入 AtomGit 上的 CANN ops-nn 仓库，揭秘华为昇腾 NPU 是如何通过 TopK 路由 与 Grouped MatMul 等黑科技算子，驾驭这种复杂的稀疏计算模式的。

如果说 Dense（稠密）大模型是一支整齐划一的正规军，那么 MoE（混合专家）模型就是一支由无数特种部队组成的联军。

在推理每一个 Token 时，MoE 模型不需要调动所有的神经元，而是通过一个 Router（路由器），从成百上千个“专家（Experts）”中挑选出最擅长的 Top-2 或 Top-K 个来干活。

这种机制虽然极大地节省了计算量，却给底层硬件带来了巨大的挑战：

1. 动态路由：每个 Token 选择的专家不同，数据流不再是静态的。
2. 显存碎片：不同专家的权重分散在显存各处，带宽利用率低。
3. 负载不均：有的专家忙死，有的专家闲死。

华为昇腾 CANN 的 ops-nn 仓库，在 AtomGit 上开源了一套针对 MoE 架构的“交通指挥系统”。通过深入阅读代码，我们可以看到 NPU 是如何通过软硬结合，实现 Token 的极速分发与聚合。

MoE 核心阵地

• CANN 组织主页： https://atomgit.com/cann
• ops-nn 源码仓库： https://atomgit.com/cann/ops-nn

---

一、 MoE 的心脏：TopK Router 算子

MoE 的第一步是“选人”。给定一个 Token 的特征向量，Router 需要计算它与所有专家的匹配度（Score），然后选出分数最高的 K 个专家。

这是一个典型的 Vector-Bound 任务。在 ops-nn 中，这不仅仅是一次排序，而是结合了 Softmax、TopK 和 Scatter/Gather 索引生成的复合算子。

昇腾 NPU 的 Vector 单元拥有强大的排序指令，可以在 时间内并行找出 TopK，而无需像 CPU 那样进行全排序。

---

二、 代码实战：构建高效的 Expert Router

我们将展示一个简化版的 Router 核心逻辑：输入路由分数（Logits），输出选中的专家索引（Indices）和对应的权重（Weights）。

Ascend C 核心代码逻辑

#include "kernel_operator.h"

using namespace AscendC;

constexpr int32_t NUM_EXPERTS = 64;  // 假设有64个专家
constexpr int32_t TOP_K = 2;         // 每个Token选2个专家
constexpr int32_t BLOCK_LEN = 128;   // 处理的Token数量

class KernelMoERouter {
public:
    __aicore__ inline KernelMoERouter() {}

    __aicore__ inline void Init(GM_ADDR logits, GM_ADDR indices, GM_ADDR weights) {
        logitsGm.SetGlobalBuffer((__gm__ half *)logits);
        indicesGm.SetGlobalBuffer((__gm__ int32_t *)indices);
        weightsGm.SetGlobalBuffer((__gm__ half *)weights);

        pipe.InitBuffer(inQueue, 1, BLOCK_LEN * NUM_EXPERTS * sizeof(half));
        pipe.InitBuffer(outQueueIdx, 1, BLOCK_LEN * TOP_K * sizeof(int32_t));
        pipe.InitBuffer(outQueueWt, 1, BLOCK_LEN * TOP_K * sizeof(half));
    }

    __aicore__ inline void Process() {
        // 1. 搬入 Logits [Block, Num_Experts]
        LocalTensor<half> inputLocal = inQueue.AllocTensor<half>();
        DataCopy(inputLocal, logitsGm, BLOCK_LEN * NUM_EXPERTS);
        inQueue.EnQue(inputLocal);

        Compute();

        // 3. 搬出结果
        // ... (省略搬出逻辑)
    }

private:
    __aicore__ inline void Compute() {
        LocalTensor<half> logits = inQueue.DeQue<half>();
        LocalTensor<int32_t> topkIdx = outQueueIdx.AllocTensor<int32_t>();
        LocalTensor<half> topkVal = outQueueWt.AllocTensor<half>();

        // 临时 Buffer
        LocalTensor<half> softmaxOut = outQueueWt.AllocTensor<half>(); // 复用空间

        // --- Step 1: Softmax 归一化 ---
        // 使得路由得分变成概率分布
        // Softmax(softmaxOut, logits, ...); 
        
        // --- Step 2: TopK 选择 ---
        // 这是 MoE 的核心。Ascend C 提供了专门的 TopK 指令或者 Sort 指令
        // 这里的逻辑是对每一行（每个Token）在 64 个专家中找最大的 2 个
        
        // 伪代码：对 Tensor 的最后一维进行 TopK
        // TopK(topkVal, topkIdx, softmaxOut, TOP_K);
        
        // 在 ops-nn 的真实实现中，为了极致性能，可能会使用 Bitonic Sort (双调排序) 网络
        // 或者利用 Vector 单元的 "Compare & Select" 并行指令
        
        // --- Step 3: 归一化权重 (Normalize Weights) ---
        // 选出的 Top2 权重之和通常需要重新归一化为 1
        // sum = w1 + w2; w1 /= sum; w2 /= sum;
        // Vector Add + Div
        
        outQueueIdx.EnQue(topkIdx);
        outQueueWt.EnQue(topkVal);
        inQueue.FreeTensor(logits);
    }

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, 1> inQueue;
    TQue<QuePosition::VECOUT, 1> outQueueIdx, outQueueWt;
    GlobalTensor<half> logitsGm, weightsGm;
    GlobalTensor<int32_t> indicesGm;
};

3. 代码背后的优化哲学

在 ops-nn 仓库的 MoE 实现中，有几个关键点值得注意：

• 避免条件分支：TopK 的过程在 CPU 上通常是大量的 if (a > b)，这在 GPU/NPU 上是性能杀手。ops-nn 采用的是 Bit-mask 或 SIMD Min/Max 指令，全流水线执行，没有任何分支预测失败的惩罚。
• 索引计算：Router 算完后，需要生成 Gather 指令所需的偏移量。CANN 提供了极其高效的整数向量运算能力，能瞬间生成成千上万个内存地址偏移，为后续的数据搬运做准备。

---

三、 挑战稀疏性：Grouped MatMul (GMMA)

选完专家后，真正困难的来了：如何计算？

传统的 MatMul 要求矩阵是整块的。但在 MoE 中，Token 被打散了。例如，Expert A 分到了 10 个 Token，Expert B 分到了 3 个 Token，Expert C 分到了 0 个。

如果一个个循环调用 MatMul，Kernel 启动开销会炸裂。

在 AtomGit 的 ops-nn 仓库中，CANN 引入了 Grouped MatMul (GMMA) 的概念。它的思想是：“虽然你们形状各异，但我把你们打包在一起算。”

GMMA 的算子逻辑

1. Sort & Pack：根据 Router 的输出，将所有 Token 按照专家 ID 进行排序和重排。这样，分给 Expert A 的数据在内存中就连续了。
2. Multi-Stream Compute：利用 NPU 的多流能力，或者特定的 Grouped MatMul 指令，一次性提交多个小矩阵乘法任务。
3. Cube 利用率：ops-nn 中的代码会极其精细地处理 Padding（填充）。因为 Cube 单元通常喜欢 16x16 的倍数，对于分到 3 个 Token 这种尴尬的情况，代码会自动 Padding 到 16，算完后再切掉，以换取 Cube 单元的满载运行。

---

四、 为什么 ops-nn 是 MoE 开发者的必读库？

随着 DeepSeek 等开源 MoE 模型的流行，越来越多的开发者尝试在本地部署微调。

1. 理解 EP (Expert Parallelism) 通信：
   虽然 ops-nn 专注单卡算子，但其中的 All2All 通信衔接逻辑（即 Dispatch 和 Combine 阶段）是理解分布式 MoE 的基础。你可以在代码中看到数据是如何被“打包”准备发送的。
2. 自定义负载均衡 Loss：
   MoE 训练需要 Aux Loss 来保证负载均衡。ops-nn 中的 LogSoftmax 和 ReduceMean 组合实现，展示了如何高效计算这种辅助损失函数。
3. 从 Dense 到 Sparse 的思维转变：
   阅读 ops-nn 的 MoE 源码，是一次思维升级。你将不再把矩阵看作铁板一块，而是看作可以动态拆解、路由、重组的流体。

---

五、 结语：算力，随需而动

MoE 架构的本质，是对算力资源的极致精细化管理。它不再“大水漫灌”，而是“精准滴灌”。

华为 CANN 的 ops-nn 仓库，就是这套精细化灌溉系统的阀门控制器。通过 TopK Router 和 Grouped MatMul，昇腾 NPU 实现了在万亿参数海洋中的精准导航。

如果你想驾驭 DeepSeek-V3 这种级别的庞然大物，或者想设计自己的 MoE 架构，AtomGit 上的 ops-nn 仓库是你绕不开的“藏经阁”。

开启稀疏计算之旅：

• 加入 CANN 开发者社区： https://atomgit.com/cann
• 探索 MoE 算子源码： https://atomgit.com/cann/ops-nn