CANN ops-transformer 算子详解：MoeInitRoutingQuant 的 MoE 路由量化实现与 aclnn 调用实践
【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer
导读
MoeInitRoutingQuant是 CANN ops-transformer 仓库（moe/moe_init_routing_quant）中面向 MoE（Mixture of Experts，混合专家）大模型的关键算子，它在 MoE 路由（routing）流程中承担"按专家排序 + token 特征扩展 + 量化"三位一体的职责：基于aclnnMoeGatingTopKSoftmax输出的专家选择结果，将原始 token 特征按专家序号重排为专家连续分块的布局，同时完成 FP16/BF16/FP32 到 INT8 的在线量化，为后续专家并行计算提供紧凑、可直接喂给下游算子的输入。读完本文，你将掌握该算子的计算公式与数据流、全部入参与出参的语义与约束、两段式 aclnn 接口的调用方法，以及从 Host 侧 shape 推导、Tiling 到 Kernel 侧排序/重排/量化实现的完整源码脉络。
算子定位：MoE 路由链路中的"排序 + 扩展 + 量化"
在 MoE 大模型中，每一层通常包含多个专家子网络，输入 token 并非全部经过所有专家，而是通过门控网络挑选 Top-K 个专家。本算子位于这条链路的关键节点：
- 门控层先通过 aclnnMoeGatingTopKSoftmax 计算每个 token 的 Top-K 专家得分与索引；
MoeInitRoutingQuant拿到expertIdx（每个 token 对应的 K 个专家）与rowIdx（原始行位置）后，将全部NUM_ROWS × K个"（token, 专家）"组合按专家序号排序，使同一个专家处理的 token 连续排布；- 同时按排序后的顺序将原始 token 特征
x复制（gather）到扩展后的位置，并就地完成 INT8 量化，输出expandedXOut等三个结果。
从仓库文件布局看，该算子是一个标准的 CANN 算子工程：算子定义与 shape 推导、Tiling 计算、AscendC Kernel 实现、aclnn 调用示例 一应俱全，并配套了 UT/ST 测试。
计算公式与数据流
设输入 token 数为NUM_ROWS，每个 token 选择K个专家，expertIdx的 shape 为[NUM_ROWS, K]。算子将其展平为一行后按专家序号排序，核心公式如下：
$$ expandedExpertIdxOut,\ sortedRowIdx=keyValueSort(expertIdx,rowIdx) $$
$$ expandedRowIdxOut[sortedRowIdx[i]]=i $$
$$ expandedXOut[i]=quant(x[sortedRowIdx[i]%NUM_ROWS]) $$
三个公式分别对应三个输出：
| 公式 | 含义 | 对应输出 |
|---|---|---|
| keyValueSort | 以expertIdx为键、rowIdx为值联合排序，返回排序后的专家序号（键）与排序下标序列（值） | expandedExpertIdxOut、排序下标sortedRowIdx |
| 逆映射 | 记录"排序后位置 → 原始位置"的映射，即expandedX第i行取自x的哪一行 | expandedRowIdxOut |
| 扩展 + 量化 | 按排序结果从x中取出行（sortedRowIdx[i] % NUM_ROWS还原原始行号），并对每个元素做量化 | expandedXOut |
其中quant为线性对称/非对称量化，量化参数为scale与offset：
$$ quant(v)=round(v \times scale + offset) $$
典型量化参数取值（来自 aclnnMoeInitRoutingQuant 文档）为 $S=(R_{max}-R_{min})/(Q_{max}-Q_{min})$，其中 $R_{max}/R_{min}$ 为输入浮点数的最大/最小值，$Q_{max}/Q_{min}$ 为定点表示的最大/最小可表示值。
一个直观的小例子
以官方示例test_aclnn_moe_init_routing_quant.cpp中的数据（NUM_ROWS=3, K=2, H=4）推演：
x：shape[3, 4]，三行特征分别为0.1/0.2/0.3；expertIdx：shape[3, 2]，即[[1,2],[0,1],[0,2]]，表示第 0 行 token 去专家 1、2，第 1 行去专家 0、1，第 2 行去专家 0、2；rowIdx：shape[3, 2]，即[[0,3],[1,4],[2,5]]，其中第 1 列3/4/5表示这些位置对应的原始行是1/2/0（可视为已做过多轮 swap 后的行标记）。
对(expertIdx, rowIdx)键值对按专家序号升序排序后，expandedExpertIdxOut得到[0,0,1,1,2,2]，expandedRowIdxOut给出各排序位置对应的原始行索引，expandedXOut按该索引从x中取值并量化（示例中scale=1.0, offset=1.0，即v+1取整），最终输出 shape 为[6, 4]的 INT8 张量。这一布局使得"同一专家的 token 连续存放"，是后续专家计算高效并行的前提。
参数说明
算子共有 3 个输入、3 个属性、3 个输出（详见 README）：
| 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 |
|---|---|---|---|---|
x | 输入 | MOE 的输入，即 token 特征输入，对应公式中的x | FLOAT16、BFLOAT16、FLOAT32 | ND |
rowIdx | 输入 | 指示每个位置对应的原始行位置，对应公式中的rowIdx | INT32 | ND |
expertIdx | 输入 | aclnnMoeInitRoutingQuantSoftmax的输出，即每一行特征对应的 K 个处理专家，对应公式中的expertIdx | INT32 | ND |
activeNum | 属性 | 表示总的最大处理 row 数且大于等于 0，expandedXOut只有这么多行是有效的 | INT64 | - |
scale | 属性 | 量化计算需要 | DOUBLE | - |
offset | 属性 | 量化计算需要 | DOUBLE | - |
expandedXOut | 输出 | 根据expertIdx扩展过的特征，对应公式中的expandedX | INT8 | ND |
expandedRowIdxOut | 输出 | expandedX和x的映射关系，对应公式中的expandedRowIdx | INT32 | ND |
expandedExpertIdxOut | 输出 | 输出expertIdx排序后的结果，对应公式中的expandedExpertIdx | INT32 | ND |
各参数的 shape 与使用约束（API 文档细化）
依据 aclnnMoeInitRoutingQuant 接口文档，各 tensor 的维度要求如下：
x：2D，shape(NUM_ROWS, H)，支持空 tensor；rowIdx：2D，shape 与expertIdx一致，数值从 0 开始沿第 1 维递增，支持空 tensor；expertIdx：2D，shape(NUM_ROWS, K)，支持空 tensor；activeNum：标量，值域>= 0，expandedXOut仅有activeNum行有效；expandedXOut：2D，shape(min(NUM_ROWS, activeNum) * K, H)，支持空 tensor；expandedRowIdxOut：1D，shapeNUM_ROWS * K，支持空 tensor；expandedExpertIdxOut：1D，shapeNUM_ROWS * K，支持空 tensor。
x、rowIdx、expertIdx支持非连续 tensor（对应 API 文档中"非连续 Tensor"标记为 √），而三个输出要求连续内存。activeNum、scale、offset在 aclnn 接口中作为标量参数直接传入。
来自算子定义文件的佐证
在 moe_init_routing_quant_def.cpp 中，算子原型与 README 完全对应：
- 输入
x支持DT_FLOAT16 / DT_BF16 / DT_FLOAT，row_idx、expert_idx为DT_INT32，均使用 ND 格式并声明AutoContiguous； - 输出
expanded_x固定为DT_INT8，expanded_row_idx、expanded_expert_idx为DT_INT32； - 属性
active_num声明为 Int，scale、offset声明为 Float； - 注册了
ascend910b与ascend910_93两套 AICore 配置，并在 op_host/config 下提供了对应moe_init_routing_quant_binary.json与moe_init_routing_quant_simplified_key.ini配置文件。
产品支持情况与约束
根据算子 README 的产品支持矩阵：
| 产品 | 是否支持 |
|---|---|
| Ascend 950PR/Ascend 950DT | × |
| Atlas A3 训练系列产品/Atlas A3 推理系列产品 | × |
| Atlas A2 训练系列产品/Atlas A2 推理系列产品 | √ |
| Atlas 200I/500 A2 推理产品 | × |
| Atlas 推理系列产品 | × |
| Atlas 训练系列产品 | × |
需要说明的是，接口级文档 aclnnMoeInitRoutingQuant.md 中列出 Atlas A3 系列为"支持"，而算子 README 标记为"×"，两处口径存在差异；从算子定义注册的ascend910b（对应 A2 系列）与ascend910_93（对应 A3 系列）AICore 配置看，A3 侧代码路径已具备。实际使用时请以当前 CANN 版本配套的支持列表为准。
其余约束：README 中"约束说明"为无，API 文档则补充了一条确定性约束——aclnnMoeInitRoutingQuant默认采用确定性实现（与 确定性计算说明 相关），便于结果复现与调试。
调用说明：两段式 aclnn 接口
MoeInitRoutingQuant遵循 CANN 算子库通用的两段式接口调用范式：先调用GetWorkspaceSize接口完成入参校验、推导 workspace 大小并创建执行器；再调用执行接口真正下发计算。
函数原型
aclnnStatus aclnnMoeInitRoutingQuantGetWorkspaceSize( const aclTensor *x, const aclTensor *rowIdx, const aclTensor *expertIdx, int64_t activeNum, double scale, double offset, const aclTensor *expandedXOut, const aclTensor *expandedRowIdxOut, const aclTensor *expandedExpertIdxOut, uint64_t *workspaceSize, aclOpExecutor **executor); aclnnStatus aclnnMoeInitRoutingQuant( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream);第一段接口：aclnnMoeInitRoutingQuantGetWorkspaceSize
输入为x、rowIdx、expertIdx、activeNum、scale、offset与三个输出 tensor，输出为workspaceSize（需要在 Device 侧申请的 workspace 大小）与executor（封装了算子计算流程的执行器）。
返回aclnnStatus状态码（错误码定义见 aclnn 返回码），第一段接口完成入参校验，常见报错场景：
| 返回值 | 错误码 | 描述 |
|---|---|---|
ACLNN_ERR_PARAM_NULLPTR | 161001 | 输入和输出的 Tensor 是空指针 |
ACLNN_ERR_PARAM_INVALID | 161002 | 输入和输出的数据类型不在支持的范围内 |
ACLNN_ERR_INNER_TILING_ERROR | 561002 | x的 shape 维度不为 2 |
rowIdx的 shape 不为 2 或rowIdx与expertIdx的 shape 不相等 | ||
activeNum小于 0 | ||
expandedRowIdxOut与expandedExpertIdxOut的 shape 不相等，且不等于NUM_ROWS * K | ||
expandedXOut的 shape 不等于(min(NUM_ROWS, activeNum) * K, H) |
第二段接口：aclnnMoeInitRoutingQuant
参数为第一段接口返回的workspace内存地址、workspaceSize、executor，以及指定执行任务的stream流；同样返回aclnnStatus状态码。
完整调用示例与运行步骤
示例代码位于 examples/test_aclnn_moe_init_routing_quant.cpp，其完整编译与执行流程请参考仓库的编译与运行样例。核心流程如下：
#include "acl/acl.h" #include "aclnnop/aclnn_moe_init_routing_quant.h" #include <iostream> #include <vector> // ... CHECK_RET / LOG_PRINT / GetShapeSize / Init / CreateAclTensor 等辅助函数见示例文件 ... int main() { int32_t deviceId = 0; // 根据实际 device 填写 aclrtStream stream; auto ret = Init(deviceId, &stream); // aclInit + aclrtSetDevice + aclrtCreateStream // 1. 构造输入与输出 std::vector<int64_t> xShape = {3, 4}; // NUM_ROWS=3, H=4 std::vector<int64_t> idxShape = {3, 2}; // NUM_ROWS=3, K=2 std::vector<int64_t> expandedXOutShape = {6, 4}; // min(3, activeNum=3)*2=6 std::vector<int64_t> idxOutShape = {6}; // NUM_ROWS*K=6 int64_t activeNum = 3; double scale = 1.0; double offset = 1.0; std::vector<float> xHostData = {0.1,0.1,0.1,0.1, 0.2,0.2,0.2,0.2, 0.3,0.3,0.3,0.3}; std::vector<int> expertIdxHostData = {1,2, 0,1, 0,2}; std::vector<int> rowIdxHostData = {0,3, 1,4, 2,5}; std::vector<int8_t> expandedXOutHostData(24, 0); std::vector<int> expandedRowIdxOutHostData(6, 0); std::vector<int> expandedExpertIdxOutHostData(6, 0); // 通过 CreateAclTensor 完成 aclrtMalloc + aclrtMemcpy(H2D) + aclCreateTensor(ACL_FORMAT_ND) // x: ACL_FLOAT；rowIdx/expertIdx: ACL_INT32；expandedXOut: ACL_INT8；两个 idxOut: ACL_INT32 // 2. 两段式调用 uint64_t workspaceSize = 0; aclOpExecutor* executor; ret = aclnnMoeInitRoutingQuantGetWorkspaceSize(x, rowIdx, expertIdx, activeNum, scale, offset, expandedXOut, expandedRowIdxOut, expandedExpertIdxOut, &workspaceSize, &executor); void* workspaceAddr = nullptr; if (workspaceSize > 0) { ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } ret = aclnnMoeInitRoutingQuant(workspaceAddr, workspaceSize, executor, stream); // 3. 同步等待并取回结果（D2H） ret = aclrtSynchronizeStream(stream); // aclrtMemcpy 分别取回 expandedXData(INT8)、expandedRowIdxData、expandedExpertIdxData 并打印 // 4. 释放资源 aclDestroyTensor(x); aclDestroyTensor(rowIdx); aclDestroyTensor(expertIdx); aclDestroyTensor(expandedXOut); aclDestroyTensor(expandedRowIdxOut); aclDestroyTensor(expandedExpertIdxOut); aclrtFree(xDeviceAddr); /* ... 其余 device 内存 ... */ if (workspaceSize > 0) aclrtFree(workspaceAddr); aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }要点提示：
- 调用头文件为
aclnnop/aclnn_moe_init_routing_quant.h； expandedXOut的 shape 必须等于(min(NUM_ROWS, activeNum) * K, H)，示例中min(3,3)*2=6；- workspace 仅在
workspaceSize > 0时需要申请，结束后一并释放； - 输出结果通过
aclrtMemcpy(..., ACL_MEMCPY_DEVICE_TO_HOST)回拷到 Host 侧打印验证。
源码级原理：从 shape 推导到 Kernel 流水线
Host 侧：InferShape 与 Tiling
在 moe_init_routing_quant_infershape.cpp 中，shape 推导逻辑与 API 文档约束一一对应：
- 校验
x、rowIdx、expertIdx的维度为 2（或动态维度 -2），三者第一维必须一致，rowIdx与expertIdx第二维（K）必须一致； - 校验
activeNum >= 0； - 推导输出：
expanded_x的 shape 为(min(n, activeNum) * k, cols)，expanded_row_idx与expanded_expert_idx的 shape 为(n * k)，其中k = max(rowIdx_k, expertIdx_k)； - 输出数据类型固定为
expanded_x → INT8、两个 idx 输出 → INT32。
Tiling 数据结构 表明该算子把任务拆成了四个子阶段分别做 Tiling：QuantVBSComputeTilingData（排序核心数/每核元素数/每核循环数）、QuantVMSMiddleComputeTilingData、QuantSortOutComputeTilingData（单次循环最大元素数）与QuantGatherOutComputeTilingData（gather 阶段的每核行数、每核 K、每循环行列数、maxColsOneLoop、splitFlag等），并在MoeInitRoutingQuantTilingData中统一封装coreNum/n/cols/k/scale/offset与各子结构，体现了"按行、按 K、按列分片"的多核并行切分思路。
Kernel 侧：三段式流水线
Kernel 入口在 moe_init_routing_quant.cpp，moe_init_routing_quant函数按TILING_KEY分发，依次执行三个算子阶段：
- 排序阶段（Sort）：根据 Tiling 结果，小规模数据走
MoeSortOneCore（moe_quant_sort_one_core.h），大规模数据走MoeSortMultiCore（moe_quant_sort_multi_core.h），内部复用多路归并排序MoeMrgsort（moe_quant_mrgsort.h）。从moe_quant_common.h的常量可见其面向昇腾向量指令做了精细设计：ONE_REPEAT_SORT_NUM = 32、BLOCK_BYTES = 32（32 字节对齐，即 8 个 INT32/32 个 INT8 一个 block）、归并路数支持 2/3/4 路、并提供 256 长度的assist辅助索引表。归并排序循环内以CopyIn → MrgsortCompute → UpdateSortInfo → CopyOut流水化推进，直到所有元素排完。 - 源到目标映射阶段（SrcToDst）：
MoeSrcToDstOp（moe_quant_src_to_dst_op.h）利用排序阶段写入 workspace 的expandDstToSrcRow中间结果，反推并写出expandedRowIdxOut（即"排序后位置 → 原始行"的映射表），支持多核下按blockIdx分片、SyncAll全局同步。 - Gather + 量化阶段（GatherOut）：
MoeGatherOut<T>（moe_quant_gather_out.h）与面向少 active 行的MoeGatherOutSmallActiveRow（moe_quant_gather_out_small_activate_row.h）按expandedRowIdxOut从x中 gather 对应行，并在Compute中完成量化：对 FP16 直接Muls(scale)+Adds(offset)+Cast(INT8, CAST_RINT)；对 FP32 先转 FP16 再做同样的乘加与取整转换；对 BF16 则先Cast到 FP32 再经 FP16 中转完成量化。量化全程使用CAST_RINT舍入模式，与文档中"确定性实现"的要求一致。
值得注意的是，moe_quant_common.h中MoeInitRoutingQuant命名空间下的公共工具（Ceil、Align、AlignBytes、Min/Max及辅助常量）同时被上述各阶段头文件复用，整条流水线体现了"排序 — 映射 — 扩展量化"的清晰分层，也从源码层面印证了 README 三组公式的实现路径。
配套测试与验证
仓库为该算子提供了完整的测试保障：
- UT（单元测试）：tests/ut/op_host 下的
test_moe_init_routing_quant_infershape.cpp与test_moe_init_routing_quant_tiling.cpp覆盖 shape 推导与 Tiling 计算；tests/ut/op_kernel 下的test_moe_init_routing_quant.cpp通过moe_init_routing_quant_data/gen_data.py、gen_tiling.py生成测试数据与 Tiling 参数，验证 Kernel 计算结果。 - ST（系统测试）：tests/st/aclnnMoeInitRoutingQuant 下的
atk_aclnnMoeInitRoutingQuant.json与executor_aclnnMoeInitRoutingQuant.py提供了基于 ATK 框架的端到端用例，可直接作为回归样例参考。
延伸阅读
- 上游门控算子：aclnnMoeGatingTopKSoftmax，本算子的
expertIdx即来源于此； - 接口调用范式：两段式接口说明；
- 错误码含义：aclnn 返回码；
- 样例编译执行：编译与运行样例；
- 同系列 MoE 算子可参考 moe_init_routing（非量化版本）、moe_finalize_routing 与 moe_token_permute，对比理解 MoE 路由后 token 重排的完整家族。
【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
)
