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中面向 MoEMixture 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,\ sortedRowIdxkeyValueSort(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_ROWS3, K2, H4推演xshape[3, 4]三行特征分别为0.1/0.2/0.3expertIdxshape[3, 2]即[[1,2],[0,1],[0,2]]表示第 0 行 token 去专家 1、2第 1 行去专家 0、1第 2 行去专家 0、2rowIdxshape[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中取值并量化示例中scale1.0, offset1.0即v1取整最终输出 shape 为[6, 4]的 INT8 张量。这一布局使得同一专家的 token 连续存放是后续专家计算高效并行的前提。参数说明算子共有 3 个输入、3 个属性、3 个输出详见 README参数名输入/输出/属性描述数据类型数据格式x输入MOE 的输入即 token 特征输入对应公式中的xFLOAT16、BFLOAT16、FLOAT32NDrowIdx输入指示每个位置对应的原始行位置对应公式中的rowIdxINT32NDexpertIdx输入aclnnMoeInitRoutingQuantSoftmax的输出即每一行特征对应的 K 个处理专家对应公式中的expertIdxINT32NDactiveNum属性表示总的最大处理 row 数且大于等于 0expandedXOut只有这么多行是有效的INT64-scale属性量化计算需要DOUBLE-offset属性量化计算需要DOUBLE-expandedXOut输出根据expertIdx扩展过的特征对应公式中的expandedXINT8NDexpandedRowIdxOut输出expandedX和x的映射关系对应公式中的expandedRowIdxINT32NDexpandedExpertIdxOut输出输出expertIdx排序后的结果对应公式中的expandedExpertIdxINT32ND各参数的 shape 与使用约束API 文档细化依据 aclnnMoeInitRoutingQuant 接口文档各 tensor 的维度要求如下x2Dshape(NUM_ROWS, H)支持空 tensorrowIdx2Dshape 与expertIdx一致数值从 0 开始沿第 1 维递增支持空 tensorexpertIdx2Dshape(NUM_ROWS, K)支持空 tensoractiveNum标量值域 0expandedXOut仅有activeNum行有效expandedXOut2Dshape(min(NUM_ROWS, activeNum) * K, H)支持空 tensorexpandedRowIdxOut1DshapeNUM_ROWS * K支持空 tensorexpandedExpertIdxOut1DshapeNUM_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_FLOATrow_idx、expert_idx为DT_INT32均使用 ND 格式并声明AutoContiguous输出expanded_x固定为DT_INT8expanded_row_idx、expanded_expert_idx为DT_INT32属性active_num声明为 Intscale、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_NULLPTR161001输入和输出的 Tensor 是空指针ACLNN_ERR_PARAM_INVALID161002输入和输出的数据类型不在支持的范围内ACLNN_ERR_INNER_TILING_ERROR561002x的 shape 维度不为 2rowIdx的 shape 不为 2 或rowIdx与expertIdx的 shape 不相等activeNum小于 0expandedRowIdxOut与expandedExpertIdxOut的 shape 不相等且不等于NUM_ROWS * KexpandedXOut的 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::vectorint64_t xShape {3, 4}; // NUM_ROWS3, H4 std::vectorint64_t idxShape {3, 2}; // NUM_ROWS3, K2 std::vectorint64_t expandedXOutShape {6, 4}; // min(3, activeNum3)*26 std::vectorint64_t idxOutShape {6}; // NUM_ROWS*K6 int64_t activeNum 3; double scale 1.0; double offset 1.0; std::vectorfloat 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::vectorint expertIdxHostData {1,2, 0,1, 0,2}; std::vectorint rowIdxHostData {0,3, 1,4, 2,5}; std::vectorint8_t expandedXOutHostData(24, 0); std::vectorint expandedRowIdxOutHostData(6, 0); std::vectorint expandedExpertIdxOutHostData(6, 0); // 通过 CreateAclTensor 完成 aclrtMalloc aclrtMemcpy(H2D) aclCreateTensor(ACL_FORMAT_ND) // x: ACL_FLOATrowIdx/expertIdx: ACL_INT32expandedXOut: 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.hexpandedXOut的 shape 必须等于(min(NUM_ROWS, activeNum) * K, H)示例中min(3,3)*26workspace 仅在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 数据结构 表明该算子把任务拆成了四个子阶段分别做 TilingQuantVBSComputeTilingData排序核心数/每核元素数/每核循环数、QuantVMSMiddleComputeTilingData、QuantSortOutComputeTilingData单次循环最大元素数与QuantGatherOutComputeTilingDatagather 阶段的每核行数、每核 K、每循环行列数、maxColsOneLoop、splitFlag等并在MoeInitRoutingQuantTilingData中统一封装coreNum/n/cols/k/scale/offset与各子结构体现了按行、按 K、按列分片的多核并行切分思路。Kernel 侧三段式流水线Kernel 入口在 moe_init_routing_quant.cppmoe_init_routing_quant函数按TILING_KEY分发依次执行三个算子阶段排序阶段Sort根据 Tiling 结果小规模数据走MoeSortOneCoremoe_quant_sort_one_core.h大规模数据走MoeSortMultiCoremoe_quant_sort_multi_core.h内部复用多路归并排序MoeMrgsortmoe_quant_mrgsort.h。从moe_quant_common.h的常量可见其面向昇腾向量指令做了精细设计ONE_REPEAT_SORT_NUM 32、BLOCK_BYTES 3232 字节对齐即 8 个 INT32/32 个 INT8 一个 block、归并路数支持 2/3/4 路、并提供 256 长度的assist辅助索引表。归并排序循环内以CopyIn → MrgsortCompute → UpdateSortInfo → CopyOut流水化推进直到所有元素排完。源到目标映射阶段SrcToDstMoeSrcToDstOpmoe_quant_src_to_dst_op.h利用排序阶段写入 workspace 的expandDstToSrcRow中间结果反推并写出expandedRowIdxOut即排序后位置 → 原始行的映射表支持多核下按blockIdx分片、SyncAll全局同步。Gather 量化阶段GatherOutMoeGatherOutTmoe_quant_gather_out.h与面向少 active 行的MoeGatherOutSmallActiveRowmoe_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），仅供参考