摘要:在AI技术快速发展的今天,我们常常听到开发者们提到一个问题:AI算子开发复杂度高、周期长。有没有一种方式,能让这个过程变得更简单、更直观?
在AI技术快速发展的今天,我们常常听到开发者们提到一个问题:AI算子开发复杂度高、周期长。有没有一种方式,能让这个过程变得更简单、更直观?
最近,在CANN开源开放系列活动中,我们看到了一个新的可能——TileLang-Ascend。
什么是TileLang?
TileLang是由北京大学杨智副教授团队主导开发的一款开源AI算子编程语言。它的设计初衷很纯粹:让AI算子开发更容易。
随着AI大模型的不断深入发展和应用场景对吞吐、时延的更高要求,AI硬件加速技术也不断向前演进。为了实现更高的数据传输和计算效率,新一代AI计算体系架构设计要求以更大的Tile数据块粒度来进行指令集的设计。
这种架构演进趋势对AI编程范式也提出了新的要求,需要在更大的Tile粒度上进行编程和编译,实现Tile编程模型计算抽象和Tile硬件抽象的一致,这样才能充分挖掘硬件的极致性能。TileLang就是在这种背景下应运而生的。
它采用类Python的语法,让熟悉Python的开发者能够快速上手。它的主要设计理念是提出的Tile级语言抽象,通过将数据分块处理(Tiling)实现高效内存管理和计算调度。
TileLang设计上分为Tile-VM和和API控制原语两部分。Tile-VM主要包含两个抽象,一个是对计算任务的Tile Task,另外一个是执行Tile Task的硬件Tile抽象。Tile Task是一个针对不同数据块层次的多级Tile,对应到硬件的不同内存层级,硬件操作也按照Tile粒度去搬运和计算。
TileLang设计了一套API控制原语来实现对Tile的多种操作,例如下面列出的Tile的分配操作以及计算操作。内存分配API可以在指定的内存层级上分配存储单元,从而提升算子的性能。
计算类操作API则提供对Tile粒度的数据操作,将数据块的索引等内部细节封装起来,大大简化了编程实现,并且容易理解和维护。
TileLang在昇腾NPU上的实践
2025年9月,DeepSeek-V3.2-Exp发布并开源,昇腾也同步完成了对TileLang的适配,并开源了Sparse Flash Attention和Lightning Indexer两个算子的TileLang实现(参考:
https://gitcode.com/cann/cann-recipes-infer/blob/master/docs/models/deepseek-v3.2-exp/deepseek_v3.2_exp_tilelang_operator_guide.md)。这意味着,现在开发者可以在昇腾NPU上使用TileLang来开发算子了。
TileLang在NPU上的整个工作流程分为编译和运行两个阶段,如下图所示:
编译阶段
多级Lowering转换:TileLang算子根据NPU硬件特性进行多级降级,生成针对昇腾硬件优化的TensorIR表示。
Ascend C代码生成:基于TensorIR,使用专门的Ascend Codegen模块生成对应的Ascend C代码。
动态库编译:通过毕昇编译器(BiSheng)将Ascend C代码编译成动态链接库(.so文件)。
运行阶段
库文件加载:通过torch_npu运行时库将.so文件加载到Python环境中。
函数封装:将算子封装为可调用的Python函数对象。
执行调用:用户以普通Python函数方式调用,提供输入张量即可在昇腾NPU上执行计算并获得结果。
下面的TileLang-Ascend的整体编译运行流程图进一步展示了本轮适配中修改的三大核心组件(深绿色显示)是如何协同工作的:
TileLang API原语
算子语言原语是支撑算子开发与执行的基础操作单元,其贯穿算子运行的完整生命周期。Tilelang API原语封装了昇腾Ascend C后端接口,在此基础上提供精简高效的原语指令供用户使用,各原语通过功能协同,实现算子对硬件资源的高效调度与数据处理逻辑的可靠运行。
我们可以把TileLang API原语理解为一套预先定义好的、基础的操作指令。它们封装了昇腾NPU底层硬件的复杂接口,提供了如内存分配、数据搬运、计算等精简高效的指令。开发者使用这些像积木一样的原语,就能组合出各种复杂的算子,实现对硬件资源的高效调度。
针对昇腾后端实现的几类典型原语介绍如下表所示:
CodeGen
在TileLang面向昇腾NPU的算子开发流程中,CodeGen(代码生成)模块是实现“前端简洁开发”与“后端高效执行”衔接的关键组件,其核心作用是通过指令模板封装与指令映射两大机制,将TileLang API高层抽象语法自动转化为标准Ascend C语言,为后端算子开发提供标准化、规范化与高效化支撑。
我们可以把它当作整个流程的“翻译官”,它的任务就是把前端简洁的TileLang代码,自动转换为后端NPU能高效执行的Ascend C代码。
JIT
JIT(Just-in-time,即时编译)是一种动态编译技术,Tilelang算子开发过程中通过JIT调用CodeGen生成Ascend C代码,并对整个过程进行动态调控,解决静态编译的局限性,确保生成的Ascend C代码动态适配NPU特性,同时最大化提升算子执行效率。
使用JIT技术让整个流程变得更加智能灵活。它在算子运行时动态工作,解决静态编译的局限。
TileLang-Ascend快速入门
下面让我们通过一个matmul的例子介绍TileLang的基础语法结构以及如何快速开发一个TileLang算子。
@tilelang.jit装饰器标识函数会被jit动态编译,out_idx表示主计算函数返回参数索引。
@T.prim_func装饰器在模块内定义主计算函数。
T.Tensor声明数据类型Tensor数据缓冲区,并指定其形状和数据类型。
T.Kernel原语触发算子kernel调用,对应Ascend C的kernel调用。
T.Scope标识代码的计算单元,”C”表示在cube核运行,“V”表示在vector核上运行。
return main返回主计算函数对象供调用。
上述的算子执行流程为:
通过T.kernel调用创建并发任务执行单元,并发任务数为m_num*n_num,其中m_num和n_mum是分别沿M轴和N轴的数据块切分数量,二者乘积即为总的切分数据块数量,也即对应的任务数量。
通过T.alloc_L1分配两个输入矩阵块的存储单元,T.alloc_L0C为输出矩阵存储单元。K_L1为沿K轴的切分,循环单元通过T.gemm_v0实现对各个K轴切分块的计算和累加。
计算结束,通过T.copy将计算结果搬运对应存储单元。
可以看到使用TileLang-Ascend开发AI算子相比传统的算子开发方式,代码更加简洁,结构也更清晰。
加入我们,共同打造更丝滑的开发体验
TileLang-Ascend还处于起步阶段,很多功能还在不断完善中。比如,目前对TileLang算子的打印调试功能还在开发中,现在还是沿用的Ascend C提供的PRINT功能。
具体的使用方式可以参见:
对TileLang算子的优化,我们也会尽量将通用的优化方法融入到TileLang的编译框架中,而不被用户感知,用户只需要实现算子的计算逻辑,即可实现高性能的算子。
TileLang-Ascend开源社区:
TileLang-Ascend环境搭建指南:
来源:新浪财经