CANN ops-transformer 算子解析：aclnnMoeFinalizeRoutingV2 接口详解与 MoE 输出合并实战

CANN ops-transformer 算子解析aclnnMoeFinalizeRoutingV2 接口详解与 MoE 输出合并实战【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读本文围绕 CANN ops-transformer 仓库中 MoEMixture of Experts专家混合路由收尾算子aclnnMoeFinalizeRoutingV2展开系统讲解其功能定位、计算公式、两段式接口原型、全部入参语义、分产品约束与调用示例。本文以官方接口文档 docs/aclnnMoeFinalizeRoutingV2.md 为主体并结合 op_host 下的算子定义、Shape 推导、Tiling 与 Kernel 实现以及 tests 下的 golden 参考实现帮助读者既能在 NPU 上正确调用该算子也能从源码层面理解其内部工作机理。算子定位MoE 路由的最后一公里在 MoE 大模型推理与训练流程中Token 经过 Gating Network 选出 top-K 专家后会被路由Routing到不同专家执行 FFNFeed-Forward Network前馈网络计算最终需要把来自不同专家的输出结果按行合并Finalize Routing回原始顺序。MoeFinalizeRoutingV2正是这一流程的收尾算子合并 MoE FFN 的输出结果同时支持叠加共享专家Shared Expert结果、专家偏置bias与加权系数scales。其整体计算公式为$$ expertidexpertIdx[i,k] $$$$ out(i,j)x1_{i,j}x2_{i,j}\sum_{k0}^{K}(scales_{i,k}(expandedX_{expandedRowIdx_{iknum_rows},j}bias_{expertid,j})) $$其中各符号含义如下符号含义NUM_ROWS行数即参与计算的 Token/行数量K每个 Token 从总专家数 E 中选出的专家个数top-KHhidden size即每个 Token 序列长度为列数Eexpert num专家总数要求 E ≥ KCexpert capacity专家处理 Token 数量的能力阈值drop pad 场景从公式可见算子对第i行的输出等于共享专家输出x1、x2的直加再加上对 K 个被选专家加权scales后的expandedX行行号由expandedRowIdx给出与专家偏置bias按expertIdx索引之和。这一语义与仓库中测试 golden 参考实现 tests/assets/golden.py 完全一致。产品支持情况该算子在不同昇腾产品上的支持情况如下产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品Atlas 310P 系列推理支持Atlas 训练系列产品Atlas 910 系列训练不支持此外从仓库源码看算子定义在 moe_finalize_routing_v2_def.cpp 中还注册了kirinx90、kirin9030两款 Kirin 处理器调用GetKirinCoreConfig()统一配置且 README 中明确 Kirin X90/Kirin 9030 处理器系列产品不支持 BFLOAT16。在 README.md 中同样给出了完整支持矩阵可作为交叉印证。两段式接口与函数原型按照 CANN 算子库的统一约定每个算子采用两段式接口调用详见 两段式接口说明先调用aclnnMoeFinalizeRoutingV2GetWorkspaceSize完成入参校验、构图并获取计算所需 workspace 大小与包含算子计算流程的执行器executor再调用aclnnMoeFinalizeRoutingV2真正在指定 Stream 上执行计算。aclnnStatus aclnnMoeFinalizeRoutingV2GetWorkspaceSize( const aclTensor *expandedX, const aclTensor *expandedRowIdx, const aclTensor *x1Optional, const aclTensor *x2Optional, const aclTensor *biasOptional, const aclTensor *scalesOptional, const aclTensor *expertIdxOptional, int64_t dropPadMode, const aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnMoeFinalizeRoutingV2( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)从源码看第一段接口的实现在 op_api/aclnn_moe_finalize_routing_v2.cpp 中先执行CheckV2入参校验然后对全部输入逐一执行l0op::Contiguous保证连续性再调用l0op::MoeFinalizeRoutingV2构图最后通过uniqueExecutor-GetWorkspaceSize()返回 workspace 大小、uniqueExecutor.ReleaseTo(executor)释放执行器句柄。第二段接口则统一走CommonOpExecutorRun完成执行aclnn_moe_finalize_routing_v2.cpp。第一段接口参数详解aclnnMoeFinalizeRoutingV2GetWorkspaceSize的完整参数说明如下参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorexpandedX(aclTensor*)输入公式中的 expandedXMoE 的 FFN 输出-FLOAT16、BFLOAT16、FLOAT32NDdrop less 场景(NUM_ROWS * K, H)drop pad 场景(E, C, H)√expandedRowIdx(aclTensor*)输入公式中的 expandedRowIdx-INT32ND(NUM_ROWS * K)√x1Optional(aclTensor*)输入公式中的 x1表示第一个共享专家-与 expandedX 一致ND与 out 一致√x2Optional(aclTensor*)输入公式中的 x2表示第二个共享专家-与 expandedX 一致ND与 out 一致√biasOptional(aclTensor*)输入公式中的 bias表示偏置量-与 expandedX 一致ND(E, H)√scalesOptional(aclTensor*)输入公式中的 scales-FLOAT16、BFLOAT16、FLOAT32ND(NUM_ROWS, K)√expertIdxOptional(aclTensor*)输入公式中的 expertIdxTensor 中的值取值范围是 [0, E-1]INT32ND(NUM_ROWS, K)√dropPadMode(int64_t)输入表示是否支持丢弃模式及 expandedRowIdx 的排列方式取值范围为 [0, 3]----out(aclTensor*)输出公式中的输出-与 expandedX 一致ND(NUM_ROWS, H)×workspaceSize(uint64_t)输出返回需要在 Device 侧申请的 workspace 大小-----executor(aclOpExecutor**)输出返回 op 执行器包含了算子计算流程-----分产品差异约束不同产品对输入类型与组合的约束存在差异官方文档与源码校验逻辑op_api/moe_finalize_routing_common.h保持一致总结如下Atlas A2910B/Atlas A3910_93训练与推理系列产品对应源码IsCommonValidationChip()分支expandedX 要求是 2D/3D Tensor支持 FLOAT16、BFLOAT16、FLOAT32支持 drop less 与 drop pad 两种场景混合精度模式下支持 expandedX 为 BFLOAT16 时 scalesOptional 为 FLOAT32非混合精度模式下scalesOptional 数据类型要求与 expandedX 一致。Ascend 950PR/Ascend 950DTexpandedX 要求是 2D/3D Tensor支持 FLOAT16、BFLOAT16、FLOAT32支持 drop less 与 drop pad 场景scalesOptional 数据类型可以与 expandedX 不一致。Atlas 推理系列产品310P对应源码Is310P()分支expandedX 要求是 2D Tensor数据类型支持 FLOAT16、FLOAT32shape 要求尾轴 H 为 32 对齐x1Optional、x2Optional、biasOptional、expertIdxOptional 仅支持传入 nullptr仅支持 dropPadMode 传入 2scalesOptional 数据类型支持 FLOAT16、FLOAT32且需要与 expandedX 一致。源码层面310P 分支调用CheckParams310Pmoe_finalize_routing_common.h其数据类型白名单仅为{FLOAT16, FLOAT}DTYPE_SUPPORT_LIST_X_310P而通用分支CheckParams支持{FLOAT16, BF16, FLOAT}DTYPE_SUPPORT_LIST_X与文档描述一一对应。返回值与错误码aclnnStatus返回状态码的完整语义参见 aclnn 返回码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的必选输入、必选输出或者必选属性是空指针ACLNN_ERR_PARAM_INVALID161002输入和输出的数据类型、数据格式或维度不在支持的范围之内ACLNN_ERR_INNER_NULLPTR561103多个输入 tensor 之间的 shape 信息不匹配或输入属性和输入 tensor 之间的 shape 信息不匹配校验逻辑在 op_api/moe_finalize_routing_common.h 中分四步执行空指针检查CheckNotNull、数据类型检查CheckDtypeValid、数据格式检查CheckFormatValidexpandedX 在 drop pad 场景还额外允许 3D 的 NCL 格式、shape 维度检查CheckShapeValid要求 drop less 场景 expandedX 为 2D、drop pad 场景为 3DexpandedRowIdx 为 1D各可选输入与 out 均为 2D。第二段接口参数详解参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口aclnnMoeFinalizeRoutingV2GetWorkspaceSize获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream返回值同为aclnnStatus语义见 aclnn 返回码。约束说明确定性计算aclnnMoeFinalizeRoutingV2默认为确定性实现相同输入在不同次调用中结果可复现。符号约定NUM_ROWS 表示行数K 表示从总专家 E 中选出的专家个数H 表示 hidden size列数E 表示专家数需满足 E ≥ KC 表示专家处理 Token 数量的能力阈值expert capacity。expandedRowIdx 取值范围当dropPadMode为 0、2 时Tensor 中值取值范围是 [0, NUM_ROWS * K - 1]当dropPadMode为 1、3 时值取值范围是 [-1, E * C - 1]其中 -1 表示该位置被 pad/丢弃。x1 与 x2 的依赖关系x1Optional 未输入时x2Optional 也不能输入。scales 与 K 的关系scalesOptional 不存在时K 为 1。bias 与 expertIdx 的依赖关系biasOptional 存在时expertIdxOptional 必须同时存在。dropPadMode 取值与 expandedRowIdx 排列方式dropPadMode场景expandedRowIdx 排列对应上游算子输出格式0drop less按列排列与 aclnnMoeInitRouting 输出格式对应1drop pad按列排列与 aclnnMoeInitRouting 输出格式对应2drop less按行排列与 aclnnMoeInitRoutingV2 输出格式对应3drop pad按行排列与 aclnnMoeInitRoutingV2 输出格式对应其中“按列排列”表示展开索引计算方式为idx k * num_rows i即expandedRowIdx[i k*NUM_ROWS]而“按行排列”为idx i * K k。这一点在 golden 参考实现 golden.py 中也有明确体现读者可据此推演自己场景下应使用的取值。 8.空 Tensor 与标量输入不支持空 Tensor不支持标量 Tensor。从源码理解算子实现链路除接口文档外仓库源码完整呈现了该算子从 Shape 推导到核内计算的实现链路可作为理解算子行为的深层依据。Shape 推导InferShapeop_host/moe_finalize_routing_v2_infershape.cpp 实现了输出 Shape 推导校验dropPadMode必须在 [0, 3] 范围内校验 drop lessmode 0/2时 expandedX 为 2Ddrop padmode 1/3时 expandedX 为 3DexpandedRowIdx 为 1D其余可选输入与 out 为 2D交叉一致性校验CheckCrossInputConsistencyscales 与 expertIdx 的第 1 维必须一致同为 Kx1、x2、scales、expertIdx 的第 0 维行数必须一致x1、x2、bias 的第 1 维与 expandedX 尾维H必须一致输出 Shape 推导MoeCopyShapeInput2OutputWithIdx输出为 2D(NUM_ROWS, H)行数取自 scales 的第 0 维若 scales 为空且 k 0则行数 expandedRowIdx 长度 / k输出数据类型与 expandedX 一致。Tiling 与 Workspaceop_host/moe_finalize_routing_v2_tiling.cpp 定义了 Tiling 流程workspace 固定预留 16MBWORKSPACE_RESERVED 16 * 1024 * 1024并在TilingPrepareForMoeFinalizeRoutingV2中获取 AIV 核数与 UB 内存大小作为编译期信息。Tiling 数据结构moe_finalize_routing_v2_tiling.h中包含totalCoreNum、usedCoreNum、H 维度切分normalH/unnormalH/hSliceNum、K 维度切分normalK/unnormalK/kSliceNum、各核行数分配normalCoreHandleNum/tailCoreHandleNum等字段可以看出算子按“核间分行、核内按 H 与 K 双重切分”的策略进行计算。算子注册与 Kernel 实现算子定义moe_finalize_routing_v2_def.cpp通过OP_ADD(MoeFinalizeRoutingV2)注册输入expanded_x、expanded_row_idx为 REQUIREDx1、x2、bias、scales、expert_idx为 OPTIONAL输出y为 REQUIRED属性包括drop_pad_mode默认 0、zero_expert_range/copy_expert_range/constant_expert_range默认 [-1, -1]与k默认 1。AICore 配置按产品拆分为 membase910B/910_93、regbase950支持更多扩展输入 x/a1/a2/v、310P 与 Kirin 多套配置。Kernel 侧op_kernel针对不同数据类型与切分策略提供了大量专用实现文件如moe_finalize_routing_v2_bf16_cuth_k2.h、moe_finalize_routing_v2_fp_db.h、moe_finalize_routing_v2_bf16_all_bias.h、moe_finalize_routing_v2_unpermute_fast.h等并通过moe_finalize_routing_v2_apt.cpp完成适配Tiling Key 常量moe_finalize_routing_v2_tiling.h中也能看到 FLOAT/BF16/FLOAT16 × 大 K/普通 K/CutH K2/K4/DB/全 bias/混合精度等组合的分派维度。调用示例下面给出完整的可编译参考示例与仓库 examples/test_aclnn_moe_finalize_routing_v2.cpp 及官方文档一致具体编译与运行过程可参考 编译与运行样例。#include acl/acl.h #include aclnnop/aclnn_moe_finalize_routing_v2.h #include iostream #include vector #include cstdio #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shape_size 1; for (auto i : shape) { shape_size * i; } return shape_size; } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法资源初始化 auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret); return 0; } template typename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size GetShapeSize(shape) * sizeof(T); // 调用aclrtMalloc申请device侧内存 auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed. ERROR: %d\n, ret); return ret); // 计算连续tensor的strides std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1.固定写法device/stream初始化参考acl API手册 // 根据自己的实际device填写deviceId int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); // check根据自己的需要处理 CHECK_RET(ret 0, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出需要根据API的接口自定义构造 // 示例NUM_ROWS3K2H4E2 std::vectorint64_t expandedXShape {3 * 2, 4}; // (NUM_ROWS*K, H)drop less 列排列 std::vectorint64_t x1Shape {3, 4}; // (NUM_ROWS, H) std::vectorint64_t x2OptionalShape {3, 4}; // (NUM_ROWS, H) std::vectorint64_t biasShape {2, 4}; // (E, H) std::vectorint64_t scalesShape {3, 2}; // (NUM_ROWS, K) std::vectorint64_t expandedExpertIdxShape {3, 2}; // (NUM_ROWS, K) std::vectorint64_t expandedRowIdxShape {3 * 2}; // (NUM_ROWS*K) std::vectorint64_t outShape {3, 4}; // (NUM_ROWS, H) void* expandedXAddr nullptr; void* x1Addr nullptr; void* x2OptionalAddr nullptr; void* biasAddr nullptr; void* scalesDeviceAddr nullptr; void* expandedExpertIdxAddr nullptr; void* expandedRowIdxAddr nullptr; void* outDeviceAddr nullptr; aclTensor* expandedX nullptr; aclTensor* x1 nullptr; aclTensor* x2Optional nullptr; aclTensor* bias nullptr; aclTensor* scales nullptr; aclTensor* expandedExpertIdx nullptr; aclTensor* expandedRowIdx nullptr; aclTensor* out nullptr; std::vectorfloat expandedXHostData {0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1, 8.1, 9.1, 10.1, 11.1, 0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1, 8.1, 9.1, 10.1, 11.1}; std::vectorfloat x1HostData {0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2, 8.2, 9.2, 10.2, 11.2}; std::vectorfloat x2OptionalHostData {0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2, 8.2, 9.2, 10.2, 11.2}; std::vectorfloat biasHostData {0.2, 0.4, 0.2, 0.4, 0.2, 0.4, 0.2, 0.4}; std::vectorfloat scalesHostData {1.3, 1.6, 1.2, 1.8, 1.2, 2.3}; std::vectorint32_t expandedExpertIdxHostData {0, 1, 0, 1, 0, 1}; std::vectorint32_t expandedRowIdxHostData {2, 1, 4, 3, 0, 5}; std::vectorfloat outHostData(12, 0.0f); int64_t dropPadMode 0; // 创建expandedX aclTensor ret CreateAclTensor(expandedXHostData, expandedXShape, expandedXAddr, aclDataType::ACL_FLOAT, expandedX); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建x1 aclTensor ret CreateAclTensor(x1HostData, x1Shape, x1Addr, aclDataType::ACL_FLOAT, x1); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建x2Optional aclTensor ret CreateAclTensor(x2OptionalHostData, x2OptionalShape, x2OptionalAddr, aclDataType::ACL_FLOAT, x2Optional); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建bias aclTensor ret CreateAclTensor(biasHostData, biasShape, biasAddr, aclDataType::ACL_FLOAT, bias); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建scales aclTensor ret CreateAclTensor(scalesHostData, scalesShape, scalesDeviceAddr, aclDataType::ACL_FLOAT, scales); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建expandedExpertIdx aclTensor ret CreateAclTensor(expandedExpertIdxHostData, expandedExpertIdxShape, expandedExpertIdxAddr, aclDataType::ACL_INT32, expandedExpertIdx); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建expandedRowIdx aclTensor ret CreateAclTensor(expandedRowIdxHostData, expandedRowIdxShape, expandedRowIdxAddr, aclDataType::ACL_INT32, expandedRowIdx); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建Out aclTensor ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT, out); CHECK_RET(ret ACL_SUCCESS, return ret); // 3.调用CANN算子库API需要修改为具体的算子接口 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnMoeFinalizeRoutingV2第一段接口 ret aclnnMoeFinalizeRoutingV2GetWorkspaceSize(expandedX, expandedRowIdx, x1, x2Optional, bias, scales, expandedExpertIdx, dropPadMode, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnMoeFinalizeRoutingV2GetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret); } // 调用aclnnMoeFinalizeRoutingV2第二段接口 ret aclnnMoeFinalizeRoutingV2(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnMoeFinalizeRoutingV2 failed. ERROR: %d\n, ret); return ret); // 4.固定写法同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 5. 获取输出的值将device侧内存上的结果拷贝至Host侧需要根据具体API的接口定义修改 auto size GetShapeSize(outShape); std::vectorfloat resultData(size, 0.0f); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr, size * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy result from device to host failed. ERROR: %d\n, ret); return ret); for (int64_t i 0; i size; i) { LOG_PRINT(result[%ld] is: %f\n, i, resultData[i]); } // 6. 释放aclTensor和aclScalar需要根据具体API的接口定义修改 aclDestroyTensor(expandedX); aclDestroyTensor(x1); aclDestroyTensor(x2Optional); aclDestroyTensor(bias); aclDestroyTensor(scales); aclDestroyTensor(expandedExpertIdx); aclDestroyTensor(expandedRowIdx); aclDestroyTensor(out); // 7. 释放device资源需要根据具体API的接口定义修改 aclrtFree(expandedXAddr); aclrtFree(x1Addr); aclrtFree(x2OptionalAddr); aclrtFree(biasAddr); aclrtFree(scalesDeviceAddr); aclrtFree(expandedExpertIdxAddr); aclrtFree(expandedRowIdxAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例要点说明该示例构造的是 drop less、按列排列dropPadMode0的最小可运行场景NUM_ROWS3、K2、H4、E2。expandedRowIdxHostData {2, 1, 4, 3, 0, 5}为 6 个展开行的目标行号列排列下第 i 行第 k 个专家对应下标i k*3CreateAclTensor模板函数统一完成了 Device 内存申请、Host→Device 拷贝、连续 strides 计算与aclCreateTensor创建实际使用时可按需替换为aclrtMallocAsync、aclrtMemcpyAsync等异步接口示例中所有输入均以非空 Tensor 传入若某可选输入不参与计算对应位置传nullptr即可如只使用共享专家 x1 时x2Optional 需传 nullptr详见“约束说明”第 4 条若想验证 drop pad 场景dropPadMode1/3需将 expandedX 改为 3D(E, C, H)并将 expandedRowIdx 中 pad 位置置为 -1。测试与验证golden 参考实现仓库在 tests/assets/golden.py 中提供了算子行为的 Python golden 参考实现完整复刻了前述计算公式可作为结果比对与自测依据输出初始化为x1 x2若传入遍历i ∈ [0, num_rows)、kk ∈ [0, K)按 dropPadMode 计算展开索引0/1 为列排列kk * num_rows i2/3 为行排列i * K kkdrop pad 场景mode 1/3下expandedRowIdx -1的行直接跳过取expandedX[expandedRowIdx[idx]]作为专家输出若传入 bias 与 expertIdx 则叠加bias[expertIdx[i, kk]]再乘scales[i, kk]累加进输出容差配置按stat_rel_err相对误差统计校验 float32/float16/bfloat16。此外tests/st/aclnnMoeFinalizeRoutingV2 下提供了基于 ATK 框架的 st 用例atk_aclnnMoeFinalizeRoutingV2.json数据用例与executor_aclnnMoeFinalizeRoutingV2.py执行器tests/ut 下包含 infershape、tiling、kernel 与 onnx 插件的单测感兴趣可结合 测试用例 CSV 覆盖更多边界组合。小结aclnnMoeFinalizeRoutingV2是 CANN ops-transformer 中 MoE 计算链路的收尾算子负责将路由展开后的专家 FFN 输出按expandedRowIdx合并回原始行序并叠加共享专家、偏置与缩放系数。使用时的关键决策点集中在三处dropPadMode 的选择决定 expandedX 是 2D 还是 3D、expandedRowIdx 排列方式与取值区间、可选输入的依赖关系x2 依赖 x1、bias 依赖 expertIdx、产品差异约束尤其是 310P 的 H 对齐与模式限制。本文结合 官方接口文档 与仓库源码给出了从公式语义、参数约束到完整调用示例的闭环说明可作为基于 CANN 开发 MoE 推理/训练算子的直接参考。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考