MoeComputeExpertTokens

产品支持情况

产品 是否支持
Ascend 950PR/Ascend 950DT
Atlas A3 训练系列产品/Atlas A3 推理系列产品
Atlas A2 训练系列产品/Atlas A2 推理系列产品
Atlas 200I/500 A2 推理产品 ×
Atlas 推理系列产品 ×
Atlas 训练系列产品 ×
Kirin X90 处理器系列产品
Kirin 9030 处理器系列产品

功能说明

  • 算子功能:MoE计算中,通过二分查找的方式查找每个专家处理的最后一行的位置。

  • 计算公式

    outi=BinarySearch(sortedExperts,numExperts)out_{i}=BinarySearch(sortedExperts,numExperts)

参数说明

参数名 输入/输出/属性 描述 数据类型 数据格式
sortedExperts 输入 公式中的sortedExperts。 INT32 ND
numExperts 属性 总专家数。 INT64 -
out 输出 公式中的输出。 INT32 ND

约束说明

  • sortedExperts的shape大小需要小于2(24)2^(24)
  • numExperts的输入常值需要大于0,但不能超过2048。
  • 输入shape大小不要超过device可分配的内存上限,否则会导致异常终止。

调用说明

调用方式 调用样例 说明
aclnn调用 test_aclnn_moe_compute_expert_tokens 通过接口方式调用MoeComputeExpertTokens算子。