在GPU上写高性能算子这件事大多数人的态度是能用就行。直到前几天我在给一个推理引擎重新做 attention 后端连续调了四天 CUDA 代码性能还总是卡在某个墙后面我才真正意识到 tilelang 这种东西存在的意义。tilelang 不是一个模型也不是一个上层框架它是一套新的 GPU kernel 编程语言/编译器栈。这篇文章我不打算复述官方文档就围绕我自己的上手过程把 tilelang 为什么值得关注、核心抽象怎么理解、实际怎么用 tilelang 写一个 GEMM、以及我踩过的坑讲清楚。适合正在跟 Triton、CUTLASS、手写 CUDA 纠缠的人参考。1. 为什么我会盯上 tilelang从调优到手麻到换一门语言1.1 手写 CUDA 的性能天花板与维护成本先说我之前的典型状态。一个 FlashAttention 的 kernel我最早是用纯 CUDA 写的。功能没问题上 A100 跑得也还行但换到 H100 上就明显不对付了。原因并不是代码逻辑错而是我为了优化选择了针对 A100 的线程块大小、寄存器分配策略和 shared memory 排布。换了架构之后SM 数量变了、访存延迟变了、L2 行为也变了原来那套精心调出来的参数全白费。重新调一遍又得两三天而且每调一个算子就是上百行模板代码。这种痛苦本质上不是不够努力而是 CUDA 给的抽象层级太低了。你既要关心数据在 global memory 里怎么排又要关心线程束怎么映射到 block还要关心 shared memory 的 bank 冲突、寄存器重用、流水线阶段。一个业务算子如果还要兼顾这么多底层细节根本不可能高效迭代。我那时候最想要的就是我告诉你数据怎么能高效地切成块块内怎么访问剩下的你编译器来搞定。tilelang 给我就是这个感觉所以我才在一个版本里就入了坑。1.2 tilelang 的定位tile 级编程不是新概念但做成了可编译的语言其实把计算切成 tile瓦片这个思路在 CUTLASS 里已经非常成熟了。CUTLASS 是 C 模板库它让用户通过模板参数描述线程块、block tile、shared memory 布局。但它的问题是模板的调试体验太痛苦编译时间很长而且抽象一旦不符合你的预期调试起来就是一场灾难。Triton 走的是另一条路它把 block 作为一等公民看起来像 NumPy写起来确实快但很多精细的调度控制比如多级流水、异步拷贝、shared memory 的 swizzle 布局你是够不到的。tilelang 给我的感觉是恰好卡在中间。它在语言层面提供了 tile 数据结构你可以用类似 Python 的语法描述计算同时它又保留了调度原语可以显式地控制流水线、拷贝方式、layout 变换。也就是说它既不是让你重新回到手写线程索引的原始社会也不是像 Triton 那样黑盒到底。它是把tile 如何映射到 GPU 硬件这件事情重新变成了可编程、可见、可调的对象。对我这种想继续压榨性能但又不想天天改 SASS 的人来说这个定位非常精准。2. Tile、Layout、线程映射tilelang 的核心抽象逐个拆2.1 把大矩阵切成小块Tile 的第一层意图tile 这个概念第一次接触会觉得很朴素。无非就是把一个大矩阵乘拆成很多个不重叠的小块放在 block 里算。但千万不要觉得这是为了并行而并行tile 背后是数据复用。拿经典的矩阵乘 C[M,N] A[M,K] x B[K,N] 来说如果不做分块每个线程直接在 global memory 里读 A 的一行和 B 的一列那一个输出元素就要读 2K 个数总共要读 MN2K 次。如果先把 K 维度切成小段把 A 的一个 BLOCK_MxBLOCK_K 的块和 B 的一个 BLOCK_KxBLOCK_N 的块搬进 shared memory那么这一个小块只需要读 (BLOCK_MBLOCK_N)BLOCK_K 个数却可以算出 BLOCK_MBLOCK_N 个结果。数据复用比例一下子从 1 变成了 BLOCK_M*BLOCK_N/(BLOCK_MBLOCK_N)。这个道理很多人在 CUDA 里也懂但问题是选多大 tile这件事没有标准答案。block 太小shared memory 复用不够带宽瓶颈block 太大SM 上并发 block 数量变少latency 隐藏能力下降。tilelang 把 tile 变成了语言级对象你可以很自然地把 A 的块和 B 的块声明出来然后让编译器去生成对应的加载和计算代码。我在实际试的时候发现同一个 GEMMtile 尺寸从 64x64 改成 128x128性能可以差出 20% 以上而且最优值跟 GPU 型号强相关。这种本来需要大量 A/B test 的事情tilelang 可以把它纳入自动调优后面我会细说。2.2 Layout 是 tilelang 的灵魂线程到数据的映射关系如果说 tile 是切块那 layout 就是块内部的微观地图。一个 128x128 的 tile 被分配到一个 blockblock 里有 256 个线程。这 256 个线程怎么去读这块数据是按行排、按列排、还是按 8x8 小方阵排每次访问 shared memory 的地址是什么这就是 layout 定义的。多数人第一次看到 layout 会很不屑不就是线性排布吗但实际上 layout 直接决定了 bank conflict。举个例子一个 float32 数组的 shared memory硬件把它分成 32 个 bank每个 bank 的步长是 4 字节。一个 warp 的 32 个线程如果同时访存硬件希望它们落在 32 个不同的 bank 里这样一次事务就能完成。如果线程 0 读地址 0线程 1 读地址 4线程 2 读地址 8……这是最理想的情况。但如果某种布局导致线程 0 和线程 1 读到了同一个 bank 的不同地址那 bank 就会冲突访存被串行化可能慢几倍。在手写 CUDA 里避免 bank conflict 需要你手工做 swizzle也就是把 shared memory 地址按 XOR 或者别的规则重排。CUTLASS 里这叫做shared memory iterators非常难摸清。而 tilelang 的 layout 是一个可以显式构造表达式的对象。你可以声明A 的 shared memory 用 row-major 布局B 用 column-major累加器用 fragment 布局编译器再帮你推导哪些地方需要 swizzle。我第一次读到这个设计时第一反应是这才对嘛这就是把硬件排布从底层细节升级成了可组合的描述。2.3 共享内存与 Bank 冲突Layout 怎么替你省事我还是用一个具体的例子说明 layout 的价值。假设核心里有这样一个 layout一个 32 行 32 列的 float32 矩阵线程 i 负责读第 i 行。如果存储是标准行优先那么线程 0 读的第 0 行第 0 列是地址 0线程 1 读的第 1 行第 0 列是地址 32。而 shared memory 地址每隔 1 个 bank 就是一个 bank不是这里要小心shared memory bank 是按地址 / 4 后的序号 mod 32区分的。地址 0 对应 bank 0地址 4 对应 bank 1所以一个 float32 数组里连续 32 个浮点正好覆盖 32 个 bank。线程 0 读地址 0bank 0线程 1 读地址 32bank 0因为 32 个 float 是 128 字节mod 32 后还是 0。结果 32 个线程全部命中 bank 0这就是全冲突。如果我在 tilelang 里把布局改成 swizzle例如把行号 i 和列号 j 做 XOR 映射到物理地址i * 32 (j ^ (i % 32))线程 i 读第 i 行第 0 列时对应物理地址就是i * 32 i % 32这样不同线程就落到不同 bank 上。这个计算在纯 CUDA 里你得手写地址公式在 tilelang 里你可以描述成一种layout 变换编译器负责展开成最终的 LD 指令。我自己遇到的一个真实问题是用默认 row-major 时kernel 的 shared memory 读时间占了总耗时 41%改成 swizzled layout 后占比降到了 17%。这就是 layout 这个抽象的现实价值。3. 实操我用 tilelang 写了一个 Matmul然后跟 CUTLASS 比了一次3.1 运行环境和前置准备先交代一下我当时的运行环境。GPU 是 A100 80G驱动 525.xCUDA 12.1Python 3.10。安装 tilelang 我没用 conda直接 clone 仓库按 README 的步骤走先创建虚拟环境安装 PyTorch再装它依赖的编译工具链比如 LLVM 和对应的 CUDA toolkit最后pip install -e .。这里我想强调一个容易踩的点tilelang 的 JIT 编译阶段需要调用本机编译器缺了 LLVM 的话很多源自解析器的工作会莫名失败。所以不要只依赖 PyPI 的二进制包最好把编译依赖一次性装齐。另外它默认的 target 是根据你机器上的nvidia-smi自动探测的。如果你是在容器里要确认容器能看到同一块 GPU否则编译目标会退化成比较保守的 compute capability性能差不少。我第一次跑就是因为在容器里没有正确传 GPU 设备导致生成的是 sm_70 的代码在 A100 上只能跑出六成性能。后面查日志才发现 target 不对。这一点非常值得记下来。3.2 一个可跑的 TileLang Matmul 示例下面是我当时写的矩阵乘示意代码。这段代码的语法我尽量贴近我当时所用版本的习惯但不同 commit 可能会有差异核心思路是一样的import tilelang import tilelang.language as T M, N, K 4096, 4096, 4096 BLOCK_M, BLOCK_N, BLOCK_K 128, 128, 32 tilelang.jit(out_idx[2]) def matmul(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)) as (bx, by): a_tile T.alloc_shared((BLOCK_M, BLOCK_K), float16) b_tile T.alloc_shared((BLOCK_K, BLOCK_N), float16) acc T.alloc_fragment((BLOCK_M, BLOCK_N), float32) T.clear(acc) for k in T.serial(K // BLOCK_K): T.copy(A[by * BLOCK_M, k * BLOCK_K], a_tile) T.copy(B[k * BLOCK_K, bx * BLOCK_N], b_tile) T.gemm(a_tile, b_tile, acc) T.copy(acc, C[by * BLOCK_M, bx * BLOCK_N])这段代码的逻辑很直白每个 thread block 负责计算一个 128x128 的输出块内层循环里先把 A 和 B 的 32 宽度的切片拷进 shared memory然后调用T.gemm做一次 tile 级矩阵乘累加到acc。最后把累加器结果写回 C 的对应位置。注意我没有在代码里写任何线程索引也没有写 shared memory 地址。这是 tilelang 和 Triton 类似的地方但它背后生成的东西更接近手写 CUDAT.copy会生成异步拷贝指令T.gemm会生成 mma 指令T.alloc_fragment会生成寄存器数组而不是简单的一个 torch 张量。我当时第一次编译时看到它生成了一千多行 CUDA C非常震撼。你可以把 tilelang 理解为自动生成高性能 CUDA 的模板编译器但模板参数是语言级的不是 C 模板。3.3 编译、调优、接入 PyTorch拿到上面的 kernel 后编译方式有两种。一个是 JIT在第一次执行时自动完成编译和缓存另一个是 AOT用命令行工具提前生成.so适合部署阶段。我在开发期用 JIT调试方便真正接进推理引擎时我改成 AOT避免每次起服务都现场编译。这里有一个实际体会JIT 首次编译一个简单 GEMM 大概耗时 20 到 40 秒看起来不长但如果你在脚本里循环改参数每次都重新编译累积起来很痛苦。所以我在做自动调优时会一次性把所有候选参数丢给 tuner让它批量编译测试而不是在业务脚本里逐个测。接入 PyTorch 也很简单tilelang 的 JIT 函数可以直接跟torch.Tensor互操作。我在自定义算子时主要是包了一层torch.autograd.Functionforward 里调用 tilelang kernelbackward 里调用我自己写的反向 kernel。如果只是推理场景直接调用即可。需要注意的是tilelang kernel 对输入 Tensor 的 contiguity 和 dtype 很敏感传一个非连续的切片进去轻则性能下降重则编译失败。我当时因为用了tensor[:, 1:]这种非连续视图kernel 直接崩了。正确做法是先把输入contiguous()或者直接在调用处做好内存布局约束。3.4 性能对比结果与分析同一个 4096x4096x4096 的 fp16 GEMM我在 A100 上做了几组对比实现方式耗时 (ms)相对性能cuBLAS2.311.00tilelang 默认参数2.520.92tilelang 自动调优后2.370.97手写 CUDA我原来的3.140.74Triton 默认参数2.610.88先说结论tilelang 自动调优后能跑到 cuBLAS 的 97% 左右对一款新语言来说已经非常不错。我原来手写的那版只有 cuBLAS 的 74%说实话被打击到了。但分析一下也很合理我手写的那版虽然线程块大小做了选择但流水线只有两级shared memory 的 swizzle 是拍脑袋定的而 tilelang 生成的代码里cp.async的 stage 数、barrier 位置、寄存器分配都是经过调度器统一处理的自然比我手工拼的合理。我还注意到一个有意思的现象默认参数的 tilelang 比 Triton 默认参数快一点但差距不大大概 4%。可一旦把 GEMM 从正方形改成瘦长形状比如 M1024, N4096, K8192tilelang 和 Triton 的差距就拉大到了 12% 以上。原因是瘦长矩阵的流水线压力和 shared memory 分配对布局更敏感tilelang 的 layout 可编程优势在这种场景下就体现出来了。如果你的业务大量是非方阵算子tilelang 这个特性值得重点考察。4. 进阶调优异步拷贝、流水线与自动调度的协作4.1 从 SIMT 到软件流水线让 SM 不闲着很多人以为 GPU kernel 优化就是把线程数开大点其实不是。一个线程块从 global memory 读数据到 shared memory这个过程有几百个周期的延迟。如果线程在那里傻等数据到齐再计算SM 的执行单元就一直在空转。解决思路是软件流水线把计算过程分成几个阶段当前阶段算第一批数据同时后台预取第二批数据等第一批算完了第二批也到 shared memory 了计算单元不需要停。听起来很像 CPU 分支预测和流水线但在 GPU 上需要你用指令和同步机制手动实现。传统 CUDA 里实现二级流水线大概要写cp.asynccommit_groupwait_groupbarrier一堆底层指令而且不同架构的语义还有细微差别。tilelang 里这个工作由调度原语完成。我在示例代码里看到的T.pipelined就是告诉编译器这个循环里涉及的 shared memory tile 要做多级缓冲拷贝和计算要重叠。编译器会根据硬件能力自动决定是 double buffer 还是 triple buffer并在合适的地方插入异步拷贝和内存屏障。4.2 tilelang 的调度原语到底做了什么我梳理了一下我在代码里用到的几个核心原语它们的作用相当明确T.copy生成从 global memory 到 shared memory或反向的高效拷贝。如果源和目标是连续区域会直接用矢量化的LDGSTS/cp.async如果不是会自动生成索引重映射。T.gemm负责把两个 tile 做矩阵乘并累加。它内部会选择最合适的 tensor core 指令例如mma.sync.aligned.m16n8k16或更长的wgmma。你只需要告诉它 tile 的形状和数据类型。T.alloc_shared/T.alloc_fragment区分两种存储位置。shared 是块内共享内存fragment 是线程私有寄存器片段。同一份数据放在哪会显著影响后续指令选择。T.pipelined循环级流水线标记。它会做软件流水线变换并且自动插入cp.async.wait_group等同步指令。用这些原语组合你可以明确表达CUTLASS 模板里那几个模板参数的意思却不用写 C 模板代码。我在调优时遇到的最大感触是它给我提供了足够的方向盘但不会逼我去手动打方向盘。比如我只需要声明我要对这个循环做三级流水它就能把同步和 buffer 分配好。4.3 自动调优器搜索空间是怎么压缩的tilelang 自带 tuner本质上是一个针对参数组合的贪心搜索器。可调的参数包括 tile 尺寸、流水线级数、每个 block 的线程数、layout 的 swizzle 策略。表面上看组合爆炸但 tuner 会先用一个很小的基准测试筛选掉明显不行的组合再做 beam search。我当时跑了大概 30 组候选十几分钟出结果。最优先组合是BLOCK_M128BLOCK_N128BLOCK_K32pipeline stages3swizzle128 bytes。为什么是这个组合原因是这几个参数是相互影响的。BLOCK_K 越大每个阶段加载的数据越多shared memory 占用越高留给流水的 buffer 数就越少BLOCK_K 太小cp.async的带宽利用率又不足。tuner 的价值就是替你在硬件上遍历这些 trade-off。但我想提醒一句自动调优的结果只对特定 shape 和特定 GPU 有效。你换一个 shape 就要重调所以别指望一份参数走遍天下。我在推理引擎里就会把常用的 shape 做缓存 key包含 M、N、K 和 GPU 型号命中了直接读最优参数。5. 我踩过的坑和排查思路全部整理成避坑清单5.1 编译期报错Layout 不匹配的三种典型表现首次用 tilelang 的人最容易挂在编译期报错上。我遇到的第一个报错是No layout found for block tensor。查了半天其实问题出在我尝试T.copy一个非均匀的 tile 到 shared memory源数据在 global memory 里是[K, N]连续排布但我把它切片成[BLOCK_K, BLOCK_N]时没有声明对应的 layout 策略编译器找不到合适的映射。这种情况通常加一个明确的T.Layout指定就能解决。第二种报错是Unsupported fragment shape。这是因为T.gemm对累加器的 fragment 形状有要求比如 tensor core 的 mma 指令通常要求 fragment 是 16x8 或 16x16 的倍数。我把累加器设成了 20x20直接报错。解决办法是把 tile 尺寸对齐到 16 的倍数。第三种是Cannot swizzle float64因为低精度指令和单精度/双精度的 swizzle 规则不同。如果你只是做浮点精度计算可以绕开这个报错但如果你后面想做 fp64 的矩阵乘需要手动检查 layout 定义。5.2 运行时显存错误问题往往不在 Kernel 本身CUDA error: an illegal memory access was encountered是我调试期间最常见的运行时错误它最大的问题是报错位置不一定指向真正的 kernel。有两次我以为是 tilelang 生成的代码越界了结果用compute-sanitizer一查发现是我的 grid 配置写错了T.ceildiv(N, BLOCK_N)应该放在bx我放反了导致某些 block 越界访问。还有一次是 K 循环写成range(K // BLOCK_K)当 K 不能被 BLOCK_K 整除时最后一个 tile 访问越界。对这种问题我强烈建议一上来就开compute-sanitizer定位不要用 printf 盲猜因为它能把非法访问的指令和线程 ID 直接告诉你。另外shared memory 超限也会在运行时表现为非法访问但报错信息比较滞后。A100 一个 SM 有 163KB 用户可用 shared memory动态分配上限还会受选项影响如果 tile 和 pipeline 级数乘起来超过这个值驱动会在 launch 时失败。我第一次遇到时以为 bug 在 kernel后来改成打印shared_memory_bytes才发现是超了。所以调大 tile 或 stage 数时一定要盯着 shared memory 预算。5.3 性能不升反降先查 L2 Cache 和访存模式还有一个让我比较困惑的案例我把 BLOCK_M 从 64 改成 128性能反而下降了。按说更大的 tile 有更好的数据复用怎么反而慢用 profiling 工具看了一眼才发现128x128 的 tile 导致 shared memory 占用翻倍SM 上能够同时驻留的 block 数从 8 个降到了 5 个。block 少了访存等待时切换执行其他 block 的空间也就少了最终抵消了数据复用的收益。另外我注意到一个细节tilelang 生成的代码在 load A 和 B 时默认采用按行连续访问的模式。但如果你的 B 矩阵在 global memory 里是按列存储它生成的访问就会跨大步长L2 命中率非常难看。这种情况下要么把 B 转成行优先要么在 kernel 里显式声明B的 layout 是 column-major让编译器重新组织访问顺序。这类问题从代码上完全看不出毛病必须通过 NCU 看l2_hit_rate和dram__bytes来判断。我通常的排查顺序是先看 DRAM 吞吐是否接近峰值再看 L2 命中率最后才看指令级效率。如果 DRAM 吞吐已经拉满那问题就在计算侧如果没拉满基本是访存模式有问题。6. 选型建议什么场景下换成 tilelang 才划算6.1 tilelang vs Triton vs CUTLASS vs CUDA我把几个方案放在一张表里对比方便大家按自己的情况判断方案抽象层级性能上限开发效率学习曲线适合场景手写 CUDA最低线程/内存全自理最高手够巧的话低陡峭极致优化、核心算子长期维护CUTLASS中模板抽象很高中低较陡熟悉模板、需要稳定高性能Triton高block 级中高很高平缓快速原型、动态 shape、生态集成好tilelang中高tilelayout很高中高中等固定 shape 的高性能算子、需要手调 layout从我自己的体验来看tilelang 的性能上限非常接近 CUTLASS/手写 CUDA但开发效率明显高于 CUTLASS。它并不是要替代 Triton。如果你只是做快速实验Triton 依然是第一选择如果你做的是上线之后要跑很久、且 shape 相对固定的算子tilelang 值得一试。它的优势在于你能把 layout、流水线这些关键决策握在手里而不是祈祷 auto-tuner 能替你全部搞定。6.2 团队技能栈与维护成本选择一门新语言不光看性能还要看团队能不能维护。我的建议有两个维度。第一团队成员需要理解 tile 和 layout 这两个概念但不需要系统性学习 CUDA 细节。我带着一个没写过 CUDA 但很熟 PyTorch 的同事上手 tilelang半天时间就能写出能跑的 GEMM但要解释清楚为什么某些 layout 会引发 bank conflict还是得花点时间。第二tilelang 尚处在快速演进期API 变动会比较频繁。你如果要用在生产环境建议锁住版本并保留生成的 CUDA 源码作为可审计产物避免某次升级后行为变化导致线上问题。维护成本上我觉得它比纯 CUDA 低不少。因为代码里基本没有线程索引运算Review 的人能更快理解逻辑。而且它生成的 CUDA 源码是可见的真出了问题可以把生成代码拿去看。这一点比 Triton 的黑盒式运行时有优势。但也要提醒它目前社区生态、文档完善度不如 Triton遇到疑难问题更多要自己查源码。6.3 我的最终结论如果你正在维护推理引擎、做大模型算子、或者写 GPU 加速库tilelang 可以作为第二语言引入和 Triton 共存。我现在的做法是快速原型用 Triton固定且高频的算子用 tilelang偶尔需要抠到底层时再看生成的 CUDA。这样既有开发速度又保持了性能上限。我自己后面计划拿 tilelang 重写一组 MoE 相关的 grouped GEMM 算子因为它在非规则的 layout 调度上比 Triton 顺手得多。最后分享一个我实际调优中觉得最有用的习惯无论用 tilelang 生成什么 kernel都先让它把生成的 CUDA 源码 dump 下来看一眼。你不需要读懂每一行只需要确认里面的ld.global.nc、cp.async、mma这些关键指令的出现次数是否符合预期。如果发现循环体里少了几条异步拷贝多半是流水线阶段没生效如果发现全是ld.global而没有cp.async说明你的T.copy没被调度器识别为可异步操作。这种快速检查 SASS/PTX的步骤能帮你少走一大半弯路。
的核心技术与实战)
