最近一直在折腾 GPU 算子的高性能实现，手写 CUDA kernel 写了一年多以后，越来越觉得常规的“算子开发—手工调优”路线有点撑不住场面。直到把 tilelang 这个项目翻出来，花了两周时间把几个核心算子重新实现了一遍，才确定这条路是值得认真投入的。

tilelang 是一个面向 GPU 的、以 tile 为基本抽象的高性能算子编译框架。它做的事情简单说就是：你用一套贴近手写 CUDA 思维的高层语言描述计算，编译器和自动调度器负责把 tile 映射到线程块、共享内存、寄存器，最终生成可以直接调用的 CUDA kernel。对于做推理引擎、训练框架、算子库的人来说，它解决的是“既要性能又要开发效率”的老大难问题。即使你之前完全没搞过编译器，把它当作一个“能生成高效 CUDA 代码的嵌入式 DSL”来用，学习曲线也比想象中平缓。

这篇文章我会从设计思路、语言形态、编译流程一路讲到实际跑通的完整例子，最后再把我踩过的坑和性能排查经验整理出来。想用它替换手写 CUDA、或者想找 Triton 之外另一种路线的人，这篇应该能帮你省不少时间。

1. tilelang 到底是什么：一个重新思考算子编译的项目

1.1 手写高性能算子到底难在哪

以前手写一个 CUDA GEMM，第一反应就是切 tile：blockIdx对应输出矩阵的一个大块，threadIdx对应块内的元素映射，数据先从 global memory 搬到 shared memory，再围绕寄存器分段和计算流水做循环。听起来不复杂，但真正要写得快，细节多得让人头皮发麻：shared memory 的 bank conflict 要怎么避开、加载用cp.async还是普通拷贝、双缓冲还是多级流水线、每个线程持有多少个寄存器值才算均衡、向量化宽度怎么对齐float4、L2 局部性如何利用……

这些问题每一项都够折腾大半天，而且调好一个 shape 之后，换一个规模又得重新来一遍。

tilelang 的切入点就在这里：它不试图取代你思考这些优化，而是把“切块 + 数据搬运 + 计算映射”这套思维本身变成语言结构。你在代码里写清楚的切分方式、流水线阶段、向量化偏好，会被编译器理解并转换成真正的 CUDA 实现。换句话讲，你还是在写 CUDA 风格的代码，但不再需要事无巨细地手工管理线程和寄存器的每一根线。

1.2 tile 作为一等公民：核心抽象的概念

很多编译框架会先定义算子的数学语义，再做循环变换，最后才谈得上 tile 的概念。tilelang 反过来：tile 从一开始就是一级公民，整个语言就是围绕“把一个大的计算切成 tile，在每个 tile 内做数据搬运和计算”来设计的。

描述一个算子时，你通常会先声明一个Kernel，告诉你编译器这个 kernel 被划分成多少个线程块；然后在 kernel 内部，用alloc_fragment声明属于当前线程块的寄存器/共享内存数据块，用copy和gemm这类的高层原语完成数据和计算操作。编译器知道这些原语的语义，所以它有机会在背后做寄存器级的分段、双缓冲、流水线重排，而不需要你在代码里显式写出每个线程的寄存器索引。

这种抽象的一个直接好处是代码读起来像伪代码，但跑起来又接近手写性能。我第一次看到生成的 CUDA 源码时仔细检查过，它的 shared memory 分段、循环展开、向量化负载都做得有板有眼，有些地方甚至让我觉得“如果是手写，我也会这么写”。

1.3 和 TVM、Triton 的定位差异

如果你用过 TVM，会发现 tilelang 的 Python 前端和 TVM 的 TensorIR 长得很像。这并不奇怪，tilelang 本身就和 TVM 生态关系密切，在底层 IR 上做了大量复用。但它和 TVM 的定位不同：TVM 是个全栈编译器，目标是从高层模型一路 down 到各种硬件后端，因此它的调度系统非常庞大，学习成本也高；tilelang 则聚焦在“算子级的高性能代码生成”这个更窄的场景，把复杂度和灵活性都往这个方向集中。

