最近一直在折腾 GPU 算子的高性能实现手写 CUDA kernel 写了一年多以后越来越觉得常规的“算子开发—手工调优”路线有点撑不住场面。直到把 tilelang 这个项目翻出来花了两周时间把几个核心算子重新实现了一遍才确定这条路是值得认真投入的。tilelang 是一个面向 GPU 的、以 tile 为基本抽象的高性能算子编译框架。它做的事情简单说就是你用一套贴近手写 CUDA 思维的高层语言描述计算编译器和自动调度器负责把 tile 映射到线程块、共享内存、寄存器最终生成可以直接调用的 CUDA kernel。对于做推理引擎、训练框架、算子库的人来说它解决的是“既要性能又要开发效率”的老大难问题。即使你之前完全没搞过编译器把它当作一个“能生成高效 CUDA 代码的嵌入式 DSL”来用学习曲线也比想象中平缓。这篇文章我会从设计思路、语言形态、编译流程一路讲到实际跑通的完整例子最后再把我踩过的坑和性能排查经验整理出来。想用它替换手写 CUDA、或者想找 Triton 之外另一种路线的人这篇应该能帮你省不少时间。1. tilelang 到底是什么一个重新思考算子编译的项目1.1 手写高性能算子到底难在哪以前手写一个 CUDA GEMM第一反应就是切 tileblockIdx对应输出矩阵的一个大块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 kernel2.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), threads256) as (bx, by): acc T.alloc_fragment((block_M, block_N), float32) for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages3): 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_stages3表示数据预取和计算重叠三级。这一句话在实际生成中可能就对应着双缓冲、异步拷贝和循环展开的整套组合。手动写过 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 兼容性会带来额外便利因为你可以在同一个后端里做更多编译期分析。表格化对比会更直观维度手写 CUDATritontilelang性能上限最高高高开发效率低最高中高显式控制力最强弱中强学习成本高低中代码可读性看功力好很好嵌入集成难度手工封装中等较低多 shape 适配每 shape 手工调靠编译器靠自动调优4. 实操用 tilelang 写出并调优一个 GEMM4.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), threads256) as (bx, by): acc T.alloc_fragment((block_M, block_N), float32) for ko in T.Pipelined(T.ceildiv(K, block_K), num_stagesnum_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), dtypetorch.float16, devicecuda) b torch.randn((K, N), dtypetorch.float16, devicecuda) c torch.empty((M, N), dtypetorch.float16, devicecuda) kernel(a, b, c)你可能会注意到block_K这个参数不能乱选它通常应该是 16 或 32 的倍数这样才能和 FP16 的 8 字节/16 字节向量化加载对齐。太少会导致循环开销过大太多会浪费 shared memory一般 32 到 64 比较常用。关于自动调优实际使用时可以把上面的 config 列表扩充成更细的网格然后调用框架的搜索接口让它挨个编译并跑 benchmark。这里我建议第一次不要搜索太大先把几组最直观的配置跑完看看性能水平在什么量级再决定要不要扩大搜索空间。4.3 性能验证与正确性检查调优完了不能直接上线先做两件事正确性对比和性能对比。正确性对比的做法是用 PyTorch 做参考实现对生成 kernel 的输出做误差检查。FP16 GEMM 这种算子用atol1e-2, rtol1e-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 四个最容易让性能“卡死”的典型问题我调试过程中遇到最多的问题几乎都集中在下面几个方向。现象典型原因排查与解决性能始终上不去换什么配置都差不多tile 尺寸太小或太大导致并行度不足或 shared memory 占用过多先固定 block_M/block_N 在 64~128 之间再微调 block_K 和 num_stages小 shape 比手写慢很多循环和启动开销占比过高流水线优势发挥不出来小 shape 减少 num_stages甚至关掉 Pipelined 直接跑shared memory 莫名超限block_K 或 num_stages 过大估算 shared memory 占用每级流水线的 tile 数据 fragment 数据控制在硬件限额内结果精度反复不对累加顺序/精度问题或数据搬移错位先用小规模测试逐 tile 对比确认 block 索引与 copy 地址的映射正确第一类问题是最常见的也是新手最容易浪费时间的。调参时不要一上来就几百个组合全跑手工调优最有效的办法是“保持其他变量固定一次只改一个维度”。比如先固定 block_Mblock_N128、num_stages3把 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 的时候就决定把手头一部分手写算子迁过来了。