CANN ops-transformer 算子解析aclnnMoeInitRoutingV3 两段式接口、路由重排与全量化模式实战【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer本文是 CANN ops-transformer 仓库中moe_init_routing_v3算子的深度技术指南。该算子负责 MoEMixture of Experts混合专家模型的 routing 计算将 Gating 网络产出的专家索引expertIdx转换为可直接送入专家计算的扩展特征expandedXOut与索引映射expandedRowIdxOut并支持不量化、静态量化、动态量化及多种 FP8/FP4/INT4/HIFLOAT8 量化输出。阅读本文后你将掌握该算子的两段式 aclnn 接口用法、全部 18 种 quantMode 量化语义、Dropless/DropPad/Active 三种输出模式、按产品的支持差异与性能模板准入条件并能参照完整示例代码在 NPU 上独立跑通该算子。算子定位MoE 推理链路中的路由核心在 MoE 大模型中一个 token 并不会被所有专家处理而是由 Gating 网络例如aclnnMoeGatingTopKSoftmaxV2为每个 token 选出 Top-K 个专家并给出权重。aclnnMoeInitRoutingV3正是承接这一结果、完成路由动作的算子它根据 aclnnMoeGatingTopKSoftmaxV2 计算出的 expertIdx 结果对 token 特征做重排Gather/Scatter与量化输出专家侧可直接消费的expandedXOut及配套索引是整个 MoE 前向链路上承上启下的关键算子。该算子的产品支持情况如下见 README 与 接口文档产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持相对 V2 的功能变更aclnnMoeInitRoutingV3是针对 V2 接口 aclnnMoeInitRoutingV2 的演进版本核心变更可以归纳为三点量化能力增强新增动态与静态量化功能支持expandedXOut的 INT8 量化输出在 Ascend 950 系列上进一步扩展出 MXFP8quantMode2/3、FP8 PerGroup4/5、HIF86/7/8、MXFP49、FP8 PerBlock11/12、INT4 动态量化13、FP8 PerGroup Amax14/15、MXFP8 RoundScaleAmax16/17等多种量化模式。输出重定义删除输出expertTokensBeforeCapacityOut新增输出expertTokensCountOrCumsumOut并新增activeExpertRangeOptional参数支持筛选有效范围内的 expertId从 README 的功能说明可见。属性重定义原有expertTokensBeforeCapacityFlag(bool)和expertTokensCountOrCumsumFlag(int)被重新定义为expertsTokensNumFlag(bool)与expertTokensNumType(int)通过DropPadMode、expertsTokensNumFlag、expertTokensNumType三者组合统一了 cumsum / count / key_value 三种直方图输出格式对应关系如下DropPadModeexpertsTokensNumFlagexpertTokensNumType输出格式说明0true0cumsum 模式expertTokensCountOrCumsumOut表示按排序后各专家处理 token 的计数前缀和直方图0true1count 模式expertTokensCountOrCumsumOut表示按排序后各专家处理 token 的单独计数直方图0true2key_value 模式输出 shape 为[expert_num, 2]表示每个专家和该专家处理非零 token 数量的累计值1true1输出模式为 count 模式不开启false不开启不输出expertTokensCountOrCumsumOut工作原理排序、索引映射与量化公式算子的计算过程在 接口文档 中有完整的数学描述整体分为六步第 1 步对 expertIdx 做 key-value 排序。输入专家索引expertIdxshape 为(NUM_ROWS, K)按专家 id 排序得到排序后的sortedExpertIdx与对应的原始位置序号sortedRowIdx$$sortedExpertIdx,\ sortedRowIdxkeyValueSort(expertIdx,rowIdx)$$第 2 步以sortedRowIdx做位置映射得到expandedRowIdxOut。具体输出哪种索引由属性rowIdxType决定rowIdxType1输出 scatter 索引$expandedRowIdxOut[i]sortedRowIdx[i]$rowIdxType0输出 gather 索引$expandedRowIdxOut[sortedRowIdx[i]]i$。第 3 步统计直方图。对排序后的sortedExpertIdx统计各专家 token 数得到expertTokensCountOrCumsumOutOptional$$expertTokensCountOrCumsumOutOptional[i]Histogram(sortedExpertIdx)$$第 4 步量化计算quantMode ≠ -1 时。静态量化与动态量化的公式如下静态量化$quantResultround((x*scaleOptional)offsetOptional)$动态量化不输入 scale$dynamicQuantScaleOutOptional row_max(abs(x)) / 127$$quantResult round(x / dynamicQuantScaleOutOptional)$动态量化输入 scale$dynamicQuantScaleOutOptional row_max(abs(x * scaleOptional)) / 127$$quantResult round(x / dynamicQuantScaleOutOptional)$当quantMode13INT4 动态量化时对称量化范围变为 [-8, 7]scale 计算中的分母为 7量化结果沿 H 维每两个 INT4 值打包为 1 个字节因此要求 H 为偶数。第 5 步按索引搬运 token 得到expandedXOut。若活跃 expert 范围为全专家范围即activeExpertRangeOptional覆盖全部专家则按 Scatter 索引搬运否则按 Gather 索引搬运Scatter 搬运$expandedXOut[i]x[scatterRowIdx[i] // K]$量化场景为 $expandedXOut[i]quantResult[scatterRowIdx[i] // K]$Gather 搬运$expandedXOut[gatherRowIdx[i]]x[i // K]$量化场景为 $expandedXOut[gatherRowIdx[i]]quantResult[i // K]$。在dropPadMode1时还会将每个专家需要处理的 token 个数对齐为expertCapacity个超过容量的 token 被 Drop不足的用 0 填充Pad。第 6 步计算有效元素数量availableIdxNum。即expertIdx中落在activeExpertRangeOptional[expertStart, expertEnd)范围内元素的个数-1表示无效专家不参与路由计算$$availableIdxNum |{x\in expertIdx| expert_start \le x expert_end }|$$两段式接口与函数原型与 CANN aclnn 系列算子一致aclnnMoeInitRoutingV3采用两段式接口详见 两段式接口说明必须先调用aclnnMoeInitRoutingV3GetWorkspaceSize完成入参校验、计算 workspace 大小并生成包含算子计算流程的执行器再调用aclnnMoeInitRoutingV3真正执行计算。aclnnStatus aclnnMoeInitRoutingV3GetWorkspaceSize( const aclTensor *x, const aclTensor *expertIdx, const aclTensor *scaleOptional, const aclTensor *offsetOptional, int64_t activeNum, int64_t expertCapacity, int64_t expertNum, int64_t dropPadMode, int64_t expertTokensNumType, bool expertTokensNumFlag, int64_t quantMode, const aclIntArray *activeExpertRangeOptional, int64_t rowIdxType, const aclTensor *expandedXOut, const aclTensor *expandedRowIdxOut, const aclTensor *expertTokensCountOrCumsumOut, const aclTensor *expandedScaleOut, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnMoeInitRoutingV3( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)接口头文件在 op_host/op_api/aclnn_moe_init_routing_v3.h 中声明extern C导出头文件位于aclnnop/aclnn_moe_init_routing_v3.h实现位于 op_host/op_api/aclnn_moe_init_routing_v3.cpp。参数详解输入参数Tensor参数名输入/输出描述使用说明数据类型shapex输入MoE 的 token 特征输入shape 为(NUM_ROWS, H)quantMode9/13 的 MXFP4/INT4 动态量化场景以及 quantMode-1 且 x 为 FLOAT4_E2M1 的非量化透传场景H 要求为偶数见下文量化模式表格2expertIdx输入每行特征对应的 K 个处理专家元素专家 id 不能超过专家数-1 表示无效专家该位置不参与路由会被过滤shape 为(NUM_ROWS, K)INT322scaleOptional输入可选用于计算量化结果的参数非量化场景可选静态量化quantMode0必须输入且 shape 为[1,]动态量化quantMode1可选输入时要求 2D(expertEnd-expertStart, H)INT4 动态量化quantMode13可选输入时 shape 为(1, H)表示按 H 维广播的 smooth scaleMXFP8/HIF8 直转/FP8 PerGroup/MXFP4/FP8 PerBlock/MXFP8 RoundScaleAmax 场景不输入HIF8 PERTENSORquantMode7必须输入 shape[1,]FLOAT32、FLOAT8_E8M01-3offsetOptional输入可选用于计算 quant 结果的偏移值非量化、动态量化quantMode1、MXFP8/HIF8/MXFP4/FP8 PerGroup/FP8 PerBlock/INT4/MXFP8 RoundScaleAmax 场景均不输入静态量化quantMode0必须输入shape 为[1,]FLOAT32-属性参数int64_t / bool / aclIntArray参数名类型描述与取值范围activeNumINT64总的最大处理 row 数。A2/A3dropPadMode0 时支持 ≥-1-1/0 表示不限制0 时最多处理min(activeNum, NUM_ROWS*K)行。Ascend 950不用于限制处理行数仅接受 -1、0 或NUM_ROWS*KexpertCapacityINT64每个专家能处理的 token 数。Dropless 场景不使用DropPad 场景必须校验且取值范围为(0, NUM_ROWS]expertNumINT64专家数。key_value 模式expertTokensNumType2时范围为 [1, 5120]其它模式 [1, 10240]dropPadModeINT640 表示 Dropless 场景1 表示 DropPad 场景expertTokensNumTypeINT640cumsum 模式1count 模式2key_value 模式expertTokensNumFlagBOOL是否输出 expertTokensCountOrCumsumOutfalse/truequantModeINT64量化场景取值 -1、0、1、2、3、4、5、6、7、8、9、11、12、13、14、15、16、17详见量化模式小节activeExpertRangeOptionalaclIntArray活跃 expert 范围[expertStart, expertEnd)左闭右开要求 ≥0 且 expertEnd ≤ expertNumDropPad 场景下必须为[0, expertNum]rowIdxTypeINT640gather 类型索引1scatter 类型索引输出参数Tensor参数名描述与 shapeexpandedXOut根据 expertIdx 扩展过的特征。Dropless[NUM_ROWS*K, H]Active[min(activeNum, NUM_ROWS*K), H]DropPad3D[expertNum, expertCapacity, H]。数据类型随量化模式变化expandedRowIdxOutexpandedXOut 与 x 的索引映射shape(NUM_ROWS*K,)前 availableIdxNum 个元素有效rowIdxType0 时无效数据填 -1rowIdxType1 时无效数据未初始化。INT32expertTokensCountOrCumsumOut各专家处理 token 数量的统计或累加值。type0/1 时 shape 为[expertEnd-expertStart]type2 时 shape 为[expertNum, 2]。INT64expandedScaleOut量化过程中 scale 的中间值。非量化 输入 scale 时 shape[NUM_ROWS*K]FLOAT32x 为 FP8/FP4 时 shape[NUM_ROWS*K, CeilDiv(H,64), 2]FLOAT8_E8M0MXFP82/3/16/17时[NUM_ROWS*K, M]MCeilAlign(CeilDiv(H,32),2)FLOAT8_E8M0HIF8 PERTOKEN 时[NUM_ROWS*K]FLOAT32MXFP4 时[NUM_ROWS*K, CeilDiv(H,64), 2]FLOAT8_E8M0FP8 PerGroup4/5/14/15时[NUM_ROWS*K, CeilDiv(H,128)]FLOAT32FP8 PerBlock11/12时[NUM_ROWS*K, CeilDiv(H,256), 2]FLOAT32DropPad 场景为[expertNum*expertCapacity]FLOAT32参数校验在 op_host/moe_init_routing_v3_infershape.cpp 中实现例如validQuantModes集合定义了全部 18 个合法 quantMode第 82-99 行CheckScaleShapeForStaticQuant要求静态量化 scale 为 1D 且 shape 为 1CheckScaleShapeForDynamicQuant要求动态量化 scale 为 2D(activeExpertRange, H)GetAndCheckAttrRowIdxType在 DropPad 模式下强制rowIdxType0CalculateOutputDims中activeNum0||-1时取n*k否则取min(activeNum, n*k)。这些源码级约束与文档描述一一对应。量化模式quantMode全解量化模式是 V3 相比 V2 最核心的增强expandedXOut的输出数据类型随 quantMode 变化quantMode场景说明expandedXOut 数据类型-1不量化透传x 支持 FLOAT16/BFLOAT16/FLOAT32/INT8/HIFLOAT8/FLOAT4_E2M1/FLOAT8_E4M3FN/FLOAT8_E5M2同 x0静态量化必须输入 scale/offsetINT81动态量化INT8scale 可选INT82 / 3MXFP8 量化FLOAT8_E5M2 / FLOAT8_E4M3FN4 / 5FP8 PerGroupGroupSize128RoundScaleFLOAT8_E5M2 / FLOAT8_E4M3FN6HIF8 直转量化HIFLOAT87HIF8 PERTENSOR 量化必须输入 scaleHIFLOAT88HIF8 PERTOKEN 量化HIFLOAT89MXFP4 量化FLOAT4_E2M111 / 12FP8 PerBlockBlockSize128expandedScaleOut 为 FLOAT32 三维布局FLOAT8_E5M2 / FLOAT8_E4M3FN13INT4 动态量化x 为 FLOAT32/BFLOAT16INT414 / 15FP8 PerGroupGroupSize128RoundScaleAmaxFLOAT8_E5M2 / FLOAT8_E4M3FN16 / 17MXFP8 RoundScaleAmaxFLOAT8_E5M2 / FLOAT8_E4M3FN各量化模式下 x 的数据类型约束quantMode0/1 支持 FLOAT16/BFLOAT16/FLOAT32quantMode2、3、4、5、6、7、8、9、11、12、14、15、16、17 支持 FLOAT16/BFLOAT16quantMode13 支持 FLOAT32/BFLOAT16。这些 dtype 推导逻辑在 infershape 的DetermineOutputDtypes中实现同时在 op_host/moe_init_routing_v3_def.cpp 的OpAICoreConfigascend950 配置中注册了对应 25 组输入/输出数据类型组合。产品差异Atlas A2/A3 系列仅支持 -1、0、1且 quantMode-1 时 x 仅支持 FLOAT16/BFLOAT16/FLOAT32/INT8其余全部量化模式仅 Ascend 950PR/950DT 支持。Ascend 950 上expertTokensNumFlag仅支持 true。约束说明与性能模板确定性计算aclnnMoeInitRoutingV3为默认确定性实现保证相同输入多次运行的输出一致。空 Tensor 处理A2/A3NUM_ROWS0时进入空 Tensor 处理路径Ascend 950NUM_ROWS0或K0时没有路由元素进入空 Tensor 处理路径专家计数为 0输出 shape 仍需满足相应模式约束NUM_ROWS*K0且H0时仍走正常路由流程生成expandedRowIdxOut和expertTokensCountOrCumsumOut。输入值域不支持输入为 inf/-inf/nanactiveExpertRangeOptional需满足expertStart expertEnd、expertStart ≥ 0、expertEnd ≤ 10240infershape 中GetAndCheckAttrActiveExpertRange的校验。性能模板准入条件A2/A3 产品上算子根据输入特征自动选择性能模板不满足条件时退化为通用模板性能模板类型准入条件低时延性能模板x/expertIdx/scaleOptional shape 分别为 (1, 7168)/(1, 8)/(256, 7168)x 为 BFLOAT16属性activeExpertRangeOptional[0, 256]、quantMode1、expertTokensNumType2、expertNum256大batch性能模板NUM_ROWS ∈ [384, 8192]K8属性expertNum256、expertEnd-expertStart≤32、quantMode-1、rowIdxType1、expertTokensNumType1全载性能模板dropPadMode0Ascend 950 非量化场景还支持 dropPadMode1此时 quantMode-1、rowIdxType0、activeExpertRange[0, expertNum]。该模板在输入 shape 较小时将搬入、排序、计算放入同一个 kernel 内完成避免多核同步开销README 中还补充了 950 上的计数排序性能模板FullLoad 模板要求 x 为 BFLOAT16/FLOAT16/FLOAT32/INT8、expertNum≤1024、expertEnd-expertStart≤32、quantMode-1、dropPadMode0CutOrigin 模板要求 expertNum≤1024、expertEnd-expertStart≤128、quantMode 为 -1 或 0、NUM_ROWS*H*x.dtype 1.5*totalUbSize、NUM_ROWS*K ≥ 8192。性能模板的实际调度由 op_host/moe_init_routing_v3_tiling.cpp 及 op_host/moe_init_routing_v3_tiling_arch35.cpp 完成tiling 数据字段coreNum、expertStart/expertEnd、quantMode、dropPadMode、各子任务分核参数等定义在 op_host/moe_init_routing_v3_tiling.h 中op_host/moe_init_routing_v3_tiling_base.cpp 负责在 Tiling 准备阶段获取 AIV 核数与 UB 大小TilingPrepareForMoeInitRountingV3为分核决策提供硬件信息。950 DropPad 模式特殊约束Ascend 950 上 dropPadMode1 时rowIdxType 仅支持 0gather 索引activeExpertRangeOptional 必须为[0, expertNum]expertTokensNumType 仅支持 1count 模式quantMode 仅支持 -1非量化且 x 数据类型仅支持 FLOAT16/BFLOAT16/FLOAT32/INT8/HIFLOAT8。返回码与错误处理两段式接口均返回aclnnStatus状态码具体参见 aclnn 返回码说明。第一段接口GetWorkspaceSize完成入参校验典型错误如下返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001计算输入和计算输出是空指针ACLNN_ERR_PARAM_NULLPTR161002输入和输出的数据类型不在支持的范围内ACLNN_ERR_INNER_TILING_ERROR561002输入、输出 Tensor 的 shape 不在支持的范围内或输入的属性不在支持的范围内完整调用示例下面是在单卡上跑通aclnnMoeInitRoutingV3的完整 C 示例该示例也以 examples/test_aclnn_moe_init_routing_v3.cpp 形式存在于仓库中编译与运行的整体流程可参考 编译与运行样例#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_moe_init_routing_v3.h #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初始化 int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret 0, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出 std::vectorint64_t xShape {3, 2}; std::vectorint64_t expertIdxShape {3, 4}; std::vectorint64_t scaleShape {3}; std::vectorint64_t offsetShape {1}; std::vectorint64_t expandedXOutShape {12, 2}; std::vectorint64_t expandedRowIdxOutShape {12}; std::vectorint64_t expertTokensCountOrCumsumOutOptionalShape {4}; std::vectorint64_t expandedScaleOutOptionalShape {12}; std::vectorint64_t activeExpertRangeArray {0, 4}; void *xDeviceAddr nullptr; void *expertIdxDeviceAddr nullptr; void *scaleDeviceAddr nullptr; void *offsetDeviceAddr nullptr; void *expandedXOutDeviceAddr nullptr; void *expandedRowIdxOutDeviceAddr nullptr; void *expertTokensCountOrCumsumOutOptionalDeviceAddr nullptr; void *expandedScaleOutOptionalDeviceAddr nullptr; aclTensor *x nullptr; aclTensor *expertIdx nullptr; aclTensor *scale nullptr; aclTensor *offset nullptr; int64_t activeNum 12; int64_t expertCapacity 0; int64_t expertNum 256; int64_t dropPadMode 0; int64_t expertTokensNumType 1; bool expertTokensNumFlag true; int64_t quantMode -1; aclIntArray *activeExpertRange aclCreateIntArray(activeExpertRangeArray.data(), activeExpertRangeArray.size()); int64_t rowIdxType 1; aclTensor *expandedXOut nullptr; aclTensor *expandedRowIdxOut nullptr; aclTensor *expertTokensCountOrCumsumOutOptional nullptr; aclTensor *expandedScaleOutOptional nullptr; std::vectorfloat xHostData {0.1, 0.1, 0.2, 0.2, 0.3, 0.3}; std::vectorint expertIdxHostData {1, 2, 0, 3, 0, 2, 1, 3, 0, 1, 3, 2}; std::vectorfloat scaleHostData {0.3423, 0.1652, 0.2652}; std::vectorfloat offsetHostData {1.8369}; std::vectorfloat expandedXOutHostData {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; std::vectorint expandedRowIdxOutHostData {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; std::vectorint64_t expertTokensCountOrCumsumOutOptionalHostData {0, 0, 0, 0}; std::vectorfloat expandedScaleOutOptionalHostData {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; // 创建输入 aclTensor ret CreateAclTensor(xHostData, xShape, xDeviceAddr, aclDataType::ACL_FLOAT, x); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(expertIdxHostData, expertIdxShape, expertIdxDeviceAddr, aclDataType::ACL_INT32, expertIdx); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(scaleHostData, scaleShape, scaleDeviceAddr, aclDataType::ACL_FLOAT, scale); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(offsetHostData, offsetShape, offsetDeviceAddr, aclDataType::ACL_FLOAT, offset); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建输出 aclTensor ret CreateAclTensor(expandedXOutHostData, expandedXOutShape, expandedXOutDeviceAddr, aclDataType::ACL_FLOAT, expandedXOut); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(expandedRowIdxOutHostData, expandedRowIdxOutShape, expandedRowIdxOutDeviceAddr, aclDataType::ACL_INT32, expandedRowIdxOut); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(expertTokensCountOrCumsumOutOptionalHostData, expertTokensCountOrCumsumOutOptionalShape, expertTokensCountOrCumsumOutOptionalDeviceAddr, aclDataType::ACL_INT64, expertTokensCountOrCumsumOutOptional); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(expandedScaleOutOptionalHostData, expandedScaleOutOptionalShape, expandedScaleOutOptionalDeviceAddr, aclDataType::ACL_FLOAT, expandedScaleOutOptional); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用第一段接口获取workspaceSize和executor uint64_t workspaceSize 0; aclOpExecutor *executor; ret aclnnMoeInitRoutingV3GetWorkspaceSize( x, expertIdx, scale, offset, activeNum, expertCapacity, expertNum, dropPadMode, expertTokensNumType, expertTokensNumFlag, quantMode, activeExpertRange, rowIdxType, expandedXOut, expandedRowIdxOut, expertTokensCountOrCumsumOutOptional, expandedScaleOutOptional, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnMoeInitRoutingV3GetWorkspaceSize 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); } // 4. 调用第二段接口执行计算 ret aclnnMoeInitRoutingV3(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnMoeInitRoutingV3 failed. ERROR: %d\n, ret); return ret); // 5. 同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 6. 将device侧结果拷贝回host并打印 auto expandedXSize GetShapeSize(expandedXOutShape); std::vectorfloat expandedXData(expandedXSize, 0); ret aclrtMemcpy(expandedXData.data(), expandedXData.size() * sizeof(expandedXData[0]), expandedXOutDeviceAddr, expandedXSize * 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 expandedXSize; i) { LOG_PRINT(expandedXData[%ld] is: %f\n, i, expandedXData[i]); } auto expandedRowIdxSize GetShapeSize(expandedRowIdxOutShape); std::vectorint expandedRowIdxData(expandedRowIdxSize, 0); ret aclrtMemcpy(expandedRowIdxData.data(), expandedRowIdxData.size() * sizeof(expandedRowIdxData[0]), expandedRowIdxOutDeviceAddr, expandedRowIdxSize * sizeof(int32_t), 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 expandedRowIdxSize; i) { LOG_PRINT(expandedRowIdxData[%ld] is: %d\n, i, expandedRowIdxData[i]); } // 7. 释放aclTensor和device资源 aclDestroyTensor(x); aclDestroyTensor(expertIdx); aclDestroyTensor(scale); aclDestroyTensor(offset); aclDestroyTensor(expandedXOut); aclDestroyTensor(expandedRowIdxOut); aclDestroyTensor(expertTokensCountOrCumsumOutOptional); aclDestroyTensor(expandedScaleOutOptional); aclrtFree(xDeviceAddr); aclrtFree(expertIdxDeviceAddr); aclrtFree(scaleDeviceAddr); aclrtFree(offsetDeviceAddr); aclrtFree(expandedXOutDeviceAddr); aclrtFree(expandedRowIdxOutDeviceAddr); aclrtFree(expertTokensCountOrCumsumOutOptionalDeviceAddr); aclrtFree(expandedScaleOutOptionalDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例中xShape{3,2}、expertIdxShape{3,4}即 3 行 token、每行选 4 个专家NUM_ROWS3K4非量化模式quantMode-1下输出expandedXOutshape 为[12, 2]。若需要验证 950 上的量化路径可将quantMode调整为对应取值并同步替换expandedXOut/expandedScaleOut的数据类型与 shape。源码结构导读moe_init_routing_v3目录按 CANN 算子标准布局组织详见 目录清单 与 接口文档接口层op_host/op_api/aclnn_moe_init_routing_v3.h/.cpp暴露 aclnn 两段式 APImoe_init_routing_v3.h/.cpp为算子内部 API 封装算子注册op_host/moe_init_routing_v3_def.cpp通过OpDef声明输入/输出/属性及ascend910b、ascend910_93、ascend950三个 AICore 配置shape/dtype 推导op_host/moe_init_routing_v3_infershape.cpp全部属性与 shape 校验、输出 shape/dtype 推导tilingop_host/moe_init_routing_v3_tiling*.cpp/h分核、分块与性能模板调度kernelop_kernel/按场景拆分为 sort单核/多核/合并排序、gather含 droppad、mxfp4/mxfp8、hif8、fp8 group/perblock 量化、full_load、expert_tokens_count、mrgsort 等子模块arch35 子目录为 950 平台实现产品配置op_host/config/ascend910_93、ascend910b、ascend950 三套binary.json与simplified_key.ini测试tests/ut/op_host/op_api/test_aclnn_moe_init_routing_v3.cpp覆盖两段式接口的参数校验与执行路径ut/op_host/下还有 tiling 与 infershape 单测st/下含 aclnn 的算子级系统测试用例与golden.py数据生成脚本另有 tests/ut/op_kernel/ 的 kernel 级测试gen_data.py/gen_tiling.py生成输入与 tiling 数据。小结aclnnMoeInitRoutingV3是 CANN ops-transformer 中面向 MoE 路由场景的高阶算子它以两段式 aclnn 接口承接 Gating 结果通过排序 直方图统计 Gather/Scatter 搬运完成 token 到专家的路由分发并在搬运的同时完成从透传到 INT8/FP8/FP4/INT4/HIFLOAT8 的多种量化支持 Dropless、Active、DropPad 三种输出布局。理解其参数语义、量化模式与产品差异是将其正确接入自有 MoE 推理/训练链路、并利用性能模板获得更优执行效率的前提。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
