WeightQuantBatchMatmulExperiment 算子深度解析A16W4 PerGroup 伪量化 MatMul 的 MSD 算法与双模板流水设计【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nnWeightQuantBatchMatmulExperiment 是 CANN ops-nn 仓库experimental/matmul/weight_quant_batch_matmul_experiment下的一个自定义算子工程样例用于在 NPU 上实现A16W4 PerGroup 场景的 MatMul 伪量化计算。本算子以 MSDMulti-Step Decomposition多步分解算法为核心将 float16 激活矩阵按组展开为多张 int4 矩阵再通过 Cube 单元完成整数矩阵乘、Vector 单元完成后处理合并并同时提供基础流水与计算解耦 Preload 流水两套 kernel 模板配合 msprof 性能采集工具验证流水优化效果。阅读本文后你将掌握该算子的数学原理、算子规格、从编译安装到 aclnn 单算子调用与性能采集的完整实操流程以及基础/Preload 两套流水模板的同步机制与性能差异。一、算子概述本样例算子通过自定义算子工程实现位于仓库目录 experimental/matmul/weight_quant_batch_matmul_experiment。其 kernel 包含两个模板基础模板BASIC_MSD串行执行前处理 → 矩阵计算 → 后处理模块间通过全核同步保证数据不踩踏计算解耦流水模板PRELOAD_MSD让前处理计算单元连续执行两轮实现当前轮矩阵计算、上一轮后处理与下一轮前处理并行执行。工程内的 example 用例分别使用两个模板运行精度均正常示例同时通过 msprof 工具采集了两个模板的性能数据用于对比流水优化的收益。二、支持的 AI 处理器产品是否支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√从算子定义看weight_quant_batch_matmul_experiment_def.cpp 中通过this-AICore().AddConfig(ascend910b)注册了 AICore 配置算子本身基于 AIC AIV 异构流水ASCENDC_TPL_MIX_AIC_1_2即 1 个 AIC 搭配 2 个 AIV 的混合模板因此适用于具备独立 Cube 与 Vector 计算单元的 A2/A3 系列 NPU。三、目录结构介绍├── weight_quant_batch_matmul_experiment │ ├── examples // 样例工程aclnn 单算子调用用例 │ ├── op_host // tiling 与算子定义 │ └── op_kernel // 算子 kernel 实现进一步展开仓库中的实际文件布局examplessrc/main.cpp单算子 API 执行入口、run.sh一键编译/运行/校验/采集脚本、scripts/gen_data.py输入与真值生成、scripts/verify_result.py结果比对op_hostweight_quant_batch_matmul_experiment_def.cpp算子定义、weight_quant_batch_matmul_experiment_tiling.cppTiling 计算、weight_quant_batch_matmul_experiment_infershape.cppshape 推导op_kernelweight_quant_batch_matmul_experiment.cppkernel 入口、msd/目录下weight_quant_batch_matmul_experiment_msd_controller.h流水控制器、weight_quant_batch_matmul_experiment_cube_compute.hCube 矩阵计算、weight_quant_batch_matmul_experiment_vec_compute.hVector 前/后处理、weight_quant_batch_matmul_experiment_tool.h常量与工具figuresbasic_msd_flow.png与preload_msd_flow.png两张流水示意图。四、功能说明4.1 算子功能MSD 伪量化算法本样例算子实现 MatMul伪量化 A16W4 PerGroup场景并使用MSD 算法完成伪量化的计算过程。核心思路是对每个 group 内的激活矩阵用A_max rowMax(|A_group|)做归一化后通过乘 7.49 → 取整迭代 3 次把一组 float16 数据分解为 3 张 int4 残差矩阵分别对应高位、中位、低位信息使得伪量化误差逐级递减随后用 3 张 int4 激活矩阵分别与 int4 权重的 MatMul 结果按权重系数合并乘以A_max与反量化 scale 还原出高精度结果。其数学表达式如下计算 $A_{max}$$$A_{max} rowMax(|A_{group}|)$$计算 $tmp_{1}$$$tmp_{1} \frac{7.49 * A_{group}}{A_{max}}$$计算 $A_1$$$A_1 round(tmp_{1})$$计算 $tmp_{2}$$$tmp_{2}(tmp_{1}-A_{1})*14.98$$计算 $A_{2}$$$A_{2}round(tmp_{2})$$计算 $tmp_{3}$$$tmp_{3}(tmp_{2}-A_{2})*14.98$$计算 $A_{3}$$$A_{3}round(tmp_{3})$$构造矩阵 $A_{int}$$$ A_{int} \begin{bmatrix} A_{1} \ A_{2} \ A_{3} \ \end{bmatrix} $$计算 $Y_{int}$$$ Y_{int} \begin{bmatrix} Y_{1} \ Y_{2} \ Y_{3} \ \end{bmatrix} A_{int} \cdot Weight_{group} $$计算 $Y_{group}$$$ Y_{group} [(\frac{Y_{1}}{7.49}\frac{Y_{2}}{7.4914.98}\frac{Y_{3}}{7.4914.98*14.98})*A_{max}] * scale_{group} $$累加输出 $Y^i$$$Y^{i} Y_{group} Y^{i-1}$$其中scale_group即反量化 scaleantiquant_scale第 11 步表示按 group 依次累加最终得到完整的输出矩阵。仓库中的真值生成脚本 gen_data.py 完整实现了上述 11 步公式f1 7.49、f2 14.98逐 group 计算a_max、a1/a2/a3三个残差层的 MatMul 结果并加权累加可作为理解公式与验证算子的参照实现。kernel 侧 weight_quant_batch_matmul_experiment_vec_compute.h 中同样以multiFactors_[UNFLOD_TIMES] {7.49f, 14.98f, 14.98f}保存 MSD 展开系数与公式一一对应。4.2 算子规格算子类型(OpType)WeightQuantBatchMatmul算子输入nameshapedata typeformatx1M * Kfloat16NDx2K * Nint4NDantiquant_scaleGroupNum * Nfloat16ND算子输出yM * Nfloat16ND核函数名WeightQuantBatchMatmulExperiment算子定义源码 weight_quant_batch_matmul_experiment_def.cpp 与规格表一致xfloat16/ND、weightint4/ND、antiquant_scalefloat16/ND三个输入输出yfloat16/NDTiling 侧还会进一步校验输入必须为 2 维、x为 float16、weight为 int4、y为 float16否则报错退出。4.3 样例的默认运行规格examples 中 main.cpp 的CreateOpDesc与 gen_data.py 均以如下规格构造数据M 1N 12288K 8192groupSize 128即 GroupNum K / 128 64antiquant_scale 的 shape 为(64, 12288)激活x在[-3, 3]均匀采样float16权重在[-7, 7]采样int4 表示打包为 uint8 后每字节存两个 int4scale 在[-3, 3]采样float16。权重输入在样例中是以每字节两个 int4 打包的形式从input_weight.bin读入的与算子对 int4 权重存储格式的要求一致。真值golden.bin以 float16 写出供verify_result.py与算子输出output_z.bin比对。五、编译运行5.1 配置环境变量根据当前环境上 CANN 开发套件包toolkit 包 ops 包的安装方式选择对应配置环境变量的命令默认路径root 用户安装 CANN 软件包export ASCEND_INSTALL_PATH/usr/local/Ascend/cann默认路径非 root 用户安装 CANN 软件包export ASCEND_INSTALL_PATH$HOME/Ascend/cann指定路径 install_path安装 CANN 软件包export ASCEND_INSTALL_PATH${install_path}/cann5.2 编译与安装自定义算子包# 切换到工程根目录 cd ${git_clone_path} # 编译样例算子 run 包 bash build.sh --pkg --socascend910b --vendor_namecustom --opsweight_quant_batch_matmul_experiment --experimental # 安装自定义算子 run 包 ./build_out/cann-ops-nn-${vendor_name}-${arch}_linux.run其中--opsweight_quant_batch_matmul_experiment指定只编译该算子算子位于experimental/目录因此需加--experimental参数--socascend910b指定目标 SoC与算子定义中注册的 AICore 配置对应--vendor_namecustom指定 vendor 名安装后算子包位于$ASCEND_INSTALL_PATH/opp/vendors/custom_nn/op_api/下。5.3 编译 执行 aclnn 接口样例采集样例性能# 切换 weight_quant_batch_matmul_experiment aclnn 执行用例目录 cd ${git_clone_path}/experimental/matmul/weight_quant_batch_matmul_experiment/examples # 编译 执行 aclnn 接口 采集性能数据 bash run.sh # 切换 aclnn 用例性能数据目录 cd ${git_clone_path}/experimental/matmul/weight_quant_batch_matmul_experiment/examples/output/msprof_resultrun.sh的执行流程如下可对照脚本 run.sh 查看清理遗留的input/*.bin、output/*.bin与日志文件运行python3 scripts/gen_data.py生成输入数据input_a.bin、input_weight.bin、input_antiquant_scale_fp16.bin与真值数据output/golden.bin在build目录下cmake ../src并make编译出可执行文件execute_weight_quant_batch_matmul_experiment_op设置LD_LIBRARY_PATH指向$ASCEND_INSTALL_PATH/opp/vendors/custom_nn/op_api/lib执行算子并输出日志运行python3 scripts/verify_result.py output/output_z.bin output/golden.bin比对真值验证精度使用msprof --output./msprof_result ./execute_weight_quant_batch_matmul_experiment_op采集算子性能数据。5.4 aclnn 两段式接口调用自定义算子编译部署后会自动生成单算子 API两段式接口无需单算子描述文件即可直接调用// 第一段获取算子使用的 workspace 空间大小 aclnnStatus aclnnWeightQuantBatchMatmulExperimentGetWorkspaceSize(const aclTensor *a, const aclTensor *b, const aclTensor *bias, bool transposeX1, bool transposeX2, const aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor); // 第二段执行算子 aclnnStatus aclnnWeightQuantBatchMatmulExperiment(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream);其中第一段接口用于计算本次 API 调用所需的 workspace 内存大小获取后按workspaceSize申请 Device 侧内存再调用第二段接口执行计算。examples/src/main.cpp中InitResourceaclInit→aclrtSetDevice→aclrtGetRunMode与DestroyResourceaclrtResetDevice→aclFinalize展示了完整的 ACL 资源生命周期管理。六、流水设计6.1 基础流水模板该模板实现如下流水基础模板按前处理 → 矩阵计算 → 后处理串行推进前处理AIV 计算单元多个核并行产生后处理模块依赖的 $A_{max}$ 与矩阵计算模块依赖的 $A_{int}$即 UnfoldA 展开矩阵计算AIC/Cube 计算单元产生后处理模块依赖的 $Y_{int}$后处理AIV 计算单元读取 $Y_{int}$ 与 $A_{max}$、antiquant_scale按 MSD 公式完成合并并写回输出。由于上述模块之间存在数据依赖这些模块执行结束后都需要执行一次全核同步保证下个模块处理时上一个模块已完全使用完数据避免数据踩踏。在 weight_quant_batch_matmul_experiment_msd_controller.h 的BasicMsd中可以看到该模板的同步编排AIV 侧CrossCoreWaitFlag(SYNC_AIV_ONLY_AMAX_FLAG)等待 $A_{max}$ 写入UnfoldA完成后CrossCoreSetFlag(SYNC_AIV_ONLY_A_UNFOLD_FLAG)随后 AIC 侧等待SYNC_AIV_AIC_FLAG后执行LaunchMatmulAIV 再等待SYNC_AIC_AIV_FLAG执行MergeY全程通过CrossCoreSetFlag/CrossCoreWaitFlag成对的原语完成跨核同步且每个循环迭代内前处理、矩阵计算、后处理严格按序执行。6.2 Preload 流水模板该模板实现如下流水Preload计算解耦模板的核心思想是让前处理计算单元可以连续执行两轮从而实现当前轮的矩阵计算模板处理AIC 执行 MatMul上一轮的后处理模块AIV 执行 MergeY下一轮的前处理模块AIV 执行 UnfoldA。三者并行计算缩短关键路径上的等待时间。从源码看该模板通过**双缓冲double buffer**实现Process()中if constexpr (msdMode PRELOAD_MSD) { bufferNum DOUBLE_BUFFER_NUM; }对 $A_{max}$ 工作区、$Y_{int}$ 累加区均按cvLoopIdx_ % DOUBLE_BUFFER_NUM轮转展开后的激活矩阵aUnfold使用3 份缓冲cvLoopIdx_ % 3轮转覆盖当前轮 AIC 正在消费、上一轮 AIV 刚写入、下一轮 AIV 即将写入三个角色的同时存在PreloadMsd中 AIV 每轮先做本轮的UnfoldA若cvLoopIdx_ 0再合并上一轮的Y_intAIC 侧LaunchMatmul与等待SYNC_AIC_ONLY_AUNFOLD_FLAG的顺序保证上一轮 AIC 消费完毕后才复用展开缓冲。TilingKey 注册处weight_quant_batch_matmul_experiment_tiling_key.h定义了BASIC_MSD 0与PRELOAD_MSD 1两个模板MSD_MODE以 4 bit 编码进 TilingKeyHost 侧 weight_quant_batch_matmul_experiment_tiling.cpp 的PostTiling中默认使用BASIC_MSD注释明确可切换至 PRELOAD_MSD 流水通过GET_TPL_TILING_KEY(msdTemplate)生成对应模板的 TilingKey再context-SetTilingKey(tilingKey)实现模板的编译期选择。因此若要跑 Preload 模板只需将PostTiling中的msdTemplate改为PRELOAD_MSD重新编译算子包即可。七、性能分析Preload 流水模板通过流水优化可以在相同的计算逻辑下实现显著的性能提升。基于 Atlas A2 进行测试可以得到如下数据模板类型基础流水模板(us)preload流水模板(us)算子耗时396303即在相同的计算逻辑与样例规格M1、N12288、K8192、groupSize128下Preload 流水模板将算子耗时从基础模板的396 us降低到303 us提升约 23%。需要说明的是该数据为当前样例在 Atlas A2 上的测试结果实际收益会随矩阵规模、核数、groupSize 及 MSD 展开次数变化。复现方法分别将 weight_quant_batch_matmul_experiment_tiling.cpp 中的msdTemplate设为BASIC_MSD/PRELOAD_MSD编译算子包再执行bash run.sh即可在examples/output/msprof_result中获取两套模板的 msprof 性能数据并对比。八、关键实现细节补充8.1 Tiling 参数解析weight_quant_batch_matmul_experiment_tiling.cpp 中DoOpTiling的核心 Tiling 逻辑如下usedCoreNum aicNum按可用 AIC 核数设置 BlockDimcontext-SetBlockDim(...)singleCoreM 3 * M由于 MSD 把激活展开为 3 张 int4 矩阵实际执行 MatMul 的 shape 在 M 轴展开 3 倍singleCoreN 512示例代码中 N 默认按 512 切分nGmOffset baseN * cubeBlockIdx表示各核在 N 方向均分baseK 128groupSize 固定为 128与样例数据规格一致baseM CeilAlign(singleCoreM, 16)、depthA1/depthB1 2分别做对齐与流水深度设置dbL0C 1开启 L0C 双缓冲。此外GetWorkspaceSize为算子预留 workspace系统侧通过GetLibApiWorkSpaceSize获取用户侧统一预留 64 MB合计作为算子运行时可用的全局工作区。8.2 Kernel 入口与 workspace 布局kernel 入口 通过WQBMM_EXP_IMPL_CLASS宏实例化WqbmmExpMsdControllerxType, wType, MSD_MODE模板参数MSD_MODE由 TilingKey 在编译期决定从而让同一份 kernel 源码编译出基础/Preload 两套模板实现。Init中对 workspace 的布局对应 weight_quant_batch_matmul_experiment_msd_controller.haMaxWorkspaceGlobal_存放每行 $A_{max}$按 512B 对齐Preload 模式下分配DOUBLE_BUFFER_NUM份aUnfoldS4Global_ / aUnfoldS8Global_同一 GM 地址的两种视图存放展开后的 int4 激活S8 视图用于 Cube 读取按 3 份轮转yS32Gm_存放 Cube 输出的 int32 中间结果 $Y_{int}$Preload 模式下同样按双缓冲轮转。8.3 精度验证闭环样例的精度验证是端到端闭环的gen_data.py用 numpy 严格按 MSD 公式生成 float16 真值golden.bin算子输出output_z.bin经verify_result.py与真值比对。仓库 README 中明确说明example 的用例分别使用两个模板的精度均正常因此可以在切换模板后通过同一套数据与脚本验证两套模板的数值正确性再配合 msprof 数据对比性能。总结WeightQuantBatchMatmulExperiment 是一个结构完整、可直接编译运行的 A16W4 PerGroup 伪量化 MatMul 自定义算子样例它以 MSD 三级分解算法为核心将浮点激活按 group 展开为多张 int4 残差矩阵参与 Cube 整数矩阵乘再由 Vector 单元按权重系数合并还原同一套代码通过 TilingKey 模板机制同时提供基础流水与 Preload 流水两套实现前者以全核同步保证正确性、结构清晰后者以双缓冲 三缓冲实现当前轮矩阵计算、上一轮后处理、下一轮前处理三级并行在 Atlas A2 上将算子耗时从 396 us 优化至 303 us。对于希望研究 NPU 上伪量化算子实现、Cube/Vector 异构流水设计或 aclnn 单算子调用流程的开发者本样例是一份可直接复用的工程参考。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考