aclnnGroupedMatmulFinalizeRouting

产品支持情况

产品	是否支持
[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动作

函数原型

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

• aclnnStatus aclnnGroupedMatmulFinalizeRoutingGetWorkspaceSize(const aclTensor *x, 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 aclnnGroupedMatmulFinalizeRouting(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)

aclnnGroupedMatmulFinalizeRoutingGetWorkspaceSize

• 参数说明：

  • x（aclTensor*，计算输入）：Device侧的aclTensor，输入x(左矩阵)，undefined支持ND，数据类型支持INT8，shape支持2维，维度为(m, k)，维度m的取值范围为[1,16*1024*8]，k支持2048，不支持非连续的Tensor。
  • w（aclTensor*，计算输入）：Device侧的aclTensor，输入weight(右矩阵)，undefined支持ND，数据类型支持INT32（传入数据为INT32，实际按照8个INT4来处理），shape支持三维，当输入为INT32时维度为(e, k, n / 8)，输入转为INT4时维度为(e, k, n)，e取值范围[1,256]，k支持2048，n支持7168，不支持非连续的Tensor。
  • scaleOptional（aclTensor*，可选计算输入）：Device侧的aclTensor，量化参数中的缩放因子，undefined支持ND，数据类型支持INT64，shape支持三维，维度为(e, 1, n)，e、n和w的e、n一致，不支持非连续的Tensor。
  • biasOptional（aclTensor*，可选计算输入）：Device侧的aclTensor，矩阵的偏移，undefined支持ND，数据类型支持FLOAT32，shape支持二维，维度为(e, n)，e、n和w的e、n一致，不支持非连续的Tensor。
  • pertokenScaleOptional（aclTensor*，可选计算输入）：Device侧的aclTensor，矩阵计算的反量化参数，undefined支持ND，数据类型支持FLOAT32，shape支持一维，维度为(m)，m和x的m一致，不支持非连续的Tensor。
  • groupListOptional（aclTensor*，可选计算输入）：Device侧的aclTensor类型，代表输入和输出分组轴方向的matmul大小分布，undefined支持ND，数据类型支持INT64，shape支持一维，维度为(e)，e和w的e一致，不支持非连续的Tensor。
  • sharedInputOptional（aclTensor*，可选计算输入）：Device侧的aclTensor类型，moe计算中共享专家的输出，需要与moe专家的输出进行combine操作undefined支持ND，数据类型支持BFLOAT16，shape支持二维，维度(batch/dp,n)，batch/dp取值范围[1,2*1024]，batch取值范围[1,16*1024]。
  • logitOptional（aclTensor*，可选计算输入）：Device侧的aclTensor类型，moe专家对各个token的logit大小，矩阵乘的计算输出与该logit做乘法，然后索引进行combine，undefined支持ND，数据类型支持FLOAT32，shape支持一维，维度为(m)，m和x的m一致，不支持非连续的Tensor。
  • rowIndexOptional（aclTensor*，可选计算输入）：Device侧的aclTensor类型，moe专家输出按照该rowIndex进行combine，其中的值即为combine做scatter add的索引，undefined支持ND，数据类型支持FLOAT32，shape支持一维，维度为(m)，m和x的m一致，不支持非连续的Tensor。
  • 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，shape支持二维。
  • workspaceSize（uint64_t*，出参）：返回需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。

• 返回值：

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

  [object Object]

aclnnGroupedMatmulFinalizeRouting

• 参数说明：

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

• 返回值：

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

约束说明

伪量化场景支持类型 输入和输出支持以下数据类型组合：

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

• 在该场景中，scaleOptional代表per-channel和per-group离线融合的结果。
• 在该场景中，biasOptional代表离线计算的辅助结果，值要求为$8 \times w \times scaleOptional$，并在第一维累加。

调用示例

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