【CANN全新升级】CATLASS算子模板库分层设计,助力GEMM类算子开发效率翻倍
发表于 2025-06-09 11:42:51

Transformer架构中矩阵乘法(GEMM, General Matrix Multiplication)计算占据较高比重,其性能优化对提升整体计算效率至关重要。针对GEMM类算子编程,由于不同应用场景和优化目标存在差异,其实现变体呈现多样化特征;同时,随着算法持续演进和创新,还会不断产生新的定制化开发需求,这些需求往往难以通过预先枚举的方式全面覆盖。直接基于硬件底层能力进行GEMM类算子的定制开发,不仅技术门槛较高,还存在开发周期过长的问题。为此,昇腾CANN推出CATLASS算子模板库,采用分层模块化设计,通过将GEMM计算解耦为数据分块策略、计算单元配置等多个可灵活组合的功能组件,构建了一套支持快速搭建与组合的开发范式。

CATLASS算子模板库通过提供可复用的模板、基础组件和典型算子实践案例,面向昇腾硬件进行了深度优化。这种设计使得开发者能够基于模块化组装方式,快速完成计算流水线编排。同时,开发者可以根据具体硬件特性和计算需求,灵活定制计算内核实现方案,在确保高性能的同时,显著提升开发效率。

算子实现分层模块化设计

CATLASS采用分层抽象的设计理念,通过分析硬件架构特性和GEMM计算需求,将整体实现按照计算粒度自上而下分层设计。Device层是算子在Host侧的调用接口,提供完整算子能力;Kernel层则体现算子在NPU上的完整实现,涉及多个计算核的并行计算;Block层包含单个AI核的计算过程;Tile层则由数据搬入、数据计算、数据搬出等步骤组成,可通过最底层Basic层的指令组装实现。

该设计通过模板化方式提取各层共性逻辑,同时保留必要的差异化扩展能力,使得不同层级的软件抽象能够精准对应到特定硬件结构和计算流水阶段。算法框架中的特定步骤会延迟到子类实现,使得子类能够在不改变算法整体结构的情况下,灵活重定义其中的某些关键步骤。

图片1.png

算子流水自定义灵活配置

CATLASS提供了灵活的开发方式,开发者可以复用预置的范式来快速实现基础功能,也能够针对特定需求修改模块进行定制开发,还可以通过更换组件来实现自定义的流水组合。这种设计在保证计算性能的同时,为开发者提供了充分的灵活性和可拓展空间。

图片2.png

高性能算子实现样例

CATLASS针对不同算子应用场景,提供了高性能的算子基础组件和算子模板实现样例。样例中包含了在昇腾硬件上可参考的实现,如高效数据搬运模式,异步流水编排技巧,高性能接口使用示例。基于这些示例,开发者可以充分学习理解昇腾硬件的性能优化技术,从而针对各种优化实现场景高效完成自定义算子开发。相关样例在从上至下完全开源可参考。

高性能算子组件复用,开发效率显著提升

CATLASS针对Matmul和FA等算子的kernel性能优化场景,提供了高效的编程实现方案。在保持计算性能的同时,能够提升代码的复用程度,缩短开发周期。当前代码仓已经提供了高性能的Block和Tile层级代码模块50+,开发者可以通过复用已有模块快速拼装出自定义高性能算子,从而显著缩短开发周期。在算子样例中在Matmul算子的K轴优化实现中,开发周期从原先的2人周缩短至1人周;在MLA相关优化特性开发中,周期从8人周缩减至4人周。

定制场景算子性能提升

模板库针对定制场景使能快速开发能力,提供不同场景下的性能优化模块供开发者组装定制,算子在定制shape下的模型性能相比相应算子标杆性能提升5%-10%。

图片3.png

图片4.png

CATLASS算子模板库使用流程

以GroupedMatmul为例,以下代码示例将展示如何基于CATLASS算子模板库快速开发实现算子,教您三步完成QuantGroupedMatmul算子定制组装。(注:以下为示例代码,完整代码请参考开源代码仓内的算子模板样例)

第一步:指定输入大小,申请设备空间

1.png

第二步:指定输入信息及模板算子所需切分信息

2.1 指定硬件类型、切块大小、输入类型

2.png

2.2 指定算子计算需要的分核、切分策略等基础组件,并完成拼装

3.png

第三步:指定模板Kernel,即可调用

4.png

模板库代码仓开源,未来持续演进

CATLASS算子模板库在Gitee已正式上线,地址为https://gitee.com/ascend/catlass。当前支持A2/A3平台,提供OptmizedMatmul等20多个代码样例。未来,昇腾会持续升级,在近期会上线更多高性能的算子样例,如FA和Conv等。算子模板库社区的繁荣,离不开广大开发者的支持,欢迎各位下载使用,并在自定义开发之后将开发成果贡献至开源仓,社区会定期评选优秀开发者并给予对应激励,感谢广大开发者的支持。


「免责声明」:以上页面展示信息由第三方发布,目的在于传播更多信息,与本网站立场无关。我们不保证该信息(包括但不限于文字、数据及图表)全部或者部分内容的准确性、真实性、完整性、有效性、及时性、原创性等。相关信息并未经过本网站证实,不对您构成任何投资建议,据此操作,风险自担,以上网页呈现的图片均为自发上传,如发生图片侵权行为与我们无关,如有请直接微信联系g1002718958。

CSDN官方微信
扫描二维码,向CSDN吐槽
微信号:CSDNnews
微博关注
【免责声明:CSDN本栏目发布信息,目的在于传播更多信息,丰富网络文化,稿件仅代表作者个人观点,与CSDN无关。其原创性以及文中陈述文字和文字内容未经本网证实,对本文以及其中全部或者部分内容、文字的真实性、完整性、及时性本网不做任何保证或者承诺,请读者仅作参考,并请自行核实相关内容。您若对该稿件有任何怀疑或质疑,请立即与CSDN联系,我们将迅速给您回应并做处理。】