CANN/ops-math权重量化预处理算子 aclnnWeightQuantPreprocess【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math 查看源码产品支持情况产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品×Atlas A2 训练系列产品/Atlas A2 推理系列产品×Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×功能说明完成伪量化 Matmul包括 QuantBatchMatmulV5、GroupedMatmul-伪量化的参数预处理主要将 weight 从 ND 格式转换为 FRACTAL_NZ 格式并在需要时对 weightScale、weightOffsetOptional、biasOptional 进行同步处理。函数原型每个算子分为两段式接口调用aclnnWeightQuantPreprocessGetWorkspaceSize获取 workspace 大小及执行器调用aclnnWeightQuantPreprocess执行计算。注意用户需自行构造输出张量参考约束说明中的 shape 计算公式。aclnnStatus aclnnWeightQuantPreprocessGetWorkspaceSize( const aclTensor *weight, const aclTensor *weightScale, const aclTensor *weightOffsetOptional, const aclTensor *biasOptional, aclDataType xDtype, aclDataType xScaleDtype, int64_t kGroupSize, aclTensor *outWeight, aclTensor *outWeightScale, aclTensor *outWeightOffsetOptional, aclTensor *outBiasOptional, uint64_t *workspaceSize, aclOpExecutor **executor) aclnnStatus aclnnWeightQuantPreprocess( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)aclnnWeightQuantPreprocessGetWorkspaceSize参数说明参数名输入/输出描述使用说明数据类型数据格式维度shape非连续Tensorweight(aclTensor *)输入Matmul的权重矩阵不支持空 tensorfloat4_e2m1ND2-3仅转置场景支持weightScale(aclTensor *)输入权重的反量化 scale 参数不支持空 tensorfloat8_e8m0ND/NCL/NCHW3-4仅转置场景支持weightOffsetOptional(aclTensor *)可选输入权重的反量化 offset 参数当前 MM_MX_A8W4/GMM_MX_A8W4 数据流不支持必须为 nullptr-ND1-2仅转置场景支持biasOptional(aclTensor *)可选输入Matmul 的偏置矩阵不支持空 tensor必须 contiguousfloat16/bfloat16ND1-2不支持xDtype(aclDataType)输入Matmul的激活矩阵的数据类型-aclDataType---xScaleDtype(aclDataType)输入激活的量化 scale 参数的数据类型-aclDataType---kGroupSize(int64_t)输入权重在 per-group 量化时 K 维度的 group的大小-int64---outWeight(aclTensor *)输出预处理后的 weight-int8/int4/fp8_e4m3/hif8/fp4_e2m1NZ2-5仅转置场景支持outWeightScale(aclTensor *)输出预处理后的 weightScale-float16/bfloat16/fp8_e8m0ND/NCL/NCHW3-4仅转置场景支持outWeightOffsetOptional(aclTensor *)输出预处理后的 weightOffset当前 MM_MX_A8W4/GMM_MX_A8W4 数据流不支持必须为 nullptrfloat16/bfloat16ND1-2仅转置场景支持outBiasOptional(aclTensor *)输出预处理后的 bias必须 contiguousfloat16/bfloat16ND1-2不支持workspaceSize(uint64_t *)输出计算所需的workspace大小-uint64*---executor(aclOpExecutor **)输出包含算子计算流程的执行器-aclOpExecutor**---返回值aclnnStatus返回状态码具体参见aclnn返回码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001weight、weightScale、outWeight或outWeightScale是空指针或biasOptional非空但outBiasOptional是空指针。ACLNN_ERR_PARAM_INVALID161002输入的数据类型组合不支持无法匹配当前支持的MM_MX_A8W4/GMM_MX_A8W4数据流。weight、weightScale、outWeight或outWeightScale是空tensor或biasOptional/outBiasOptional在提供时为空tensor。weight、weightScale、biasOptional、outWeight、outWeightScale或outBiasOptional的数据类型和数据格式不在支持的范围之内。weight、weightScale、biasOptional、outWeight、outWeightScale或outBiasOptional的shape或storage shape不满足校验条件。weight或weightScale的stride不满足转置要求或biasOptional/outBiasOptional在提供时不连续。weightOffsetOptional或outWeightOffsetOptional非空或kGroupSize不等于32。ACLNN_ERR_RUNTIME_ERROR361001产品型号不支持。ACLNN_ERR_INNER_CREATE_EXECUTOR561101内部错误执行器创建失败。ACLNN_ERR_INNER_NULLPTR561103workspaceSize或executor是空指针或API内部构图接口返回空指针。aclnnWeightQuantPreprocess参数说明参数名输入/输出描述workspace输入在Device侧申请的workspace内存地址。workspaceSize输入在Device侧申请的workspace大小由第一段接口aclnnWeightQuantPreprocessGetWorkspaceSize获取。executor输入op执行器包含了算子计算流程。stream输入指定执行任务的Stream。返回值aclnnStatus返回状态码具体参见aclnn返回码。约束说明MM_MX_A8W4 数据流MM 表示 MatmulMX_A8W4 表示 x 的数据类型为 FLOAT8_E4M3FNweight 的数据类型为 FLOAT4_E2M1Mx量化模式weight数据类型FLOAT4_E2M1格式NDK % kGroupSize 0view shape2-D{K, N}storage shape{N, K}transposedstride[1, K]最后两维 transposed不支持空 tensorweightScale数据类型FLOAT8_E8M0格式ND/NCLview shape3-D{ceildiv(K, 64), N, 2}storage shape{N, ceildiv(K, 64), 2}transposedstride[2, 2*ceildiv(K,64), 1]维度0和1交换不支持空 tensorweightOffsetOptional当前不支持必须为 nullptroutWeightOffsetOptional 也必须为 nullptrbiasOptional数据类型float16/bfloat16格式ND必须为 contiguous不支持空 tensor若提供kGroupSize必须等于 32xDtypeFLOAT8_E4M3FNxScaleDtypeFLOAT8_E8M0outWeight数据类型与 weight 相同格式FRACTAL_NZ_C0_32view shape与 weight view shape 相同{K, N}storage shape4-D{ceildiv(K, 32), ceildiv(N, 16), 16, 32}outWeightScale数据类型与 weightScale 相同格式NDview shape与 weightScale view shape 相同storage shape与 weightScale storage shape 相同outBiasOptional数据类型与 biasOptional 相同格式ND必须为 contiguousview shape与 biasOptional 相同storage shape与 biasOptional 相同GMM_MX_A8W4 数据流GMM 表示 GroupedMatmulMX_A8W4 表示 x 的数据类型为 FLOAT8_E4M3FNweight 的数据类型为 FLOAT4_E2M1Mx量化模式weight数据类型FLOAT4_E2M1格式NDK % kGroupSize 0view shape3-D{G, K, N}storage shape{G, N, K}transposed最后两维交换stride[K*N, 1, K]维度1和2 transposed不支持空 tensorweightScale数据类型FLOAT8_E8M0格式ND/NCL/NCHWview shape4-D{G, ceildiv(K, 64), N, 2}storage shape{G, N, ceildiv(K, 64), 2}transposed维度2和3交换stride[2*ceildiv(K,64)*N, 2, 2*ceildiv(K,64), 1]维度2和3交换不支持空 tensorweightOffsetOptional当前不支持必须为 nullptroutWeightOffsetOptional 也必须为 nullptrbiasOptional数据类型float16/bfloat16格式ND必须为 contiguous不支持空 tensor若提供kGroupSize必须等于 32xDtypeFLOAT8_E4M3FNxScaleDtypeFLOAT8_E8M0outWeight数据类型与 weight 相同格式FRACTAL_NZ_C0_32view shape与 weight view shape 相同{G, K, N}storage shape5-D{G, ceildiv(K, 32), ceildiv(N, 16), 16, 32}outWeightScale数据类型与 weightScale 相同格式NDview shape与 weightScale view shape 相同storage shape与 weightScale storage shape 相同outBiasOptional数据类型与 biasOptional 相同格式ND必须为 contiguousview shape与 biasOptional 相同storage shape与 biasOptional 相同其余数据类型与 shape 组合为预留接口当前调用将返回 ACLNN_ERR_PARAM_INVALID调用示例示例代码如下仅供参考具体编译和执行过程请参考编译与运行样例。注意用户需自行计算并构造输出张量 shape参考约束说明中的公式outWeight viewShape与 weight viewShape 相同outWeight storageShape{CeilDiv(K, 32), CeilDiv(N, 16), 16, 32}outWeight formatACL_FORMAT_FRACTAL_NZ_C0_32#include iostream #include memory #include vector #include acl/acl.h #include aclnnop/aclnn_weight_quant_preprocess.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define CEIL_DIV(x, y) (((x) (y) - 1) / (y)) int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t size 1; for (auto d : shape) size * d; return size; } class AclRuntimeGuard { public: explicit AclRuntimeGuard(int32_t deviceId) : deviceId_(deviceId) {} ~AclRuntimeGuard() { if (stream_ ! nullptr) { aclrtDestroyStream(stream_); stream_ nullptr; } if (deviceSet_) { aclrtResetDevice(deviceId_); deviceSet_ false; } if (aclInited_) { aclFinalize(); aclInited_ false; } } int Init(aclrtStream* stream) { auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, return ret); aclInited_ true; ret aclrtSetDevice(deviceId_); CHECK_RET(ret ACL_SUCCESS, return ret); deviceSet_ true; ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, return ret); stream_ *stream; return ACL_SUCCESS; } private: int32_t deviceId_; aclrtStream stream_ nullptr; bool aclInited_ false; bool deviceSet_ false; }; int main() { int32_t deviceId 0; aclrtStream stream nullptr; AclRuntimeGuard aclGuard(deviceId); auto ret aclGuard.Init(stream); CHECK_RET(ret ACL_SUCCESS, std::cout Init failed std::endl; return ret); // weight: FLOAT4_E2M1, transposed (MM_MX_A8W4) int64_t k 64; int64_t n 128; int64_t C0 32; // FLOAT4_E2M1 对应 C032 std::vectorint64_t weightViewShape {k, n}; std::vectorint64_t weightStorageShape {n, k}; std::vectorint64_t weightStrides {1, k}; int64_t weightStorageSize GetShapeSize(weightStorageShape); int64_t weightBytes weightStorageSize / 2; // FP4: 4 bits 0.5 bytes per element std::vectorint8_t weightHostData(weightBytes, 0); void* weightDeviceAddr nullptr; ret aclrtMalloc(weightDeviceAddr, weightBytes, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, std::cout Malloc weight failed std::endl; return ret); std::unique_ptrvoid, aclError (*)(void*) weightDeviceAddrPtr(weightDeviceAddr, aclrtFree); ret aclrtMemcpy(weightDeviceAddr, weightBytes, weightHostData.data(), weightBytes, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, std::cout Memcpy weight failed std::endl; return ret); aclTensor* weight aclCreateTensor(weightViewShape.data(), weightViewShape.size(), ACL_FLOAT4_E2M1, weightStrides.data(), 0, ACL_FORMAT_ND, weightStorageShape.data(), weightStorageShape.size(), weightDeviceAddr); std::unique_ptraclTensor, aclnnStatus (*)(const aclTensor*) weightPtr(weight, aclDestroyTensor); CHECK_RET(weight ! nullptr, std::cout Create weight tensor failed std::endl; return ACL_ERROR_FAILURE); // weightScale: FLOAT8_E8M0, 3-D transposed (MM_MX_A8W4) // viewShape: {ceildiv(K,64), N, 2} {1, 128, 2} // storageShape: {N, ceildiv(K,64), 2} {128, 1, 2} // transposed stride: {2, 2, 1} (dim0 - dim1) std::vectorint64_t scaleViewShape {k / 64, n, 2}; std::vectorint64_t scaleStorageShape {n, k / 64, 2}; std::vectorint64_t scaleStrides {2, 2, 1}; int64_t scaleStorageSize GetShapeSize(scaleStorageShape); int64_t scaleBytes scaleStorageSize; // FP8: 1 byte per element std::vectorint8_t scaleHostData(scaleBytes, 0); void* scaleDeviceAddr nullptr; ret aclrtMalloc(scaleDeviceAddr, scaleBytes, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, std::cout Malloc weightScale failed std::endl; return ret); std::unique_ptrvoid, aclError (*)(void*) scaleDeviceAddrPtr(scaleDeviceAddr, aclrtFree); ret aclrtMemcpy(scaleDeviceAddr, scaleBytes, scaleHostData.data(), scaleBytes, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, std::cout Memcpy weightScale failed std::endl; return ret); aclTensor* weightScale aclCreateTensor(scaleViewShape.data(), scaleViewShape.size(), ACL_FLOAT8_E8M0, scaleStrides.data(), 0, ACL_FORMAT_ND, scaleStorageShape.data(), scaleStorageShape.size(), scaleDeviceAddr); std::unique_ptraclTensor, aclnnStatus (*)(const aclTensor*) weightScalePtr(weightScale, aclDestroyTensor); CHECK_RET(weightScale ! nullptr, std::cout Create weightScale tensor failed std::endl; return ACL_ERROR_FAILURE); // 用户自行构造 outWeight (FRACTAL_NZ_C0_32) // viewShape 与 weight viewShape 相同storageShape 按公式计算 std::vectorint64_t outWeightViewShape {k, n}; std::vectorint64_t outWeightStorageShape {CEIL_DIV(k, C0), CEIL_DIV(n, 16), 16, C0}; int64_t outWeightStorageSize GetShapeSize(outWeightStorageShape); int64_t outWeightBytes outWeightStorageSize / 2; // FP4 void* outWeightDeviceAddr nullptr; ret aclrtMalloc(outWeightDeviceAddr, outWeightBytes, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, std::cout Malloc outWeight failed std::endl; return ret); std::unique_ptrvoid, aclError (*)(void*) outWeightDeviceAddrPtr(outWeightDeviceAddr, aclrtFree); aclTensor* outWeight aclCreateTensor(outWeightViewShape.data(), outWeightViewShape.size(), ACL_FLOAT4_E2M1, nullptr, 0, ACL_FORMAT_FRACTAL_NZ_C0_32, outWeightStorageShape.data(), outWeightStorageShape.size(), outWeightDeviceAddr); std::unique_ptraclTensor, aclnnStatus (*)(const aclTensor*) outWeightPtr(outWeight, aclDestroyTensor); CHECK_RET(outWeight ! nullptr, std::cout Create outWeight tensor failed std::endl; return ACL_ERROR_FAILURE); // 构造 outWeightScale (viewShape 和 storageShape 都与 weightScale 相同) // 根据实现要求outWeightScale 的 viewShape 和 storageShape 必须都与 weightScale 相同 void* outScaleDeviceAddr nullptr; ret aclrtMalloc(outScaleDeviceAddr, scaleBytes, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, std::cout Malloc outWeightScale failed std::endl; return ret); std::unique_ptrvoid, aclError (*)(void*) outScaleDeviceAddrPtr(outScaleDeviceAddr, aclrtFree); aclTensor* outWeightScale aclCreateTensor(scaleViewShape.data(), scaleViewShape.size(), ACL_FLOAT8_E8M0, scaleStrides.data(), 0, ACL_FORMAT_ND, scaleStorageShape.data(), scaleStorageShape.size(), outScaleDeviceAddr); std::unique_ptraclTensor, aclnnStatus (*)(const aclTensor*) outWeightScalePtr(outWeightScale, aclDestroyTensor); CHECK_RET(outWeightScale ! nullptr, std::cout Create outWeightScale tensor failed std::endl; return ACL_ERROR_FAILURE); aclDataType xDtype ACL_FLOAT8_E4M3FN; aclDataType xScaleDtype ACL_FLOAT8_E8M0; int64_t kGroupSize 32; // 1. 获取 workspace 与执行器 uint64_t workspaceSize 0; aclOpExecutor* executor nullptr; ret aclnnWeightQuantPreprocessGetWorkspaceSize( weight, weightScale, nullptr, nullptr, // weightOffsetOptional, biasOptional xDtype, xScaleDtype, kGroupSize, outWeight, outWeightScale, nullptr, nullptr, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, std::cout GetWorkspaceSize failed std::endl; return ret); void* workspaceAddr nullptr; std::unique_ptrvoid, aclError (*)(void*) workspaceAddrPtr(nullptr, aclrtFree); if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, std::cout Malloc workspace failed std::endl; return ret); workspaceAddrPtr.reset(workspaceAddr); } // 2. 执行计算 ret aclnnWeightQuantPreprocess(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, std::cout Preprocess failed std::endl; return ret); ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, std::cout Synchronize failed std::endl; return ret); // 3. 释放资源 workspaceAddrPtr.reset(); outWeightScalePtr.reset(); outWeightPtr.reset(); weightScalePtr.reset(); weightPtr.reset(); outScaleDeviceAddrPtr.reset(); outWeightDeviceAddrPtr.reset(); scaleDeviceAddrPtr.reset(); weightDeviceAddrPtr.reset(); return 0; }【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考