昇腾社区首页
中文
注册
开发者
下载

aclnnGroupedMatmulSwigluQuant

产品支持情况

产品 是否支持
[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]
[object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]
[object Object]Atlas 200I/500 A2 推理产品[object Object] ×
[object Object]Atlas 推理系列产品 [object Object] ×
[object Object]Atlas 训练系列产品[object Object] ×

功能说明

  • 接口功能:融合GroupedMatmul 、dquant、swiglu和quant,详细解释见计算公式。

  • 计算公式:

    • 量化场景(A8W8):
      • 定义

        • 表示矩阵乘法。
        • 表示逐元素乘法。
        • x\left \lfloor x\right \rceil 表示将x四舍五入到最近的整数。
        • Z8={xZ128x127}\mathbb{Z_8} = \{ x \in \mathbb{Z} | −128≤x≤127 \}
        • Z32={xZ2147483648x2147483647}\mathbb{Z_{32}} = \{ x \in \mathbb{Z} | -2147483648≤x≤2147483647 \}
      • 输入

        • XZ8M×KX∈\mathbb{Z_8}^{M \times K}:输入矩阵(左矩阵),M是总token 数,K是特征维度。
        • WZ8E×K×NW∈\mathbb{Z_8}^{E \times K \times N}:分组权重矩阵(右矩阵),E是专家个数,K是特征维度,N是输出维度。
        • w_scaleRE×Nw\_scale∈\mathbb{R}^{E \times N}:分组权重矩阵(右矩阵)的逐通道缩放因子,E是专家个数,N是输出维度。
        • x_scaleRMx\_scale∈\mathbb{R}^{M}:输入矩阵(左矩阵)的逐 token缩放因子,M是总token 数。
        • grouplistNEgrouplist∈\mathbb{N}^{E}:前缀和的分组索引列表。
      • 输出

        • QZ8M×N/2Q∈\mathbb{Z_8}^{M \times N / 2}:量化后的输出矩阵。
        • Q_scaleRMQ\_scale∈\mathbb{R}^{M}:量化缩放因子。
      • 计算过程

        • 1.根据groupList[i]确定当前分组的 token ,i[0,Len(groupList)]i \in [0,Len(groupList)]

          例子:假设groupList=[3,4,4,6],从0开始计数。

          第0个右矩阵W[0,:,:],对应索引位置[0,3)的tokenx[0:3](共3-0=3个token),对应x_scale[0:3]w_scale[0]bias[0]offset[0]Q[0:3]Q_scale[0:3]Q_offset[0:3]

          第1个右矩阵W[1,:,:],对应索引位置[3,4)的tokenx[3:4](共4-3=1个token),对应x_scale[3:4]w_scale[1]bias[1]offset[1]Q[3:4]Q_scale[3:4]Q_offset[3:4]

          第2个右矩阵W[2,:,:],对应索引位置[4,4)的tokenx[4:4](共4-4=0个token),对应x_scale[4:4]w_scale[2]bias[2]offset[2]Q[4:4]Q_scale[4:4]Q_offset[4:4]

          第3个右矩阵W[3,:,:],对应索引位置[4,6)的tokenx[4:6](共6-4=2个token),对应x_scale[4:6]w_scale[3]bias[3]offset[3]Q[4:6]Q_scale[4:6]Q_offset[4:6]

          请注意:grouplist中未指定的部分将不会参与更新。 例如groupList=[12,14,18],X的shape为[30,:]。

          则第一个输出Q的shape为[30,:],其中Q[18:,:]的部分不会进行更新和初始化,其中数据为显存空间申请时的原数据。

          同理,第二个输出Q的shape为[30],其中Q_scale[18:]的部分不会进行更新或初始化,其中数据为显存空间申请时的原数据。

          即输出的Q[:grouplist[-1],:]和Q_scale[:grouplist[-1]]为有效数据部分。

        • 2.根据分组确定的入参进行如下计算:

          Ci=(XiWi)x_scalei BroadCastw_scalei BroadCastC_{i} = (X_{i}\cdot W_{i} )\odot x\_scale_{i\ BroadCast} \odot w\_scale_{i\ BroadCast}

          Ci,act,gatei=split(Ci)C_{i,act}, gate_{i} = split(C_{i})

          Si=Swish(Ci,act)gateiS_{i}=Swish(C_{i,act})\odot gate_{i}   其中Swish(x)=x1+exSwish(x)=\frac{x}{1+e^{-x}}

        • 3.量化输出结果

          Q_scalei=max(Si)127Q\_scale_{i} = \frac{max(|S_{i}|)}{127}

          Qi=SiQ_scaleiQ_{i} = \left \lfloor \frac{S_{i}}{Q\_scale_{i}}\right \rceil


    • MSD伪量化场景(A8W4):
      • 定义

        • 表示矩阵乘法。
        • 表示逐元素乘法。
        • x\left \lfloor x\right \rceil 表示将x四舍五入到最近的整数。
        • Z8={xZ128x127}\mathbb{Z_8} = \{ x \in \mathbb{Z} | −128≤x≤127 \}
        • Z4={xZ8x7}\mathbb{Z_4} = \{ x \in \mathbb{Z} | −8≤x≤7 \}
        • Z32={xZ2147483648x2147483647}\mathbb{Z_{32}} = \{ x \in \mathbb{Z} | -2147483648≤x≤2147483647 \}
      • 输入

        • XZ8M×KX∈\mathbb{Z_8}^{M \times K}:输入矩阵(左矩阵),M是总token 数,K是特征维度。
        • WZ4E×K×NW∈\mathbb{Z_4}^{E \times K \times N}:分组权重矩阵(右矩阵),E是专家个数,K是特征维度,N是输出维度。
        • biasRE×Nbias∈\mathbb{R}^{E \times N}:计算矩阵乘时的辅助矩阵(生成辅助矩阵的计算过程间下文)。
        • w_scaleRE×K_group_num×Nw\_scale∈\mathbb{R}^{E \times K\_group\_num \times N}:分组权重矩阵(右矩阵)的逐通道缩放因子,E是专家个数,K_group_num 是在K轴维度上的分组数,N是输出维度。
        • x_scaleRMx\_scale∈\mathbb{R}^{M}:输入矩阵(左矩阵)的逐token缩放因子,M是总token 数。
        • grouplistNEgrouplist∈\mathbb{N}^{E}:前缀和的分组索引列表。
      • 输出

        • QZ8M×N/2Q∈\mathbb{Z_8}^{M \times N / 2}:量化后的输出矩阵。
        • Q_scaleRMQ\_scale∈\mathbb{R}^{M}:量化缩放因子。
      • 计算过程

        • 1.根据groupList[i]确定当前分组的 token ,i[0,Len(groupList)]i \in [0,Len(groupList)]

          • 分组逻辑与A8W8相同。
        • 2.生成辅助矩阵(bias)的计算过程(请注意bias部分计算为离线生成作为输入,并非算子内部完成):

          • 当为per-channel量化(w_scalew\_scale为2维):

            biasi=8×weightScale×Σk=0K1weight[:,k,:]bias_{i} = 8 × weightScale × Σ_{k=0}^{K-1} weight[:,k,:]

          • 当为per-group量化(w_scalew\_scale为3维):

            biasi=8×Σk=0K1(weight[:,k,:]×weightScale[:,k/num_per_group,:])bias_{i} = 8 × Σ_{k=0}^{K-1} (weight[:,k,:] × weightScale[:, ⌊k/num\_per\_group⌋, :])

            注:num_per_group=K//K_group_numnum\_per\_group = K // K\_group\_num

        • 3.根据分组确定的入参进行如下计算:

          • 3.1.将左矩阵Z8\mathbb{Z_8},转变为高低位 两部分的Z4\mathbb{Z_4}

            X_high_4bitsi=Xi16X\_high\_4bits_{i} = \lfloor \frac{X_{i}}{16} \rfloor

            X_low_4bitsi=Xi&0x0f8X\_low\_4bits_{i} = X_{i} \& 0x0f - 8

          • 3.2.做矩阵乘时,使能per-channel或per-group量化

            per-channel:

            C_highi=(X_high_4bitsiWi)w_scaleiC\_high_{i} = (X\_high\_4bits_{i} \cdot W_{i}) \odot w\_scale_{i}

            C_lowi=(X_low_4bitsiWi)w_scaleiC\_low_{i} = (X\_low\_4bits_{i} \cdot W_{i}) \odot w\_scale_{i}

            per-group:

            C_highi=Σk=0K1((X_high_4bitsi[:,knum_per_group:(k+1)num_per_group]Wi[knum_per_group:(k+1)num_per_group,:])w_scalei[k,:])C\_high_{i} = \\ Σ_{k=0}^{K-1}((X\_high\_4bits_{i}[:, k * num\_per\_group : (k+1) * num\_per\_group] \cdot W_{i}[k * num\_per\_group : (k+1) * num\_per\_group, :]) \odot w\_scale_{i}[k, :] )

            C_lowi=Σk=0K1((X_low_4bitsi[:,knum_per_group:(k+1)num_per_group]Wi[knum_per_group:(k+1)num_per_group,:])w_scalei[k,:])C\_low_{i} = \\ Σ_{k=0}^{K-1}((X\_low\_4bits_{i}[:, k * num\_per\_group : (k+1) * num\_per\_group] \cdot W_{i}[k * num\_per\_group : (k+1) * num\_per\_group, :]) \odot w\_scale_{i}[k, :] )

          • 3.3.将高低位的矩阵乘结果还原为整体的结果

            Ci=(C_highi16+C_lowi+biasi)x_scaleiC_{i} = (C\_high_{i} * 16 + C\_low_{i} + bias_{i}) \odot x\_scale_{i}

            Ci,act,gatei=split(Ci)C_{i,act}, gate_{i} = split(C_{i})

            Si=Swish(Ci,act)gateiS_{i}=Swish(C_{i,act})\odot gate_{i}    其中Swish(x)=x1+exSwish(x)=\frac{x}{1+e^{-x}}

        • 3.量化输出结果

          Q_scalei=max(Si)127Q\_scale_{i} = \frac{max(|S_{i}|)}{127}

          Qi=SiQ_scaleiQ_{i} = \left \lfloor \frac{S_{i}}{Q\_scale_{i}}\right \rceil

