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則跑通了“硬件可感知”的路線?；赥ileLang的硬件抽象能力，用戶可以通過注解輸入更多的信息（如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，共同見證這一技術變革。