aclnnMoeUpdateExpert

产品支持情况

[object Object]undefined

功能说明

本API支持负载均衡和专家剪枝功能。经过映射后的专家表和mask可传入Moe层进行数据分发和处理。

• 负载均衡：为了解决负载不均衡的场景，该算子可以完成每个token的topK个专家逻辑专家号到物理卡号的映射。计算方法如下所示：

  负载均衡对于expert_ids中的第i个值，即第i个token：

  [object Object]

• 专家剪枝：支持根据阈值对token发送的topK个专家进行剪枝。计算方法如下所示：

  将shape为(BS,)的active_mask进行broadcast成为shape为(BS,K)的active_mask_tensor，其中BS对应为False的专家会直接被剪枝。对于active_mask_tensor为True的元素，满足条件也将被剪枝。

  [object Object]

函数原型

每个算子分为，必须先调用 “aclnnMoeUpdateExpertGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用“aclnnMoeUpdateExpert”接口执行计算。

aclnnMoeUpdateExpertGetWorkspaceSize

• 参数说明

  [object Object]

• 返回值：

  aclnnStatus：返回状态码，具体参见。

  第一段接口完成入参校验，出现以下场景时报错：

  [object Object]

aclnnMoeUpdateExpert

• 参数说明

  [object Object]

• 返回值

  返回aclnnStatus状态码，具体参见。

约束说明

1. 确定性计算：

   • aclnnMoeUpdateExpert默认确定性实现。

2. 接口配套与调用顺序：
   该接口必须与[object Object]及[object Object]/[object Object]接口配套使用，调用顺序固定为：
   [object Object] → [object Object] → [object Object]/[object Object]；

   或与[object Object]及[object Object]/[object Object]接口配套使用，调用顺序固定为：
   [object Object] → [object Object] → [object Object]/[object Object]；具体参考。

3. 参数一致性要求：
   调用过程中使用的[object Object]、[object Object]参数取值，所有卡需保持一致，网络不同层中也需保持一致，且需与[object Object]、[object Object]/[object Object]的对应参数一致。

4. 硬件相关定义：
   [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：单卡包含双DIE（“晶粒”或“裸片”），因此参数说明中的“本卡”均指单DIE。

5. 参数shape格式约束：

   • BS：本卡最终输出的token数量，取值范围 ( 0 < BS ≤ 512 )。
   • K：选取的topK个专家，取值范围 ( 0 < K ≤ 16 )，且需满足 ( 0 < K ≤ moeExpertNum )。
   • moeExpertNum：MoE专家数量，取值范围 (0, 1024]。
   • F：映射表[object Object]的列数，取值范围 [2, worldSize + 1]；第一列为逻辑专家的部署实例数（值>0），后F-1列为对应的物理卡号。
   • 实例总数限制：所有卡部署的MoE专家实例总数最多1024，即 ( place_per_rank * world_size ≤ 1024 )（[object Object]为单卡部署实例数）。
   • 单卡实例数一致性：每张卡部署的专家实例数必须相同。

调用示例

以[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]为例，调起MoeUpdateExpert，MoeDistributeDispatchV2和MoeDistributeCombineAddRmsNorm算子。

• 示例代码如下，仅供参考
  [object Object]