cann-samples分组MXFP8量化矩阵乘

Grouped Matmul MXFP8量化矩阵乘算子【免费下载链接】cann-samples算子领域高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples概述本示例展示了Grouped Matmul MXFP8量化矩阵乘算子在昇腾AI处理器上的完整实现。算子以专家数进行分组执行分组矩阵乘计算输入矩阵A在M维按组拼接权重矩阵B按组独立存储适用于MoE等包含多专家分组计算的推理场景。当前目录提供以下能力quant_grouped_matmul_mxfp8基于m轴分组、权重按ND包括NDExtLayout/DNExtLayout下面仅用ND统称逻辑组织的分组量化矩阵乘示例。quant_grouped_matmul_mxfp8_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是否一致。使用约束当前样例需要满足以下约束条件当前仅支持transAfalsetransBtrue和transAfalsetransBfalse两种场景。当transBtrue时A的形状为[M, K]B ND/NZ的形状为[E, N, K]/[E, K1, N1, N0, K0]N016K032N1ceil(N/N0)K1ceil(K/K0)。当transBfalse时A的形状为[M, K]B ND/NZ的形状为[E, K, N]/[E, N1, K1, K0, N0],K016,N032,K1ceil(K/K0)N1ceil(N/N0)。当前仅支持m轴分组。支持架构NPU ARCH 3510API参考Ascend C API文档输入参数算子执行文件与结果校验脚本的命令行参数格式一致program group_num m k n [transA transB]group_num专家数也就是分组数group_m_list表示每个专家对应的分组大小例如128,128,0m总的M大小要求满足m sum(group_m_list)k矩阵A的列数同时也是每组矩阵B的列数n每组矩阵B的行数也是输出矩阵每组结果的列数transA可选参数当前仅支持false默认值为falsetransB可选参数默认值为truetrue表示B以[E, N, K]/[E, K1, N1, N0, K0]组织false表示B以[E, K, N]/[E, N1, K1, K0, N0]组织transA和transB需要同时省略或同时指定取值支持0/1/true/false。其中实际参与计算的group_m_list由数据生成脚本gen_data.py或gen_data_weight_nz.py生成并写入input/input_groupList.bin。当前文件中保存的是每个分组各自的M大小允许某些组为0。golden输入数据由对应的数据生成脚本生成。编译安装后请在build_out下的本示例目录中执行该脚本。数据生成方式gen_data.pyND权重与gen_data_weight_nz.pyNZ权重支持以下两种调用方式仅将脚本名替换即可方式一显式指定group_m_listpython3 gen_script.py group_list group_m_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含义如下group_list显式分组模式直接传入每个专家的分组大小。group_m_list每个专家对应的分组大小例如128,128,0m总的M上限要求m sum(group_m_list)k矩阵乘的k维n矩阵乘的n维方式二按专家数和期望平均值随机生成group_m_listpython3 gen_script.py expect_m_per_group group_num expect_m_per_group m k n [transA transB]示例# ND权重 python3 gen_data.py expect_m_per_group 3 128 384 256 256 false false # NZ权重 python3 gen_data_weight_nz.py expect_m_per_group 3 128 384 256 256 false false含义如下expect_m_per_group随机分组模式按每组期望分组大小随机生成分组group_num专家数/分组数expect_m_per_group每组期望平均分组大小m总的M上限要求m sum(group_m_list)k矩阵乘的k维n矩阵乘的n维在该模式下脚本会随机生成长度为group_num的group_m_list并保证每个分组大小均在[floor(0.7 * expect_m_per_group),ceil(1.3 * expect_m_per_group)]范围内sum(group_m_list) m构建与运行在仓库根目录下执行全量编译与安装并进入安装目录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_mxfp8之后可按需执行以下命令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_mxfp8 3 384 256 256 # 运行transBfalse场景显式指定transA/transB python3 gen_data.py group_list 128,128,0 384 256 256 false false ./quant_grouped_matmul_mxfp8 3 384 256 256 false falseNZ权重# 生成数据方式一显式指定grouplist生成一组测试数据 python3 gen_data_weight_nz.py group_list 128,128,0 384 256 256 # 生成数据方式二按专家数和平均M随机生成grouplist python3 gen_data_weight_nz.py expect_m_per_group 3 128 384 256 256 # 运行可执行文件并校验结果默认transAfalse,transBtrue ./quant_grouped_matmul_mxfp8_split_m_weight_nz 3 384 256 256 # 运行transBfalse场景显式指定transA/transB python3 gen_data_weight_nz.py group_list 128,128,0 384 256 256 false false ./quant_grouped_matmul_mxfp8_split_m_weight_nz 3 384 256 256 false false# 可选手动再次校验用于调试/复核 python3 verify_result.py 3 384 256 256【免费下载链接】cann-samples算子领域高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考