CANN MoE门控TopK算子文档

CANN MoE门控TopK算子文档 custom-npu_moe_gating_top_k【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer产品支持情况产品是否支持Atlas A3 推理系列产品√Ascend 950PR/Ascend 950DT√功能说明算子功能MoE (Mixture of Experts) 计算中对输入x做Sigmoid、SoftMax或者Softplus计算对计算结果分组进行排序对结果取前K个专家或根据输入的词表进行hash。计算公式对输入做Sigmoid或者SoftMax$$ \begin{aligned} if\ normType 1: normOutSigmoid(x) \ else\ if\ normType 0: normOutSoftMax(x) \ else\ : normOutSoftplus(x) \end{aligned} $$如果bias不为空$$ normOut normOut bias $$对计算结果按照groupCount进行分组每组按照groupSelectMode取max或topk2的sum值对group进行排序取前kGroup个组$$ groupOut, groupId TopK(ReduceSum(TopK(Split(normOut, groupCount), k2, dim-1), dim-1),kkGroup) $$如果指定了input_ids和tid2eid则根据输入的词表进行hash操作 否则根据上一步的groupId获取normOut中对应的元素将数据再做TopK得到expertIdxOut的结果$$ y,expertIdxOutTopK(normOut[groupId, :],kk) $$对y按照输入的routedScalingFactor和eps参数进行计算得到yOut的结果$$ yOut y / (ReduceSum(y, dim-1)eps)*routedScalingFactor $$函数原型custom.npu_moe_gating_top_k(Tensor x, int k, *, Tensor? biasNone, Tensor? input_idsNone, Tensor? tid2eidNone, int k_group1, int group_count1, float routed_scaling_factor1., float eps9.9999999999999995e-21, int group_select_mode0, int renorm0, int norm_type0, bool out_flagFalse) - (Tensor, Tensor, Tensor)参数说明说明bbatch size表示输入样本批量大小、ssequence length表示输入样本序列长度、T表示bs合轴后的大小、e表示专家数量、k表示选取Top专家数。参数名类型描述xTensor输入张量支持2D或3Dshape为 (T, e) 或 (b, s, e)支持float16、bfloat16和float32kint选取的专家数量取值小于等于e且必须小于等于64biasTensor可选偏置张量shape为edtype与x相同 支持float16、bfloat16和float32input_idsTensor可选输入词表shape为T 仅支持int64取值范围为[0 ,n]n为tid2eid第一维的大小tid2eidTensor可选词表到专家id的映射关系表shape为nk 仅支持int32取值范围为[0 ,e]e代表专家数k_groupint可选选取的组数量默认为1group_countint可选总组数默认为1routed_scaling_factorfloat可选路由缩放因子默认为1epsfloat可选数值稳定性参数防止除零默认为1e-20group_select_modeint可选组选择模式0-使用最大值排序1-使用top2的和排序renormint可选重归一化标志仅支持0norm_typeint可选归一化标志0-Softmax, 1-Sigmoid, 2-Softplusout_flagbool可选是否输出归一化结果返回值说明返回值类型描述yOutTensor归一化、分组排序和TopK后的结果expertIdxOutTensor专家索引数据类型为int32outTensor归一化结果当out_flagTrue时有效约束说明renorm仅支持0表示先进行norm操作再计算topk。group_select_mode取值0和10表示使用最大值对group进行排序, 1表示使用topk2的sum值对group排序。norm_type取值0和10表示使用Softmax函数1表示使用Sigmoid函数2表示使用Softplus函数。outFlag取值true和falsetrue表示输出false表示不输出。input_ids和tid2eid都不为空表示hash场景都为空表示topk场景不允许只有一个为空。k_group和group_count为1时表示不分组排序。bias的dtype要和x相同。该接口支持推理场景下使用。该接口支持aclgraph入图。该接口与PyTorch配合使用时需要保证CANN相关包与PyTorch相关包的版本匹配。调用示例详见 test_npu_moe_gating_top_k.py【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考