custom-npu_moe_init_routing_group_quant【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer产品支持情况产品是否支持Ascend 950PR/Ascend 950DT√功能说明算子功能MoE的routing计算根据MoE Gating TopK的计算结果做routing处理支持不量化和动态量化模式。计算公式1.对输入expertIdx做排序得出排序后的结果sortedExpertIdx和对应的序号sortedRowIdx$$ sortedExpertIdx, sortedRowIdxkeyValueSort(expertIdx,rowIdx) $$2.以sortedRowIdx做位置映射得出expandedRowIdx$$ expandedRowIdx[sortedRowIdx[i]]i $$3.在drop模式下对sortedExpertIdx的每个专家统计直方图结果得出expertTokensCountOrCumsum$$ expertTokensCountOrCumsum[i]Histogram(sortedExpertIdx) $$4.计算quant结果动态quant若不输入scale $$ dynamicQuantScaleOut row_max(abs(x)) / 127 $$$$ quantResult round(x / dynamicQuantScaleOut) $$若输入scale: $$ dynamicQuantScaleOut row_max(abs(x * scale)) / 127 $$$$ quantResult round(x / dynamicQuantScaleOut) $$5.对quantResult取前NUM_ROWS个sortedRowIdx的对应位置的值得出expandedXOut$$ expandedX[i]quantResult[sortedRowIdx[i]%NUM_ROWS] $$6.expandedRowIdx的有效元素数量availableIdxNum计算方式为expertIdx中activeExpertRange范围内的元素的个数 $$ availableIdxNum |{x\in expertIdx| expert_start \le xexpert_end \ }| $$函数原型custom.npu_moe_init_routing_group_quant(Tensor x, Tensor expert_idx, Tensor? scaleNone, Tensor? offsetNone, int active_num-1, int expert_capacity-1, int expert_num-1, int drop_pad_mode-1, int expert_tokens_num_type-1, bool expert_tokens_num_flagFalse, int quant_mode-1, SymInt[] active_expert_range, int row_idx_type-1, int group_size128) -(Tensor, Tensor, Tensor, Tensor)参数说明参数名输入/输出/属性描述数据类型数据格式x输入MOE的输入即token特征输入对应公式中x。FLOAT32、FLOAT16、BFLOAT16、INT8NDexpert_idx输入每一行特征对应的K个处理专家里面元素专家id不能超过专家数。对应公式中 expertIdx。INT32NDscale可选输入表示用于计算quant结果的参数。如果不输入表示计算时不使用scale对应公式中scale。FLOAT32NDoffset可选输入表示用于计算quant结果的偏移值。在非量化场景下和动态quant场景下不输入对应公式中offset。FLOAT32NDactive_num属性表示总的最大处理row数输出expanded_x只有这么多行是有效的。INT-expert_capacity属性表示每个专家能够处理的tokens数取值范围大于等于0。INT-expert_num属性表示专家数expert_tokens_num_type为key\_value模式时取值范围为[0, 5120], 其它模式取值范围[0, 10240]。INT-drop_pad_mode属性表示是否为 DropPad 场景取值为 0 和 1。0表示 Dropless 场景该场景下不校验 expert_capacity。1表示 DropPad 场景。INT-expert_tokens_num_type属性取值为0、1和2 。0表示 comsum 模式。1表示 count 模式即输出的值为各个专家处理的 token 数量的累计值。2表示 key\_value 模式即输出的值为专家和对应专家处理 token 数量的累计值。INT-expert_tokens_num_flag属性取值为false和true。false表示不输出 expert_tokens_count_or_cumsum。true表示输出 expert_tokens_count_or_cumsum。Bool-quant_mode属性取值为-1、0、1、2、3、4、5。-1表示不量化场景。0表示静态 quant 场景。1表示动态 quant 场景。2表示MXFP8量化场景expanded_x量化到FLOAT8_E5M2。3表示MXFP8量化场景expanded_x量化到FLOAT8_E4M3FN。4表示PerGroup量化group_size固定为128expanded_x量化到FLOAT8_E5M2scale的dtype为float。5表示PerGroup量化group_size固定为128expanded_x量化到FLOAT8_E4M3FNscale的dtype为float。INT-active_expert_range可选属性长度为2数组内的值为[expertStart, expertEnd], 表示活跃的expert范围在expertStart和expertEnd之间左闭右开。要求值大于等于0并且expertEnd不大于expertNum。不输入则不使能专家筛选功能。ListInt-row_idx_type属性表示expanded_row_idx使用的索引类型取值为0、1。性能模板仅支持10表示gather类型的索引。1表示scatter类型的索引。INT-expanded_x输出根据expert_idx进行扩展过的特征。非量化场景下数据类型同x量化场景quant_mode为0、1时数据类型支持INT8quant_mode为2、3时数据类型分别支持FLOAT8_E5M2、FLOAT8_E4M3FN。FLOAT32、FLOAT16、BFLOAT16、INT8、FLOAT8_E5M2、FLOAT8_E4M3FNNDexpanded_row_idx输出expanded_x和x的索引映射关系前available_idx_num\*H个元素为有效数据其余无效数据当row_idx_type为0时无效数据由-1填充当row_idx_type为1时无效数据未初始化。INT32NDexpert_tokens_count_or_cumsum输出在expert_tokens_num_type为1的场景下表示active_expert_range范围内expert对应的处理token的总数。在expert_tokens_num_type为2的场景下表示active_expert_range范围内token总数为非0的expert以及对应expert处理token的总数。INT64NDexpanded_scale输出输出量化计算过程中scale的中间值。FLOAT32ND约束说明输入值域限制active_num 当前未使用校验需等于NUM_ROWS*K。expert_capacity 当前未使用仅校验非空。offset 当前未使用。drop_pad_mode 当前只支持0代表 Dropless 场景。expert_tokens_num_type 当前只支持 1 和 2分别代表 count 模式和 key_value 模式。expert_tokens_num_flag 只支持 true代表输出 expert_tokens_count_or_cumsum。quant_mode: 支持-1、0、1、2、3、4、5其中-1表示不量化0表示静态量化输出INT81表示动态量化输出INT82表示MXFP8量化到FLOAT8_E5M23表示MXFP8量化到FLOAT8_E4M3FN。4模式是groupSize固定为128的PerGroup量化expand_x量化到float8_e5m2类型scale为float类型。5模式是groupSize固定为128的PerGroup量化expand_x量化到float8_e4m3类型scale为float类型。输入shape约束x shape为 (N, H)expert_id : shape为(N, topK)scale : shape为(N, )或者(expert, H)输出shape约束expanded_x : shape为 (N*K, H)expanded_row_idx : shape为 (N*K)expert_tokens_count_cumsum : expert_tokens_num_type为1时shape为 (expert, ) expert_tokens_num_type为2时shape为 (expert, 2)expanded_scale : quant_mode 4或者quant_mode 5时shape为 (N*K, M)其中MCeil(H, 128)该接口支持推理场景下使用。该接口支持图模式。该接口与PyTorch配合使用时需要保证CANN相关包与PyTorch相关包的版本匹配。调用示例详见 test_npu_moe_init_routing_group_quant.py【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考