AscendAntiQuant【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit产品支持情况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支持功能说明按元素做伪量化计算比如将int8_t数据类型伪量化为half数据类型计算公式如下PER_CHANNEL场景按通道量化不开启输入转置groupSize src.shape[0] / offset.shape[0]dst[i][j] scale[i / groupSize][j] * (src[i][j] offset[i / groupSize][j])开启输入转置groupSize src.shape[1] / offset.shape[1]dst[i][j] scale[i][j / groupSize] * (src[i][j] offset[i][j / groupSize])PER_TENSOR场景按张量量化dst[i][j] scale * (src[i][j] offset)PER_TOKEN场景按token量化PER_GROUP场景按组量化根据输入数据类型的不同当前PER_GROUP分为两种场景fp4x2_e2m1_t/fp4x2_e1m2_t场景后续内容中简称为float4场景和int8_t/hifloat8_t/fp8_e5m2_t/fp8_e4m3fn_t场景后续内容中简称为b8场景。fp4x2_e2m1_t/fp4x2_e1m2_t场景float4场景groupSize可配置接口定义group的计算方向为k方向src在k方向上每groupSize个元素共享一组scale。src的shape为[m, n]时如果kDim0表示k是m方向scale的shape为[(m groupSize - 1) / groupSize, n]如果kDim1表示k是n方向scale的shape为[m(n groupSize - 1) / groupSize]。isTranspose为True表示srcscaledst都是转置的矩阵。k为m方向即公式中i轴为group的计算方向kDim0同时isTransposeFalse或者kDim1同时isTransposeTruedst[i][j] scale[i / groupSize][j] * src[i][j]k为n方向即公式中j轴为group的计算方向kDim0同时isTransposeTrue或者kDim1同时isTransposeFalsedst[i][j] scale[i][j / groupSize] * src[i][j]groupSize固定为32isTranspose为True表示srcscaledst都是转置的矩阵。不开启输入转置isTransposeFalsedst[i][j] scale[i / groupSize][j] * src[i][j]开启输入转置isTransposeTruedst[i][j] scale[i][j / groupSize] * src[i][j]int8_t/hifloat8_t/fp8_e5m2_t/fp8_e4m3fn_t场景b8场景定义group的计算方向为k方向src在k方向上每groupSize个元素共享一组scale和offset。src的shape为[m, n]时如果kDim0表示k是m方向scale和offset的shape为[(m groupSize - 1) / groupSize, n]如果kDim1表示k是n方向scale和offset的shape为[m(n groupSize - 1) / groupSize]。offset是可选输入。isTranspose为True表示srcscaledst都是转置的矩阵。k为m方向即公式中i轴为group的计算方向kDim0同时isTransposeFalse或者kDim1同时isTranspose Truek为n方向即公式中j轴为group的计算方向kDim0同时isTransposeTrue或者kDim1同时isTranspose False实现原理图1AscendAntiQuant算法框图如上图所示为AscendAntiQuant的典型场景算法框图计算过程分为如下几步均在Vector上进行精度转换将输入src转换为half类型计算offset当offset为向量时做Add计算当offset为scalar时做Adds计算计算scale当scale为向量时做Mul计算当scale为scalar时做Muls计算。图2isTranspose为False且输出为bfloat16的AscendAntiQuant算法框图在Atlas A2 训练系列产品/Atlas A2 推理系列产品上当输出为bfloat16时计算过程分为如下几步src精度转换将输入的src转换为half类型再转换为float类型存放到tmp1offset精度转换当输入的offset为向量时转换为float类型存放到tmp2为scalar时做Cast转换为float类型计算offset当输入的offset为向量时与tmp2做Add计算为scalar时做Adds计算scale精度转换当输入的scale为向量时转换为float类型存放到tmp2为scalar时做Cast转换为float类型计算scale当输入的scale为向量时用tmp2做Mul计算为scalar时做Muls计算dst精度转换将tmp1转换为bf16类型。图3AscendAntiQuant PER_TOKEN/PER_GROUP算法框图PER_TOKEN/PER_GROUP b8/float4场景的计算逻辑如下读取数据连续读取输入src根据不同的场景对输入scale和offset采用不同的读取方式例如PER_TOKEN场景做Broadcast处理PER_GROUP场景做Gather处理精度转换根据不同输入的数据类型组合对src/scale/offset进行相应的数据类型转换计算对类型转换后的数据做加乘操作精度转换将上述加乘操作得到的计算结果转换成dstT类型得到最终输出。函数原型通过sharedTmpBuffer入参传入临时空间PER_CHANNEL场景按通道量化template typename InputDataType, typename OutputDataType, bool isTranspose __aicore__ inline void AscendAntiQuant(const LocalTensorOutputDataType dst, const LocalTensorInputDataType src, const LocalTensorOutputDataType offset, const LocalTensorOutputDataType scale, const LocalTensoruint8_t sharedTmpBuffer, const uint32_t k, const AntiQuantShapeInfo shapeInfo {})PER_CHANNEL场景按通道量化不带offsettemplate typename InputDataType, typename OutputDataType, bool isTranspose __aicore__ inline void AscendAntiQuant(const LocalTensorOutputDataType dst, const LocalTensorInputDataType src, const LocalTensorOutputDataType scale, const LocalTensoruint8_t sharedTmpBuffer, const uint32_t k, const AntiQuantShapeInfo shapeInfo {})PER_TENSOR场景按张量量化template typename InputDataType, typename OutputDataType, bool isTranspose __aicore__ inline void AscendAntiQuant(const LocalTensorOutputDataType dst, const LocalTensorInputDataType src, const OutputDataType offset, const OutputDataType scale, const LocalTensoruint8_t sharedTmpBuffer, const uint32_t k, const AntiQuantShapeInfo shapeInfo {})PER_TENSOR场景按张量量化不带offsettemplate typename InputDataType, typename OutputDataType, bool isTranspose __aicore__ inline void AscendAntiQuant(const LocalTensorOutputDataType dst, const LocalTensorInputDataType src, const OutputDataType scale, const LocalTensoruint8_t sharedTmpBuffer, const uint32_t k, const AntiQuantShapeInfo shapeInfo {})PER_GROUP float4场景按组量化仅支持Ascend 950PR/Ascend 950DTtemplate typename InputDataType, typename OutputDataType, bool isTranspose __aicore__ inline void AscendAntiQuant(const LocalTensorOutputDataType dst, const LocalTensorInputDataType src, const LocalTensorfp8_e8m0_t scale, const LocalTensoruint8_t sharedTmpBuffer, const uint32_t k, const AntiQuantShapeInfo shapeInfo {})PER_TOKEN/PER_GROUP b8/float4场景按token量化/按组量化仅支持Ascend 950PR/Ascend 950DTtemplate typename dstT, typename srcT, typename scaleT, const AscendAntiQuantConfig config, const AscendAntiQuantPolicy policy __aicore__ inline void AscendAntiQuant(const LocalTensordstT dstTensor, const LocalTensorsrcT srcTensor, const LocalTensorscaleT scaleTensor, const LocalTensorscaleT offsetTensor,const LocalTensoruint8_t sharedTmpBuffer, const AscendAntiQuantParam para)接口框架申请临时空间PER_CHANNEL场景template typename InputDataType, typename OutputDataType, bool isTranspose __aicore__ inline void AscendAntiQuant(const LocalTensorOutputDataType dst, const LocalTensorInputDataType src, const LocalTensorOutputDataType offset, const LocalTensorOutputDataType scale, const uint32_t k, const AntiQuantShapeInfo shapeInfo {})PER_TENSOR场景template typename InputDataType, typename OutputDataType, bool isTranspose __aicore__ inline void AscendAntiQuant(const LocalTensorOutputDataType dst, const LocalTensorInputDataType src, const OutputDataType offset, const OutputDataType scale, const uint32_t k, const AntiQuantShapeInfo shapeInfo {})PER_GROUP float4场景groupSize固定为32仅支持Ascend 950PR/Ascend 950DTtemplate typename InputDataType, typename OutputDataType, bool isTranspose __aicore__ inline void AscendAntiQuant(const LocalTensorOutputDataType dst, const LocalTensorInputDataType src, const LocalTensorfp8_e8m0_t scale, const uint32_t k, const AntiQuantShapeInfo shapeInfo {})PER_TOKEN/PER_GROUP b8/float4场景groupSize可配置仅支持Ascend 950PR/Ascend 950DTtemplate typename dstT, typename srcT, typename scaleT, const AscendAntiQuantConfig config, const AscendAntiQuantPolicy policy __aicore__ inline void AscendAntiQuant(const LocalTensordstT dstTensor, const LocalTensorsrcT srcTensor, const LocalTensorscaleT scaleTensor, const LocalTensorscaleT offsetTensor,const AscendAntiQuantParam para)由于该接口的内部实现中涉及复杂的数学计算需要额外的临时空间来存储计算过程中的中间变量。临时空间支持接口框架申请和开发者通过sharedTmpBuffer入参传入两种方式。接口框架申请临时空间开发者无需申请但是需要预留临时空间的大小。通过sharedTmpBuffer入参传入使用该tensor作为临时空间进行处理接口框架不再申请。该方式开发者可以自行管理sharedTmpBuffer内存空间并在接口调用完成后复用该部分内存内存不会反复申请释放灵活性较高内存利用率也较高。接口框架申请的方式开发者需要预留临时空间通过sharedTmpBuffer传入的情况开发者需要为sharedTmpBuffer申请空间。临时空间大小BufferSize的获取方式如下通过GetAscendAntiQuantMaxMinTmpSize中提供的接口获取需要预留空间的范围大小。参数说明表1模板参数说明参数名描述InputDataType输入的数据类型。OutputDataType输出的数据类型。isTranspose是否开启输入数据转置。表2PER_TOKEN/PER_GROUP b8/float4场景模板参数说明参数名描述dstT目的操作数的数据类型。srcT源操作数的数据类型。scaleT缩放因子scale参数的数据类型。config量化接口配置参数定义如下方代码所示其中参数的含义如下。hasOffset量化参数offset是否参与计算。True表示offset参数参与计算。False表示offset参数不参与计算。isTranspose表示是否开启输入src转置。True表示输入src转置。False表示输入src不转置。kDimgroup的计算方向即k方向。仅在PER_GROUP场景有效支持的取值如下。0k轴是第0轴即m方向为group的计算方向1k轴是第1轴即n方向为group的计算方向。policy量化策略配置参数枚举类型可取值如下PER_TENSOR预留参数暂不支持。PER_CHANNEL预留参数暂不支持。PER_TOKEN配置为PER_TOKEN场景。PER_GROUP配置为PER_GROUP场景。PER_CHANNEL_PER_GROUP预留参数暂不支持。PER_TOKEN_PER_GROUP预留参数暂不支持。struct AscendAntiQuantConfig { bool hasOffset; bool isTranspose; int32_t kDim 1; };表3PER_TOKEN/PER_GROUP b8/float4场景支持的数据类型组合srcDtypescaleDtype/offsetDtypedstDtypeint8_thalfhalfint8_tbfloat16_tbfloat16_tint8_tfloatfloatint8_tfloathalfint8_tfloatbfloat16_thifloat8_thalfhalfhifloat8_tbfloat16_tbfloat16_thifloat8_tfloatfloathifloat8_tfloathalfhifloat8_tfloatbfloat16_tfp8_e5m2_t/fp8_e4m3fn_thalfhalffp8_e5m2_t/fp8_e4m3fn_tbfloat16_tbfloat16_tfp8_e5m2_t/fp8_e4m3fn_tfloatfloatfp8_e5m2_t/fp8_e4m3fn_tfloathalffp8_e5m2_t/fp8_e4m3fn_tfloatbfloat16_tfp4x2_e1m2_t/fp4x2_e2m1_t当前均只支持PER_GROUP场景fp8_e8m0_thalffp4x2_e1m2_t/fp4x2_e2m1_t当前均只支持PER_GROUP场景fp8_e8m0_tbfloat16_t表4接口参数说明参数名输入/输出描述dst输出目的操作数。不同型号支持的数据类型请参考dst支持的数据类型。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。src输入源操作数。不同型号支持的数据类型请参考src支持的数据类型。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。offset输入输入数据反量化时的偏移量。不同型号支持的数据类型请参考offset支持的数据类型。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。scale输入输入数据反量化时的缩放因子。不同型号支持的数据类型请参考scale支持的数据类型。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。sharedTmpBuffer输入临时缓存。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。临时空间大小BufferSize的获取方式请参考GetAscendAntiQuantMaxMinTmpSize。k输入isTranspose为true时src的shape为[N,K]isTranspose为false时src的shape为[K,N]。参数k对应其中的K值。shapeInfo输入设置参数offset和scale的shape信息仅PER_CHANNEL场景按通道量化需要配置。可选参数。在PER_CHANNEL场景如果未传入该参数或者结构体中数据设置为0将从offset和scale的ShapeInfo中获取offset和scale的shape信息。AntiQuantShapeInfo类型定义如下方代码所示。struct AntiQuantShapeInfo { uint32_t offsetHeight{0}; // offset的高 uint32_t offsetWidth{0}; // offset的宽 uint32_t scaleHeight{0}; // scale的高 uint32_t scaleWidth{0}; // scale的宽 };表5PER_TOKEN/PER_GROUP b8/float4场景接口参数说明参数名输入/输出描述dstTensor输出目的操作数。支持的数据类型为half、bfloat16_t、float。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。srcTensor输入源操作数。支持的数据类型为int8_t、fp8_e4m3fn_t、fp8_e5m2_t、hifloat8_t、fp4x2_e1m2_t、fp4x2_e2m1_t。注意对于fp4x2_e1m2_t、fp4x2_e2m1_t数据类型仅在PER_GROUP场景下支持。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。sharedTmpBuffer输入临时缓存。支持的数据类型为uint8_t。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。临时空间大小BufferSize的获取方式请参考GetAscendAntiQuantMaxMinTmpSize。scaleTensor输入量化参数scale。支持的数据类型为half、float、bfloat16_t、fp8_e8m0_t。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。offsetTensor输入量化参数offset。支持的数据类型和scaleTensor保持一致。对于float4场景offsetTensor不生效。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。para输入量化接口的参数AscendAntiQuantParam类型具体定义如下方代码所示其中参数的含义如下。mm方向元素个数。nn方向元素个数。n值对应的数据大小需满足32B对齐的要求即shape最后一维为n的输入输出均需要满足该维度上32B对齐的要求。calCount参与计算的元素个数。calCount必须是n的整数倍。groupSizePER_GROUP场景有效表示groupSize行/列数据共用一个scale/offset。groupSize的取值必须大于0且是32的整倍数。struct AscendAntiQuantParam { uint32_t m; uint32_t n; uint32_t calCount; uint32_t groupSize 0; };返回值说明无约束说明不支持源操作数与目的操作数地址重叠。操作数地址对齐要求请参见通用地址对齐约束。输入输出操作数参与计算的数据长度要求32B对齐。输入带转置场景k需要32B对齐。调用接口前确保输入数据的size正确offset和scale的size和shape正确。PER_TOKEN/PER_GROUP b8/float4场景支持情况如下Ascend 950PR/Ascend 950DT支持。Atlas A3 训练系列产品/Atlas A3 推理系列产品不支持。Atlas A2 训练系列产品/Atlas A2 推理系列产品不支持。Atlas 推理系列产品AI Core不支持。PER_TOKEN/PER_GROUP b8/float4场景连续计算方向即n方向的数据量要求32B对齐。非PER_TOKEN/PER_GROUP b8/float4场景dst支持的数据类型Ascend 950PR/Ascend 950DT支持的数据类型为half、bfloat16_t。Atlas A3 训练系列产品/Atlas A3 推理系列产品支持的数据类型为half、bfloat16_t。Atlas A2 训练系列产品/Atlas A2 推理系列产品支持的数据类型为half、bfloat16_t。Atlas 推理系列产品AI Core支持的数据类型为half。Kirin X90支持的数据类型为half。Kirin 9030支持的数据类型为half。非PER_TOKEN/PER_GROUP b8/float4场景src支持的数据类型Ascend 950PR/Ascend 950DTPER_CHANNEL和PER_TENSOR场景下支持的数据类型为int8_t、fp8_e4m3fn_t、fp8_e5m2_t、hifloat8_tPER_GROUP float4场景下支持的数据类型为fp4x2_e2m1_t、fp4x2_e1m2_t。Atlas A3 训练系列产品/Atlas A3 推理系列产品支持的数据类型为int8_t、int4b_t。Atlas A2 训练系列产品/Atlas A2 推理系列产品支持的数据类型为int8_t、int4b_t。Atlas 推理系列产品AI Core支持的数据类型为int8_t。Kirin X90支持的数据类型为int8_t。Kirin 9030支持的数据类型为int8_t。非PER_TOKEN/PER_GROUP b8/float4场景offset支持的数据类型Ascend 950PR/Ascend 950DT支持的数据类型为half、bfloat16_t。Atlas A3 训练系列产品/Atlas A3 推理系列产品支持的数据类型为half、bfloat16_t。Atlas A2 训练系列产品/Atlas A2 推理系列产品支持的数据类型为half、bfloat16_t。Atlas 推理系列产品AI Core支持的数据类型为half。Kirin X90支持的数据类型为half。Kirin 9030支持的数据类型为half。非PER_TOKEN/PER_GROUP b8/float4场景scale支持的数据类型Ascend 950PR/Ascend 950DTPER_CHANNEL和PER_TENSOR场景下支持的数据类型为half、bfloat16_tPER_GROUP float4场景下支持的数据类型为fp8_e8m0_t。Atlas A3 训练系列产品/Atlas A3 推理系列产品支持的数据类型为half、bfloat16_t。Atlas A2 训练系列产品/Atlas A2 推理系列产品支持的数据类型为half、bfloat16_t。Atlas 推理系列产品AI Core支持的数据类型为half。Kirin X90支持的数据类型为half。Kirin 9030支持的数据类型为half。调用示例// dstLocal结果张量 // srcLocal量化输入 // offsetLocal偏移参数 // scaleLocal缩放参数 // sharedTmpBuffer开发者管理的临时缓冲区用于存放内部计算中的中间变量 // kk轴长度 // shapeInfooffsetLocal和scaleLocal张量的shape信息 AscendC::AntiQuantShapeInfo shapeInfo {1, elementCountOfOffset, 1, elementCountOfOffset}; AscendC::AscendAntiQuantInputType, OutType, false( dstLocal, srcLocal, offsetLocal, scaleLocal, sharedTmpBuffer, k, shapeInfo);结果示例如下输入数据srcshape为[2,64]非转置场景: [1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1] offsetshape为[1,64]: [2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2.] scaleshape为[1,64]: [3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3.] 输出数据dstLocalshape为[2,64]: [9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9.]PER_TOKEN/PER_GROUP b8场景调用示例如下。// 注意m,n需从外部传入 constexpr static bool isReuseSource false; constexpr static AscendAntiQuantConfig config {has_offset, has_transpose, -1}; constexpr static AscendAntiQuantPolicy policy AscendAntiQuantPolicy::PER_TOKEN; AscendAntiQuantParam para; para.m m; para.n n; para.calCount calCount; AscendAntiQuantdstType, srcType, scaleType, config, policy(dstLocal, srcLocal, scaleLocal, offsetLocal, para);【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考