CANN/AscendC MmadMx矩阵计算API

# MmadMx【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit产品支持情况不传入bias的原型Ascend 950PR/Ascend 950DT支持Atlas A3 训练系列产品/Atlas A3 推理系列产品不支持Atlas A2 训练系列产品/Atlas A2 推理系列产品不支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品AI Core不支持Atlas 推理系列产品Vector Core不支持Atlas 训练系列产品不支持Kirin X90不支持Kirin 9030不支持传入bias的原型Ascend 950PR/Ascend 950DT支持Atlas A3 训练系列产品/Atlas A3 推理系列产品不支持Atlas A2 训练系列产品/Atlas A2 推理系列产品不支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品AI Core不支持Atlas 推理系列产品Vector Core不支持Atlas 训练系列产品不支持Kirin X90不支持Kirin 9030不支持功能说明头文件路径为basic_api/kernel_operator_mm_intf.h。MmadMx全称Microscaling Mmad为带有量化系数的矩阵乘法即左矩阵和右矩阵均有对应的量化系数矩阵左量化系数矩阵scaleA和右量化系数矩阵scaleB。MmadMx场景中左量化系数矩阵与左矩阵乘积右量化系数矩阵与右矩阵乘积对两个乘积的结果做矩阵乘法。MmadMx的计算公式为C(scaleA⊗A)*(scaleB⊗B)Bias“⊗”表示广播乘法左/右矩阵与左/右量化系数矩阵做乘积时K方向上每32个元素共享一个量化因子如图1所示其中以AB矩阵均为fp4x2_e2m1_t数据类型为例。A、scaleA、B、scaleB为源操作数。A为左矩阵形状为[M,K]scaleA为左量化系数矩阵形状为[M,K/32]B为右矩阵形状为[K,N]scaleB为右量化系数矩阵形状为[K/32,N]。C为目的操作数存放矩阵乘结果的矩阵形状为[M,N]。Bias为矩阵乘偏置形状为[1,N]。对(scaleA⊗A)*(scaleB⊗B)结果矩阵的每一行都采用该Bias进行偏置。图1MmadMx接口矩阵乘分形示意图 表1矩阵计算矩阵A、B、C解释说明矩阵计算逻辑矩阵计算物理位置维度输入/输出数据格式数据类型AL0A BufferM x K[Nz]数据类型scaleAL0A_MX BufferM x K/32[大Z小z]小z shape为216BL0B BufferK x N[Zn]scaleBL0B_MX BufferK/32 x N[大N小n]小n shape为216CL0C BufferM x N可支持使用偏置矩阵Bias进行初始化维度为1 x N[Nz]函数原型不传入biastemplate typename T, typename U, typename S __aicore__ inline void MmadMx(const LocalTensorT dst, const LocalTensorU fm, const LocalTensorS filter, const MmadParams mmadParams)传入biastemplate typename T, typename U, typename S, typename V __aicore__ inline void MmadMx(const LocalTensorT dst, const LocalTensorU fm, const LocalTensorS filter, const LocalTensorV bias, const MmadParams mmadParams)参数说明表2参数说明参数名称输入/输出含义dst输出目的操作数结果矩阵C类型为LocalTensor支持的物理存储位置为L0C BufferTPosition:CO1。LocalTensor的起始地址需要按照1024字节对齐。fm输入源操作数左矩阵A类型为LocalTensor支持的物理存储位置为L0A BufferTPosition:A2。左矩阵A对应的scale矩阵起始地址为A矩阵起始对应地址/16。对于fp4场景LocalTensor的起始地址需要按照512字节对齐。对于fp8场景LocalTensor的起始地址需要按照1024字节对齐。filter输入源操作数右矩阵B类型为LocalTensor支持的物理存储位置为L0B BufferTPosition:B2。右矩阵B对应的scale矩阵起始地址为B矩阵起始对应地址/16。对于fp4场景LocalTensor的起始地址需要按照512字节对齐。对于fp8场景LocalTensor的起始地址需要按照1024字节对齐。bias输入源操作数Bias矩阵类型为LocalTensor支持的物理存储位置为BT BufferTPosition:C2。LocalTensor的起始地址需要按照64字节对齐。mmadParams输入矩阵乘相关参数。该参数类型的具体定义请参考${INSTALL_DIR}/asc/include/basic_api/interface/kernel_struct_mm.h${INSTALL_DIR}请替换为CANN软件安装后文件存储路径。MmadParams参数说明请参考MmadParams结构体内参数说明。数据类型表3MmadMx接口左矩阵、右矩阵、Scale矩阵、Bias矩阵、结果矩阵支持的精度类型组合Ascend 950PR/Ascend 950DT左矩阵fm右矩阵filterScale矩阵偏置Bias结果矩阵dstfp4x2_e1m2_tfp4x2_e1m2_tfp8_e8m0_tfloatfloatfp4x2_e2m1_tfp4x2_e1m2_tfp8_e8m0_tfloatfloatfp4x2_e1m2_tfp4x2_e2m1_tfp8_e8m0_tfloatfloatfp4x2_e2m1_tfp4x2_e2m1_tfp8_e8m0_tfloatfloatfp8_e4m3fn_tfp8_e4m3fn_tfp8_e8m0_tfloatfloatfp8_e4m3fn_tfp8_e5m2_tfp8_e8m0_tfloatfloatfp8_e5m2_tfp8_e4m3fn_tfp8_e8m0_tfloatfloatfp8_e5m2_tfp8_e5m2_tfp8_e8m0_tfloatfloat返回值说明无约束说明不同矩阵对于存储位置的约束结果矩阵C只支持位于物理存储位置为L0C BufferTPosition:CO1大小256KB。左矩阵A只支持位于物理存储位置为L0A BufferTPosition:A2大小64KB。右矩阵B只支持位于物理存储位置为L0B BufferTPosition:B2大小64KB。Bias矩阵只支持位于物理存储位置为BT BufferTPosition:C2大小4KB。地址约束说明请参考表2。当M、K、N中的任意一个值为0时表示指令不会执行该接口将被视为NOP空操作。K需要是64的倍数。对于fp4场景A/B矩阵的起始地址需要按照512字节对齐对于fp8场景A/B矩阵的起始地址需要按照1024字节对齐。左矩阵A/B对应的scale矩阵起始地址为A/B矩阵起始对应地址/16。其他特殊场景约束可参考Mmad接口约束说明。调用示例完整使用样例请参考MmadMx最佳实践。【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考