和 Triton 的差异更有意思。Triton 也推崇 tile 抽象，但走的是“编译器帮你决定 block 内部怎么做”的路线，你提供的是程序指针式的 tile 语义，编译器负责把 tile 映射到线程和寄存器。tilelang 则更强调“让写 kernel 的人保留更多控制力”，你可以在 tile 内显式地控制数据是放在 fragment（寄存器/共享）里还是本地内存里，也能更直接地表达类似 flash attention 这种带复杂循环和归约结构的算子。

选型上的直觉是：如果你只想快速写出能跑的算子，不想关心底层，Triton 很合适；如果你已有一定 CUDA 功底，希望代码在被编译的前提下依然能精确控制关键优化点，并且需要一个能嵌入现有代码库的编译后端，tilelang 更对你胃口。

2. 核心设计：从 tile 语言到 CUDA kernel

2.1 一个 tile 语言写出来的 GEMM 长什么样

先不上理论，拿矩阵乘法来开刀。下面是 GEMM 在 tilelang 里非常典型的写法（示意代码，略去了 import 和配置细节）：

@T.prim_func def matmul_kernel( A: T.Tensor((M, K), "float16"), B: T.Tensor((K, N), "float16"), C: T.Tensor((M, N), "float16"), ): with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=256) as (bx, by): acc = T.alloc_fragment((block_M, block_N), "float32") for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3): a_tile = T.alloc_fragment((block_M, block_K), "float16") b_tile = T.alloc_fragment((block_K, block_N), "float16") T.copy(A[by * block_M, ko * block_K], a_tile) T.copy(B[ko * block_K, bx * block_N], b_tile) T.gemm(a_tile, b_tile, acc) T.copy(acc, C[by * block_M, bx * block_N])

先说T.Kernel(...) as (bx, by)：这一行声明了线程块的网格。bx和by分别对应 N 方向和 M 方向的块索引，维度大小用ceildiv算出来，看上去就是 CUDA 里gridDim和blockIdx的翻版。

T.alloc_fragment是这个框架里最核心的抽象之一。fragment 可以理解为一段与当前线程块强相关的数据空间，编译器会把它分配到寄存器或共享内存，并且根据后续的计算模式自动选择布局。举个例子：acc作为累加器，如果后面跟的是T.gemm，那么编译器大概率会让它驻留在寄存器中，并按线程私有布局切分；而a_tile和b_tile是要从 global memory 搬运进来的中间数据，编译器会更倾向于放进 shared memory 做缓存和重用。

T.Pipelined则在循环上开启多级流水线，num_stages=3表示数据预取和计算重叠三级。这一句话在实际生成中可能就对应着双缓冲、异步拷贝和循环展开的整套组合。手动写过 CUDA 流水线的朋友应该知道这些操作写起来有多啰嗦，在这里只是一个循环参数问题。

最后T.copy和T.gemm就是高阶原语：copy 负责 global 到 shared（或 shared 到 fragment）的数据搬运，gemm负责执行 block 内的矩阵乘累加。这几个原语你都不要自己去实现循环，编译器会帮你把搬运向量化、把计算映射到具体的线程。

2.2 从 Python 到可执行 CUDA kernel，编译流水线里发生了什么

第一次接触编译框架的人可能觉得“编译器很神秘”，其实它的流水线是可以拆开的。tilelang 的路径大致是：Python 前端把上面的代码构造成 IR，然后经过一轮轮 lowering pass 把它往下变成越来越接近硬件的表示。

粗略分几步：先是把T.Kernel、T.alloc_fragment、T.copy这些高层原语展开成具体的循环、内存分配和线程绑定关系。接着是循环变换，重点处理T.Pipelined标志，把它变成带预取和双缓冲的控制流结构。再往下是存储重写，决定 fragment 到底落在寄存器还是 shared memory，并安排好地址映射。最后是 CUDA 代码生成，把数据搬运规范化成cp.async或普通拷贝，计算部分生成真实的乘加指令，输出一段包含完整__global__函数的 C++ 代码。

由于底层继承了 TVM 的 IR 基础设施，很多成熟的 pass 比如存储规划、线程绑定、向量化识别都可以直接复用。这也就是为什么 tilelang 能够做到“前端看起来偏教学，后端生成却相当工程化”——它站在了大量编译器工作的肩膀上。

