aclnnGroupedMatmulFinalizeRoutingWeightNz

产品支持情况

产品	是否支持
[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]	√
[object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]	√
[object Object]Atlas 200I/500 A2 推理产品[object Object]	×
[object Object]Atlas 推理系列产品 [object Object]	×
[object Object]Atlas 训练系列产品[object Object]	×

功能说明

算子功能：GroupedMatmul和MoeFinalizeRouting的融合算子，GroupedMatmul计算后的输出按照索引做combine动作，支持w为昇腾亲和数据排布格式(NZ)

函数原型

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

• aclnnStatus aclnnGroupedMatmulFinalizeRoutingWeightNzGetWorkspaceSize(const aclTensor *x, const aclTensor *w, const aclTensor *scaleOptional, const aclTensor* biasOptional, const aclTensor *pertokenScaleOptional, const aclTensor *groupListOptional, const aclTensor *sharedInputOptional, const aclTensor* logitOptional, const aclTensor *rowIndexOptional, int64_t dtype, float sharedInputWeight, int64_t sharedInputOffset, bool transposeX, bool transposeW, int64_t groupListType, aclTensor *y, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnStatus aclnnGroupedMatmulFinalizeRoutingWeightNz(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)

aclnnGroupedMatmulFinalizeRoutingWeightNzGetWorkspaceSize

• 参数说明：

  • x（aclTensor*，计算输入）：Device侧的aclTensor，输入x(左矩阵)，undefined支持ND，shape支持2维，维度为(m, k)， 数据类型支持INT8，维度m的取值范围为[1,16*1024*8]；k支持256、512、1024、1408、2048。
  • w（aclTensor*，计算输入）：Device侧的aclTensor，输入weight(右矩阵)，undefined支持昇腾亲和数据排布格式(NZ)。数据类型支持INT8，shape支持5维。维度为(e, n1, k1, k0, n0)，其中k0 = 16，n0 = 32， x shape中的k和w shape中的k1需要满足以下关系：ceilDiv（k,16） = k1。可使用aclnnCalculateMatmulWeightSizeV2接口以及aclnnTransMatmulWeight接口完成输入Format从ND到昇腾亲和数据排布格式（NZ）的转换。e取值范围[1,256]。
  • scaleOptional（aclTensor*，可选计算输入）：Device侧的aclTensor，代表per-channel量化参数中的缩放因子，undefined支持ND，数据类型支持FLOAT32，shape是2维(e, n)，n = n1 * n0，e和w的e一致，n支持2048、7168、7680。
  • biasOptional（aclTensor*，可选计算输入）：Device侧的aclTensor，矩阵偏移bias，undefined支持ND，数据类型支持FLOAT32，当前为预留参数，暂不生效，传入空指针即可。
  • pertokenScaleOptional（aclTensor*，可选计算输入）：Device侧的aclTensor，矩阵计算的反量化参数，undefined支持ND，对应x矩阵，维度为(m)，数据类型支持FLOAT32，不支持非连续的Tensor。
  • groupListOptional（aclTensor*，可选计算输入）：Device侧的aclTensor，代表输入和输出分组轴方向的matmul大小分布，数据类型支持INT64，undefined支持ND，维度为(e)，grouplist的值总和小于等于m。
  • sharedInputOptional（aclTensor*，可选计算输入）：Device侧的aclTensor，moe计算中共享专家的输出，需要与moe专家的输出进行combine操作，数据类型支持BFLOAT16，不支持非连续的Tensor。undefined支持ND，维度(batch/dp, n)，batch/dp取值范围[1, 2*1024]，batch取值范围[1, 16*1024]。
  • logitOptional（aclTensor*，可选计算输入）：Device侧的aclTensor，moe专家对各个token的logit大小，矩阵乘的计算输出与该logit做乘法，然后索引进行combine，不支持非连续的Tensor。undefined支持ND，维度(m)，数据类型支持FLOAT32。
  • rowIndexOptional（aclTensor*，可选计算输入）：Device侧的aclTensor，moe专家输出按照该rowIndex进行combine，其中的值即为combine做scatter add的索引，不支持非连续的Tensor。undefined支持ND，维度为(m)，数据类型支持INT64、INT32。
  • dtype（int64_t，计算输入）：GroupedMatmul计算的输出类型，0：FLOAT32；1：FLOAT16；2：BFLOAT16，仅支持取0。
  • sharedInputWeight（float，计算输入）：共享专家与moe专家进行combine的系数，sharedInput先于该参数乘，然后在和moe专家结果累加。
  • sharedInputOffset（int64_t，计算输入）：共享专家输出的在总输出中的偏移。
  • transposeX（bool，计算输入）：左矩阵是否转置，仅支持false。
  • transposeW（bool，计算输入）：右矩阵是否转置，仅支持false。
  • groupListType(int64_t，计算输入)：GroupedMatmul的分组模式：配置为0：cumsum模式，即为前缀和；配置为1：count模式。
  • y（aclTensor*，计算输出）：2D的Tensor，不支持非连续的Tensor，输出的数据类型固定为FLOAT32。undefined支持ND，维度(batch, n)。
  • workspaceSize（uint64_t*，出参）：返回需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。

• 返回值：

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

  [object Object]

aclnnGroupedMatmulFinalizeRoutingWeightNz

• 参数说明：

  • workspace（void*，入参）：在Device侧申请的workspace内存地址。
  • workspaceSize（uint64_t*，入参）：在Device侧申请的workspace大小，由第一段接口aclnnGroupedMatmulFinalizeRoutingWeightNzGetWorkspaceSize获取。
  • executor（aclOpExecutor*，入参）：op执行器，包含了算子计算流程。
  • stream（aclrtStream，入参）：指定执行任务的Stream。

• 返回值：

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

约束说明

输入和输出支持以下数据类型组合：

x	w	scale	bias	pertokenScale	groupList	sharedInput	logit	rowIndex	out
INT8	INT8	FLOAT32	null	FLOAT32	INT64	BFLOAT16	FLOAT32	INT64	FLOAT
INT8	INT8	FLOAT32	null	FLOAT32	INT64	null	null	INT64	FLOAT

调用示例

示例代码如下，仅供参考，具体编译和执行过程请参考undefined。