• 若[object Object]不为[object Object]，即非动态量化场景：

  

• 若[object Object]为[object Object]，即动态量化场景：

  

x([object Object])：表示计算使用的token数据，需根据[object Object]来发送给其他卡。要求2维张量，shape为(BS, H)，表示有BS(batch size)个token，数据类型支持[object Object]、[object Object]，数据格式为$ND$，支持非连续的Tensor。

expert_ids([object Object])：表示每个token的topK个专家索引，决定每个token要发给哪些专家。要求2维张量，shape为(BS, K)，数据类型支持[object Object]，数据格式为$ND$，支持非连续的Tensor。对应的[object Object]输入，张量里value取值范围为[0, moe_expert_num)，且同一行中的K个value不能重复。

group_ep([object Object])：EP通信域名称，专家并行的通信域。字符串长度范围为[1,128)。[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]时不能和[object Object]相同。

ep_world_size([object Object])：EP通信域size。

• [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：取值支持16、32、64。
• [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：取值支持8、16、32、64、128、144、256、288。

ep_rank_id([object Object])：EP通信域本卡ID，取值范围[0, ep_world_size)，同一个EP通信域中各卡的[object Object]不重复。

moe_expert_num([object Object])：MoE专家数量，取值范围[1, 512]，并且满足以下条件：moe_expert_num%(ep_world_size - shared_expert_rank_num)=0。

• [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：还需满足moe_expert_num/(ep_world_size - shared_expert_rank_num) <= 24。

[object Object]*[object Object]：必选参数，代表其之前的变量是位置相关的，必须按照顺序输入；之后的变量是可选参数，位置无关，需要使用键值对赋值，不赋值会使用默认值。

scales([object Object])：可选参数，表示每个专家的权重，非量化场景不传入，动态量化场景可传可不传。若传值要求为2维张量，如果有共享专家，shape为(shared_expert_num+moe_expert_num, H)，如果没有共享专家，shape为(moe_expert_num, H)，数据类型支持[object Object]，数据格式为$ND$，不支持非连续的Tensor。

x_active_mask([object Object])：预留参数，暂未使用，使用默认值即可。

expert_scales([object Object])：可选参数，表示每个token的topK个专家权重。

• [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：要求2维张量，shape为(BS, K)，数据类型支持[object Object]，数据格式为$ND$，支持非连续的Tensor。
• [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：暂不支持该参数，使用默认值即可。

group_tp([object Object])：可选参数，TP通信域名称，数据并行的通信域。若有TP域通信需要传参，若无TP域通信，使用默认值""即可。

• [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：eager模式使用默认值即可，图模式传入与[object Object]相同。
• [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：字符串长度范围为[1, 128)，不能和[object Object]相同。

tp_world_size([object Object])：可选参数，TP通信域size。有TP域通信才需要传参。

• [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：不支持TP域通信，使用默认值0即可。
• [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：当有TP域通信时，取值范围[0, 2]，0和1表示无TP域通信，2表示有TP域通信。

tp_rank_id([object Object])：可选参数，TP通信域本卡ID。有TP域通信才需要传参。

• [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：不支持TP域通信，使用默认值即可。
• [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：当有TP域通信时，取值范围[0, 1]，默认为0，同一个TP通信域中各卡的[object Object]不重复。无TP域通信时，传0即可。

expert_shard_type([object Object])：表示共享专家卡排布类型。

• [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：暂不支持该参数，使用默认值即可。
• [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：当前仅支持0，表示共享专家卡排在MoE专家卡前面。

shared_expert_num([object Object])：表示共享专家数量，一个共享专家可以复制部署到多个卡上。

• [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：暂不支持该参数，使用默认值即可。
• [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：仅支持1，默认值为1。

shared_expert_rank_num([object Object])：可选参数，表示共享专家卡数量。

• [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：不支持共享专家，传0即可。
• [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：取值范围[0, ep_world_size)。取0表示无共享专家，不取0时需满足ep_world_size%shared_expert_rank_num=0。

quant_mode([object Object])：可选参数，表示量化模式。支持取值：0表示非量化（默认），2表示动态量化。当[object Object]为2，[object Object]不为None；当[object Object]为0，[object Object]为None。

global_bs([object Object])：可选参数，表示EP域全局的BS大小。

• [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：当每个rank的BS不同时，支持传入max_bs*ep_world_size或者256*ep_world_size，其中max_bs表示单rank BS最大值，建议按max_bs*ep_world_size传入，固定按256*ep_world_size传入，在后续版本BS大于256的场景下会无法支持；当每个rank的BS相同时，支持取值0或BS*ep_world_size。
• [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：当每个rank的BS不同时，支持传入max_bs*ep_world_size，其中max_bs表示单rank BS最大值；当每个rank的BS相同时，支持取值0或BS*ep_world_size。

expert_token_nums_type([object Object])：可选参数，表示输出[object Object]的值类型，取值范围[0, 1]，0表示每个专家收到token数量的前缀和，1表示每个专家收到的token数量（默认）。

expand_x([object Object])：表示本卡收到的token数据，要求2维张量，shape为(max(tp_world_size, 1) *A, H)，A表示在EP通信域可能收到的最大token数，数据类型支持[object Object]、[object Object]、[object Object]。量化时类型为[object Object]，非量化时与[object Object]数据类型保持一致。数据格式为$ND$，支持非连续的Tensor。

dynamic_scales([object Object])：表示计算得到的动态量化参数。当[object Object]非0时才有该输出，要求1维张量，shape为(A,)，数据类型支持[object Object]，数据格式支持$ND$，支持非连续的Tensor。

expand_idx([object Object])：表示给同一专家发送的token个数，要求1维张量，shape为(BS * K, )。数据类型支持[object Object]，数据格式为$ND$，支持非连续的Tensor。对应的[object Object]输入。

expert_token_nums([object Object])：本卡每个专家实际收到的token数量，要求1维张量，shape为(local_expert_num,)，数据类型[object Object]，数据格式支持$ND$，支持非连续的Tensor。

ep_recv_counts([object Object])：表示EP通信域各卡收到的token数（token数以前缀和的形式表示），要求1维张量，数据类型[object Object]，数据格式支持$ND$，支持非连续的Tensor。对应的[object Object]输入。

• [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：要求shape为(moe_expert_num+2*global_bs*K*server_num, )，前[object Object]个数表示在EP通信域内，该卡上每个专家收到来自其他各卡的token数（以前缀和的形式表示），2*global_bs*K*server_num用于存储机间和机内通信前，combine可提前做reduce操作的token个数和通信区偏移量，[object Object]传入0时此处按照bs*ep_world_size计算。
• [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：要求shape为(ep_world_size*max(tp_world_size, 1)*local_expert_num, )。

tp_recv_counts([object Object])：表示TP通信域各卡收到的token数量。对应的[object Object]输入。

• [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：不支持TP通信域，暂无该输出，
• [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：支持TP通信域，要求一个1维张量，shape为(tp_world_size, )，数据类型支持[object Object]，数据格式为$ND$，支持非连续的Tensor。

expand_scales([object Object])：表示[object Object]与[object Object]一起进行alltoallv之后的输出。

• [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：要求一个1维张量，shape为(A, )，数据类型支持[object Object]，数据格式要求为$ND$，支持非连续的Tensor。
• [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：暂不支持该输出，返回None。

该接口支持推理场景下使用。

该接口支持静态图模式，[object Object]和[object Object]必须配套使用。

在不同产品型号、不同通信算法或不同版本中，[object Object]的Tensor输出[object Object]、[object Object]、[object Object]、[object Object]中的元素值可能不同，使用时直接将上述Tensor传给[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]、[object Object]、[object Object]、[object Object]网络中不同层中也需保持一致，且和对应参数也保持一致。

[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：该场景下单卡包含双DIE（简称为“晶粒”或“裸片”），因此参数说明里的“本卡”均表示单DIE。

参数里Shape使用的变量如下：

• A：表示本卡接收的最大token数量，取值范围如下

  • 对于共享专家，当global_bs为0时，要满足A=BS*shared_expert_num/shared_expert_rank_num；当global_bs非0时，要满足A=global_bs*shared_expert_num/shared_expert_rank_num。
  • 对于MoE专家，当global_bs为0时，要满足A>=BS*ep_world_size*min(local_expert_num, K)；当global_bs非0时，要满足A>=global_bs* min(local_expert_num, K)。

• H：表示hidden size隐藏层大小。

  • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：取值范围(0,7168]，且保证是32的整数倍。
  • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：仅支持 7168。

• BS：表示待发送的token数量。

  • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：取值范围为0<BS≤256。
  • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：取值范围为0<BS≤512。

• K：表示选取topK个专家，需满足0<K≤moe_expert_num。

  • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：保证取值范围为0<K≤16。
  • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：保证取值范围为0<K≤8。

• server_num：表示服务器的节点数，取值只支持2、4、8。

  • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：仅该场景的shape使用了该变量。

• local_expert_num：表示本卡专家数量。

  • 对于共享专家卡，local_expert_num=1
  • 对于MoE专家卡，local_expert_num=moe_expert_num/(ep_world_size-shared_expert_rank_num)，当local_expert_num>1时，不支持TP域通信。

调用本接口前需检查HCCL_BUFFSIZE环境变量取值是否合理：

[object Object]

• [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：要求>=2*(BS*ep_world_size*min(local_expert_num, K)*H*sizeof(uint16)+2MB)。
• [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：要求>=2且满足1024^2*(HCCL_BUFFSIZE-2)/2>=BS*2*(H+128)*(ep_world_size*local_expert_num+K+1)，local_expert_num需使用MoE专家卡的本卡专家数。

[object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：配置环境变量HCCL_INTRA_PCIE_ENABLE=1和HCCL_INTRA_ROCE_ENABLE=0可以减少跨机通信数据量，提升算子性能。此时要求HCCL_BUFFSIZE>=moe_expert_num*BS*(H*sizeof(dtype_x)+4*((K+7)/8*8)*sizeof(uint32))+4MB+100MB。并且，对于入参moe_expert_num，只要求moe_expert_num%(ep_world_size - shared_expert_rank_num)=0，不要求moe_expert_num/(ep_world_size - shared_expert_rank_num) <= 24。

本文公式中的“/”表示整除。

通信域使用约束：

• 一个模型中的[object Object]和[object Object]算子仅支持相同EP通信域，且该通信域中不允许有其他算子。

• 一个模型中的[object Object]和[object Object]算子仅支持相同TP通信域或都不支持TP通信域，有TP通信域时该通信域中不允许有其他算子。

版本配套约束：

静态图模式下，从Ascend Extension for PyTorch 8.0.0版本开始，Ascend Extension for PyTorch框架会对静态图中最后一个节点输出结果做Meta推导与inferShape推导的结果强校验。当图中只有一个Dispatch算子，若CANN版本落后于Ascend Extension for PyTorch版本，会出现Shape不匹配报错，建议用户升级CANN版本，详细的版本配套关系参见《Ascend Extension for PyTorch 版本说明》中“[object Object]相关产品版本配套说明[object Object]”。

单算子模式调用

图模式调用