WeightQuantBatchMatmulExperiment 算子深度解析：A16W4 PerGroup 伪量化 MatMul 的 MSD 算法与双模板流水设计

WeightQuantBatchMatmulExperiment 算子深度解析：A16W4 PerGroup 伪量化 MatMul 的 MSD 算法与双模板流水设计

【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn

WeightQuantBatchMatmulExperiment 是 CANN ops-nn 仓库（experimental/matmul/weight_quant_batch_matmul_experiment）下的一个自定义算子工程样例，用于在 NPU 上实现A16W4 PerGroup 场景的 MatMul 伪量化计算。本算子以 MSD（Multi-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 处理器

从算子定义看，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 实现

进一步展开仓库中的实际文件布局：

• examples：src/main.cpp（单算子 API 执行入口）、run.sh（一键编译/运行/校验/采集脚本）、scripts/gen_data.py（输入与真值生成）、scripts/verify_result.py（结果比对）；
• op_host：weight_quant_batch_matmul_experiment_def.cpp（算子定义）、weight_quant_batch_matmul_experiment_tiling.cpp（Tiling 计算）、weight_quant_batch_matmul_experiment_infershape.cpp（shape 推导）；
• op_kernel：weight_quant_batch_matmul_experiment.cpp（kernel 入口）、msd/目录下weight_quant_batch_matmul_experiment_msd_controller.h（流水控制器）、weight_quant_batch_matmul_experiment_cube_compute.h（Cube 矩阵计算）、weight_quant_batch_matmul_experiment_vec_compute.h（Vector 前/后处理）、weight_quant_batch_matmul_experiment_tool.h（常量与工具）；
• figures：basic_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 还原出高精度结果。其数学表达式如下：

1. 计算 $A_{max}$：

$$A_{max} = rowMax(|A_{group}|)$$

2. 计算 $tmp_{1}$：

$$tmp_{1} = \frac{7.49 * A_{group}}{A_{max}}$$

3. 计算 $A_1$：

$$A_1 = round(tmp_{1})$$

4. 计算 $tmp_{2}$：

$$tmp_{2}=(tmp_{1}-A_{1})*14.98$$

5. 计算 $A_{2}$：

$$A_{2}=round(tmp_{2})$$

6. 计算 $tmp_{3}$：

$$tmp_{3}=(tmp_{2}-A_{2})*14.98$$

7. 计算 $A_{3}$：

$$A_{3}=round(tmp_{3})$$

8. 构造矩阵 $A_{int}$：

$$ A_{int} = \begin{bmatrix} A_{1} \ A_{2} \ A_{3} \ \end{bmatrix} $$

9. 计算 $Y_{int}$：

$$ Y_{int} = \begin{bmatrix} Y_{1} \ Y_{2} \ Y_{3} \ \end{bmatrix} = A_{int} \cdot Weight_{group} $$

10. 计算 $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} $$

11. 累加输出 $Y^i$：

$$Y^{i} = Y_{group} + Y^{i-1}$$

其中scale_group即反量化 scale（antiquant_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 算子规格

算子定义源码 weight_quant_batch_matmul_experiment_def.cpp 与规格表一致：x（float16/ND）、weight（int4/ND）、antiquant_scale（float16/ND）三个输入，输出y（float16/ND）；Tiling 侧还会进一步校验输入必须为 2 维、x为 float16、weight为 int4、y为 float16，否则报错退出。

4.3 样例的默认运行规格

examples 中 main.cpp 的CreateOpDesc与 gen_data.py 均以如下规格构造数据：

• M = 1，N = 12288，K = 8192；
• groupSize = 128，即 GroupNum = K / 128 = 64，antiquant_scale 的 shape 为(64, 12288)；
• 激活x在[-3, 3]均匀采样（float16），权重在[-7, 7]采样（int4 表示，打包为 uint8 后每字节存两个 int4），scale 在[-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}/cann

5.2 编译与安装自定义算子包

# 切换到工程根目录 cd ${git_clone_path} # 编译样例算子 run 包 bash build.sh --pkg --soc=ascend910b --vendor_name=custom --ops=weight_quant_batch_matmul_experiment --experimental # 安装自定义算子 run 包 ./build_out/cann-ops-nn-${vendor_name}-${arch}_linux.run

其中：

• --ops=weight_quant_batch_matmul_experiment指定只编译该算子（算子位于experimental/目录，因此需加--experimental参数）；
• --soc=ascend910b指定目标 SoC，与算子定义中注册的 AICore 配置对应；
• --vendor_name=custom指定 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_result

run.sh的执行流程如下（可对照脚本 run.sh 查看）：

1. 清理遗留的input/*.bin、output/*.bin与日志文件；
2. 运行python3 scripts/gen_data.py生成输入数据（input_a.bin、input_weight.bin、input_antiquant_scale_fp16.bin）与真值数据（output/golden.bin）；
3. 在build目录下cmake ../src并make，编译出可执行文件execute_weight_quant_batch_matmul_experiment_op；
4. 设置LD_LIBRARY_PATH指向$ASCEND_INSTALL_PATH/opp/vendors/custom_nn/op_api/lib，执行算子并输出日志；
5. 运行python3 scripts/verify_result.py output/output_z.bin output/golden.bin比对真值，验证精度；
6. 使用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中InitResource（aclInit→aclrtSetDevice→aclrtGetRunMode）与DestroyResource（aclrtResetDevice→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后执行LaunchMatmul，AIV 再等待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_int；AIC 侧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 编码进 TilingKey；Host 侧 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 进行测试，可以得到如下数据：

即在相同的计算逻辑与样例规格（M=1、N=12288、K=8192、groupSize=128）下，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 核数设置 BlockDim，context->SetBlockDim(...)）；
• singleCoreM = 3 * M：由于 MSD 把激活展开为 3 张 int4 矩阵，实际执行 MatMul 的 shape 在 M 轴展开 3 倍；
• singleCoreN = 512：示例代码中 N 默认按 512 切分，nGmOffset = baseN * cubeBlockIdx表示各核在 N 方向均分；
• baseK = 128：groupSize 固定为 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宏实例化WqbmmExpMsdController<xType, wType, MSD_MODE>，模板参数MSD_MODE由 TilingKey 在编译期决定，从而让同一份 kernel 源码编译出基础/Preload 两套模板实现。

Init中对 workspace 的布局（对应 weight_quant_batch_matmul_experiment_msd_controller.h）：

• aMaxWorkspaceGlobal_：存放每行 $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