NPU算子“智能编译”：TileLang Developer模式来了

‍
在大模型算力竞赛的下半场，硬件性能只是入场券，软件生态的易用性是决胜的关键决定因素之一。昇腾（Ascend）NPU 凭借独特的Cube/Vector（矩阵/向量）存算分离架构，在硬件性能上表现强劲。然而，对于开发者而言，要想完全释放这份性能，往往面临着“易用性”与“高性能”的艰难权衡：

使用PyTorch等上层框架开发固然简单，自定义算子性能优化却难以触达NPU特有的存储层级红利；若直接使用底层语言进行编程，则需要开发者深入理解Cube与Vector的并行关系、手动管理复杂的多级内存搬运及核间同步，开发门槛极高。

近日，TileLang AscendNPU IR Developer模式的发布，为开发者打破了这堵“高墙”。

作为TileLang-Ascend提供的两大代码生成后端之一（另一个是Ascend C），TileLang AscendNPU IR支持Expert和Developer开发模式。如果说此前推出的Expert模式是为追求100%极致性能的专家打造的“手术刀”，那么随着Developer模式的演进，TileLang-Ascend正基于AscendNPU IR强大的编译推断能力，致力于让NPU算子开发进入“智能编译”时代。

Developer模式：算子开发的“智能编译”

Developer模式的核心目标，是实现从GPU到NPU跨架构无缝迁移。对于习惯了GPU编程思维的开发者，TileLang AscendNPU IR带来了一个平滑的过渡路径。

如下所示，在将一个经典的矩阵乘（GEMM）算子从GPU迁移至昇腾NPU时，开发者面临的代码改动微乎其微。得益于TileLang统一的T.Kernel和T.Parallel抽象，算子的核心切分逻辑（Tiling）和计算逻辑几乎保持一致。开发者仅需在编译阶段调整Backend Target，TileLang编译器后端即可自动将通用的并行原语映射为昇腾NPU特有的Cube/Vector指令，屏蔽了底层的架构差异。

NPU算子代码如下：

@tilelang.jit(out_idx=[-1], target="npuir")

def matmul(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float32"):

@T.prim_func

def gemm(

A: T.Tensor((M, K), dtype),

B: T.Tensor((K, N), dtype),

C: T.Tensor((M, N), dtype),

with T.Kernel(T.ceildiv(N, block_N) * T.ceildiv(M, block_M), is_npu=True) as (cid, _):

by = cid // T.ceildiv(N, block_N)

bx = cid % T.ceildiv(N, block_N)

A_shared = T.alloc_shared((block_M, block_K), dtype)

B_shared = T.alloc_shared((block_K, block_N), dtype)

C_local = T.alloc_fragment((block_M, block_N), accum_dtype)

for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=2):

T.copy(A[by * block_M, k * block_K], A_shared)

T.copy(B[k * block_K, bx * block_N], B_shared)

T.gemm(A_shared, B_shared, C_local, initC=(k == 0))

T.copy(C_local, C[by * block_M, bx * block_N])

return gemm

来源：https://github.com/tile-ai/tilelang-ascend/blob/npuir/examples/gemm/example_gemm.py

GPU算子代码如下：

@tilelang.jit(out_idx=[-1])

def matmul(M, N, K, block_M, block_N, block_K, dtype=T.float16, accum_dtype=T.float32):

@T.prim_func

def gemm(

A: T.Tensor((M, K), dtype),

B: T.Tensor((K, N), dtype),

C: T.Tensor((M, N), dtype),

with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):

A_shared = T.alloc_shared((block_M, block_K), dtype)

B_shared = T.alloc_shared((block_K, block_N), dtype)

C_local = T.alloc_fragment((block_M, block_N), accum_dtype)

T.clear(C_local)

for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):

T.copy(A[by * block_M, k * block_K], A_shared)

T.copy(B[k * block_K, bx * block_N], B_shared)

T.gemm(A_shared, B_shared, C_local)

T.copy(C_local, C[by * block_M, bx * block_N])

return gemm

来源：https://github.com/tile-ai/tilelang/blob/main/examples/gemm/example_gemm.py

