CATLASS Quant Matmul Full LoadA 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输出矩阵样例实现CATLASS44_quant_matmul_full_loadA_tla样例算子是基于CATLASS Gemm API实现的昇腾亲和Matmul算子针对大尺寸矩阵计算场景优化设计关键算子组件包括以下几部分:Example组装quant_matmul_full_loadA_tla.cppKernel实现主Kernel文件quant_matmul_full_loadA_tla.hppBlock组件包含全载专门的mmad组件block_mmad_pingpong_full_loadA_tla.hpp反量化后处理组件block_epilogue_per_token_dequant_tla.hpp多核全载用的swizzleGemmIdentityBlockSwizzleL1FullLoad方案设计如下图所示模板库矩阵有关键参数L1TileShapeM, N, KC矩阵按照L1TileShape::M和L1TileShape::N参数切分基本块并分核而后普通Matmul模板会将A矩阵中L1TileShape::M * L1TileShape::K大小的矩阵块载入L1而A矩阵全载模板会直接将L1TileShape::M * K大小的矩阵块载入L1而对于B矩阵普通模板和A矩阵全载模板都是载入L1TileShape::K * L1TileShape::N大小的矩阵块L1。A矩阵全载模板相较普通模板的流水变动如下图所示使用全载方案在计算一开始就经MTE2完整搬入A矩阵的数据块而后pingpong搬入B矩阵执行计算。采用A矩阵全载模板时需要一半的L1空间以放入L1TileShape::M * problemShape.K的数据若L1空间不够A矩阵全载则返回报错。A矩阵全载时N轴越大单核越能多次复用L1中的A矩阵、无需再从GM或L2Cache搬运A矩阵性能收益就越大。A矩阵全载时N轴较小无法复用A矩阵性能收益较00_basic_matmul可能会出现劣化。若problemShape.M L1TileShape::M可使用GemmIdentityBlockSwizzleL1FullLoadSwizzleOffset, SwizzleDirection, AicCoreNum策略使得每个核需要处理的基本块尽可能地连在一起提升A矩阵分核全载时的块间复用率。若problemShape.M L1TileShape::M即M方向不切块分核此时使用GemmIdentityBlockSwizzle策略即可适用常用的策略参数选取可参考swizzle_explanation。以20个cube核为例常用的GemmIdentityBlockSwizzle策略的基本块分核顺序为0-1-2-...-18-19-0-1-2...-18-19-0-1-2...每个核需要处理的基本块跳跃分布而GemmIdentityBlockSwizzleL1FullLoad策略的基本块分核顺序为0-0...-0-1-1...-1-2-2...-19每个核需要处理的基本块连续分布。性能收益使用相同tileShape和swizzle参数相比标杆样例12_quant_matmul具备A矩阵全载特性的样例性能平均提升5%~15%这是由于全载实现减少了计算过程中的A矩阵搬运提升了全载矩阵的数据复用率可参考下表。[M,N,K]12_quant_matmul44_quant_matmul_full_loadA_tla[512, 4096, 1024]33.26 us31.73 us[128,16384, 1024]38.68 us37.05 us[1024, 4096, 1024]42.66 us41.96 us[512,8192,1024]52.36 us44.05 us[128,16384,2048]53.77 us52.75 us说明标杆为QuantMatmul算子。统计耗时均为核函数总耗时使用msprof工具得到。上述测试例中A、B及输出矩阵均为layout::RowMajor排布方式。测试环境说明NPU型号为910B1CANN包版本为9.0.0。【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考