CANN cann-samples weightnz 样例实战:用 FRACTAL_NZ 权重预处理消除 Matmul 随路格式转换开销 CANN cann-samples weightnz 样例实战用 FRACTAL_NZ 权重预处理消除 Matmul 随路格式转换开销【免费下载链接】cann-samplesCANN高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples本篇技术文章基于 cann-samples 仓库中 Samples/1_Features/instruction_optimization/weightnz/README.md 样例文档展开围绕昇腾 NPU 的 FRACTAL_NZNZ数据排布特性讲解其背后的硬件原理、权重预处理转换方法以及如何在 Matmul 算子中通过GM 按 NZ 排布存放权重 L1 直接按 NZ 布局消费的方式消除随路格式转换、提升数据搬运效率。读完本文后你将能够理解 NZ 排布与 ND 排布在块级组织上的差异掌握 ND 转 NZ 的转换函数与维度对齐规则完整编译运行 weightnz 样例并通过 msprof 观察优化前后的流水收益。一、背景为什么推理场景要预转 FRACTAL_NZ 排布昇腾 NPU 内部 Cube 计算单元亲和的权重排布形式是FRACTAL_NZ分形 NZ 排布。当 Matmul 算子以传统的 ND行主序排布传入权重时数据通路需要在搬运过程中随路完成 ND → FRACTAL_NZ 的格式转换这个过程会引入额外的带宽损耗使 MTE搬运引擎的带宽无法被充分利用。因此在推理场景下一个典型的优化手段是对权重进行离线预处理在模型部署前就把权重从 ND 排布转换为 FRACTAL_NZ 排布使 GMGlobal Memory片外内存中存放的权重已经是 Cube 单元可直接消费的格式。这样在算子运行时MTE2 只需将 NZ 排布数据按基本块粒度搬运到 L1不再承担格式转换的开销带宽利用率得以充分发挥。需要强调的是该方法仅适用于推理场景。在训练场景中权重需要在反向传播中不断更新无法提前做离线预处理此时若把 ND 转 FRACTAL_NZ 放在算子内部随路完成反而会引入额外耗时得不偿失。二、原理FRACTAL_NZ 的基本块与块内/块间组织方式NPU 内部 FRACTAL_NZ 的基本块大小随数据类型变化这是理解对齐规则的关键数据类型基本块大小M × N 维度FP16 / BF1616 × 16FP3216 × 8即保持 M 维度块大小 16 不变FP32 时由于元素宽度是 FP16/BF16 的两倍N 维度块大小减半K 维度的分块策略保持不变。排布组织方式上ND 排布数据在内存中按行主序连续存储FRACTAL_NZ 排布以基本块为单位组织数据基本块之间按 N 字形列优先顺序排列基本块内部按 Z 字形Morton 顺序排列。这种N 间 Z 内的组织方式与 Cube 计算单元逐块取数的访存模式对齐使得搬运数据可以拿来即用无需在片上做块内重排。三、样例代码剖析权重按 NZ 排布创建与存放weightnz 样例的主体代码位于 Samples/1_Features/instruction_optimization/weightnz/main.asc实现了一个带 GM→L1、L1→L0 双缓冲流水的 Matmul 内核其中权重 B 在 GM 与 L1 两侧均按 NZ 排布描述。3.1 核心改动GM→L1 搬运的排布描述文档给出的关键改动是将原本按 ND 排布存储的权重矩阵 B改为按 FRACTAL_NZ 排布创建和存放。文档中的改动示意如下// 为矩阵 B 创建 FRACTAL_NZ 排布K 维度上的分块大小为 cube_size即 32/sizeof(T)N 维度上的分块大小为 16 auto layoutB AscendC::Te::MakeNzLayoutT( tool::CeilAlign(k, tool::CUBE_SIZE / sizeof(T)), // 将 K 维度对齐到 cube_size 的倍数 tool::CeilAlign(n, 16) // 将 N 维度对齐到 16 的倍数 ); auto tensorBgm AscendC::Te::MakeTensor(AscendC::Te::MakeGMmemPtr(reinterpret_cast__gm__ T*(bGm)), layoutB);在仓库实际源码中这一改动对应 main.asc 中的布局描述auto layoutA AscendC::Te::MakeFrameLayoutAscendC::Te::NDExtLayoutPtn(m, k); // Create NZ layout for matrix B: tile size cube_size (32/sizeof(T)) in K-dim, 16 in N-dim auto layoutB AscendC::Te::MakeFrameLayoutAscendC::Te::NZLayoutPtn(tool::CeilAlign(k, tool::CUBE_SIZE / sizeof(T)), tool::CeilAlign(n, 16)); auto layoutC AscendC::Te::MakeFrameLayoutAscendC::Te::NDExtLayoutPtn(m, n);两个关键改动点排布创建方式权重矩阵从 GM 搬运至 L1 时存储排布由默认的行主序 ND 排布NDExtLayoutPtn切换为 FRACTAL_NZ 排布NZLayoutPtn。由于 GM 中的数据已经是离线转好的 NZ 格式MTE2 的CopyGM2L1只需按块搬运无需随路转换nd 转 nz 函数转换本身不在算子内核中完成而是由主机侧脚本离线执行见下文 3.2 节。此外从源码结构看样例在 L1→L0 的搬运中也做了配套处理L0A 使用NZLayoutPtn描述L0B 使用ZNLayoutPtn即 NZ 的转置视角描述分别对应 Cube 计算时 A/B 矩阵在 L0 的亲和取数形式与Mmad指令的输入要求相匹配参见 main.asc。3.2 主机侧权重 ND → NZ 的离线转换nd 转 nz 函数在样例中由 scripts/gen_data.py 的index_nd_to_nz_torch实现转换逻辑可概括为四步cube_size 32 // element_size # 元素大小决定 K 方向块宽bf16 时为 16 ceil_k ((k cube_size - 1) // cube_size) * cube_size # K 对齐到 cube_size 倍数 ceil_n ((n 15) // 16) * 16 # N 对齐到 16 的倍数 padded torch.zeros((ceil_k, ceil_n), ...) # 零填充对齐 padded[:k, :n] matrix num_k_tiles ceil_k // cube_size num_n_tiles ceil_n // 16 reshaped padded.reshape(num_k_tiles, cube_size, num_n_tiles, 16) shuffled reshaped.permute(2, 0, 1, 3) # [16, cube_size, num_k_tiles, num_n_tiles] return shuffled.flatten()即先将 K、N 两维做零值 Padding 对齐再按基本块内部 Z 字形、基本块之间 N 字形的顺序重排后展平。这与第二节描述的排布组织方式一一对应也印证了块间 N 字形先 N 后 K 的块序、块内 Z 字形的存储顺序。主机侧主程序在分配内存时同样以对齐后尺寸为准避免越界访问参见 main.ascint64_t alignedSize CUBE_SIZE / sizeof(bfloat16_t); int64_t alignedK (k alignedSize - 1) / alignedSize * alignedSize; int64_t alignedN (n 16 - 1) / 16 * 16; std::vectoruint16_t hostInput(m * k, 0); std::vectoruint16_t hostWeightNz(alignedK * alignedN, 0); // 权重按对齐后尺寸分配随后样例调用gen_data.py生成随机输入A 为 ND 排布、B 已转为 NZ 排布并写入input/input_b.bin经aclrtMemcpy拷贝到设备后以内核MatmulKernelbfloat16_t下发执行最后由 scripts/verify_result.py 将 NPU 输出与 CPU golden 结果比对。3.3 修改注意点按文档总结落地该优化时有三个必须注意的工程细节维度对齐处理原始权重矩阵的 K 维和 N 维可能不是基本块大小的整数倍需在转换前进行 Padding 对齐K 对齐到cube_size倍数N 对齐到 16 的倍数否则块边界处的取数会错位数据类型适配不同数据类型对应的 cube_size 不同FP16/BF16 为 16FP32 为 8需通过32 / sizeof(T)动态计算而不是写死常量。样例中CUBE_SIZE 32字节BF16 下实际块宽为32 / 2 16内存大小变化FRACTAL_NZ 排布因维度对齐会导致实际存储空间大于原始矩阵分配 GM 空间时应以对齐后的尺寸为准如 3.2 节中alignedK * alignedN避免越界访问。四、编译与执行4.1 编译样例从项目根目录启动构建整体构建流程参考仓库根目录 README.md。weightnz 样例要求NPU_ARCHdav-3510见 CMakeLists.txt 中的cann_sample_check_arch(dav-3510)检查。在仓库根目录下完成编译和安装后进入当前样例目录cmake -S . -B build -DNPU_ARCHdav-3510 cmake --build build --parallel cmake --install build --prefix ./build_out cd ./build_out/1_Features/instruction_optimization/weightnz/如需单独编译当前样例可使用以下指令cmake --build build --target weightnz cp ./Samples/1_Features/instruction_optimization/weightnz/scripts/* ./build/Samples/1_Features/instruction_optimization/weightnz/ cd ./build/Samples/1_Features/instruction_optimization/weightnz/4.2 运行样例使用可执行文件直接执行算子用例需要指定矩阵乘维度m k n程序会随机生成输入数据./weightnz 1024 2048 4096运行成功后终端将打印如下类似信息NPU 输出与 CPU golden 一致Data generated successfully! [verify] shape(1024, 4096), elements4194304 - summary (large matrix, full tensors omitted) abs_err: max2.560000e02, mean7.385254e-03, rmse1.375000e00 rel_err: max6.410256e-03 count(|abs_err| 0.001): 121 / 4194304 cpu golden (top-left 4x4): tensor([[42496., 42496., 41728., 41728.], [41728., 41984., 41216., 40960.], [41984., 41984., 41216., 41216.], [41216., 41472., 40960., 40960.]], dtypetorch.bfloat16) npu out (top-left 4x4): tensor([[42496., 42496., 41728., 41728.], [41728., 41984., 41216., 40960.], [41984., 41984., 41216., 41216.], [41216., 41472., 40960., 40960.]], dtypetorch.bfloat16) max abs diff: 256.0 point error count(0.1): 0/4194304 ratio error count(0.001): 121/4194304, error ratio: 0.000029 [PASS] NPU results are consistent with CPU.从 verify_result.py 的判据可以看到校验通过需同时满足单点相对误差abs_diff/abs(golden)不超过 0.1且绝对误差超过 0.001 的点占比不超过 0.001。如果存在精度问题则会打印错误数据并显示如下结果[ERROR] NPU results differ from CPU.4.3 测试性能运行性能测试脚本scripts/profile_matmul.py指定矩阵乘法的维度后执行。该脚本通过msprof对当前安装目录下的可执行文件进行 Profiling解析op_summary_*.csv中的 kernel、mac、scalar、mte1、mte2、fixpipe 及 icache_miss 指标并输出对比表python3 profile_matmul.py 1024 2048 4096打印如下执行结果证明样例性能测试成功[Profile Breakdowm] --------------------------------------------------------------------------------------------- | candidate | kernel(us) | mac(us) | scalar(us) | mte1(us) | mte2(us) | fixpipe(us) | icache_miss(%) | | weightnz | 53.598 | 40.511 | 2.664 | 10.659 | 35.529 | 3.216 | 1.100 | ---------------------------------------------------------------------------------------------与相同输入规模下的基础开启 DBDouble Buffer的 matmul 算子相比[Profile Breakdowm] --------------------------------------------------------------------------------------------- | candidate | kernel(us) | mac(us) | scalar(us) | mte1(us) | mte2(us) | fixpipe(us) | icache_miss(%) | | n_buffer | 66.000 | 40.810 | 2.558 | 12.681 | 37.595 | 1.980 | 1.200 | ---------------------------------------------------------------------------------------------可以看到在m1024, k2048, n4096的规模下weightnz 版本 kernel 总耗时约 53.6us对比 n_buffer 基线的 66.0us整体计算时间显著缩短性能有所提升其中 mac 耗时基本持平40.5us vs 40.8us收益主要来自搬运链路mte1 10.7us vs 12.7us、mte2 35.5us vs 37.6us与整体流水的提前结束。该基线即同目录下的 n_buffer 样例已开启双缓冲流水优化的 Matmul。由上述仿真流水图可以看出将权重矩阵预先转换为 FRACTAL_NZ 排布后MTE2 搬运不再承担随路格式转换可有效提升数据搬运效率整体流水提前收敛。五、适用场景与结论适用场景权重矩阵固定适合权重矩阵固定的推理场景可以离线一次性完成 NZ 转换并随模型部署大矩阵场景适合大 weight 矩阵场景权重搬运量大时带宽优化收益更高。结论通过将权重矩阵预转换为 FRACTAL_NZ 排布消除 Matmul 算子中的随路格式转换开销可提升整体流水效率。落地时需牢记三点对齐与内存规则K 维对齐到32 / sizeof(T)的倍数、N 维对齐到 16 的倍数、GM 内存按对齐后尺寸分配。六、支持架构weightnz 样例当前支持NPU ARCH 3510Ascend 950 平台构建参数NPU_ARCHdav-3510。若使用其他架构如dav-2201全量构建工程该样例会在配置阶段被自动跳过属预期行为。延伸阅读同属指令优化系列的 n_buffer双缓冲流水、unit_flagMMAD 与 Fixpipe 流水并行、mte2_preloadMTE2 预加载样例可与 weightnz 组合使用进一步压缩 Matmul 流水空闲时间样例目录总览见 Samples/1_Features/instruction_optimization/README.md。【免费下载链接】cann-samplesCANN高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考