算子编程这件事,过去几年一直是个“少数人游戏”。写 CUDA C++ 的人要同时懂硬件架构、懂并行模型、还得懂编译器怎么把你的代码翻译成 SASS,门槛高到让大部分算法工程师望而却步。后来 Triton 出来了,用 Python DSL 把 tile 级别的并行抽象出来,一下子把门槛砍掉一大截。但 Triton 是 OpenAI 主导的,生态和话语权都在外面。TileLang 的出现,让我看到了一条不太一样的路——它不只是“又一个 Triton 替代品”,而是从算子编程语言这个切入点,去撬动整个国产开源生态的底层叙事。这篇文章我会从 TileLang 的设计思路、核心机制、实操写法、性能调优、生态影响几个维度展开,尽量把我知道的、踩过的、想明白的都写出来。
1. 算子编程语言到底在解决什么问题
1.1 从 CUDA 到 Triton 再到 TileLang 的演进逻辑
要理解 TileLang 的价值,得先搞清楚算子编程语言这个品类是怎么来的。最早写 GPU 算子就是纯 CUDA C++,你得手动管理 thread block、shared memory、register 分配,还得考虑 bank conflict、warp divergence 这些底层细节。一个 GEMM 算子写下来,几百行代码是常态,调优周期以周为单位。这种方式灵活度最高,但人力成本极高,而且代码几乎不可移植——换一代硬件架构,可能就得重写。
Triton 的思路是把“一个 thread block 处理一块 tile”这个模式抽象出来,让开发者用 Python 语法描述 tile 级别的计算,编译器负责往下映射到具体的线程和内存层级。这个抽象层次选得很巧妙:它足够高,让你不用管线程调度;又足够低,让你能控制 shared memory 的使用和计算的分块策略。所以 Triton 在学术界和工业界都铺得很快。
TileLang 走的是类似的路,但它在几个关键点上做了不同的取舍。第一,它更强调对硬件特性的显式表达,比如你可以直接指定某个 buffer 放在 shared memory 还是 fragment(寄存器),这在 Triton 里是编译器自动决定的。第二,它的调度原语更细粒度,支持软件流水(software pipeline)、warp specialization 这些高级优化手段。第三,也是最重要的,它是国产团队主导的开源项目,从设计之初就考虑了国产硬件后端的适配问题。
1.2 TileLang 的核心定位与目标用户
TileLang 的定位很明确:面向高性能算子开发的领域特定语言(DSL),基于 TVM 的 TIR 基础设施构建,用 Python 作为前端语法。它的目标用户不是刚入门深度学习的调参侠,而是需要手写高性能 kernel 的算子工程师、编译器工程师、以及做推理框架底层优化的那批人。
我自己的感受是,TileLang 的学习曲线介于 CUDA 和 Triton 之间。如果你有 CUDA 经验,理解它的内存层级和调度原语会很快;如果你只有 PyTorch 经验,那需要补一下 GPU 执行模型的基础知识。但一旦上手,写一个高性能 GEMM 或者 Flash Attention 的效率,比手写 CUDA 快五到十倍不止。
它解决的问题可以归纳成三个层面。第一层是开发效率:用几十行 Python 代码替代几百行 CUDA,且性能不输手写。第二层是可移植性:同一份算子描述,可以通过不同的后端编译到不同硬件上,包括国产加速卡。第三层是生态自主:在算子编程这个底层环节,不再完全依赖外部主导的框架和工具链。
1.3 为什么“算子编程语言”是国产开源生态的关键拼图
很多人觉得算子编程语言是个很窄的领域,跟“生态”这种大词扯不上关系。但我的判断恰恰相反:算子编程语言是连接上层框架和底层硬件的咽喉要道。你想想,PyTorch 的算子库、推理引擎的 kernel 实现、训练框架的通信算子,最终都要落到某个具体的编程模型上。如果这个编程模型是别人定义的,那你的硬件适配、性能优化、甚至功能支持,都得跟着别人的节奏走。
国产开源生态这些年发展很快,框架层有 PaddlePaddle、MindSpore,推理层有 ncnn、MNN,但算子编程这一层一直是空白。大家要么直接用 CUDA,要么用 Triton,要么自己造一套内部 DSL 但不开源。TileLang 填补的就是这个空白。它的意义不在于技术本身有多颠覆,而在于它提供了一个公共的、开放的、可演进的算子编程基础设施,让国产硬件厂商、框架开发者、算子工程师有一个共同的协作界面。
2. TileLang 的核心设计拆解
2.1 基于 TVM TIR 的分层架构
TileLang 的架构可以粗略分成三层。最上层是 Python DSL,开发者用@T.prim_func装饰器定义算子,用T.Kernel描述 grid 和 block 的划分,用T.alloc_shared、T.alloc_fragment这些原语管理内存。中间层是 TileLang 自己的 IR,它把 Python AST 转换成一种带 tile 语义的中间表示,然后做一系列优化 pass,比如 layout inference、pipeline scheduling、memory planning。最下层是代码生成,通过 TVM 的 codegen 后端输出 CUDA C、HIP、或者国产硬件的 DSL。
这个分层设计的好处是解耦。前端语法可以独立演进,中间优化可以针对不同硬件做特化,后端 codegen 可以按需扩展。我实测下来,从 Python 到 CUDA C 的编译时间在秒级,比 TVM 原生的 TE 调度方式要快不少,因为 TileLang 的 IR 更贴近硬件,优化 pass 的搜索空间小了很多。
注意:TileLang 依赖 TVM 的特定版本,安装时建议用官方提供的 wheel 包或者从源码编译,不要混用不同版本的 TVM,否则会出现 IR 不兼容的报错。
2.2 Tile 级抽象与内存层级显式管理
TileLang 最核心的抽象是 tile。一个 tile 就是一块数据,可以是 shared memory 里的一块矩阵,也可以是寄存器里的一组向量。开发者用 tile 级别的操作来描述计算,比如T.gemm(A_shared, B_shared, C_fragment)就表示从 shared memory 读两块矩阵,做矩阵乘,结果写到寄存器 fragment 里。
这种抽象的好处是,它把“数据在哪”和“怎么算”分开了。你可以先决定 A 和 B 放在 shared memory,C 放在 fragment,然后描述计算逻辑。编译器会根据你的内存分配和计算描述,自动推导出线程映射和指令调度。这比 CUDA 里手动算 thread index、手动做 swizzle 要省心太多。
但 TileLang 没有完全把内存管理藏起来,这是它和 Triton 的一个关键区别。在 Triton 里,你写tl.load和tl.store,编译器决定数据放哪。在 TileLang 里,你可以显式地T.alloc_shared和T.alloc_fragment,甚至可以用T.annotate_layout指定数据布局。这种显式性在调优时非常有用,因为你可以精确控制 shared memory 的使用量,避免编译器做出你不想要的决策。
2.3 调度原语:软件流水与 Warp Specialization
TileLang 提供了一组调度原语,让你能表达复杂的流水线策略。最常用的是T.Pipelined,它可以把一个循环体拆成多个 stage,让数据加载和计算重叠执行。比如在 GEMM 里,你可以把 K 维度的循环标记为 pipelined,编译器会自动插入 double buffer 或者 multi-buffer,让下一轮的 A、B 加载和当前轮的矩阵乘并行。
Warp specialization 是另一个高级特性。它允许你把不同的 warp 分组,一组专门做数据搬运,一组专门做计算,通过 shared memory 做生产者-消费者同步。这个模式在 Hopper 架构上特别有效,因为 Hopper 有 TMA(Tensor Memory Accelerator)可以做异步大块数据搬运。TileLang 对 TMA 的支持是通过T.copy原语加上 pipeline 注解来实现的,写起来比手写 CUDA 的 mbarrier 简单得多。
我试过用 TileLang 写一个带 warp specialization 的 GEMM,代码量大概是 CUDA 版本的三分之一,性能能达到手写 CUTLASS 的 90% 左右。对于大部分非极致场景,这个性价比已经很高了。
2.4 多后端支持与国产硬件适配路径
TileLang 的后端架构是插件式的。目前官方支持 CUDA 和 HIP(AMD),社区在推进国产加速卡的适配。适配一个新后端的工作量主要在三块:一是 codegen,把 TileLang IR 翻译成目标硬件的编程接口;二是 runtime,处理内存分配、kernel launch、同步这些运行时逻辑;三是 intrinsic 映射,把T.gemm、T.copy这些高层原语映射到硬件提供的加速指令上。
国产硬件适配的难点在于,很多国产加速卡的编程模型和 CUDA 不完全一样。比如有的卡没有独立的 shared memory,有的卡的矩阵乘指令形状是固定的,有的卡的异步拷贝机制不同。TileLang 的 tile 抽象层提供了一个缓冲,让这些差异可以在后端消化,而不是暴露给算子开发者。这是它作为“公共基础设施”的价值所在。
3. 上手实操:从零写一个高性能 GEMM
3.1 环境搭建与依赖安装
先把环境搞起来。TileLang 的安装方式有几种,我推荐用 pip 装官方 wheel,最省事:
pip install tilelang如果你需要最新特性或者要改编译器,那就从源码编译:
git clone https://github.com/tile-ai/tilelang.git cd tilelang pip install -e .编译依赖 CMake、LLVM(建议 16 以上)、以及 CUDA Toolkit(如果要跑 NVIDIA 后端)。我踩过的坑是 LLVM 版本不匹配导致 TVM 编译失败,后来统一用 LLVM 17 就稳了。另外,如果你在容器里跑,记得把CUDA_VISIBLE_DEVICES设对,TileLang 的 runtime 会读这个环境变量。
验证安装是否成功:
import tilelang print(tilelang.__version__)能打印出版本号就说明基础环境 OK 了。
3.2 第一个 TileLang 算子:向量加法
别一上来就写 GEMM,先用一个向量加法把流程跑通。下面是一个最简单的 TileLang kernel:
import tilelang import tilelang.language as T @tilelang.jit(out_idx=[2]) def vector_add(N, block_N, dtype="float16"): @T.prim_func def main( A: T.Tensor((N,), dtype), B: T.Tensor((N,), dtype), C: T.Tensor((N,), dtype), ): with T.Kernel(T.ceildiv(N, block_N), threads=128) as bx: for i in T.Parallel(block_N): idx = bx * block_N + i if idx < N: C[idx] = A[idx] + B[idx] return main这段代码的结构很清晰。@tilelang.jit是编译装饰器,out_idx=[2]告诉编译器第三个参数是输出,会自动分配内存。T.Kernel定义 grid 维度,这里是一维 grid,每个 block 128 个线程。T.Parallel表示这个循环会被并行化到线程上。
调用方式:
import torch N = 1024 a = torch.randn(N, dtype=torch.float16, device="cuda") b = torch.randn(N, dtype=torch.float16, device="cuda") c = vector_add(N, 256)(a, b) print(c)跑通这个例子,你就理解了 TileLang 的基本骨架:Kernel 定义 grid,Parallel 做线程级并行,Tensor 描述数据形状。
3.3 分块 GEMM 的完整实现与参数选择
现在上硬菜,写一个分块 GEMM。目标计算 C = A @ B,其中 A 是 M×K,B 是 K×N,C 是 M×N。分块策略是经典的 shared memory tiling:
@tilelang.jit(out_idx=[2]) def matmul(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float32"): @T.prim_func def main( 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 ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3): T.copy(A[by * block_M, ko * block_K], A_shared) T.copy(B[ko * 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 main这段代码有几个关键点值得展开。第一,T.Kernel用了二维 grid,bx 对应 N 维度,by 对应 M 维度,这是 GEMM 的标准划分方式。第二,T.alloc_shared分配了两块 shared memory 分别存 A 和 B 的 tile,T.alloc_fragment分配了累加器。第三,T.Pipelined把 K 维度的循环做了软件流水,num_stages=3表示三缓冲,让数据加载和计算重叠。
参数选择上,block_M、block_N、block_K 的取值直接影响性能。我的经验是:对于 fp16 输入,block_M=128、block_N=128、block_K=32 是一个比较稳的起点。block_K 不宜太大,因为 shared memory 容量有限;也不宜太小,否则流水线填充开销占比高。num_stages 一般取 2 到 4,取决于 shared memory 剩余容量和计算密度。
3.4 编译产物检查与性能验证
写完 kernel 后,第一件事是看编译出来的 CUDA 代码长什么样:
kernel = matmul(1024, 1024, 1024, 128, 128, 32) print(kernel.get_kernel_source())这会打印出生成的 CUDA C 代码。你可以检查 shared memory 分配是否符合预期、pipeline 是否真的插入了 double buffer、有没有多余的同步。我一般会重点看__shared__数组的大小和cp.async指令的数量。
性能验证用 torch 做基准对比:
import torch M = N = K = 4096 a = torch.randn(M, K, dtype=torch.float16, device="cuda") b = torch.randn(K, N, dtype=torch.float16, device="cuda") kernel = matmul(M, N, K, 128, 128, 32) c = kernel(a, b) ref = a @ b print("max diff:", (c - ref).abs().max().item()) # 测速 import triton.testing ms = triton.testing.do_bench(lambda: kernel(a, b)) tflops = 2 * M * N * K / (ms * 1e-3) / 1e12 print(f"{ms:.3f} ms, {tflops:.1f} TFLOPS")在 A100 上,这个配置大概能跑到 250-300 TFLOPS(fp16),接近 cuBLAS 的水平。如果跑不到,先检查 block 配置和 num_stages,再检查生成的代码里有没有 bank conflict。
4. 性能调优与常见问题排查
4.1 内存布局与 Bank Conflict 处理
Bank conflict 是 shared memory 性能的头号杀手。TileLang 默认会做 layout inference,自动给 shared memory 加 padding 或者做 swizzle,但有时候推断结果不是最优的。如果你发现 GEMM 性能明显低于预期,可以手动指定 layout:
T.annotate_layout({ A_shared: tilelang.layout.make_swizzled_layout(A_shared), B_shared: tilelang.layout.make_swizzled_layout(B_shared), })Swizzled layout 的原理是把 shared memory 的地址做异或重排,让同一个 warp 内的线程访问不同 bank。对于 fp16 的 128×32 tile,默认行优先布局下,同一列的访问会落到同一个 bank,造成 8-way conflict。Swizzle 之后可以降到无冲突。
我实测过一个 4096³ 的 GEMM,不加 swizzle 是 180 TFLOPS,加了之后到 270 TFLOPS,差距非常明显。所以如果你在调 GEMM,swizzle 是第一个要检查的点。
4.2 流水线深度与 Shared Memory 容量权衡
num_stages不是越大越好。每增加一个 stage,就要多分配一份 A_shared 和 B_shared。以 block_M=128、block_N=128、block_K=32、fp16 为例,一份 A_shared 是 128×32×2 = 8KB,一份 B_shared 是 32×128×2 = 8KB,加起来 16KB。num_stages=3 就是 48KB,num_stages=4 就是 64KB。A100 的 shared memory 上限是 164KB(可配置),所以理论上可以开到 8 以上,但实际收益在 3-4 之后就递减了。
原因是流水线的收益来自“加载延迟被计算掩盖”,当计算时间已经大于加载时间时,再加深流水线只是浪费 shared memory。你可以用 nsight compute 看smsp__warp_issue_stalled_long_scoreboard这个指标,如果它很低,说明加载不是瓶颈,加 stage 没用。
提示:TileLang 编译时会检查 shared memory 是否超限,如果超了会报错。你可以通过
T.annotate设置动态 shared memory 大小,但要注意不同硬件的上限不同。
4.3 常见编译错误与运行时问题速查
下面这张表是我在实际使用中整理出来的常见问题:
| 问题现象 | 可能原因 | 解决方法 |
|---|---|---|
| 编译报 IR 不兼容 | TVM 版本不匹配 | 用官方 wheel 或统一源码编译 |
| kernel launch 失败 | grid/block 配置超限 | 检查 threads 是否超过 1024,grid 维度是否超限 |
| 结果不正确 | 边界处理缺失 | 检查 M/N/K 是否能被 block 整除,加边界判断 |
| 性能远低于预期 | bank conflict 或流水线未生效 | 加 swizzle layout,检查 num_stages |
| shared memory 超限 | block 配置过大 | 减小 block_K 或 num_stages |
| 编译时间过长 | 优化 pass 搜索空间大 | 简化 kernel 结构,减少动态 shape |
还有一个容易忽略的点:TileLang 的 JIT 编译是带缓存的,第一次编译慢,后续调用会命中缓存。如果你改了 kernel 代码但发现行为没变,可能是缓存没失效,清一下~/.tilelang/cache目录。
4.4 调优实战:从 180 到 290 TFLOPS 的优化记录
记录一次完整的调优过程。初始版本:block_M=128、block_N=128、block_K=32、num_stages=2、无 swizzle,跑出来 180 TFLOPS。
第一步,加 swizzle layout,涨到 230 TFLOPS。第二步,num_stages 从 2 调到 3,涨到 260 TFLOPS。第三步,block_K 从 32 调到 64,涨到 275 TFLOPS,但 shared memory 用量翻倍,num_stages 得降到 2。第四步,试了 block_M=256、block_N=128,涨到 290 TFLOPS,但寄存器压力变大,occupancy 下降。
最后的配置是 block_M=256、block_N=128、block_K=32、num_stages=3、swizzle 开启,稳定在 285-290 TFLOPS。这个过程中,nsight compute 帮了大忙,主要看三个指标:shared memory 的 bank conflict 次数、warp 的 stall 原因分布、以及 tensor core 的利用率。
5. TileLang 对国产开源生态的影响
5.1 填补算子编程层的自主空白
国产开源生态在框架层和推理层都有拿得出手的项目,但算子编程层一直是短板。PaddlePaddle 有自己的算子库,但那是框架内部的,不对外提供编程接口。各家芯片厂商有自己的 DSL,但都是私有的,互不兼容。TileLang 的出现,第一次在开源社区提供了一个厂商中立的算子编程语言。
这个意义在于,它让算子开发者有了一个共同的语言。以前你给某家国产卡写算子,得学它那套私有 DSL,换一家就得重学。现在如果大家都适配 TileLang 后端,那算子开发者只需要写一份 TileLang 代码,就能跑在不同硬件上。这大大降低了跨硬件迁移的成本,也让算子库的复用成为可能。
5.2 降低国产硬件适配门槛的路径分析
适配一个新硬件后端,传统方式是从零写一套 codegen 和 runtime,工作量大、周期长。TileLang 提供了一条更轻量的路径:复用 TVM 的基础设施,只需要实现目标硬件的 codegen 和 intrinsic 映射。具体来说,硬件厂商需要做三件事。
第一,实现 TileLang IR 到目标硬件编程接口的 codegen。如果目标硬件的编程模型和 CUDA 接近,可以基于 CUDA codegen 改;如果差异大,就得重写。第二,实现 runtime 层,包括内存分配、kernel launch、流同步。第三,把T.gemm、T.copy、T.Pipelined这些高层原语映射到硬件的加速指令上。如果硬件没有对应的加速指令,就用软件模拟,但性能会打折扣。
我了解到的情况是,已经有多家国产加速卡厂商在推进 TileLang 后端适配,进度不一。有的已经能跑通基础算子,有的还在 codegen 阶段。这个生态一旦成型,对国产硬件的软件生态是很大的补强。
5.3 社区协作模式与生态位分析
TileLang 的社区协作模式和传统的框架项目不太一样。它更像是一个“标准+实现”的组合:TileLang 定义了一套算子编程的抽象和 IR 标准,各家厂商基于这个标准做后端实现。这种模式在编译器领域有先例,比如 LLVM 定义 IR 标准,各家做前端和后端。
TileLang 在生态中的位置,可以理解为“算子层的 LLVM”。它不直接面向最终用户,而是面向算子开发者和硬件厂商。它的成功不取决于自己有多流行,而取决于有多少硬件后端和算子库愿意基于它构建。从这个角度看,TileLang 的生态位是基础设施,而不是应用框架。
5.4 对开发者的实际影响与技能建议
对算子工程师来说,TileLang 带来的变化是实实在在的。以前写一个高性能算子,得同时懂 CUDA、懂硬件架构、懂编译器。现在用 TileLang,你主要需要懂的是算法逻辑和 tile 级别的优化策略,底层的线程映射和指令调度交给编译器。这降低了入门门槛,但也意味着你需要理解编译器的行为,才能写出高性能代码。
我的建议是,如果你在做推理框架或训练框架的底层优化,TileLang 值得花时间学。学习路径可以是:先跑通官方 example,然后自己写几个经典算子(GEMM、Flash Attention、LayerNorm),再尝试调优和自定义调度。如果你在做国产硬件适配,那更需要深入理解 TileLang 的 IR 和后端接口,因为这是你对接生态的入口。
6. 一些实操心得与后续扩展方向
6.1 我踩过的几个坑
第一个坑是版本管理。TileLang 迭代很快,不同版本的 API 有变化。我有一次用旧版本的写法在新版本上跑,报了一堆莫名其妙的错。后来养成习惯,每个项目固定一个版本,用 requirements.txt 锁死。
第二个坑是动态 shape。TileLang 对动态 shape 的支持还在完善中,如果你的 M/N/K 是运行时才确定的,可能需要编译多个 kernel 变体,或者用 padding 把 shape 对齐到 block 的整数倍。padding 的方式简单但浪费计算,多 variant 的方式高效但编译开销大。
第三个坑是调试。TileLang 编译出来的 CUDA 代码可读性一般,直接看汇编更痛苦。我的做法是先用小 shape 跑正确性,再用 nsight compute 看性能指标,定位到具体问题后再去看生成的代码。不要一上来就啃生成的 CUDA,效率太低。
6.2 后续可以深入的方向
如果你已经把基础算子写熟了,可以往这几个方向深入。一是自定义调度原语,TileLang 允许你注册自己的 IR pass 和调度策略,针对特定硬件做极致优化。二是多算子融合,把多个算子合并成一个 kernel,减少 kernel launch 开销和内存往返。三是自动调优,结合 autotuning 框架,让编译器自动搜索最优的 block 配置和流水线参数。
还有一个值得关注的方向是 TileLang 和推理引擎的结合。现在大部分推理引擎的算子库还是手写的,如果能把 TileLang 集成进去,用 JIT 编译的方式生成算子,就能根据实际 shape 做特化,性能可能比预编译的通用算子更好。这个思路在业界已经有实践,TileLang 提供了一个不错的工具基础。
6.3 给不同阶段读者的学习建议
如果你是刚接触 GPU 编程的新手,建议先补一下 CUDA 的基础知识,理解 thread、block、shared memory、warp 这些概念,再来看 TileLang 的抽象,会顺畅很多。如果你有 CUDA 经验,直接上手写 GEMM 和 Flash Attention,遇到问题查官方 example 和文档,基本能解决。
如果你是从业多年的算子工程师,TileLang 对你来说最大的价值是提效。你可以把以前手写 CUDA 的算子用 TileLang 重写一遍,对比性能和开发时间,感受一下抽象层次提升带来的收益。同时,你也可以参与到 TileLang 的社区建设中,比如贡献后端适配、优化 pass、或者算子库,这对个人成长和生态建设都有好处。
最后分享一个小技巧:TileLang 的 example 目录里有大量高质量的算子实现,包括 GEMM、Flash Attention、MoE、卷积等。与其从零开始写,不如先读一遍这些 example,理解它们的调度策略和参数选择,然后基于它们改。这是上手最快的路径,也是我自己的做法。