需要记住的一点是：tilelang 生成的不是 PTX 或者二进制，而是可读的 CUDA C++。这意味着你可以在nvcc编译之前把生成代码打开来看，检查它到底把你的意图实现成了什么样。这种可观测性对调试非常重要，我后面会展开讲。

2.3 自动调优：性能要接近手写，只是“会编译”不够

能生成代码只是第一步，生成“快代码”是另一回事。同样的算法，block 大小取 64 还是 128，流水线 depth 取 2 还是 4，中间数据到底放寄存器还是共享内存，性能可能差出一倍以上。tilelang 没有把这个环节留给你拍脑袋，而是内置了一个自动调优器：给定搜索空间，它会在真实设备上跑出一组基准，选出一套最优配置。

搜索空间一般是开发者自己定义的。以 GEMM 为例，常见维度包括block_M/block_N的尺寸、block_K循环步长、流水线阶段数、线程数，以及是否开启某些向量化标志。用起来不复杂：定义好配置列表，调用调优器，它会编译、测量、记录并比较结果，最后返回最优参数。

我自己的体会是，自动调优不一定每次都能超越你最精细的手写版本，但它的下限非常高。绝大多数情况下，只要能给出合理的搜索空间，得到的结果就是“极好”的水平，而且它能帮你覆盖到多个 shape。对做算子库的人来说，这省掉的精力是巨大的——过去要为十几种 shape 各手工调一遍，现在在搜索空间里一跑就好了。

3. 选型对比：tilelang、手写 CUDA 与 Triton 到底怎么选

3.1 会手写 CUDA 的人，为什么还要看 tilelang

我接触过的不少工程师对手写 CUDA 有感情，觉得“编译器生成的东西总归不如手撸的透明”。这种顾虑不是没道理，但实际要分场景。如果你是在给单一模型写一个固定的 kernel，而且这个 kernel 未来几年都不变，手写确实可以做到极致，因为你有无限时间打磨循环结构和 shared memory 布局。

但如果你身处算子库、推理引擎或者框架开发，情况完全不同：算子形态多、shape 分布广、需求变更快。这时候“一个人手写一个算子”的成本是线性增长的，而 tilelang 的价值在于“一个模式写出来，调优器帮你覆盖各种 shape”。我并不是说它可以完全替代手写 CUDA，而是说它能把 80% 的“重复造轮子”工作自动化，让你把精力留给真正需要手工精调的那 20%。

这里还有一层心理上的优势：tilelang 生成的是 CUDA C++，不是封装好的黑盒。你觉得它哪里做得不好，可以直接看生成的代码，再微调你的描述或搜索空间重新生成。这种透明度是很多高层框架给不了的。

3.2 Triton 与 tilelang：两条“tile 路线”的分歧

Triton 这两年在社区里也很火，它和马赛克式的 tile 编程模型非常契合 GPU 硬件。Triton 的哲学是“用指令式的 tile 操作描述计算，编译器自动分配线程和寄存器和数据布局”，开发者通常不会去干预一个 tile 内的具体布局，而是把注意力放在算法结构上。

tilelang 在抽象层级上虽然也是 tile，但它更强调开发者对布局和流水线的显式控制。比如前面示例里的alloc_fragment明确告诉你有一份数据以 fragment 形式存在，你可以在语义上区分“这是拿来复用的数据”和“这是累加器专用的数据”，而不是统统交给编译器去猜。

这种差异在 flash attention 这类算子上表现得非常明显。Flash attention 的核心是“按块加载 K/V，在寄存器里做 online softmax，再继续融入下一个块”，这种结构需要你对循环内数据的生命周期有精确理解。tilelang 的 fragment + pipelined 组合能比较自然地表达这种结构，而 Triton 虽然也能写，但你会感觉到编译器替你决策的成分更多，有些关键布局不一定如你所愿。

不是说谁好谁差，而是分工思路不同。需要把“程序员的显式意图”尽可能变成性能的时候，tilelang 的表达细致度更高；希望快速把算法跑起来、不关心底层细节的时候，Triton 更省事。

3.3 按场景选型的实际建议

做选型时，可以按下面这套逻辑来判断：

