昇腾 CANN 分组矩阵乘量化算子样例集实战指南：MXFP4 / MXFP8 / HiFloat8 / MXA8W4

昇腾 CANN 分组矩阵乘量化算子样例集实战指南MXFP4 / MXFP8 / HiFloat8 / MXA8W4【免费下载链接】cann-samplesCANN高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples本篇文章围绕 CANN 高性能实战样例仓库cann-samples中grouped_matmul_recipes目录展开系统介绍面向 MoE 等多专家推理/训练场景的分组矩阵乘Grouped Matmul量化算子在昇腾 AI 处理器NPU ARCH 3510上的典型实现样例从目录结构、样例能力、使用约束、输入参数与数据生成方式到构建、运行、结果校验的完整实操流程并延伸到 MX 量化分组矩阵乘的性能建模与优化策略帮助读者快速上手并深入理解昇腾平台上的分组量化矩阵乘编程实践。一、样例集概览什么是分组矩阵乘分组矩阵乘Grouped Matmul是面向 MoEMixture of Experts等包含多专家分组计算场景的算子形态。与普通矩阵乘不同Grouped Matmul 按专家数进行分组每组执行一次独立的矩阵乘计算输入矩阵A在M维按组拼接权重矩阵B按组独立存储每个分组的分组值由输入group_list指定允许分组值为0表示该专家未被选中。grouped_matmul_recipes目录汇总了分组矩阵乘在昇腾 AI 处理器上的典型量化实现样例每个样例提供完整的算子代码.asc、运行脚本与说明文档可直接编译运行并做性能对比。该目录位于 Samples/2_Performance/grouped_matmul_story/grouped_matmul_recipes是 grouped_matmul_story 主题下算子实现与示例代码部分的主体。目录结构grouped_matmul_recipes/ ├── CMakeLists.txt ├── README.md ├── common/ # 公共工具host/kernel ├── include/ # 共享头文件block、kernel、tile、tiling 等 └── examples/ ├── quant_grouped_matmul_mxfp4/ # MXFP4 分组量化矩阵乘样例 ├── quant_grouped_matmul_mxfp8/ # MXFP8 分组量化矩阵乘样例 ├── quant_grouped_matmul_hif8/ # HiFloat8 分组量化矩阵乘样例 └── weight_quant_grouped_matmul_mxfp8fp4/ # MXA8W4 权重量化分组矩阵乘样例common/下分为host_utils/acl_utils.h、common_utils.h、io_utils.h与kernel_utils/common_utils.h、layout_utils.h、tensor_utils.h提供 host 侧算子启动与 kernel 侧通用工具include/下按block/MMA 计算块与调度器、kernel/kernel 主体、policy/dispatch 策略、prologue/prologue 处理、tile/GM↔UB↔L1 数据搬运与格式转换、tiling/tiling 结构与 tiling 数据、utils/常量与布局结构体分层组织共享头文件。二、样例列表四种量化数据类型样例数据类型说明quant_grouped_matmul_mxfp4MXFP4MXFP4 分组量化矩阵乘quant_grouped_matmul_mxfp8MXFP8MXFP8 分组量化矩阵乘quant_grouped_matmul_hif8HiFloat8HiFloat8 分组量化矩阵乘weight_quant_grouped_matmul_mxfp8fp4MXA8W4MXA8W4 权重量化分组矩阵乘从 性能优化指南 中的对比可知MX 量化家族的主要差异在于对比项MXFP8MXFP4A/B 数据类型float8_e4m3fnfloat4_e2m1量化方式GroupSize32 的 per-group 量化同左scaleA/Bfloat8_e8m0同左分组M 轴分组A NDB ND/DN/NZ/ZNK 轴分组A DNB NDM 轴分组A NDB ND/DN/NZ/ZN显存压缩比相比 FP16/BF16 内存占用减少约 50%相比 FP16/BF16 内存占用减少约 75%典型用途模型训推兼顾速度与精度的平衡模型推理侧重极致显存效率与推理速度其中 MX 量化的本质是 Microscaling 量化即 pergroup-pergroup 量化模式G-G 量化是量化参数类型为FLOAT8_E8M0且 group size 为 32 的特例详见 量化介绍。三、样例能力与实现变体每个样例目录下都包含若干可执行程序变体与配套脚本覆盖不同分组轴、不同权重存储布局与不同流水深度MXFP4quant_grouped_matmul_mxfp4quant_grouped_matmul_mxfp4_split_m基于 m 轴分组、权重按 ND包括 NDExtLayout/DNExtLayout统称 ND逻辑组织的分组量化矩阵乘示例quant_grouped_matmul_mxfp4_split_m_weight_nz基于 m 轴分组、权重按 GM 上 NZ包括 NZLayout/ZNLayout统称 NZ存储的分组量化矩阵乘示例gen_data.py生成 ND 权重输入数据和 CPU golden 结果gen_data_weight_nz.py生成 NZ 权重输入数据和 CPU golden 结果verify_result.py校验 NPU 输出与 CPU golden 是否一致。MXFP8quant_grouped_matmul_mxfp8MXFP8 样例提供最丰富的变体集合quant_grouped_matmul_mxfp8_split_m基于 m 轴分组、权重按 ND 逻辑组织的分组量化矩阵乘示例quant_grouped_matmul_mxfp8_split_m_3buffer基于 m 轴分组、权重按 ND 逻辑组织的 3buffer 分组量化矩阵乘示例其余示例均为 2bufferquant_grouped_matmul_mxfp8_split_m_weight_nz基于 m 轴分组、权重按 GM 上 NZ 存储的分组量化矩阵乘示例quant_grouped_matmul_mxfp8_split_m_scale_b_nz基于 m 轴分组、权重按 ND 逻辑组织、ScaleB 按 GM 上 scaleB 私有格式仅包括 NNLayout形状为(N1, Kg, 16, 2)存储的分组量化矩阵乘示例其中N1ceil(N/16)Kgceil(K/64)ScaleB 私有格式与 transB 无关只有上述一种quant_grouped_matmul_mxfp8_split_k基于 K 轴分组、权重按 ND 逻辑组织的分组量化矩阵乘示例仅支持transAtrue, transBfalse场景gen_data.py/gen_data_weight_nz.py/gen_data_scale_b_nz.py分别生成 ND 权重、NZ 权重、ND 权重 NZ ScaleB 的输入数据与 CPU golden 结果verify_result.py校验 NPU 输出与 CPU golden 是否一致。HiFloat8quant_grouped_matmul_hif8quant_grouped_matmul_hif8_split_m_tt基于 M 轴分组、采用 T-T 量化模式的分组量化矩阵乘示例每个 group 一组 FP32 标量pertoken_scale.binscale.binquant_grouped_matmul_hif8_split_m_tc基于 M 轴分组、采用 T-C 量化模式的分组量化矩阵乘示例每个 group 的 N 维uint64per-channelscale.binscripts/gen_data_tt.py/scripts/gen_data_tc.py分别生成 T-T、T-C 量化模式输入数据和 CPU golden 结果scripts/verify_result.py校验 NPU 输出与 CPU golden 是否一致。T-T 与 T-C 是量化介绍中常见组合量化的两种模式T-T 即 pertensor-pertensor 量化全量化的一种T-C 即 pertensor-perchannel 量化两者分别以组内共享标量或 per-channel 向量作为量化参数。MXA8W4weight_quant_grouped_matmul_mxfp8fp4weight_quant_grouped_matmul_mxfp8fp4基于 m 轴分组的权重量化 grouped matmul 执行程序A 为 FP8(E4M3)、B 为 FP4(E2M1)、输出 BF16scripts/gen_data.py生成输入数据和 CPU golden 结果scripts/verify_result.py校验 NPU 输出与 CPU golden 是否一致scripts/batch_test_accuracy.py批量随机精度回归测试脚本。四、输入参数与命令行约定四个样例的算子执行文件与结果校验脚本的命令行参数格式一致MXA8W4 样例仅支持默认转置组合不包含 transA/transB 参数program group_num m k n [transA transB]参数含义group_num专家数也就是分组数group_value_list由数据生成脚本生成表示每个专家对应的分组大小例如128,128,0m总的M大小要求满足m sum(group_value_list)K 轴分组时改为对k的约束k矩阵A的列数同时也是每组矩阵B的列数K 轴分组时要求k sum(group_value_list)n每组矩阵B的行数也是输出矩阵每组结果的列数transA可选参数默认值为falsefalse表示 A 以[M, K]组织true表示 A 以[K, M]组织MXFP8 split_k 场景transB可选参数默认值为truetrue表示 B 以[E, N, K]NZ 为[E, K1, N1, N0, K0]组织false表示 B 以[E, K, N]NZ 为[E, N1, K1, K0, N0]组织。transA和transB需要同时省略或同时指定取值支持0/1/true/false。实际参与计算的group_value_list由数据生成脚本生成写入input/input_groupList.bin该文件保存每个分组各自的M大小允许某些组为0。从数据生成脚本 scripts/gen_data.py 的源码可以印证这些约定的底层实现脚本定义了GROUP_LIST_MODE group_list与EXPECT_M_PER_GROUP_MODE expect_m_per_group两种模式以及DEFAULT_TRANS_A False、DEFAULT_TRANS_B True默认值parse_group_m_list()逐项解析逗号分隔的分组大小并拒绝负值build_random_group_m_list()则校验group_num 0并在m group_num * floor(0.7 * expect_m_per_group)时直接报错。五、数据生成方式gen_data.pyND 权重、gen_data_weight_nz.pyNZ 权重、gen_data_scale_b_nz.pyND 权重 NZ ScaleBMXFP8以及 HiFloat8 样例的gen_data_tt.py/gen_data_tc.py均支持以下两种调用方式仅将脚本名替换即可方式一显式指定group_value_listpython3 gen_script.py group_list group_value_list m k n [transA transB]示例# ND权重 python3 gen_data.py group_list 128,128,0 384 256 256 false false # NZ权重 python3 gen_data_weight_nz.py group_list 128,128,0 384 256 256 false false # NZ ScaleB仅支持M轴分组权重仍为ND python3 gen_data_scale_b_nz.py group_list 128,128,0 384 256 256 false falsegroup_list显式分组模式直接传入每个专家的分组大小group_value_list每个专家对应的分组大小例如128,128,0m矩阵乘的m维要求m sum(group_value_list)k矩阵乘的k维n矩阵乘的n维。方式二按专家数和期望平均值随机生成group_value_listpython3 gen_script.py expect_m_per_group group_num expect_m_per_group m k n [transA transB]示例python3 gen_data.py expect_m_per_group 3 128 384 256 256 false falseexpect_m_per_group随机分组模式按每组期望分组大小随机生成分组group_num专家数 / 分组数expect_m_per_group每组期望平均分组大小m/k/n同上。在该模式下脚本会随机生成长度为group_num的group_value_list并保证每个分组大小均在[floor(0.7 * expect_m_per_group), ceil(1.3 * expect_m_per_group)]范围内sum(group_value_list) m源码中还会提前校验m group_num * floor(0.7 * expect_m_per_group)不满足时直接抛出异常避免生成无法满足总M上限的分组。golden 输入数据由对应的数据生成脚本生成编译安装后请在build_out下的本示例目录中执行该脚本ND 权重用gen_data.pyNZ 权重用gen_data_weight_nz.py依此类推。六、使用约束与形状约定MXFP4仅支持transAfalse, transBtrue和transAfalse, transBfalse两种场景当transBtrue时A 形状为[M, K]B ND/NZ 形状为[E, N, K]/[E, K1, N1, N0, K0]其中N016、K064、N1ceil(N/N0)、K1ceil(K/K0)当transBfalse时A 形状为[M, K]B ND/NZ 形状为[E, K, N]/[E, N1, K1, K0, N0]其中K016、N064、K1ceil(K/K0)、N1ceil(N/N0)仅支持 m 轴分组MXFP4 内轴K必须为偶数当transBfalse时N也必须为偶数。MXFP8支持transAfalse, transBtrue、transAfalse, transBfalse和transAtrue, transBfalse三种转置组合以及 M/K 轴分组对应约束为M 轴分组且transAfalse, transBtrueA 形状[M, K]B ND/NZ 形状[E, N, K]/[E, K1, N1, N0, K0]N016、K032ScaleB ND/NZ 形状[E, N, Kg, SCALE_C02]/[E, N1, Kg, N0, SCALE_C02]N016、SCALE_C02、Kgceil(K/64)、N1ceil(N/16)M 轴分组且transAfalse, transBfalseA 形状[M, K]B ND/NZ 形状[E, K, N]/[E, N1, K1, K0, N0]K016、N032ScaleB ND/NZ 形状[E, Kg, N, SCALE_C02]/[E, N1, Kg, N0, SCALE_C02]K 轴分组即transAtrue, transBfalseA 形状[K, M]B ND 形状[K, N]。HiFloat8仅支持 M 轴分组split-M 路径transA必须为false支持transBtrueB 形状[E, N, K]与transBfalseB 形状[E, K, N]两种场景A 形状均为[M, K]输入数据类型为hifloat8输出数据类型为bfloat16。MXA8W4仅支持A不转置、B转置场景A 形状[M, K]B 形状[E, N, K]仅支持 m 轴分组k需为64的正整数倍数据生成与校验脚本按该约束实现n需为32的正整数倍精度回归脚本按该约束实现m需满足m sum(group_m_list)。支持架构以上样例均面向NPU ARCH 3510构建时通过-DNPU_ARCHdav-3510指定。七、构建与运行在仓库根目录下执行全量编译与安装并进入安装目录cmake -S . -B build -DNPU_ARCHdav-3510 cmake --build build --parallel cmake --install build --prefix ./build_out cd build_out/2_Performance/grouped_matmul_story/grouped_matmul_recipes/quant_grouped_matmul_mxfp4构建配置由 grouped_matmul_recipes/CMakeLists.txt 定义add_grouped_matmul_recipe为每个.asc源文件生成可执行目标链接m、dl、platform、tiling_api与cann_samples::tensor_api并按样例子目录安装到grouped_matmul_recipes/${variant}下同时会通过install(DIRECTORY ... FILES_MATCHING PATTERN *.py PATTERN *.sh)将每个样例的scripts/辅助脚本安装到与可执行文件相同的目录保证安装目录下可直接运行数据生成与结果校验脚本。MXFP4 运行示例ND 权重# 生成数据方式一显式指定grouplist生成一组测试数据 python3 gen_data.py group_list 128,128,0 384 256 256 # 生成数据方式二按专家数和平均M随机生成grouplist python3 gen_data.py expect_m_per_group 3 128 384 256 256 # 运行可执行文件并校验结果默认transAfalse, transBtrue ./quant_grouped_matmul_mxfp4_split_m 3 384 256 256 # 运行transBfalse场景显式指定transA/transB python3 gen_data.py group_list 128,128,0 384 256 256 false false ./quant_grouped_matmul_mxfp4_split_m 3 384 256 256 false falseNZ 权重python3 gen_data_weight_nz.py group_list 128,128,0 384 256 256 ./quant_grouped_matmul_mxfp4_split_m_weight_nz 3 384 256 256 python3 gen_data_weight_nz.py group_list 128,128,0 384 256 256 false false ./quant_grouped_matmul_mxfp4_split_m_weight_nz 3 384 256 256 false falseMXFP8 运行示例ND 权重覆盖全部变体# 默认场景 transAfalse, transBtrue python3 gen_data.py group_list 128,128,0 384 256 256 ./quant_grouped_matmul_mxfp8_split_m 3 384 256 256 # M轴分组: transAfalse, transBfalse python3 gen_data.py group_list 128,128,0 384 256 256 false false ./quant_grouped_matmul_mxfp8_split_m 3 384 256 256 false false # M轴分组3buffer场景 ./quant_grouped_matmul_mxfp8_split_m_3buffer 3 384 256 256 false false # K轴分组: transAtrue, transBfalse python3 gen_data.py group_list 128,128,0 384 384 256 true false ./quant_grouped_matmul_mxfp8_split_k 3 384 384 256 true falseNZ 权重与 NZ ScaleB仅支持 M 轴分组python3 gen_data_weight_nz.py group_list 128,128,0 384 256 256 ./quant_grouped_matmul_mxfp8_split_m_weight_nz 3 384 256 256 python3 gen_data_scale_b_nz.py group_list 128,128,0 384 256 256 ./quant_grouped_matmul_mxfp8_split_m_scale_b_nz 3 384 256 256HiFloat8 运行示例# T-T量化模式 python3 scripts/gen_data_tt.py group_list 64,80,96 256 128 256 ./quant_grouped_matmul_hif8_split_m_tt 3 256 128 256 # T-C量化模式transBfalse 需显式指定 python3 scripts/gen_data_tc.py group_list 64,80,96 256 128 256 false false ./quant_grouped_matmul_hif8_split_m_tc 3 256 128 256 false falseMXA8W4 运行示例python3 gen_data.py group_list 128,128,0 384 256 256 ./weight_quant_grouped_matmul_mxfp8fp4 3 384 256 256MXA8W4 的程序在执行完成后会自动调用verify_result.py进行结果校验。四个样例均支持可选的手动复核python3 verify_result.py 3 384 256 256gen_data.py生成的文件清单以 MXA8W4 为例其余样例类似CPU golden: output/output_cpu.bin input/input_a.bin input/input_b.bin input/input_scaleA.bin input/input_scaleB.bin input/input_groupList.bin运行可执行文件后额外生成NPU output: output/output_npu.bin八、性能优化指南分组矩阵乘的性能建模与优化样例的模板实现与优化策略详见 MX 量化分组矩阵乘算子性能优化指南。推荐先从quant_grouped_matmul_mxfp4开始便于快速验证脚本与可执行文件的配套流程。分组原理M 轴分组与 K 轴分组GroupedMatmul 由 E 个单 Matmul 组成按分组更新A/B/ScaleA/ScaleB/C的 GM 基址偏移M 轴分组将当前组的(M_e, N, K)视作单个 MX 矩阵乘以行维度 M 为划分依据将整体矩阵乘拆分为多个行维度 M 分片子矩阵乘A(M, K)、B(E, N, K)或(E, K, N)各组 B 矩阵大小一致各组输出沿 M 维度拼接。K 轴分组将当前组的(M, N, K_e)视作单个 MX 矩阵乘以内积收缩维度 K 为划分依据A(K, M)、B(K, N)将 K 维度切分为多个子通道块各组输出大小一致各组在对应 K 分片上独立计算并拼接结果。分组内的 MX 矩阵乘实现与约束和 quant matmul MX 矩阵乘一致可参考 quant_matmul_mx_performance.md。性能瓶颈分析单 MX Matmul 的性能瓶颈分为两类CUBE Bound算子性能受限于硬件算力规格本身已实现连续的 MMAD 计算此时需重点关注多核计算负载是否均衡避免单核 Cube Bound 而整体 Cube 利用率偏低Memory Bound算子性能受限于数据搬运能力按瓶颈所在流水进一步区分为MTE2 Bound、MTE1 Bound与FIXPIPE Bound。在单 MX Matmul 的 Bound 分类之上分组场景额外引入组间切换与负载不均每个分组的数据不能组间复用各分组值不同可能导致每个分组 MX Matmul 的性能瓶颈不同。若每个 Matmul 都是同种 Bound如都是 CUBE Bound则 GroupedMatmul 是该种 Bound否则性能瓶颈需按各组时间加权估算而非由单次 Matmul 的(M,N,K)决定。优化手段清单Double Buffer双缓冲两个缓冲区交替工作用计算与数据准备的叠加隐藏内存访问延迟。2-Buffer 覆盖基本重叠关系、资源开销小3-Buffer 仅在 L1 侧增加第三份阶段缓冲A/B 数据L0 层仍为 2-Buffer用于减少搬运抖动、带宽瞬时不足造成的断流风险。3-Buffer 需满足 L1 容量约束A1 B1 scaleA scaleB A3 HALF_L1_SIZE且A2 B2 scaleA scaleB B3 HALF_L1_SIZE不满足时建议回退 2-Buffer 并结合 SWAT/Bank 冲突优化。UnitFlag单元标志为 MMAD 计算指令与 FIXPIPE 数据搬运指令提供 512B 粒度的基于内存访问的细粒度同步在无法开启 L0C Double-Buffer 时有效提高计算与搬出流水并行度。SWAT自适应滑动窗口模板通过提升多核单次访问的 L2 命中率提高 MTE2 搬运效率在 M 轴上设定固定窗口、根据尾块大小灵活调整并沿 N 方向Z型滑动使首轮搬运即可做到 MMAD 指令不断流。L1 Bank 冲突优化L1 缓冲区以 256KB 粒度分为两个 Bank开启 L1 Double Buffer 时将两份缓存放置于不同 Bank避免读写冲突导致的 MTE1 带宽下降。Scale 缓存优化Scale 数据量仅为输入矩阵的 1/32输入矩阵较小时可提前载入后续所需 Scale 并在 L1 上缓存减少搬运次数、缓解带宽利用率低的问题。WeightNZ 优化NZ/ZN 格式的 weight 数据可通过普通DataCopy搬运相比带格式转换的ND2NZ指令带宽利用率更高能显著降低 MTE2 耗时。双页表当分组 matmul 的 m 较小baseM时weight 数据只使用一次可将 weight 从 HBM 直接加载进 L1 缓冲区避免 L2 有效数据回写 HBM 造成的额外 MTE2 耗时。group 间核负载均衡GroupedMatmul 独有下一个分组 Matmul 接着上一个分组结束的核数往下分配核避免每个分组均从 0 核开始计算。M 轴分组 M 计算负载均衡GroupedMatmul 独有group_list 是 device tensorM 轴分组时每分组M_e在 host 侧未知需避免某些核 M 轴计算量偏大导致的快慢核性能劣化K 轴分组不做该均衡实现模板暂不切 K。尾轮负载均衡将最后一轮未完全分配的基本块进行二次切分均匀分配到多核消除算力浪费。优化策略选择指南Bound 类型推荐优化策略优先级CUBE BoundSWAT 多核负载均衡高MTE2 BoundScale 缓存 WeightNZ 搬运优化 双页表高MTE1 BoundL1 Bank 冲突优化中FIXPIPE BoundUnitFlag中流水停顿Double Buffer高调优实践遵循性能分析Profiling 性能建模定位 Bound→ 瓶颈识别 → 策略选择 → 参数调优tiling 参数与缓冲区使用→ 效果验证 → 迭代优化的闭环流程。从性能优化指南的模板归纳看SWAT 模板作为基础模板覆盖了上述全部优化手段其具体实现对应本样例集中的 MXFP4 splitM、MXFP4 splitM WeightNZ、MXFP8 splitM、MXFP8 splitM 3Buffer、MXFP8 splitM WeightNZ、MXFP8 splitK 等.asc文件。九、总结grouped_matmul_recipes提供了昇腾 AI 处理器上分组矩阵乘量化算子的完整实战样例矩阵从数据维度覆盖 MXFP4 / MXFP8 / HiFloat8 / MXA8W4 四种量化数据类型从实现维度覆盖 M 轴 / K 轴分组、ND / NZ / ScaleB-NZ 权重布局、2-Buffer / 3-Buffer 流水与 T-T / T-C 量化模式从工具链维度覆盖数据生成、编译构建、运行校验与精度回归。配合 量化介绍 与 性能优化指南开发者可以按选定数据类型样例 → 理解参数与约束 → 生成数据 → 构建运行 → 性能建模定位 Bound → 按策略选择指南迭代调优的路径系统掌握昇腾平台上 MoE 等分组计算场景的高性能量化矩阵乘开发方法。【免费下载链接】cann-samplesCANN高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考