TileLang：国产算子编程语言如何撬动开源生态

算子编程这件事过去几年一直是个“少数人游戏”。写 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 memoryC 放在 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 SpecializationTileLang 提供了一组调度原语让你能表达复杂的流水线策略。最常用的是T.Pipelined它可以把一个循环体拆成多个 stage让数据加载和计算重叠执行。比如在 GEMM 里你可以把 K 维度的循环标记为 pipelined编译器会自动插入 double buffer 或者 multi-buffer让下一轮的 A、B 加载和当前轮的矩阵乘并行。Warp specialization 是另一个高级特性。它允许你把不同的 warp 分组一组专门做数据搬运一组专门做计算通过 shared memory 做生产者-消费者同步。这个模式在 Hopper 架构上特别有效因为 Hopper 有 TMATensor Memory Accelerator可以做异步大块数据搬运。TileLang 对 TMA 的支持是通过T.copy原语加上 pipeline 注解来实现的写起来比手写 CUDA 的 mbarrier 简单得多。我试过用 TileLang 写一个带 warp specialization 的 GEMM代码量大概是 CUDA 版本的三分之一性能能达到手写 CUTLASS 的 90% 左右。对于大部分非极致场景这个性价比已经很高了。2.4 多后端支持与国产硬件适配路径TileLang 的后端架构是插件式的。目前官方支持 CUDA 和 HIPAMD社区在推进国产加速卡的适配。适配一个新后端的工作量主要在三块一是 codegen把 TileLang IR 翻译成目标硬件的编程接口二是 runtime处理内存分配、kernel launch、同步这些运行时逻辑三是 intrinsic 映射把T.gemm、T.copy这些高层原语映射到硬件提供的加速指令上。国产硬件适配的难点在于很多国产加速卡的编程模型和 CUDA 不完全一样。比如有的卡没有独立的 shared memory有的卡的矩阵乘指令形状是固定的有的卡的异步拷贝机制不同。TileLang 的 tile 抽象层提供了一个缓冲让这些差异可以在后端消化而不是暴露给算子开发者。这是它作为“公共基础设施”的价值所在。3. 上手实操从零写一个高性能 GEMM3.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 kernelimport tilelang import tilelang.language as T tilelang.jit(out_idx[2]) def vector_add(N, block_N, dtypefloat16): 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), threads128) 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, dtypetorch.float16, devicecuda) b torch.randn(N, dtypetorch.float16, devicecuda) c vector_add(N, 256)(a, b) print(c)跑通这个例子你就理解了 TileLang 的基本骨架Kernel 定义 gridParallel 做线程级并行Tensor 描述数据形状。3.3 分块 GEMM 的完整实现与参数选择现在上硬菜写一个分块 GEMM。目标计算 C A B其中 A 是 M×KB 是 K×NC 是 M×N。分块策略是经典的 shared memory tilingtilelang.jit(out_idx[2]) def matmul(M, N, K, block_M, block_N, block_K, dtypefloat16, accum_dtypefloat32): 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), threads128) 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_stages3): 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用了二维 gridbx 对应 N 维度by 对应 M 维度这是 GEMM 的标准划分方式。第二T.alloc_shared分配了两块 shared memory 分别存 A 和 B 的 tileT.alloc_fragment分配了累加器。第三T.Pipelined把 K 维度的循环做了软件流水num_stages3表示三缓冲让数据加载和计算重叠。参数选择上block_M、block_N、block_K 的取值直接影响性能。我的经验是对于 fp16 输入block_M128、block_N128、block_K32 是一个比较稳的起点。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, dtypetorch.float16, devicecuda) b torch.randn(K, N, dtypetorch.float16, devicecuda) 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 TFLOPSfp16接近 cuBLAS 的水平。如果跑不到先检查 block 配置和 num_stages再检查生成的代码里有没有 bank conflict。4. 性能调优与常见问题排查4.1 内存布局与 Bank Conflict 处理Bank conflict 是 shared memory 性能的头号杀手。TileLang 默认会做 layout inference自动给 shared memory 加 padding 或者做 swizzle但有时候推断结果不是最优的。如果你发现 GEMM 性能明显低于预期可以手动指定 layoutT.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差距非常明显。所以如果你在调 GEMMswizzle 是第一个要检查的点。4.2 流水线深度与 Shared Memory 容量权衡num_stages不是越大越好。每增加一个 stage就要多分配一份 A_shared 和 B_shared。以 block_M128、block_N128、block_K32、fp16 为例一份 A_shared 是 128×32×2 8KB一份 B_shared 是 32×128×2 8KB加起来 16KB。num_stages3 就是 48KBnum_stages4 就是 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 是否超过 1024grid 维度是否超限结果不正确边界处理缺失检查 M/N/K 是否能被 block 整除加边界判断性能远低于预期bank conflict 或流水线未生效加 swizzle layout检查 num_stagesshared memory 超限block 配置过大减小 block_K 或 num_stages编译时间过长优化 pass 搜索空间大简化 kernel 结构减少动态 shape还有一个容易忽略的点TileLang 的 JIT 编译是带缓存的第一次编译慢后续调用会命中缓存。如果你改了 kernel 代码但发现行为没变可能是缓存没失效清一下~/.tilelang/cache目录。4.4 调优实战从 180 到 290 TFLOPS 的优化记录记录一次完整的调优过程。初始版本block_M128、block_N128、block_K32、num_stages2、无 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_M256、block_N128涨到 290 TFLOPS但寄存器压力变大occupancy 下降。最后的配置是 block_M256、block_N128、block_K32、num_stages3、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理解它们的调度策略和参数选择然后基于它们改。这是上手最快的路径也是我自己的做法。