函数原型

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

  • aclnnStatus aclnnGroupedMatmulSwigluQuantGetWorkspaceSize(const aclTensor *x, const aclTensor *weight, const aclTensor *bias, const aclTensor *offset, const aclTensor *weightScale, const aclTensor *xScale, const aclTensor *groupList, aclTensor *output, aclTensor *outputScale, aclTensor *outputOffset, uint64_t *workspaceSize, aclOpExecutor **executor)
  • aclnnStatus aclnnGroupedMatmulSwigluQuant(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

aclnnGroupedMatmulSwigluQuantGetWorkspaceSize

  • 参数说明:

    • x(aclTensor*,计算输入):左矩阵,公式中的XX,Device侧的aclTensor。shape支持2维,假设shape为[M,K],则K必须小于65536,数据类型支持INT8,支持ND,支持

    • weight(aclTensor*,计算输入):权重矩阵,公式中的WW,Device侧的aclTensor。shape支持3维,数据类型支持INT8、INT4、INT32(INT32为适配用途,实际1个INT32会被解释为8个INT4数据),支持FRACTAL_NZ,支持

    • bias(aclTensor*,计算输入):计算矩阵乘时的辅助矩阵,公式中的biasbias,shape支持2维,数据类型支持fp32,仅A8W4场景生效,A8W8场景需传空指针。

    • offset(aclTensor*,计算输入):per-channel非对称反量化的偏移,公式中的offsetoffset,shape支持2维,数据类型支持Float,预留输入,暂不支持,需要传空指针。

    • weightScale(aclTensor*,计算输入):右矩阵的量化因子,公式中的w_scalew\_scale,Device侧的aclTensor。首轴长度需与weight的首轴维度相等,尾轴长度需要与weight还原为ND格式的尾轴相同,支持ND,支持

      • A8W4场景:shape支持2或3维,数据类型支持UINT64。
      • A8W8场景:shape支持2维,数据类型支持FLOAT、FLOAT16、BFLOAT16。
    • xScale(aclTensor*,计算输入):左矩阵的量化因子,公式中的x_scalex\_scale,Device侧的aclTensor。shape支持1维,长度需与x的首轴维度相等,数据类型支持FLOAT,支持ND,支持

    • groupList(aclTensor*,计算输入):指示每个分组参与计算的Token个数,公式中的grouplistgrouplist,Device侧的aclTensor。shape支持1维,长度需与weight的首轴维度相等,数据类型支持INT64,支持ND,支持,grouplist中的最后一个值约束了输出数据的有效部分,详见功能说明中的计算过程部分。

    • output(aclTensor*,计算输出):输出的量化结果,公式中的QQ,Device侧的aclTensor。数据类型支持INT8,shape支持2维,Device侧的aclTensor。支持ND,支持

    • outputScale(aclTensor*,计算输出):输出的量化因子,公式中的Q_scaleQ\_scale,Device侧的aclTensor。数据类型支持FLOAT,shape支持1维,Device侧的aclTensor。支持ND,支持

    • outputOffset(aclTensor*,计算输出):输出的非对称量化的偏移,公式中的Q_offsetQ\_offset,Device侧的aclTensor,shape支持1维,数据类型支持FLOAT,预留输入,暂不支持,需要传空指针。

    • workspaceSize(uint64_t*,出参):返回用户需要在npu device侧申请的workspace大小。

    • executor(aclOpExecutor**,计算输出):返回op执行器,包含了算子计算流程。

  • 返回值:

    aclnnStatus:返回状态码,具体参见

    [object Object]

aclnnGroupedMatmulSwigluQuant

  • 参数说明:

    • workspace(void*,入参):在Device侧申请的workspace内存地址。
    • workspaceSize(uint64_t,入参):在Device侧申请的workspace大小,由第一段接口aclnnGroupedMatmulSwigluQuantGetWorkspaceSize获取。
    • executor(aclOpExecutor*,入参):op执行器,包含了算子计算流程。
    • stream(aclrtStream,入参):指定执行任务的Stream。
  • 返回值:

    aclnnStatus:返回状态码,具体参见

约束说明

  • 仅支持内部使用,支持特定shape

调用示例

  • aclnn单算子调用方式

通过aclnn单算子调用示例代码如下,仅供参考,具体编译和执行过程请参考

[object Object]