这充分体现了跨架构无缝迁移的魅力——原有的Pythonic算子逻辑、Tiling策略以及循环结构均可直接复用，开发者无需重写核心算法，即可低成本完成跨硬件的算子迁移。

达成这一核心目标的背后，是TileLang AscendNPU IR通过编译器的智能分析与推断，将原本需要开发者手动管理的底层硬件细节实现了“自动化”。具体而言，Developer模式以如下四大核心特性，大幅降低了NPU编程的心智负担：

1. 消除显式CV分核，用户零感知

在昇腾NPU的异构架构中，计算任务被分为Cube（矩阵计算）和Vector（向量计算）两类核心。在传统开发模式下，开发者必须显式地指定哪段代码运行在Cube核，哪段运行在 Vector核，并处理两者间的交互。而在Developer模式下，编译器具备了自动推断计算任务类型的能力。它能根据计算逻辑，自动将任务分派给最合适的计算单元，开发者无需再手动编写T.Scope进行分核指定，真正实现了对底层异构核心的“零感知”。

2. 消除显式内存管理，自动搬运与同步

数据在Global Memory、L1、L0A/L0B、Unified Buffer (UB) 等多级存储间流转，是NPU编程中最繁琐的部分。TileLang AscendNPU IR通过引入自动内存规划，让编译器接管了这一切。这主要体现在如下两方面：

跨硬件架构的内存抽象：通过T.alloc_shared、T.alloc_fragment等跨硬件架构的两层抽象，开发者可以在不同的硬件内存级别中灵活分配缓冲区。这既屏蔽了底层硬件的差异，又能确保性能的提升。

数据跨层自动搬运：利用T.copy接口，开发者不再需要手动管理复杂的跨层数据搬运细节，编译器会自动处理数据在不同层级间的流动，并插入必要的同步指令。

3. 自动流水并行（T.Pipelined）

为了掩盖内存访问延迟，流水线并行（Pipeline Parallelism）是高性能算子的标配。在Developer模式中，这一高阶技巧被封装得极为简洁。本质上，通过T.Pipelined原语，编译器能够自动实现Double Buffer和缓冲区重用。它不仅处理常规的数据依赖，更能自动实现CV流水并行，将计算与数据搬运重叠执行，在不增加开发者负担的前提下，榨干硬件的每一分算力。

4. 向量化计算的自动映射（T.Parallel）

针对Vector核心的计算，TileLang AscendNPU IR引入了T.Parallel原语，使得开发者可以使用通用的数学符号（如a+b）或TileLang内置函数（如T.exp, T.max）在T.Parallel作用域内编写代码，AscendNPU IR后端会自动将其Lowering为昇腾硬件原生的向量指令，并自动处理数据切分和Mask掩码逻辑。这不仅保持了代码的Pythonic风格，更确保了在NPU Vector单元上的高效执行。

实战演示：CV融合算子的“智能编译”

如下代码展示了一个典型的CV融合算子开发实例。在该场景中，我们需要先执行矩阵乘法（Cube密集计算），随后紧接一个Softmax（Vector向量计算）。

@tilelang.jit(target="npuir")

def minicv(M, N, K, block_M, block_N):

m_num = M // block_M

n_num = N // block_N

@T.prim_func

