- 接口功能:MoE计算中,对x的输出做Softmax计算,取topK操作。其中yOut为softmax的topK结果;expertIdxOut为topK的值的索引结果,即对应的专家序号;rowIdxOut为与expertIdxOut相同shape的列取值结果,指示每个位置对应的原始行位置。如果expertIdxOut中对应的行的finished为True,则直接填num_expert值(即x的最后一个轴大小)。
- 计算公式:
每个算子分为,必须先调用“aclnnMoeGatingTopKSoftmaxGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器,再调用“aclnnMoeGatingTopKSoftmax”接口执行计算。
[object Object]
[object Object]
确定性计算:
- aclnnMoeGatingTopKSoftmax默认确定性实现。
Atlas 350 加速卡:
- x最后一维的大小(即专家数)取值范围为[1, 2048]。
[object Object]