• 如果目标硬件很固定，算子 shape 也非常固定，又有足够的时间去手工优化：手写 CUDA 依然是最优解，没有任何框架能完全替代你针对特定 shape 做的极致手工调整。
• 如果属于模型开发阶段，需要快速验证各种新算子的性能上限：Triton 很合适，开发速度快，学习成本低，性能也足够好。
• 如果你在维护算子库或推理引擎，有大量 shape 要覆盖、希望生成的代码能嵌入到现有编译链路中、同时你又希望保留对关键优化点的话语权：tilelang 是更值得投入的方向。
• 如果项目已经基于 TVM 做了不少基础设施：tilelang 的 IR 兼容性会带来额外便利，因为你可以在同一个后端里做更多编译期分析。

表格化对比会更直观：

4. 实操：用 tilelang 写出并调优一个 GEMM

4.1 前置环境准备：一套最省心的组合

要跑通 tilelang，硬件上需要一张支持 CUDA 的 NVIDIA GPU，建议显存 8GB 以上，架构在 Turing 及之后会体验更好，因为只有这些架构才支持cp.async一类特性无缝衔接。软件方面需要 CUDA Toolkit 11.8 或更高版本，Python 3.9 以上，最好再有一个干净的 conda 环境。

安装其实很简单，直接 pip 装就行（示意命令）：

pip install tilelang

如果是第一次用，建议先跑一下它自带的测试用例，确认你的 GPU 驱动和 CUDA 环境没问题。源码编译也不是必须的，除非你想改框架内部代码或者需要针对特定硬件做定制。

验证环境最快的办法是跑一个最小例子：定义一个极简 kernel，编译，调用，检查返回值。能跑通就说明工具链没问题，再开始玩复杂的算子。

4.2 完整实现一个可运行的 GEMM

我以 FP16 输入、FP32 累加、输出 FP16 的 GEMM 为例。M、N、K 可以设成通用变量，代码和之前示意基本一致，这里再加上自动调优的调用：

import tilelang import tilelang.language as T @T.prim_func def matmul_kernel( A: T.Tensor((M, K), "float16"), B: T.Tensor((K, N), "float16"), C: T.Tensor((M, N), "float16"), ): with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=256) as (bx, by): acc = T.alloc_fragment((block_M, block_N), "float32") for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): a_tile = T.alloc_fragment((block_M, block_K), "float16") b_tile = T.alloc_fragment((block_K, block_N), "float16") T.copy(A[by * block_M, ko * block_K], a_tile) T.copy(B[ko * block_K, bx * block_N], b_tile) T.gemm(a_tile, b_tile, acc) T.copy(acc, C[by * block_M, bx * block_N]) # 配置搜索空间，取其中一组或跑调优器 configs = [ {"block_M": 64, "block_N": 64, "block_K": 32, "num_stages": 2}, {"block_M": 128, "block_N": 128, "block_K": 32, "num_stages": 3}, {"block_M": 128, "block_N": 128, "block_K": 64, "num_stages": 4}, ] # 编译 kernel = tilelang.compile(matmul_kernel, **best_config) # 调用 a = torch.randn((M, K), dtype=torch.float16, device="cuda") b = torch.randn((K, N), dtype=torch.float16, device="cuda") c = torch.empty((M, N), dtype=torch.float16, device="cuda") kernel(a, b, c)

你可能会注意到block_K这个参数不能乱选：它通常应该是 16 或 32 的倍数，这样才能和 FP16 的 8 字节/16 字节向量化加载对齐。太少会导致循环开销过大，太多会浪费 shared memory，一般 32 到 64 比较常用。

关于自动调优，实际使用时可以把上面的 config 列表扩充成更细的网格，然后调用框架的搜索接口，让它挨个编译并跑 benchmark。这里我建议第一次不要搜索太大，先把几组最直观的配置跑完，看看性能水平在什么量级，再决定要不要扩大搜索空间。

4.3 性能验证与正确性检查

调优完了不能直接上线，先做两件事：正确性对比和性能对比。

正确性对比的做法是用 PyTorch 做参考实现，对生成 kernel 的输出做误差检查。FP16 GEMM 这种算子，用atol=1e-2, rtol=1e-2通常足够，因为累加发生在 FP32，误差主要来自输入本身的有限精度。如果误差过大，先检查是不是block_K对齐有问题，再看是不是 shared memory 复用导致的陈旧数据没有清干净。