def main(

A: T.Tensor((M, K), dtype),

B: T.Tensor((K, N), dtype),

C: T.Tensor((M, N), inner_dtype),

D: T.Tensor((M, N), inner_dtype),

with T.Kernel(m_num * n_num, is_npu=True) as (cid, _):

blockx = cid // n_num

bx = blockx * block_M

blocky = cid % n_num

by = blocky * block_N

A_BUF = T.alloc_shared((block_M, K), dtype)

B_BUF = T.alloc_shared((K, block_N), dtype)

C_BUF = T.alloc_fragment((block_M, block_N), inner_dtype)

D_BUF = T.alloc_fragment((block_M, block_N), inner_dtype)

T.copy(A[bx, 0], A_BUF, [block_M, K])

T.copy(B[0, by], B_BUF, [K, block_N])

T.gemm(A_BUF, B_BUF, C_BUF, [block_M, K, block_N], initC = True)

T.exp (C_BUF, D_BUF)

T.copy(C_BUF, C[bx, by], [block_M, block_N])

T.copy(D_BUF, D[bx, by], [block_M, block_N])

return main

来源：https://github.com/tile-ai/tilelang-ascend/blob/npuir/examples/mixcv/example_mixcv.py

我们可以看到，在Developer模式下，代表Cube矩阵计算的T.gemm与代表Vector向量计算的T.exp被自然地书写在同一段逻辑中。

首先是零感知分核与同步：代码中完全没有显式的SetFlag/WaitFlag同步原语，也无需手动划分Cube与Vector的执行Scope，编译器自动识别出T.gemm归属Cube核，T.exp归属Vector核，并自动插入同步指令处理数据依赖。

同时也实现了自动流水与内存管理：在整个计算过程中，开发者无需管理复杂的L1/L0/UB内存搬运，编译器后端自动实现了Cube计算与Vector搬运的流水线重叠（Double buffer）。

TileLang在昇腾生态中的独特价值

在当前的AI编译器领域，TileLang-Ascend并非在重复造轮子，而是占据了一个不可替代的独特生态位。与业界主流工具相比，它展现出了鲜明的差异化优势。

1. 硬件可感知，留住性能红利

现有技术如Triton，其设计哲学是“硬件不可知”，试图用一套逻辑屏蔽所有硬件细节。然而，这一注重通用性的技术路线，在拥有显式存储层级和存算分离（Cube/Vector）架构的专用硬件上，往往难以触达最底层的性能红利。

TileLang AscendNPU IR则跑通了“硬件可感知”的路线。基于TileLang的硬件抽象能力，用户可以通过注解输入更多的信息（如T.alloc_shared等），这些信息帮助AscendNPU IR编译器后端深度理解开发者的意图与NPU的硬件特性，从而做出更优的指令调度与内存分配。这种“懂硬件”的高级抽象，让TileLang既能保持易用性，又能触达通用编译器难以触及的底层优化空间。

此外，由于AscendNPU IR是基于MLIR框架构建的，它天然继承了MLIR生态的灵活性。对于熟悉MLIR的开发者而言，凭借既有经验即可无缝上手，轻松扩展自定义算子。同时，开发者还可以根据需求，在AscendNPU IR编译栈的任意层级进行细粒度的优化与调试。

2. 分层抽象，完善昇腾工具链

TileLang AscendNPU IR通过多级AscendNPU IR对接，构建了灵活的开发体系。它支持Developer和Expert等不同层级的开发模式，既能满足大众开发者对于算子编程易用性的诉求，又能满足专家用户对于极致性能的渴望，为昇腾生态提供了一个覆盖全场景的算子开发解决方案。

生态新范式：技术孵化与开源共建

TileLang Ascend的快速成熟，不仅仅是一个技术项目的成功，更是昇腾“开源开放” 平台战略与“产学研协同”模式的缩影。

首先，该项目得到了“北京大学鲲鹏昇腾科教创新卓越中心”（简称卓越中心）的强力支持。卓越中心致力于将高校在MLIR、编译技术领域的前沿探索，迅速转化为工业界的实战工具，TileLang-Ascend正是这种“学术界前沿技术+产业界硬件底座”化学反应的产物。

其次是开源共建放大成果。昇腾作为一个开放生态，AscendNPU IR为学术界及其他有创新需求的开发者提供了基于昇腾平台的创新研发机会，TileLang-Ascend从第一行代码开始就坚持开源，它不仅服务于华为内部，更深度协同PyTorch、vLLM等开源社区。

目前，社区已经基于TileLang实现了FlashAttention、MatMul等核心算子，并正在吸引越来越多的开发者参与共建，形成“技术孵化-开源共建-生态反哺”的良性循环。

结语：让算力触手可及

随着Developer模式基础特性的陆续开源，TileLang-Ascend正在将国产NPU算子开发带入一个自动化、智能化的新阶段。

对于开发者而言，这意味着不再需要深陷于流水线编排和内存管理的细节泥潭，只需关注算法逻辑本身。TileLang-Ascend正通过AscendNPU IR的技术突破，让算力触手可及。

想要了解技术路线的更多细节和代码样例的开发者，请移步代码仓tile-ai/tilelang-ascend，关注npuir分支：https://github.com/tile-ai/tilelang-ascend/tree/npuir，共同见证这一技术变革。