aclnnFlatQuantV2 完整指南从两级矩阵乘到 FP4 E2M1 量化融合【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nnaclnnFlatQuantV2 是 CANN ops-nn 库里一个计算 压缩二合一的融合算子它先对输入张量做两次小矩阵乘法先右乘 kroneckerP2、再左乘 kroneckerP1紧接着对乘积做 per-token 或 per-group 量化直接吐出 INT4、INT32 或 FP4 E2M1FLOAT4_E2M1低精度数据。如果你想在 Ascend 950 系列上把随机旋转矩阵变换和量化打包成一步完成就是它了。30 秒选版本V2 和 V1 的能力几乎完全相同唯一的分水岭是dstTypeMax这个参数对比维度V1aclnnFlatQuantV2aclnnFlatQuantV2矩阵乘 量化融合支持支持per-token 量化INT4/INT32 输出支持支持per-group 量化FLOAT4_E2M1 输出支持但目标精度最大值不可调固定取该精度能表示的最大值支持且可通过dstTypeMax自定义目标数据类型最大值dstTypeMax取值无此参数0.0关闭或 6.0~12.0仅 FLOAT4_E2M1 场景生效怎么选给你三条直接的建议需要在 FP4 E2M1 量化时把目标最大值压到 6.0~12.0 之间的某个值 —— 只能用 V2只关心 INT4/INT32 输出或者 per-group 场景下用默认的最大值就够 —— V1、V2 随便挑目标是更精细地控制量化范围、在精度和性能之间做权衡 —— 推荐 V2。工作原理从矩阵乘到量化整个算子可以拆成两段前半段是两次小矩阵乘后半段是按模式分叉的量化逻辑。两级小矩阵乘先把输入张量x的 shape 记为 [K, M, N]。第一级乘法的公式是$$ x x kroneckerP2 $$用大白话说kroneckerP2 是一个 [N, N] 的方阵从右边作用到x上相当于对每一行向量做一次随机旋转式的线性变换变换后的中间结果就是x。第二级乘法紧跟其后$$ x kroneckerP1x $$这一步里 kroneckerP1 是 [M, M] 的方阵从左边作用对x的行方向再做一次同样的变换。至此x的形状仍是 [K, M, N]量化阶段的所有计算都基于它。per-token 模式INT4 / INT32 输出per-token 的思路很朴素x的第 0 维即 K 维上每一个token一个 [M, N] 的矩阵共享一个量化因子。第一步沿 0 维对每个 token 求最大绝对值再除以 (7 / clipRatio) 得到 quantScale$$ quantScale [max(abs(x[0,:,:])),max(abs(x[1,:,:])),...,max(abs(x[K,:,:]))]/(7 / clipRatio) $$这一步在做什么把每个 token 的数据动态地压进 INT4 的表示区间。除以 (7 / clipRatio) 而不是除以 15本质是留了裁剪余量——clipRatio越小动态范围越保守越不容易溢出。第二步直接用量化因子做除法得到输出$$ out x / quantScale $$结果就是每个 token 一组缩放后的低精度数据quantScale 本身作为 [K] 形状的 FLOAT32 张量单独输出供后续反量化使用。per-group 模式FLOAT4_E2M1 输出per-group 量化对应 FP4 E2M1 输出。这个模式更复杂行为完全由dstTypeMax的取值决定分三种情况。情况一dstTypeMax 0.0默认 mx_quantize 路径先把x的 [K, M, N] reshape 成 [K, M*N]记为 x2在第二维上按 groupSize32 分组每组 32 个元素 {e0, e1, ..., e31}算出组内最大值$$ emax max(e0,e1....e31) $$这一步相当于给每组定一个标杆值后面所有缩放都围绕它展开。计算 reduceMaxValue 和 sharedExp$$ reduceMaxValue log2(reduceMax(x2),groupSize32) $$$$ sharedExp[K,M*N/32] reduceMaxValue -emax $$在 log2 域里做差得到每组的共享指数这是把每组的最大值对齐到 2 的整数次幂的关键。由共享指数还原出量化因子$$ quantScale[K,M*N/32] 2 ^ {sharedExp[K,M*N/32]} $$每组共享同一个 quantScale逐组做除法得到输出$$ out x2 / quantScale $$此时 quantScale 以 FLOAT8_E8M0 格式、[K, ceilDiv(M*N,64), 2] 的形状输出是典型的 MXmicro-scaling布局。情况二dstTypeMax 6.0 / 7.0这一档走的是带尾数修正的动态量化路径把输入沿 axis 维按 k blocksize 分组每组 k 个数 ${V_i}{i1}^{k}$动态量化为 ${mxscale, {P_i}{i1}^{k}}$其中 k blocksize组共享指数的计算规则带一个尾数检查的分支$$ shared_exp \begin{cases} ceil(log_2(max_i(|V_i|))) - emax, \text{如果} 尾数位的高比特前一/两位 \text{为1且尾数不全为0} \ floor(log_2(max_i(|V_i|))) - emax, \text{其他} \end{cases} $$为什么要分 ceil 和 floor 两种情况当最大值落在 FP4 尾数容易向上冲的区间时用 ceil 可以防止量化后溢出其余情况用 floor 更省动态范围。组内每个元素按各自位置量化$$ P_i cast_to_dst_type(V_i/mxscale, round_mode), \space i\space from\space 1\space to\space blocksize $$量化后的 $P_i$ 按对应 $V_i$ 的位置拼成输出 yOutmxscale 按 axis 维上的分组拼成 mxscaleOut。情况三dstTypeMax 取 8.0~12.0 等其他值这是把 FP32 数据映射到 FP8 表示范围内的块缩放路径步骤如下把长向量按 k 个元素一块切开每块独立计算缩放因子 $S_{fp32}^b$再用同一个因子把块内所有元素映射到低精度类型。最后一块不足 k 个时缺失值视为 0按完整块处理先找块内最大绝对值$$ Amax(D_{fp32}^b)max({|d_{i}|}_{i1}^{k}) $$这一步给出该块最大能有多冲。再把它归一化到目标类型可表示的范围。当 dst_max_value 0 时$Amax(DType)$ 取目标精度能表示的最大值当 dst_max_value ! 0 时直接取你传入的 dstTypeMax$$ S_{fp32}^b \frac{Amax(D_{fp32}^b)}{Amax(DType)} $$这就是 V2 新增dstTypeMax的意义所在把分母换成你指定的值量化范围立即跟着变。接下来把浮点缩放因子 $S_{fp32}^b$ 转成目标格式下能表示的缩放值 $S_{ue8m0}^b$先提取它的无偏指数 $E_{int}^b$ 和尾数 $M_{fixp}^b$然后为了保证量化时不溢出对指数做有条件的向上取整$$ E_{int}^b \begin{cases} E_{int}^b 1, \text{如果} S_{fp32}^b \text{为正规数且} E_{int}^b 254 \text{且} M_{fixp}^b 0 \ E_{int}^b, \text{否则} \end{cases} $$由取整后的指数定出块缩放因子与块转换因子$$ S_{ue8m0}^b 2^{E_{int}^b} $$$$ R_{fp32}^b \frac{1}{fp32(S_{ue8m0}^b)} $$最后落到每个元素上$d^i DType(d_{fp32}^i \cdot R_{fp32}^n)$最终输出是 $\left(S^b, [d^i]{i1}^k\right)$ 这样的块缩放因子 块内量化数据组合其中 $S^b$ 即 $S{ue8m0}^b$$[d^i]_{i1}^k$ 是块内量化结果。量化后的 $P_i$ 按对应 $V_i$ 的位置组成输出 yOutmxscale 按 axis 维分组组成 mxscaleOut。硬件支持情况这个算子的硬件覆盖面比较聚焦目前只有 Ascend 950 系列支持产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练 / 推理系列产品×Atlas A2 训练 / 推理系列产品×Atlas 200I / 500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×如果你不在 950 系列上可以看看同目录下功能相似的 aclnnFlatQuantV1其硬件覆盖面更广。两段式调用接口CANN 的 aclnn 接口统一采用两段式设计原理详见 两段式接口。为什么要拆两段因为第一段要完成入参校验、shape 推导和 tiling 规划这些信息的产物是一个 executor 对象和一个 workspace 大小你拿到它们、把 Device 侧内存准备好之后第二段才真正调度 kernel 干活。这样设计也让同一次调用可以复用 executor避免重复规划。先调用第一段接口拿到 workspace 大小和执行器aclnnStatus aclnnFlatQuantV2GetWorkspaceSize( const aclTensor *x, const aclTensor *kroneckerP1, const aclTensor *kroneckerP2, double clipRatio, double dstTypeMax, aclTensor *out, aclTensor *quantScale, uint64_t *workspaceSize, aclOpExecutor **executor)再调用第二段接口执行计算aclnnStatus aclnnFlatQuantV2( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)两个接口都返回aclnnStatus状态码各码值含义参见 aclnn 返回码。参数速查aclnnFlatQuantV2GetWorkspaceSize参数名方向说明与约束数据类型shape 维度x输入原始数据对应公式中的 x• 不允许空 Tensor• shape 为 [K, M, N]K ∈ [1, 262144]M、N ∈ [1, 256]• out 为 INT32 时 N 必须是 8 的倍数8 个数打包进一个 32 位字• out 为 INT4 或 FLOAT4_E2M1 时 N 必须是偶数INT4 两两打包进字节FLOAT16、BFLOAT163kroneckerP1输入左乘矩阵对应公式中的 kroneckerP1• 不允许空 Tensor• shape 为 [M, M]M 与 x 的 M 维一致• 数据类型必须与 x 相同FLOAT16、BFLOAT162kroneckerP2输入右乘矩阵对应公式中的 kroneckerP2• 不允许空 Tensor• shape 为 [N, N]N 与 x 的 N 维一致• 仅 per-group 场景允许 [0, 0]含义是跳过右乘、只做左乘• 数据类型必须与 x 相同FLOAT16、BFLOAT162clipRatio输入量化裁剪比例• 取值范围 (0, 1]超出即报错--dstTypeMax输入目标数据类型最大值对应公式中的 Amax(DType)• 0.0 表示不使用• 6.0~12.0 表示自定义最大值仅 FLOAT4_E2M1 场景生效DOUBLE-out输出量化后的低精度数据• 不允许空 Tensor• INT32尾轴为 x 尾轴的 1/8其余维一致• INT4shape 与 x 一致• FLOAT4_E2M1shape 为 [K, M*N]INT4、INT32、FLOAT4_E2M13 或 2quantScale输出量化因子• 不允许空 Tensor• INT4/INT32 场景shape [K]FLOAT32• FLOAT4_E2M1 场景shape [K, ceilDiv(M*N,64), 2]FLOAT8_E8M0FLOAT32、FLOAT8_E8M01 或 3workspaceSize输出返回需要在 Device 侧申请的 workspace 大小--executor输出返回包含算子计算流程的执行器--所有 Tensor 均要求 ND 数据格式且支持非连续 Tensor。aclnnFlatQuantV2执行接口参数名方向说明workspace输入Device 侧申请的 workspace 内存地址workspaceSize输入workspace 大小取自第一段接口的返回值executor输入第一段接口生成的 op 执行器stream输入指定执行任务的 aclrtStream错误码速查第一段接口在入参校验阶段就可能返回错误全部集中在两个码上返回码错误码触发场景ACLNN_ERR_PARAM_NULLPTR161001必选输入 x、kroneckerP1、kroneckerP2 或必选输出 out、quantScale 是空指针ACLNN_ERR_PARAM_INVALID161002以下任一校验不通过• x、kroneckerP1、kroneckerP2、out、quantScale 的数据类型或数据格式不在支持范围• kroneckerP1、kroneckerP2 与 x 的数据类型不一致• x 的维度不为 3• x 第一维超出 [1, 262144]或第二维/第三维超出 [1, 256]• kroneckerP1 维度不为 2或其两维与 x 的 M 维不一致• per-token 场景下 kroneckerP2 维度不为 2或其两维与 x 的 N 维不一致• INT4/INT32 场景下 quantScale 维度不为 1或第一维与 x 的 K 维不一致• FLOAT4_E2M1 场景下 quantScale 维度不为 3或第一维与 x 的 K 维不一致• clipRatio 超出 (0, 1]• out 为 INT4 时 x 尾轴不是偶数或 x 与 out 的 shape 不一致• out 为 INT32 时 x 尾轴不是 out 尾轴的 8 倍或两者非尾轴大小不一致完整示例下面是一个可直接运行的 INT32 输出per-token 模式完整样例shape 取 KMN16。编译与运行方法参见 编译与运行样例完整版本可对照 matmul/flat_quant/examples/ 下的test_aclnn_flat_quant_v2.cpp。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_flat_quant_v2.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t size 1; for (auto i : shape) { size * i; } return size; } // 固定写法设备与流初始化 int Init(int32_t deviceId, aclrtStream* stream) { auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, return ret); return 0; } // 通用工具device 内存分配 host 数据拷贝 创建 aclTensor 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, return ret); ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, 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]; } *tensor aclCreateTensor( shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1. 初始化 device/stream int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret 0, return ret); // 2. 构造输入输出x[16,16,16] float16INT32 输出尾轴 16/82 std::vectorint64_t xShape {16, 16, 16}; std::vectorint64_t kroneckerP1Shape {16, 16}; std::vectorint64_t kroneckerP2Shape {16, 16}; std::vectorint64_t outShape {16, 16, 2}; std::vectorint64_t quantScaleShape {16}; void* xDeviceAddr nullptr; void* kroneckerP1DeviceAddr nullptr; void* kroneckerP2DeviceAddr nullptr; void* outDeviceAddr nullptr; void* quantScaleDeviceAddr nullptr; aclTensor* x nullptr; aclTensor* kroneckerP1 nullptr; aclTensor* kroneckerP2 nullptr; aclTensor* out nullptr; aclTensor* quantScale nullptr; double clipRatio 1.0; double dstTypeMax 0.0; std::vectoraclFloat16 xHostData(16 * 16 * 16, aclFloatToFloat16(1)); std::vectoraclFloat16 kroneckerP1HostData(16 * 16, aclFloatToFloat16(1)); std::vectoraclFloat16 kroneckerP2HostData(16 * 16, aclFloatToFloat16(1)); std::vectorint32_t outHostData(16 * 16 * 2, 1); std::vectorfloat quantScaleHostData(16, 0); // 依次创建五个 aclTensor ret CreateAclTensor(xHostData, xShape, xDeviceAddr, aclDataType::ACL_FLOAT16, x); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor( kroneckerP1HostData, kroneckerP1Shape, kroneckerP1DeviceAddr, aclDataType::ACL_FLOAT16, kroneckerP1); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor( kroneckerP2HostData, kroneckerP2Shape, kroneckerP2DeviceAddr, aclDataType::ACL_FLOAT16, kroneckerP2); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_INT32, out); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor( quantScaleHostData, quantScaleShape, quantScaleDeviceAddr, aclDataType::ACL_FLOAT, quantScale); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 第一段接口入参校验 计算 workspace 大小生成 executor uint64_t workspaceSize 0; aclOpExecutor* executor; ret aclnnFlatQuantV2GetWorkspaceSize( x, kroneckerP1, kroneckerP2, clipRatio, dstTypeMax, out, quantScale, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, return ret); // 4. 按第一段返回的大小申请 device workspace void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, return ret); } // 5. 第二段接口真正执行两级矩阵乘 量化 ret aclnnFlatQuantV2(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, return ret); // 6. 同步流等待 kernel 结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, return ret); // 7. 把结果拷回 host 查看 auto size GetShapeSize(outShape); std::vectorint32_t resultData(size, 0); ret aclrtMemcpy( resultData.data(), size * sizeof(int32_t), outDeviceAddr, size * sizeof(int32_t), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, return ret); for (int64_t i 0; i size; i) { std::cout result[ i ] is: resultData[i] std::endl; } auto quantScaleSize GetShapeSize(quantScaleShape); std::vectorfloat quantScaleResultData(quantScaleSize, 0); ret aclrtMemcpy( quantScaleResultData.data(), quantScaleSize * sizeof(float), quantScaleDeviceAddr, quantScaleSize * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, return ret); for (int64_t i 0; i quantScaleSize; i) { std::cout quantScale[ i ] is: quantScaleResultData[i] std::endl; } // 8. 释放 aclTensor 与 device 内存 aclDestroyTensor(x); aclDestroyTensor(kroneckerP1); aclDestroyTensor(kroneckerP2); aclDestroyTensor(out); aclDestroyTensor(quantScale); aclrtFree(xDeviceAddr); aclrtFree(kroneckerP1DeviceAddr); aclrtFree(kroneckerP2DeviceAddr); aclrtFree(outDeviceAddr); aclrtFree(quantScaleDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }⚠️ 约束与常见坑确定性计算aclnnFlatQuantV2 默认就是确定性实现同一输入多次运行结果一致不需要额外配置N 的偶数约束INT4 和 FLOAT4_E2M1 都要求 x 的 N 维为偶数——每 2 个数才够打包进一个字节奇数维没法对齐第一段接口会直接返回 161002N 的 8 倍约束INT32 输出时 x 尾轴必须是 out 尾轴的 8 倍因为 8 个 4 位整数恰好装进一个 32 位字K、M、N 上限K ≤ 262144M、N ≤ 256。超出这个范围要回到 [1, 262144] / [1, 256] 区间内重新设计形状数据类型必须对齐x、kroneckerP1、kroneckerP2 三个输入必须同为 FLOAT16 或同为 BFLOAT16混着用会报 161002dstTypeMax 只服务 FP4 E2M1只有 FLOAT4_E2M1 输出时设置才有意义取值 0.0 或 6.0~12.0INT4/INT32 场景传非零值属于误用per-group 场景的 [0, 0] 技巧只想做左乘时可以把 kroneckerP2 传成 [0, 0] 的张量来跳过右乘但这一特性只在 per-group 场景下合法per-token 场景下 kroneckerP2 必须是完整的 [N, N] 矩阵。总结aclnnFlatQuantV2 把两次小矩阵乘 动态量化压进一个 kernel是 950 系列上做 FP4 E2M1 权重/激活压缩的核心工具V2 相对 V1 的增量是dstTypeMax让你能自定义目标精度最大值换取更精细的量化范围控制。算子源码位于 matmul/flat_quant/ 目录官方接口文档见 matmul/flat_quant/docs/aclnnFlatQuantV2.md如需拉取整个仓库可执行git clone https://gitcode.com/cann/ops-nn。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
