CANN/ops-nn分组量化SwiGLU算子

📅 2026/7/4 8:01:24
CANN/ops-nn分组量化SwiGLU算子
aclnnSwigluGroupQuant【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn 查看源码产品支持情况产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品×Atlas A2 训练系列产品/Atlas A2 推理系列产品×Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×功能说明接口功能在SwiGLU激活后执行分组低比特量化支持FP8和FP4输出。计算公式 令输入x的最后一维为D左半部分为A右半部分为B计算$$ y_{tmp}silu(A) \times B $$若传入weight则量化前执行$$ y_{tmp}y_{tmp} \times weight $$进行量化$$ scalerow_max(abs(y_{tmp}))/dstTypeScale $$$$ y Cast(Mul(y_{tmp}, 1/scale)) $$quant_mode为0时输出FP8类型的yOut和FLOAT32类型的yScaleOutquant_mode为1时输出FP8/FP4类型的yOut和FLOAT8_E8M0类型的yScaleOutquant_mode为2、3时输出HIFP8类型的yOut和FLOAT32类型的yScaleOut。函数原型每个算子分为两段式接口必须先调用“aclnnSwigluGroupQuantGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器再调用“aclnnSwigluGroupQuant”接口执行计算。aclnnStatus aclnnSwigluGroupQuantGetWorkspaceSize( const aclTensor *x, const aclTensor *weightOptional, const aclTensor *groupIndexOptional, const aclTensor *scaleOptional, int64_t dstType, int64_t quantMode, int64_t blockSize, bool roundScale, float clampLimit, float dstTypeMax, bool outputOrigin, const aclTensor *yOut, const aclTensor *yScaleOut, const aclTensor *yOriginOut, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnSwigluGroupQuant( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)aclnnSwigluGroupQuantGetWorkspaceSize参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorxaclTensor*输入SwiGLU输入。shape为[...,D]。D必须大于等于256且能被256整除。不支持空Tensor。FLOAT、FLOAT16、BFLOAT16ND1-7×weightOptionalaclTensor*输入量化前按token乘到SwiGLU输出上的权重。可选参数不支持空Tensor。不为空时数据类型为FLOAT32元素个数需等于x除最后一维外的元素个数之积。FLOAT32ND1-2×groupIndexOptionalaclTensor*输入count模式的group token数。可选参数不支持空Tensor。不为空时数据类型为INT64shape为[groupNum]。INT64ND1×scaleOptionalaclTensor*输入静态量化输入的scale张量。可选参数仅quant_mode为2时使用。quant_mode为3时不使用可传空。FLOAT32ND1×dstTypeint64_t输入目标量化类型。支持取值35、36、40、41分别表示FLOAT8_E5M2、FLOAT8_E4M3FN、FLOAT4_E2M1、FLOAT4_E1M2。dstType为40或41时quantMode必须为1。----quantModeint64_t输入量化模式。支持取值0、1、2、3。0表示Block FP8模式。1表示MX模式。2表示HIFP8静态量化模式。3表示HIFP8动态量化模式。----blockSizeint64_t输入量化块大小。0表示使用当前量化模式的默认block大小。quantMode为0时支持0或128。quantMode为1时支持0或32。quantMode为2或3时该参数不生效默认0。----roundScalebool输入是否将scale取整为2的幂。quantMode为1时roundScale必须为true。quantMode为2或3时该参数不生效。----clampLimitfloat输入SwiGLU计算前的clamp阈值。-1.0表示不启用clamp。启用clamp时clampLimit必须大于0。----dstTypeMaxfloat输入目标量化类型的最大有限值。quant_mode为2或3时用于计算scale amax / dstTypeMax。默认值为15.0HIFLOAT8的最大值。----outputOriginbool输入是否输出量化前的SwiGLU结果。MX FP4模式下yOriginOut仅作占位。quant_mode为2或3时支持输出原始激活值。----yaclTensor*输出量化输出。dstType为35或36时shape为[...,D/2]。dstType为40或41时shape为[...,D/4]。dstType为27时shape为[...,D/2]。数据类型需与dstType一致。不支持空Tensor。HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FN、FLOAT4_E2M1、FLOAT4_E1M2ND1-7×yScaleaclTensor*输出量化scale输出。quantMode为0时shape为[...,ceil((D/2)/128)]数据类型为FLOAT32。quantMode为1时shape为[...,ceil(ceil((D/2)/32)/2),2]数据类型为FLOAT8_E8M0。quantMode为2或3时无groupIndex时shape为[1]有groupIndex时shape为[G]数据类型为FLOAT32。不支持空Tensor。FLOAT32、FLOAT8_E8M0ND1-8×yOriginaclTensor*输出量化前的SwiGLU结果。shape为[...,D/2]。数据类型需与x一致。不支持空Tensor。FLOAT、FLOAT16、BFLOAT16ND1-7×workspaceSizeuint64_t*输出返回需要在Device侧申请的workspace大小。-----executoraclOpExecutor**输出返回op执行器包含了算子计算流程。-----返回值aclnnStatus返回状态码具体参见aclnn返回码。第一段接口完成入参校验出现以下场景时报错返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001x、yOut、yScaleOut、yOriginOut、workspaceSize或executor存在空指针。ACLNN_ERR_PARAM_INVALID161002输入或输出的数据类型不在支持范围内。输入或输出的shape不满足约束。dstType、quantMode、blockSize、roundScale或clampLimit不符合当前支持的值。dstType、quantMode、blockSize、roundScale和yScaleOut数据类型组合不匹配。weightOptional或groupIndexOptional不满足可选输入约束。ACLNN_ERR_INNER_TILING_ERROR561002多个输入tensor之间的shape信息不匹配、输入属性不在取值范围详见参数说明。aclnnSwigluGroupQuant参数说明参数名输入/输出描述workspace输入在Device侧申请的workspace内存地址。workspaceSize输入在Device侧申请的workspace大小由第一段接口aclnnSwigluGroupQuantGetWorkspaceSize获取。executor输入op执行器包含了算子计算流程。stream输入指定执行任务的Stream。返回值aclnnStatus返回状态码具体参见aclnn返回码。约束说明quantMode为0时仅支持FP8输出blockSize支持0或128。quantMode为1时支持FP8/FP4输出blockSize支持0或32roundScale必须为true。quantMode为2时支持HIFP8静态量化输出dstType, blockSize和roundScale不生效。quantMode为3时支持HIFP8动态量化输出dstType, blockSize和roundScale不生效。dstType为FLOAT4_E2M1或FLOAT4_E1M2时必须使用quantMode1。yScale的数据类型必须与quantMode匹配Block FP8为FLOAT32MX为FLOAT8_E8M0HIFP8为FLOAT32。quantMode为2或3时groupIndexOptional可用于MoE场景的分组量化yScaleOut的shape为[G]G为group数量。确定性计算aclnnSwigluGroupQuant默认确定性实现。调用示例示例代码如下仅供参考具体编译和执行过程请参考编译与运行样例。#include cstdint #include cstring #include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_swiglu_group_quant.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto dim : shape) { shapeSize * dim; } return shapeSize; } int Init(int32_t deviceId, aclrtStream* stream) { auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret); return ACL_SUCCESS; } bool CheckHardwareSupport() { const char* socName aclrtGetSocName(); if (socName nullptr) { LOG_PRINT(Warning: Cannot get SOC name, skip hardware check\n); return true; } LOG_PRINT(Current SOC: %s\n, socName); if (strstr(socName, Ascend950) ! nullptr || strstr(socName, ascend950) ! nullptr) { return true; } LOG_PRINT(Warning: SwigluGroupQuant only supports Ascend950, current SOC %s is not supported. Skip test.\n, socName); return false; } void Finalize(int32_t deviceId, aclrtStream stream) { (void)aclrtDestroyStream(stream); (void)aclrtResetDevice(deviceId); (void)aclFinalize(); } template typename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size GetShapeSize(shape) * sizeof(T); auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed. ERROR: %d\n, ret); return ret); std::vectorint64_t strides(shape.size(), 1); for (int64_t i static_castint64_t(shape.size()) - 2; i 0; --i) { strides[i] shape[i 1] * strides[i 1]; } *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return ACL_SUCCESS; } int main() { int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); if (!CheckHardwareSupport()) { LOG_PRINT(\n Test SKIPPED (hardware not supported) \n); Finalize(deviceId, stream); return ACL_SUCCESS; } std::vectorint64_t xShape {2, 256}; std::vectorint64_t yShape {2, 128}; std::vectorint64_t yScaleShape {2, 1}; std::vectorint64_t yOriginShape {2, 128}; std::vectoruint16_t xHostData(GetShapeSize(xShape), 0); for (size_t i 0; i xHostData.size(); i) { xHostData[i] static_castuint16_t(i % 23); } std::vectoruint8_t yHostData(GetShapeSize(yShape), 0); std::vectorfloat yScaleHostData(GetShapeSize(yScaleShape), 0.0f); std::vectoruint16_t yOriginHostData(GetShapeSize(yOriginShape), 0); void* xDeviceAddr nullptr; void* yDeviceAddr nullptr; void* yScaleDeviceAddr nullptr; void* yOriginDeviceAddr nullptr; aclTensor* x nullptr; aclTensor* y nullptr; aclTensor* yScale nullptr; aclTensor* yOrigin nullptr; ret CreateAclTensor(xHostData, xShape, xDeviceAddr, ACL_FLOAT16, x); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(yHostData, yShape, yDeviceAddr, ACL_FLOAT8_E4M3FN, y); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(yScaleHostData, yScaleShape, yScaleDeviceAddr, ACL_FLOAT, yScale); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(yOriginHostData, yOriginShape, yOriginDeviceAddr, ACL_FLOAT16, yOrigin); CHECK_RET(ret ACL_SUCCESS, return ret); int64_t dstType 36; int64_t quantMode 0; int64_t blockSize 0; bool roundScale false; double clampLimit -1.0; bool outputOrigin false; uint64_t workspaceSize 0; aclOpExecutor* executor nullptr; ret aclnnSwigluGroupQuantGetWorkspaceSize(x, nullptr, nullptr, dstType, quantMode, blockSize, roundScale, clampLimit, outputOrigin, y, yScale, yOrigin, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnSwigluGroupQuantGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret); } ret aclnnSwigluGroupQuant(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnSwigluGroupQuant failed. ERROR: %d\n, ret); return ret); ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); std::vectoruint8_t resultData(GetShapeSize(yShape), 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), yDeviceAddr, resultData.size() * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy result from device to host failed. ERROR: %d\n, ret); return ret); LOG_PRINT(result[0] is: %d\n, resultData[0]); aclDestroyTensor(x); aclDestroyTensor(y); aclDestroyTensor(yScale); aclDestroyTensor(yOrigin); aclrtFree(xDeviceAddr); aclrtFree(yDeviceAddr); aclrtFree(yScaleDeviceAddr); aclrtFree(yOriginDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } Finalize(deviceId, stream); return ACL_SUCCESS; }【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考