前言昇腾NPU作为华为面向AI计算的核心硬件平台其矩阵乘法算子开发长期面临定制难度高、开发周期长的困境。CANN软件栈在此背景下推出CATLASS算子模板库将通用矩阵乘法GEMM计算逻辑进行分层抽象与模板化使得算子代码可复用、可替换、可局部修改。CATLASS的核心价值在于开发者不再需要从零手写硬件亲和的底层流水编排而是通过组装预置模板组件快速构建针对特定shape和精度需求的高性能算子实现。在定制shape场景下CATLASS实现能达到标杆性能的零点九八到一点二倍这一数据已经过大量实测验证。从社区版本迭代来看v1.0.0到v1.5.0之间CATLASS逐步覆盖了基础Matmul、Batched Matmul、Grouped Matmul、量化Matmul、Flash Attention推理、StreamK Matmul、稀疏Matmul、W4A4量化等众多算子形态同时新增了对下一代昇腾硬件Ascend 950PR和Ascend 950DT的支持。这些演进并非简单的功能堆叠而是模板库分层设计理念在更多实际场景中的持续验证与拓展。CANN生态通过CATLASS这种模板化工程方案把昇腾NPU硬件能力转化为开发者可便捷调用的算子组装能力缩短了从算法设计到算子落地的距离。CATLASS的分层架构设计哲学CATLASS的全称是CANN Templates for Linear Algebra Subroutines中文名为昇腾算子模板库。其设计出发点是对GEMM类算子的计算过程进行结构化拆解。在Transformer架构中矩阵乘法计算占据重要比重不同场景和不同优化点的实现变种众多算法演进过程中会诞生大量新的定制化开发诉求难以事先枚举覆盖。直接基于硬件能力定制开发GEMM类算子时开发者面临的最大痛点是算子代码与硬件细节高度耦合修改一处牵动全局不同优化策略之间无法共享代码逻辑。CATLASS采用分层模块化设计解决这一痛点。整个GEMM计算被解耦为多个层级Device层负责Host侧适配与Kernel调度Kernel层负责Block间循环与后处理组合Block层负责单个Process内的主循环计算Tile层负责硬件指令级的分片搬运与MMAD运算。每个层级的模板组件都可以独立替换、局部修改而不影响其他层级的逻辑。这种分层抽象并非简单地把代码拆成多个文件。它的设计依据来自昇腾NPU硬件架构本身的层级特性。昇腾NPU的AI Core内部存在多级缓冲区全局内存GM、L1 Buffer、L0A/L0B/L0C Buffer数据从GM搬运到L1再搬运到L0进行MMAD计算这一搬运路径天然对应了Block层和Tile层的职责划分。Block层管理L1 Buffer上的数据分块策略Tile层管理L0 Buffer上的硬件指令调用。硬件流水线MTE1、MTE2、FixPipe等之间的数据依赖和同步操作也在Block层和Tile层得到对应处理。模板化方式提取各层共性逻辑的同时保留了必要的差异化扩展能力。算法框架中的特定步骤延迟到子类实现子类能够在不改变算法整体结构的情况下灵活重定义关键步骤。这种模板模式Template Method Pattern在C模板元编程中的体现正是CATLASS命名的由来——它不是一套固定算子集合而是一套可组装、可特化的模板框架。分层架构带来的工程收益体现在多个方面。代码复用率大幅提高——同一个TileMmad实现可被不同DispatchPolicy的BlockMmad组合使用同一个BlockMmad可被不同Epilogue的Kernel组合使用。修改隔离性得到保证——调整L1分块策略只需修改TileShape参数不会影响Kernel层和Device层的逻辑。跨平台适配成本降低——新硬件平台接入时只需扩展Tile层的特化实现上层逻辑保持共享。GEMM三层嵌套循环与CATLASS组件映射矩阵乘法的经典算法可以用三层嵌套循环描述。外层两重循环对应Block维度的M和N方向分块这两重循环在昇腾NPU上通过多个AI Core的并行执行来实现——代码上并不显式表达为for循环而是通过BlockIdx区分不同核心处理的数据块。内层循环对应K方向上的tile迭代每一次迭代包含从GM到L1的数据搬运、从L1到L0的数据搬运、以及L0上的MMAD计算。CATLASS将这三层嵌套循环映射到四层API组件。Device层对应整个算子的Host侧入口Kernel层对应Block间循环编排Block层对应K方向主循环Tile层对应单次MMAD指令调用。用户组装内核时遵循严格的实例化顺序创建特化的Block层MMAD组件、指定Block层后处理类型、指定数据走位方式BlockScheduler、在Kernel层将MMAD和后处理组合、将Kernel放入Device适配器。// 第一步创建所需的特化Block层mmad参数usingDispatchPolicyGemm::MmadAtlasA2Pingpongtrue;usingL1TileShapeGemmShape128,256,256;usingL0TileShapeGemmShape128,256,64;usingATypeGemm::GemmTypeElementA,LayoutA;usingBTypeGemm::GemmTypeElementB,LayoutB;usingCTypeGemm::GemmTypeElementC,LayoutC;usingBlockMmadGemm::Block::BlockMmadDispatchPolicy,L1TileShape,L0TileShape,AType,BType,CType;// 第二步指定Block层的后处理类型usingBlockEpiloguevoid;// 第三步指定计算时的数据走位方式usingBlockSchedulertypenameGemm::Block::GemmIdentityBlockSwizzle;// 第四步在kernel层将mmad和后处理组合usingMatmulKernelGemm::Kernel::BasicMatmulBlockMmad,BlockEpilogue,BlockScheduler;// 第五步将kernel放入device适配器usingMatmulCatlass::Gemm::Device::DeviceGemmMatmulKernel;这五步实例化顺序严格对应了从底层硬件计算单元到上层Host调用的组装链路。BlockMmad是计算核心必须先确定其调度策略和TileShapeBlockEpilogue可选配置后处理逻辑BlockScheduler决定多核间数据分块映射Kernel层将Block计算与后处理组合为完整内核逻辑Device层封装Host侧参数解析和Kernel启动。这种自底向上的组装顺序确保每一层的参数选择都有下层约束作为依据避免参数不匹配导致的编译或运行错误。DeviceGemm作为最外层适配器承担了Host侧Tiling参数解析、GM地址管理、Kernel启动参数构造等职责。开发者在使用DeviceGemm时只需传入矩阵A、B、C的GM地址和Tiling描述DeviceGemm自动完成Kernel参数打包和调度。这种封装方式将Host侧的样板代码降到最低开发者可以将精力集中在Block层和Tile层的参数选择上。Kernel层的BasicMatmul实现了Block维度的MN循环逻辑。它通过BlockScheduler获取当前Block负责的MN分块坐标之后再调用BlockMmad执行该分块的K方向主循环再调用BlockEpilogue执行后处理。BasicMatmul的代码逻辑在不同算子形态间保持一致——无论是基础Matmul还是Grouped MatmulMN循环的框架结构不变差异通过BlockMmad和BlockScheduler的模板参数体现。DispatchPolicy调度策略深度剖析DispatchPolicy是BlockMmad最重要的模板参数之一它决定了Block层主循环的具体实现方式。CATLASS采用基于标签的调度策略类型来特例化Block层MMAD实现这避免了代码重复使通用代码更容易编写并提供了清晰的扩展点供用户插入定制实现。CATLASS目前提供四种DispatchPolicy分别针对不同性能需求场景MmadAtlasA2Pingpong在L1和L0A/B设置Pingpong Buffer机制通过双缓冲区交替使用实现数据搬运与计算并行。STAGES参数控制缓冲区片数ENABLE_UNIT_FLAG参数控制是否启用Mmad运算与L0C结果拷出全局内存的细粒度并行。这是最基础的调度策略适用于00_basic_matmul、01_batched_matmul、03_matmul_add等简单场景。MmadAtlasA2Preload在Pingpong基础上增加了ShuffleK策略和Block间预加载能力。ShuffleK策略通过重新排列K方向上的计算顺序来优化Cache利用率Block间预加载允许当前Block的计算与下一个Block的数据搬运提前并行执行。MmadAtlasA2PreloadAsync引入了nBuffer机制支持L1和L0A/L0B/L0C上配置不同数量的缓冲区片同时支持ShuffleK策略、Block间预加载和group间预加载。PRELOAD_STAGES参数控制GM到L1数据读取多少次后启动L1到L0的数据搬运和Mmad计算L1_STAGES参数控制L1缓冲区片数量。这种异步调度策略适用于Grouped Matmul等需要多group计算的场景。structMmadAtlasA2PreloadAsync{staticconstexpruint32_tPRELOAD_STAGES1;staticconstexpruint32_tL1_STAGES2;staticconstexpruint32_tL0A_STAGES2;staticconstexpruint32_tL0B_STAGES2;staticconstexpruint32_tL0C_STAGES1;staticconstexprboolENABLE_UINT_FLAGfalse;staticconstexprboolENABLE_SHUFFLE_Ktrue;};多级缓冲区片数配置让开发者能精确控制每个缓冲区层级的资源占用。L0C设为1片是因为MMAD结果在L0C上只需暂存一轮计算完成后立即拷出到GM或经FixPipe处理L0A和L0B设为2片是为了支持Pingpong交替搬运L1设为2片配合PRELOAD_STAGES等于1实现预取一轮后启动计算。这种参数化设计让同一套模板代码能适配不同的内存预算和并行需求无需为每种配置重写实现逻辑。MmadAtlasA2PreloadAsyncWithCallback在PreloadAsync基础上增加了aic和aiv之间同步命令的callback机制用户可以将自定义同步逻辑以callback形式传入Block层由Block层决定调用时机。这种设计在量化Matmul等需要aic与aiv协作的场景中尤为关键因为量化计算中aic侧的反量化操作和aiv侧的MMAD计算之间存在精确的数据依赖关系。四种DispatchPolicy形成了从简单到复杂的递进关系。Pingpong满足基础场景Preload增加了ShuffleK和Block间预取PreloadAsync进一步支持nBuffer和group间预取PreloadAsyncWithCallback在此基础上提供自定义同步机制。开发者可以根据算子复杂度选择合适的策略层级不必为简单算子承担复杂策略的配置开销。TileShape分块参数与硬件缓冲区资源适配TileShape参数由三元组表示分别是M、N、K维度方向上的分块大小。L1TileShape对应L1 Buffer上的基本块尺寸L0TileShape对应L0 Buffer上的基本块尺寸。这两个参数的选择直接影响计算性能和内存利用率。L1TileShape的选择需要考虑L1 Buffer的容量限制。昇腾NPU的AI Core上L1 Buffer大小有限Atlas A2系列约为1.5MB分块尺寸过大会导致缓冲区溢出分块尺寸过小则无法充分利用MMAD指令的计算能力。以GemmShape128, 256, 256为例当矩阵A为FP16数据类型时单个L1分块占用128乘以256乘以2字节等于64KB矩阵B同样占用128KB256乘以256乘以2字节合计192KB远小于L1 Buffer容量为Pingpong双缓冲留出了充足空间。L0TileShape的选择需要与MMAD指令的硬件支持对齐。昇腾NPU的Cube单元支持特定尺寸的MMAD指令L0上的分块尺寸必须是这些指令尺寸的整数倍。GemmShape128, 256, 64的K维度设为64对应了Cube单元一次MMAD计算的K方向计算量M和N方向则根据L1分块和L0缓冲区容量综合确定。分块参数的选择还与DispatchPolicy耦合。Pingpong策略下L1需要2片缓冲区L1TileShape占用的空间需要满足2片不超过L1总容量PreloadAsync策略下需要更多片数对L1TileShape的约束更加严格。这种参数间的耦合关系正是CATLASS模板化组装的挑战所在——每个参数都不是孤立选择而是在硬件资源、调度策略、计算需求三者之间寻找平衡。// TileShape与缓冲区容量约束的工程计算示例// L1 Buffer容量约1.5MB (Atlas A2)// Pingpong STAGES2时需要为A和B各预留2片空间// FP16数据类型下:// A矩阵单片 L1TileShape.M * L1TileShape.K * sizeof(FP16)// B矩阵单片 L1TileShape.K * L1TileShape.N * sizeof(FP16)// 总占用 STAGES * (A单片 B单片)// GemmShape128, 256, 256:// A单片 128 * 256 * 2 65536 bytes 64KB// B单片 256 * 256 * 2 131072 bytes 128KB// 总占用 2 * (64KB 128KB) 384KB 1536KB, 约束满足constexpruint32_tl1_total_bytes1536*1024;constexpruint32_ta_tile_bytes128*256*2;constexpruint32_tb_tile_bytes256*256*2;constexpruint32_ttotal_stages_occupancy2*(a_tile_bytesb_tile_bytes);static_assert(total_stages_occupancyl1_total_bytes,L1 overflow);TileShape参数与缓冲区容量之间的约束关系必须在编译期确定C模板参数的静态特性恰好满足这一需求。将TileShape作为模板参数而非运行时参数使得编译器能在编译期检查缓冲区容量约束是否满足避免了运行时缓冲区溢出导致的计算错误或设备异常。static_assert等编译期断言机制进一步强化了这一安全保障开发者在编译阶段就能发现参数配置错误而非在运行时遇到难以排查的设备异常。L0TileShape的K维度选择也具有工程讲究。K维度过大时单次MMAD计算的数据量超过L0A和L0B的容量无法一次性完成K维度过小时MMAD指令的启动开销占比过高计算效率下降。64这个数值对应了Cube单元在FP16数据类型下的最优K方向计算粒度既充分利用了L0A/L0B的存储空间又避免了过细粒度带来的指令调度开销。Swizzle策略与多核数据走位编排BlockScheduler负责决定多核间的数据分块映射方式即不同AI Core处理哪些M和N方向的矩阵分块。CATLASS提供了GemmIdentityBlockSwizzle等基础策略也支持更复杂的Swizzle编排。Swizzle策略的核心目标是在多核并行场景下优化GM访问的带宽利用率。当多个AI Core按照简单的行列顺序分配分块时相邻Core可能在同一时刻访问GM上相近的地址区域导致Bank冲突和带宽拥塞。Swizzle策略通过重新排列分块到Core的映射顺序使得同一时刻活跃的Core访问的GM地址区域尽可能分散从而减少Bank冲突。以一个8乘以8的分块网格为例Identity Swizzle按行优先顺序将分块分配给Core 0、1、2、3……同一行上的分块连续映射到相邻Core这些Core在同一时刻需要读取矩阵B的相同K方向数据段造成GM读取热点。经过Swizzle重排后同一行上的分块被分散到不相邻的Core这些Core读取矩阵B的不同数据段GM读取压力更加均匀。CATLASS设计文档中的Swizzle策略说明详细介绍了不同Swizzle模式对分块顺序的影响。Identity Swizzle保持原始行列顺序适用于分块数量较少或GM带宽充足的场景更复杂的Swizzle模式则在大规模并行计算中带来明显的带宽优化效果。Swizzle策略的选择与DispatchPolicy中的ShuffleK策略相互配合。ShuffleK影响K方向上的计算顺序Swizzle影响MN平面上的分块映射两者共同决定了数据访问模式和计算顺序的时空分布。在Flash Attention等需要精细流水编排的算子中这种双重控制尤为重要因为Flash Attention的Q和K分块计算结果需要按特定顺序传递给后续步骤Swizzle编排既要考虑带宽优化又要满足数据依赖约束。量化Matmul与FixPipe随路量化机制量化推理是当前AI部署的关键技术路径CATLASS从v1.3.0版本开始支持FixPipe随路量化并在v1.5.0版本新增了W8A8 Per-Token量化、Per-Group和Per-Block量化Matmul TLA等示例。FixPipe是昇腾NPU上的一个硬件流水线单元它能在MMAD计算结果从L0C拷出到GM的过程中同步执行量化、反量化、类型转换等操作。这种随路处理方式避免了额外的内存往返开销——计算结果不需要先写回GM再单独调用量化算子处理而是在拷出路径上一步完成。在W8A8量化Matmul场景中矩阵A和矩阵B以INT8格式存储计算时需要先将INT8数据反量化到FP16或BF16进行MMAD再将FP32累加结果量化回INT8格式输出。CATLASS通过TileCopy中的量化与反量化配置和Epilogue中的FixPipe处理将这一流程编排为计算与量化的融合流水线。TileCopy在数据从GM搬运到L1或从L1搬运到L0的过程中执行反量化FixPipe在结果从L0C拷出到GM的过程中执行量化。Per-Token量化对每个Token使用独立的量化参数scale和zero_pointPer-Group量化对每组元素使用共享的量化参数Per-Block量化则对整个矩阵分块使用统一参数。不同量化粒度的选择影响量化参数的存储开销和精度损失。Per-Token量化精度最高但参数存储开销最大Per-Block量化参数存储开销最小但精度损失相对较大。CATLASS的模板化设计让开发者能通过配置GemmType参数灵活切换量化粒度无需重写计算流水逻辑。v1.5.0版本新增的动态W8A8 Per-Token量化示例103_dynamic_optimized_quant_matmul_per_token_basic展示了在Matmul泛化工程框架下实现Per-Token量化的完整流程。这个示例中量化参数作为额外的输入张量传入KernelTileCopy在搬运矩阵A和B的数据时根据每个Token对应的scale和zero_point执行反量化MMAD计算在反量化后的高精度数据上进行。W4A4量化是量化场景的极致压缩。矩阵A以4-bit权重存储矩阵B同样4-bit计算时反量化到更高精度进行MMAD。4-bit量化将内存占用压缩到原来的四分之一但反量化操作引入了额外计算开销。CATLASS通过FixPipe随路反量化将这一开销最小化使得W4A4量化在推理场景中具备实用价值。v1.4.0版本引入的38_w4a4_matmul_per_token_per_channel_dequant示例展示了这一机制的完整实现。仓库地址https://atomgit.com/cann/catlass