CANN/catlass量化多核切K矩阵乘TLA示例
CATLASS Quant Multi Core SplitK Matmul Tla 样例介绍【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass原型设计名称/Name类型/Class数据类型/Dtype维度/Dims格式/Format描述/DescriptionmatAinTensorint8[m, k]ND左矩阵matBinTensorint8[n, k]ND右矩阵支持转置scaleinTensorfloat[n]NDperChannel量化系数perTokenScaleinTensorfloat[m]NDperToken量化系数matDoutTensorbf16[m, n]ND输出矩阵样例实现CATLASS52_quant_multi_core_splitk_matmul_tla样例算子是基于CATLASS Gemm API实现的昇腾亲和Matmul算子针对大尺寸矩阵计算场景优化设计关键算子组件包括以下几部分:Example组装quant_multi_core_splitk_matmul_tla.cppKernel实现主Kernel文件quant_multi_core_splitk_matmul_tla.hppBlock组件包含基础的mmad组件block_mmad_pingpong_tla.hpp可替换反量化后处理组件block_epilogue_per_token_dequant_tla.hpp方案设计基础的多核切K方案参考这里主要说明多核切K下量化后处理适配的两种方案且本样例根据Plan B实现。算子在AIV上的计算为两块①矩阵乘多核切K后沿K轴的ReduceAdd累加②矩阵乘结果的反量化。Plan A是AIC完成矩阵乘后、AIV先做反量化再做累加Plan B是AIC完成矩阵乘后、AIV先做累加再做反量化。对于每个核存在多轮基本块计算的场景Plan A可以有CV互相掩盖的收益若每个核仅一轮基本块计算Plan B的反量化计算量更少。性能收益性能和 12_quant_matmul 比较优化点在于C矩阵较小时通过切K实现负载均衡劣化点在于将反量化放在最后单独处理没有cv融合掩盖并增加了AIV的ReduceAdd动作完成K方向上累加| M | N | K | 标杆耗时(us) | 耗时(us) | 加速比 | | -- | -- | ---- | ----- | ----- | ----- | | 128 | 256 | 1024 | 18.64 | 18.11 | 1.03 | | 128 | 256 | 2048 | 23.23 | 20.67 | 1.12 | | 128 | 256 | 4096 | 32.97 | 21.15 | 1.56 | | 128 | 256 | 8192 | 47.73 | 20.01 | 2.39 | | 128 | 256 | 16384 | 78.52 | 23.83 | 3.30 | | 256 | 512 | 16384 | 75.39 | 34.29 | 2.20 | | 256 | 1024 | 16384 | 77.35 | 51.41 | 1.50 |标杆为QuantMatmul算子。统计耗时均为核函数总耗时使用msprof工具得到。上述测试例中A矩阵为layout::RowMajor排布方式B矩阵为layout::ColumnMajor排布方式。测试环境说明NPU型号为910B3CANN包版本为8.5.0。【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考