人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载QuantMatmulReduceSum 是 CANN ops-nn 神经网络算子库中实现「量化分组矩阵乘 跨 batch 累加」融合计算的算子它对 INT8 量化的多组矩阵分别做矩阵乘按行/列缩放后把所有组的计算结果累加为一张 (M, N) 输出矩阵属于典型的「Matmul ReduceSum」二合一融合算子。本文基于 matmul/quant_matmul_reduce_sum 目录下的官方 README、aclnn 接口文档与完整源码系统讲解其产品支持范围、数学语义、输入输出约束、两段式 aclnn 接口用法并深入 op_def、InferShape、Tiling 与 Kernel 源码帮助你在 Atlas A2/A3 系列 NPU 上正确、高效地完成 INT8 量化矩阵乘累加计算。产品支持情况QuantMatmulReduceSum 算子以及对应的 aclnnQuantMatmulReduceSumWeightNz 接口在不同产品系列上的支持情况如下表所示产品是否支持Ascend 950PR/Ascend 950DT×Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×也就是说当前该算子仅面向 Atlas A2 / Atlas A3 训练与推理系列产品Ascend 910B 及后续同源平台使用前请先确认目标设备的型号归属。这一支持范围同样可以从算子注册代码中印证在 op_host/quant_matmul_reduce_sum_def.cpp 中QuantMatmulReduceSum的 AICore 配置仅通过AddConfig(ascend910b, ...)与AddConfig(ascend910_93, ...)添加其余平台未注册。功能说明量化分组矩阵乘再求和算子语义QuantMatmulReduceSum 完成「量化的分组矩阵计算」随后把所有组的矩阵计算结果相加后输出。以 batch 组量化矩阵为输入每组内部执行一次 INT8 矩阵乘乘积累加结果经x1Scale行缩放与x2Scale列缩放双重缩放后在 batch 维度上求和最终只输出一张 (M, N) 矩阵。计算公式$$ out \sum_{i0}^{batch}(x1_i x2_i) * x1Scale * x2Scale $$其中$x1$ 是维度为 $(batch, M, K)$ 的左矩阵INT8$x2$ 是维度为 $(batch, K, N)$ 的右矩阵INT8$out$ 是维度为 $(M, N)$ 的输出矩阵$x1Scale$ 是维度为 $(batch, M)$ 的缩放向量FLOAT32对矩阵乘结果做行缩放$x2Scale$ 是维度为 $(N,)$ 的缩放向量BFLOAT16对矩阵乘结果做列缩放。从公式可以直观理解该算子的典型落地场景多 batch 的量化矩阵乘结果需要在 batch 维上归约例如多头注意力中多 head 计算结果累加、批量样本的聚合输出等在 NPU 上若拆分为「先逐组 Matmul 再 ReduceSum」会引入多次 Kernel 启动与中间张量搬移而 QuantMatmulReduceSum 将其融合为一次 Kernel 完成显著降低调用开销。广播语义两个缩放向量在实际计算时会被广播为 (batch, M, N) 维度后参与乘算x1Scale形状为 (batch, M)广播为 (batch, M, N) 后对每组的矩阵乘结果按行缩放x2Scale形状为 (N,)广播为 (batch, M, N) 后对每组的矩阵乘结果按列缩放。该语义在 docs/aclnnQuantMatmulReduceSumWeightNz.md 的参数说明中明确给出。参数说明下表是算子层IR/Graph 层的输入输出定义其中 x1、x2、dims 为必选输入bias、x1_scale、x2_scale、y_scale、x1_offset、x2_offset、y_offset、x2_table 均为可选输入当前版本下这些可选参数在实际调用时均需传空参数名输入/输出/属性描述数据类型数据格式x1输入矩阵乘运算中的左矩阵。INT8NDx2输入矩阵乘运算中的右矩阵。INT8FRACTAL_NZx1Scale输入对矩阵乘结果进行行缩放的一维向量。FLOAT32NDx2Scale输入对矩阵乘结果进行列缩放的一维向量。BFLOAT16NDout输出结果矩阵。FLOAT32ND说明上表出自算子 README其中out的数据类型为 FLOAT32 是 IR 层原型早期的描述当前仓库实际实现中输出y/out的数据类型为BFLOAT16见下方「源码级实现解析」中 op_def 与 InferShape 的分析以及 aclnn 接口文档中 out 的 BF16 说明。实际使用请以 aclnn 接口文档与当前源码为准。对应的 IR 层完整输入输出定义包含所有预留可选输入可以在 op_host/quant_matmul_reduce_sum_def.cpp 中查看必选输入x1INT8/ND、x2INT8/FRACTAL_NZ、dimsINT64/NDreduce 维度可选输入biasBF16、x1_scaleFLOAT、x2_scaleBF16、y_scaleUINT64、x1_offsetBF16、x2_offsetBF16、y_offsetBF16、x2_tableINT8输出yBF16/ND属性dtype输出数据类型必填、compute_type默认 -1、transpose_x1默认 false、transpose_x2默认 false、group_size默认 -1、keep_dims默认 false。约束说明使用 QuantMatmulReduceSum 算子时需满足以下约束不支持空 Tensorx1、x2、x1Scale、x2Scale、out 均不能为空左右矩阵不支持非连续 TensortransposeX1、transposeX2当前版本仅支持falsekeepDims当前版本仅支持false输出不保留 batch 维度恒为 (M, N)dimsreduce 维度当前版本仅支持[0]即固定在第 0 维batch 维做 ReduceSumgroupSize为预留参数当前版本需传入 0yScale、x1Offset、x2Offset、yOffset、bias为预留参数当前版本需传入 nullptr空 Tensor 或空指针确定性说明aclnnQuantMatmulReduceSumWeightNz 默认非确定性实现可通过aclrtCtxSetSysParamOpt开启确定性。aclnn 调用说明与两段式接口调用方式该算子对外提供 aclnn 单算子调用接口支持通过aclnnQuantMatmulReduceSumWeightNz等方式调用。样例代码位于 examples/test_aclnn_quant_matmul_reduce_sum_weight_nz.cpp另有小 shape 边界样例 examples/test_aclnn_quant_matmul_reduce_sum_small_case.cpp接口详细说明见 docs/aclnnQuantMatmulReduceSumWeightNz.md。两段式接口aclnn 单算子调用采用 CANN 标准的两段式接口必须先调用aclnnQuantMatmulReduceSumWeightNzGetWorkspaceSize获取计算所需的 workspace 大小以及包含算子计算流程的执行器再调用aclnnQuantMatmulReduceSumWeightNz执行计算。第一段接口获取 workspace 大小与执行器aclnnStatus aclnnQuantMatmulReduceSumWeightNzGetWorkspaceSize( const aclTensor *x1, const aclTensor *x2, const aclTensor *x1Scale, const aclTensor *x2Scale, const aclTensor *yScale, const aclTensor *x1Offset, const aclTensor *x2Offset, const aclTensor *yOffset, const aclTensor *bias, bool transposeX1, bool transposeX2, int64_t groupSize, const aclIntArray *dims, bool keepDims, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)第二段接口执行计算aclnnStatus aclnnQuantMatmulReduceSumWeightNz( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)第一段接口参数详解参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensorx1aclTensor*输入公式中的x1。不支持空Tensor。INT8ND(batch, m, k)-x2aclTensor*输入公式中的x2。不支持空Tensor。各维度表示 (batch, n1, k1, k0, n0)其中 k016、n032x1 shape 中的 k 与 x2 shape 中的 k1 需满足 ceil(k/16)k1x2 的 n1 与 out 的 n 需满足 ceil(n/n0)n1。可使用 aclnnCalculateMatmulWeightSizeV2 与 aclnnTransMatmulWeight 接口完成 ND → FRACTAL_NZ 格式转换原始 ND shape 为 (batch, k, n)。INT8NZ5维-x1ScaleaclTensor*输入公式中的x1Scale。不支持空Tensor。计算时广播为 (batch, m, n)。FLOAT32ND(batch, m)√x2ScaleaclTensor*输入公式中的x2Scale。不支持空Tensor。计算时广播为 (batch, m, n)。BFLOAT16ND(n,)√yScaleaclTensor*输入预留参数当前版本不支持。需要传入nullptr。----x1OffsetaclTensor*输入预留参数当前版本不支持。需要传入nullptr。----x2OffsetaclTensor*输入预留参数当前版本不支持。需要传入nullptr。----yOffsetaclTensor*输入预留参数当前版本不支持。需要传入nullptr。----biasaclTensor*输入预留参数当前版本不支持。需要传入nullptr。----transposeX1bool输入x1的输入shape是否包含transpose。当前版本仅支持false。bool---transposeX2bool输入x2的输入shape是否包含transpose。当前版本仅支持false。bool---groupSizeint64_t输入预留参数当前版本不支持。需要传入0。----dimsaclIntArray*输入指定reduce维度。当前版本仅支持填[0]表示在第0维batch维做ReduceSum。INT64---keepDimsbool输入是否在输出张量中保留输入张量的维度。当前版本仅支持false。----outaclTensor*输出公式中的out。-BFLOAT16ND(m, n)√workspaceSizeuint64_t*输出返回需要在Device侧申请的workspace大小。-----executoraclOpExecutor**输出返回op执行器包含了算子计算流程。-----返回值与错误码两段接口均返回aclnnStatus状态码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 x1、x2、x1Scale、x2Scale 或 out 是空指针。ACLNN_ERR_PARAM_INVALID161002x1、x2、x1Scale、x2Scale 或 out 的数据类型和数据格式不在支持范围之内或 shape 不满足校验条件或为空 Tensor。这些校验逻辑在源码中有完整对应在 op_host/op_api/aclnn_quant_matmul_reduce_sum.cpp 中CheckInputs()依次执行CheckInputExistence()空指针检查、预留参数拦截、CheckInputAttr()transposeX1/transposeX2/groupSize/keepDims 仅允许默认值、CheckDtype()x1/x2 必须 INT8、x1Scale 必须 FLOAT、x2Scale 与 out 必须 BF16、CheckFormat()x2 必须 FRACTAL_NZ其余禁止 NZ、CheckInputShape()x1/x2 必须 3 维且 batch、k 相等x1Scale 必须 (batch, m)x2Scale 必须 (n,)。调用示例完整可运行代码以下示例来自 examples/test_aclnn_quant_matmul_reduce_sum_weight_nz.cpp完整展示了从初始化、构造 Tensor、两段式调用到取回结果的全部流程。示例采用 b8、m2048、k1024、n7168 的 shape所有输入数据初始化为 1.0INT8 为 1便于快速验证结果。#include cstdint #include cstring #include iostream #include memory #include vector #include acl/acl.h #include aclnnop/aclnn_quant_matmul_reduce_sum.h #include aclnnop/aclnn_trans_matmul_weight.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define CHECK_FREE_RET(cond, return_expr) \ do { \ if (!(cond)) { \ Finalize(deviceId, stream); \ 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 i : shape) { shapeSize * i; } return shapeSize; } float Bfloat16ToFloat(uint16_t value) { // BF16是FP32的高16位低16位补0后即可得到对应的FP32位表示。 const uint32_t bits static_castuint32_t(value) 16U; float result 0.0F; std::memcpy(result, bits, sizeof(result)); return result; } 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 0; } 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); // 调用aclrtMalloc申请device侧内存 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); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 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); // 计算连续tensor的strides std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } void Finalize(int32_t deviceId, aclrtStream stream) { aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); } // 创建FRACTAL_NZ格式的x2通过aclnnCalculateMatmulWeightSizeV2计算NZ存储大小 template typename T int CreateAclTensorX2(const std::vectorT hostData, const std::vectorint64_t shape, void **deviceAddr, aclDataType dataType, aclTensor **tensor) { auto size static_castuint64_t(GetShapeSize(shape)); const aclIntArray *mat2Size aclCreateIntArray(shape.data(), shape.size()); auto ret aclnnCalculateMatmulWeightSizeV2(mat2Size, dataType, size); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnCalculateMatmulWeightSizeV2 failed. ERROR: %d\n, ret); return ret); size * sizeof(T); 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 shape.size() - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } std::vectorint64_t storageShape; storageShape.push_back(GetShapeSize(shape)); *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_FRACTAL_NZ, storageShape.data(), storageShape.size(), *deviceAddr); return 0; } int aclnnQuantMatmulWeightNzTest(int32_t deviceId, aclrtStream stream) { auto ret Init(deviceId, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出 int64_t b 8; int64_t m 2048; int64_t k 1024; int64_t n 7168; // 创建x1 aclTensor std::vectorint64_t x1Shape {b, m, k}; void *x1DeviceAddr nullptr; aclTensor *x1 nullptr; std::vectorint8_t x1HostData(b * m * k, 1); ret CreateAclTensor(x1HostData, x1Shape, x1DeviceAddr, aclDataType::ACL_INT8, x1); std::unique_ptraclTensor, aclnnStatus (*)(const aclTensor *) x1TensorPtr(x1, aclDestroyTensor); std::unique_ptrvoid, aclError (*)(void *) x1DeviceAddrPtr(x1DeviceAddr, aclrtFree); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建FRACTAL_NZ格式的x2 aclTensor std::vectorint64_t x2Shape {b, k, n}; void *x2DeviceAddr nullptr; aclTensor *x2 nullptr; std::vectorint8_t x2HostData(b * k * n, 1); ret CreateAclTensorX2(x2HostData, x2Shape, x2DeviceAddr, aclDataType::ACL_INT8, x2); std::unique_ptraclTensor, aclnnStatus (*)(const aclTensor *) x2HPTensorPtr(x2, aclDestroyTensor); std::unique_ptrvoid, aclError (*)(void *) x2HPDeviceAddrPtr(x2DeviceAddr, aclrtFree); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建x1Scale aclTensor std::vectorint64_t x1ScaleShape {b, m}; void *x1ScaleDeviceAddr nullptr; std::vectorfloat x1ScaleHostData(b * m, 1); aclTensor *x1Scale nullptr; ret CreateAclTensor(x1ScaleHostData, x1ScaleShape, x1ScaleDeviceAddr, aclDataType::ACL_FLOAT, x1Scale); std::unique_ptraclTensor, aclnnStatus (*)(const aclTensor *) x1ScaleTensorPtr(x1Scale, aclDestroyTensor); std::unique_ptrvoid, aclError (*)(void *) x1ScaleDeviceAddrPtr(x1ScaleDeviceAddr, aclrtFree); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建x2Scale aclTensorBF16数值1.0的位表示为0x3F80 std::vectorint64_t x2ScaleShape {n}; void *x2ScaleDeviceAddr nullptr; aclTensor *x2Scale nullptr; constexpr uint16_t bf16One 0x3F80U; std::vectoruint16_t x2ScaleHostData(n, bf16One); ret CreateAclTensor(x2ScaleHostData, x2ScaleShape, x2ScaleDeviceAddr, aclDataType::ACL_BF16, x2Scale); std::unique_ptraclTensor, aclnnStatus (*)(const aclTensor *) x2ScaleTensorPtr(x2Scale, aclDestroyTensor); std::unique_ptrvoid, aclError (*)(void *) x2ScaleDeviceAddrPtr(x2ScaleDeviceAddr, aclrtFree); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建out aclTensorBF16数值0.0的位表示为0x0000 std::vectorint64_t outShape {m, n}; void *outDeviceAddr nullptr; aclTensor *out nullptr; std::vectoruint16_t outHostData(m * n, 0); ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_BF16, out); std::unique_ptraclTensor, aclnnStatus (*)(const aclTensor *) outTensorPtr(out, aclDestroyTensor); std::unique_ptrvoid, aclError (*)(void *) outDeviceAddrPtr(outDeviceAddr, aclrtFree); CHECK_RET(ret ACL_SUCCESS, return ret); bool transposeX1 false; bool transposeX2 false; // 创建dims aclIntArray第0维batch维做ReduceSum std::vectorint64_t dimsData {0}; aclIntArray *dims nullptr; dims aclCreateIntArray(dimsData.data(), dimsData.size()); CHECK_RET(dims ! nullptr, return ret); // 3. 两段式调用CANN算子库API uint64_t workspaceSize 0; aclOpExecutor *executor nullptr; // 第一段接口获取workspace大小与执行器预留参数传nullptrgroupSize传0keepDims传false ret aclnnQuantMatmulReduceSumWeightNzGetWorkspaceSize( x1, x2, x1Scale, x2Scale, nullptr, nullptr, nullptr, nullptr, nullptr, transposeX1, transposeX2, 0, dims, false, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnQuantMatmulReduceSumWeightNzGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 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, LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret); workspaceAddrPtr.reset(workspaceAddr); } // 第二段接口执行计算 ret aclnnQuantMatmulReduceSumWeightNz(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnQuantMatmulReduceSumWeightNz failed. ERROR: %d\n, ret); return ret); // 4.固定写法同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 5. 获取输出的值将device侧内存上的结果拷贝至host侧 auto size GetShapeSize(outShape); // C无法直接打印BF16先用uint16_t读取其位表示再将BF16转换为FP32 std::vectoruint16_t resultData(size, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr, 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); for (int64_t i 0; i 5; i) { LOG_PRINT(result[%ld] is: %f\n, i, Bfloat16ToFloat(resultData[i])); } return ACL_SUCCESS; } int main() { // 1.固定写法device/stream初始化参考acl API手册根据自己的实际device填写deviceId int32_t deviceId 0; aclrtStream stream; auto ret aclnnQuantMatmulWeightNzTest(deviceId, stream); CHECK_FREE_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnQuantMatmulWeightNzTest failed. ERROR: %d\n, ret); return ret); Finalize(deviceId, stream); return 0; }代码要点解读x2 的 FRACTAL_NZ 格式是本次调用的关键点x2 在 aclnn 接口层要求以 5 维 NZ 排布k016、n032 分形块传入。样例通过aclnnCalculateMatmulWeightSizeV2计算 NZ 存储大小并据此申请内存再以ACL_FORMAT_FRACTAL_NZ创建 aclTensor实际业务中若输入为 ND 格式可先用aclnnTransMatmulWeight完成 ND → NZ 转换。BF16 位操作C 无原生 BF16 类型样例用uint16_t保存 BF16 位表示1.0 为0x3F800.0 为0x0000读取结果时通过Bfloat16ToFloat左移 16 位补零转为 FP32 后打印。两段式固定流程第一段接口校验参数并构造执行器 → 按workspaceSize用aclrtMalloc申请 workspace → 第二段接口在指定 stream 上执行 →aclrtSynchronizeStream同步等待完成。源码级实现解析算子注册与属性定义op_defop_host/quant_matmul_reduce_sum_def.cpp 通过OP_ADD(QuantMatmulReduceSum)注册算子原型。除上文列出的输入输出外该文件还设置了丰富的编译配置DynamicCompileStaticFlag(true)动态编译静态化DynamicFormatFlag(true)支持动态格式DynamicRankSupportFlag(true)/DynamicShapeSupportFlag(true)支持动态 rank 与动态 shapePrecisionReduceFlag(false)不做精度降低通过ExtendCfgInfo(aclnnSupport.value, support_aclnn)声明算子支持 aclnn 调用。InferShape 推导逻辑op_host/quant_matmul_reduce_sum_infershape.cpp 实现了 shape 与 dtype 推导x1、x2 必须均为 3 维dimX1 ! 3 || dimX2 ! 3报错x1 与 x2 的 batch第 0 维必须相等、kx1 第 2 维与 x2 第 1 维必须相等输出恒为 2 维(mX1, nX2)shapeOut-SetDimNum(2)第 0 维取 x1 的 m第 1 维取 x2 的 n与「在 batch 维 ReduceSum 且 keepDimsfalse」的语义完全一致输出 dtype 固定为DT_BF16。计算图构建l0 封装op_host/op_api/aclnn_quant_matmul_reduce_sum.cpp 是 aclnn 两段式接口的直接实现。第一段接口内部依次完成CheckInputs()全量参数校验空指针、预留参数、属性默认值、dtype、format、shape详见上文错误码章节PreProcess()将各输入 Tensor 的 original shape 记录为 view shapeProcessL0()组装计算图先对 x1、x1Scale、x2Scale 执行l0op::Contiguous保证连续性对应约束中「不支持非连续 Tensor」的实现再调用l0op::QuantMatmulReduceSum见 op_host/op_api/quant_matmul_reduce_sum.cpp该 l0 接口通过INFER_SHAPE与ADD_TO_LAUNCHER_LIST_AICORE将算子挂入 AICore 执行队列最后若输出 out 为非连续 Tensor通过l0op::ViewCopy将计算得到的连续结果拷贝回 out 视图*workspaceSize uniqueExecutor-GetWorkspaceSize()返回所需 workspace 大小。第二段接口则直接调用CommonOpExecutorRun(workspace, workspaceSize, executor, stream)完成在指定 stream 上的异步执行。Tiling 与 Kernel 实现AICore 侧从 Tiling 源码 op_host/op_tiling/quant_matmul_reduce_sum_tiling.cpp 可以看到该算子的底层计算特征计算核心是 INT8 输入、INT32 累加的矩阵乘inputParams_.cDtype ge::DT_INT32x1Scale 必须是 FLOAT、x2Scale 必须是 BF16、输出 y 必须是 BF16编译信息采集包含 UB/L1/L2/L0_A/L0_B/L0_C 各级缓存容量与 AIC 核数compileInfo_.aicNum其中对 L2 容量做了特殊纠正当系统上报的 L2 为 96MB 时按 168MB 实际物理大小计算L2_REAL_SIZE 168、L2_FAKE_SIZE 96以保证大 shape 场景下的 Tiling 切分正确Kernel 侧 op_kernel/quant_matmul_reduce_sum.cpp 使用MixCoreAICAIV 混合核架构实现KERNEL_TASK_TYPE(0, KERNEL_TYPE_MIX_AIC_1_2)与KERNEL_TASK_TYPE(1, KERNEL_TYPE_MIX_AIC_1_1)定义了 AIC 与 AIV 的比例模式。执行流程分为两步先用QuantMatmulReduceSumInitOutput对输出 y 做清零因为累加结果要写入输出再通过QuantMatmulReduceSumQuantMixCoreCompute执行真正的量化矩阵乘与跨 batch 累加矩阵乘使用MatmulTypeGM, CubeFormat::ND, int8_tx1与MatmulTypeGM, CubeFormat::NZ, int8_tx2模板与 IR 层 x1ND、x2FRACTAL_NZ 的格式约束一一对应。由此可见QuantMatmulReduceSum 的「分组矩阵乘 ReduceSum」融合并非在 Host 侧拆解为多个算子而是在 AICore Kernel 内部一次性完成AIC 负责矩阵乘计算AIV 负责缩放与跨 batch 累加归约中间结果驻留在 on-chip 缓存与 workspace 中避免了多算子级联带来的额外显存读写。单元测试仓库为该算子提供了完整的 Host 侧与 Kernel 侧单测tests/ut/op_host/test_aclnn_quant_matmul_reduce_sum_api.cpp验证 aclnn 两段式接口的参数校验与基本调用tests/ut/op_host/test_quant_matmul_reduce_sum_infershape.cpp验证 InferShape 推导规则tests/ut/op_host/test_quant_matmul_reduce_sum_tiling.cpp验证 Tiling 参数计算tests/ut/op_kernel/test_quant_matmul_reduce_sum.cppKernel 计算正确性测试Tiling 数据结构定义见 tests/ut/op_kernel/quant_matmul_reduce_sum_tiling_def.h。其中small_case样例examples/test_aclnn_quant_matmul_reduce_sum_small_case.cpp专门覆盖了输出小于 32 字节的极小 shape 边界输出 32B 时 Kernel 走DataCopyPad路径输出 32B 时为对齐边界输出 32B 但存在 tail 时走 tail 处理路径这些边界场景是编写自测用例时值得参考的用例集。使用建议与注意事项平台确认调用前先确认设备属于 Atlas A2/A3 训练或推理系列Ascend 950、Atlas 200I/500 A2 及其他老平台不支持本算子。x2 格式预处理x2 需要 FRACTAL_NZ 5 维格式业务侧 ND 权重务必先用aclnnCalculateMatmulWeightSizeV2计算 NZ 存储大小、aclnnTransMatmulWeight完成格式转换示例源码中已给出可直接复用的CreateAclTensorX2模板函数。预留参数必须传空yScale/x1Offset/x2Offset/yOffset/bias传 nullptrgroupSize传 0transposeX1/transposeX2/keepDims传 falsedims传[0]——传错任一均会被第一段接口以ACLNN_ERR_PARAM_INVALID161002拦截。输出累加语义算子内部先对输出清零再累加输出为跨 batch 归约后的 (M, N) 单矩阵若上层需要保留逐组结果请拆分多次调用或使用其他 Matmul 类算子。确定性开关算子默认非确定性实现对精度可复现性有要求的场景可通过aclrtCtxSetSysParamOpt开启确定性。通过本文的公式语义、接口参数与源码实现三重视角你可以在 Atlas A2/A3 平台上快速落地「INT8 量化分组矩阵乘 跨 batch 求和」的融合计算并能够独立完成样例代码的改编、参数校验问题的排查与 shape 边界用例的构造。赞分享人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载相关推荐CANN ops-transformer GroupedMatmulAdd 算子全解析分组矩阵乘与原地累加的 K 轴分组实现与 aclnn 调用指南CANN ops transformer GroupedMatmulAdd 算子全解析分组矩阵乘与原地累加的 K 轴分组实现与 aclnn 调用指南 导读 本算子库人工智能深度学习AscendCANN ops-transformer 算子详解aclnnQuantGroupedMatmulInplaceAdd 量化分组矩阵乘原位累加融合算子CANN ops transformer 算子详解aclnnQuantGroupedMatmulInplaceAdd 量化分组矩阵乘原位累加融合算子 本技术指算子库人工智能深度学习AscendCANN ops-transformer 分组矩阵乘累加算子 aclnnGroupedMatmulAdd 接口详解与实战指南CANN ops transformer 分组矩阵乘累加算子 aclnnGroupedMatmulAdd 接口详解与实战指南 aclnnGroupedMatmu算子库人工智能深度学习Ascend上一篇Buzz语音转录终极指南3步打造你的本地隐私保护工作站下一篇FinRobot 金融 AI 代理平台完整入门指南五步跑通你的第一份股票研究报告创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