性能对比就更直接了，写一个循环跑kernel，记录平均耗时，然后用 cuBLAS 的同一 shape 做对照。我实测过几次，在惯例配置下，一个 128x128 tile 的 tilelang GEMM 基本能达到 cuBLAS 的 80% 到 95% 水平；如果 shape 比较友好、数据规模大，自动调优后甚至可以持平或略超。这里我不建议追求“一定超过 cuBLAS”，能和它在一个量级就已经赢了，毕竟 cuBLAS 是 NVIDIA 花了大精力手工调出来的。

如果要在真·生产环境使用，还可以把生成好的 CUDA 源码导出，走常规的 nvcc 编译链路，并通过 runtime API 加载，这样能避开每次启动时的 Python 解析开销，也方便和现有 C++ 工程集成。

5. 实战踩坑：性能上不去的核心原因与排查实录

5.1 四个最容易让性能“卡死”的典型问题

我调试过程中遇到最多的问题，几乎都集中在下面几个方向。

第一类问题是最常见的，也是新手最容易浪费时间的。调参时不要一上来就几百个组合全跑，手工调优最有效的办法是“保持其他变量固定，一次只改一个维度”。比如先固定 block_M=block_N=128、num_stages=3，把 block_K 从 16 调到 64，看趋势；再固定 block_K 去调 stages。这样你能直观看到每个维度对性能的影响。

5.2 编译报错怎么看：从错误信息反推问题前因

tilelang 的报错信息有些来自编译器底层，看起来会比较“硬核”，最典型的问题有两类：一类是 IR 不合法，比如 copy 两边的 shape 不匹配或者张量索引越界；另一类是在 CUDA 代码生成阶段报 nvcc 错误。

遇到第一类错误时，我的经验是别急着改代码，先看报错里提示的 tensor 名。绝大多数情况是alloc_fragment的形状和后续copy的源/目标维度不一致导致的。比如你想把一个(block_M, block_K)的 tile 从A[by*block_M, ko*block_K]拷贝进来，那源张量的连续维度必须能覆盖这个 tile，而且 A 的 K 方向对齐要能整除 block_K。

遇到第二类错误，先确认你能不能在本地单独编译生成好的 CUDA 源码。如果单独用 nvcc 编译也失败，说明生成代码里确实有非法语法，通常和模板参数推导、constexpr 表达式长度或 shared memory 声明有关，这时候把生成源码拿出来看，定位会快很多。

还有一个容易被忽略的小坑：调用 kernel 时传入的张量必须在 CUDA 上且类型和 dtype 完全匹配。我用 torch 封装时经常不小心把 CPU tensor 传进去，然后得到莫名其妙的空结果而不是报错，这是最隐蔽的坑之一。

5.3 一点经验之谈：怎么把调优时间控制在合理范围

自动调优很好用，但搜索空间设计不好同样会浪费时间。以 GEMM 为例，一个比较合理的初始搜索空间是：block_M ∈ {64, 128, 256}，block_N ∈ {64, 128, 256}，block_K ∈ {32, 64}，num_stages ∈ {2, 3, 4}，这样组合数量已经上百。如果每个配置都跑几十次 benchmark，整体时间可能很夸张。

我的做法是分阶段。第一轮只测 block_M/block_N 的组合，把 block_K 和 stages 固定成默认值；第二轮用第一轮筛出的两三组最优尺寸去遍历 block_K 和 stages。这样整个搜索时间可以缩短 70% 以上，而且结果不会差太多。

另外，自动调优的结果要留档。每个 shape 的最优配置、实际吞吐量都记下来。模型更新后 shape 变了，直接查历史配置表，能省掉大量重复劳动。

回头看，tilelang 在“程序员显式控制”和“编译器自动优化”之间找到了一条很舒服的中间路线。它既没有把开发者丢进 CUDA 的细节泥潭，也没有像更上层的框架那样替你做太多决定。如果你正好在找一种“能看懂生成的每一行 CUDA，却不自己写每一行”的算子开发方案，我个人强烈建议花一个下午把上面的例子跑一遍。你可能会和我一样，在写出第一个自动调优 kernel 的时候，就决定把手头一部分手写算子迁过来了。