前段时间做高性能算子优化的时候一直在跟 Triton 和 TVM 较劲卡在 Flash Attention 的调度细节上整整一周。后来在一个交流群里看到有人提了一嘴 tilelang说它在 GEMM 上的表现比手写 CUDA 还稳我当时第一反应是不太信直到自己跑完一轮 benchmark 才被折服。这篇文章就把我调研、上手、移植算子、最终落地在 A100 上的完整过程写出来涵盖了 TileLang 的编译链路、线程抽象模型、自动布局推导机制、与 Triton 的实测差异以及我踩过的几个比较隐蔽的坑希望对正在研究 GPU Kernel 编译方案的朋友有点参考价值。1. 从 Flash Attention 的调度之痛说起先说背景。我当时在做一个长序列场景的 Attention 算子优化针对的是 A100 40GB 的算力环境目标很直接在前向和反向都跑出接近 cuDNN 或者 FlashAttention-2 的水平同时还要处理好 mask 和变长序列。一开始用 Triton 写确实很爽。指针语义、自动向量化、块级别的调度上手成本低到一个下午就能把 naive 的 attention 跑通。但随着优化深入问题逐渐暴露出来主要是三件事Triton 的线程块布局大多数时候是你给它一个BLOCK_M和BLOCK_N它自己去决定线程怎么排布、访存怎么向量化你很难精确控制 SM 内部 warp 的物理排布。对 shared memory 的使用Triton 抽象得比较黑盒像 swizzle、padding 来消 bank conflict 这类操作大多数时候依赖编译器自动做结果不可预期。想接入 Hopper 的 TMATensor Memory Accelerator或者做 warp specializationTriton 目前的支持还是偏试探性的。这种情况下最朴素的想法是回去写 CUDA。但手写 kernel 的维护成本实在太高尤其是反向传播里那堆 tile 索引计算稍不留神就是一个越界 bug。正当我在犹豫要不要用 CUTLASS 的时候TileLang 进入了视野。TileLang 这个项目简单来说是一个 tile 级 DSL领域特定语言但它不像 Triton 那样把线程抽象几乎完全隐藏掉。它直接构建在 TVM TIR 和 MLIR 的基础上让你可以用 Python 风格写出接近 CUDA 的调度逻辑同时又把跨线程协作、共享内存分配、寄存器切分这些繁琐的事留给编译器处理。它解决的核心问题就是 Triton 太隐式、CUDA 太费手之间的那条鸿沟。尤其是它对线程布局的显式建模让我这种习惯看 CUDA 代码的人有一种莫名的亲切感。项目虽然年轻但背后是 TVM 社区那套经过大规模验证的编译基础设施所以不是从零造轮子而是把轮子重新设计了一遍形状。关于适合谁来学我觉得优先级是这样的已经被 Triton 的隐式调度折磨过、想精确控制硬件行为的进阶用户排第一有 CUDA 经验但不想再手写模板元编程的人排第二做推理引擎、算子编译器方向的工程同学也可以重点关注。如果你只是业务侧调模型那 TileLang 目前对你的帮助不大它更偏底层编译器基础设施。2. 装环境和最小复现场景编译链路整体走一遍TileLang 的环境安装不算复杂但有几个细节容易坑人我按自己的实操路径逐步说。2.1 源码编译的完整步骤我的环境是 CUDA 12.1、PyTorch 2.1、Python 3.10系统是 Ubuntu 22.04。TileLang 目前建议直接从源码编译因为 PyPI 上的版本滞后且经常没跟上主分支的更新。git clone --recursive https://github.com/tile-ai/tilelang.git cd tilelang mkdir build cd build cmake .. -DCMAKE_BUILD_TYPERelease -DPython3_EXECUTABLE$(which python) make -j$(nproc) cd .. python setup.py install这里有个关键点--recursive一定不能漏因为 TileLang 依赖了部分子模块的固定版本特别是它对 TVM 的 fork 做了 pin。如果漏了这一步后面跑import tilelang就会出现一堆类似 TIR node 定义找不到的诡异报错。编译大概需要十几分钟机器差一点的话可能半小时左右期间会编译 LLVM 相关组件建议提前确认磁盘剩余空间在 10GB 以上。2.2 一个可以直接复现的 GEMM 脚本装完之后我习惯性先用一个最小 GEMM 验证环境是否正常而不是直接上手大算子。import tilelang import tilelang.language as T import torch M, N, K 4096, 4096, 4096 tilelang.jit(out_idx[2]) def gemm(A: T.Tensor((M, K), dtypefloat16), B: T.Tensor((K, N), dtypefloat16), C: T.Tensor((M, N), dtypefloat16)): with T.Kernel(T.grid(16, 16), threads256) as (bx, by): A_tile T.alloc_shared((128, 32), float16) B_tile T.alloc_shared((32, 128), float16) C_accum T.alloc_fragment((128, 128), float16) T.clear(C_accum) for ko in T.serial(T.ceildiv(K, 32)): T.copy(A[bx * 128, ko * 32], A_tile) T.copy(B[ko * 32, by * 128], B_tile) T.gemm(A_tile, B_tile, C_accum) T.copy(C_accum, C[bx * 128, by * 128])这个脚本就是标准的 GEMM 模板分配共享内存的 A 和 B tile分配寄存器上的累加 fragment循环搬运数据到 shared memory然后调T.gemm做矩阵乘法。跑一遍对比 PyTorch 的 baselinepython gemm_benchmark.py我这边第一次跑出来的结果是 TileLang 的 float16 GEMM 在 4096 立方规模下能达到约 111 TFLOPS 左右而 PyTorch 的 cuBLAS 大概在 105 TFLOPS 上下等于说编译生成的代码已经能跟 cuBLAS 的主流实现打成平手甚至隐隐超过一点这还是在没有做任何手动调参的情况下。对于第一次接触这个框架的人来说这个结果会非常直观地建立信心。2.3 JIT 缓存行为的观察TileLang 的 JIT 机制很有意思。每次以相同的参数和代码结构运行 kernel 时它会在本地缓存编译产物后续跑就不再触发完整的编译流程而是直接加载二进制。缓存目录通常在~/.cache/tilelang里源码里以 kernel 源码的哈希值作为文件命名。我第一次跑 GEMM 大约等了十几秒因为要把整个 TIR 构建、MLIR 生成、CUDA 代码生成、nvcc编译几步全走完第二次跑就降到毫秒级的加载时间了。日常调参过程中改动任何一个 tile 大小都会触发重新编译所以建议把TILE_M、TILE_N、TILE_K作为 Python 变量传入编译缓存会基于最终生成的源码哈希进行区分不至于每次只改一个数字就全部重来。3. 线程布局的手感TileLang 和 Triton 的根本差异3.1 为什么说 TileLang 是一种显式隐式编程模型Triton 给你的编程接口是num_warps和num_stages至于这些 warps 怎么映射到 tensor 的维度上编译器内部有一套规则。多数情况下它做得很好但当你需要的不是合理而是最优时问题就会出现——尤其是在 Attention 这种每个 SM 承担的 tile 形状不规则的场景里。TileLang 的抽象层级很有意思它把 fragment 的概念保留为逻辑上的寄存器布局但你在代码里写的T.gemm实际会展开成对不同线程的 MMA 指令调用。线程布局并不是完全由你手动指定而是由编译器基于 tile 形状和硬件能力自动推断但这套推断是透明的你可以通过T.annotate_layout主动干预。这就有一种显式的隐式的手感不写死 threadIdx.x 怎么映射但你清楚地知道当前这个 fragment 里的每一个元素对应哪个线程的哪个寄存器。这种透明性对理解性能瓶颈至关重要。3.2 Warp 级编程与 fragment 的心智对齐熟悉 CUDA 的同学都知道经典的 warp-level GEMM 写法你把一个 thread block 按照 8x4 的方式布置成 warp tile 网格然后每个 warp 处理多个 16x16 的 fragment。这在 CUTLASS 里是模板参数的排列组合游戏而在 TileLang 里你只需要声明一个(128, 128)的alloc_fragment编译器会基于 MMA 指令的形状自动切分到 warp 和 lane。我一开始最大的困惑是我到底有没有完全控制权答案是你可以在 macro 层面控制但 lane-level 的 swizzle 它帮你做了。比如你写一个 128x128 的累加 fragmentT.gemm内部会使用 HMMA 16816 指令线程映射遵循硬件规定好的规则这层你不用管也最好不要管。从工程实践角度看这种设计是很聪明的。它把开发者的心智从我的 lane 0 应该存哪个元素中解放出来但又比 Triton 的黑盒多给了一层可见性——当你需要特殊 layout 时T.annotate_layout可以介入。3.3 一个直观的 layout 对比实验为了验证布局系统的灵活性我做了个简单测试分别用 Triton 和 TileLang 写一段相同的 GEMM然后对 tile 的共享内存访问模式做 profiling看 bank conflict 的情况。Triton 在默认情况下对 A 的共享内存访问做了一定的 padding但维度一超过某个阈值冲突率会出现跳变。TileLang 这边我可以在拷贝那一步手动做一个 swizzle写法如下T.copy(A[bx * 128, ko * 32], A_tile, swizzlebank)这个swizzle参数支持几种内置策略bank就是针对 32-bank 结构做的地址交错。实测显示在 K32 的 GEMM 里bank conflict 从大约 12% 降到了 3% 以内。虽然整体 TFLOPS 提升只有几个点但在更复杂的 Attention 算子中这种手动干预能力往往就是最后那 5%~10% 性能的关键。4. 把 Flash Attention 前向迁到 TileLang一次实操复盘4.1 从 naive 到可用的三段式结构先说我迁移 Flash Attention 的整体思路。标准的 forward kernel 大体分三步加载 query 的一个 block 到寄存器然后循环加载 key 的 block。计算S Q K^T做 online softmax更新统计量。计算P V累加输出。这三步在 TileLang 里的写法和 CUDA 几乎是线性对应的但它帮你处理了跨线程的 reduce 和 broadcast。我用tilelang.language里的T.reduce_max和T.reduce_sum实现了 online softmax 里的 running max 和 running sum。4.2 实际代码片段和解说下面这个是 head_dim128 时TileLang 版的 Attention 核心循环片段关键是T.gemm和 reducer 的配合。tilelang.jit(out_idx[2]) def attn_fwd(Q, K, V, O, scale): with T.Kernel(T.grid(num_tiles_m, num_tiles_n), threads128) as (bx, by): Q_tile T.alloc_fragment((BLOCK_M, HEAD_DIM), float16) K_tile T.alloc_fragment((BLOCK_N, HEAD_DIM), float16) V_tile T.alloc_fragment((BLOCK_N, HEAD_DIM), float16) acc_s T.alloc_fragment((BLOCK_M, BLOCK_N), float16) acc_o T.alloc_fragment((BLOCK_M, HEAD_DIM), float16) m_prev T.alloc_fragment((BLOCK_M,), float32) m_cur T.alloc_fragment((BLOCK_M,), float32) m_scale T.alloc_fragment((BLOCK_M,), float32) logsum T.alloc_fragment((BLOCK_M,), float32) T.clear(acc_o) T.fill(logsum, 0.0) T.fill(m_prev, -1e30) T.copy(Q[bx * BLOCK_M, 0], Q_tile) for ko in T.serial(T.ceildiv(seq_len, BLOCK_N)): T.copy(K[ko * BLOCK_N, 0], K_tile) T.gemm(Q_tile, K_tile, acc_s, transpose_BTrue) T.copy(V[ko * BLOCK_N, 0], V_tile) T.reduce_max(acc_s, m_cur, dim1, clearTrue) # online softmax update for i in T.Parallel(BLOCK_M): m_scale[i] T.exp2(m_prev[i] * scale - m_cur[i] * scale) # rescale acc_o for i, j in T.Parallel(BLOCK_M, HEAD_DIM): acc_o[i, j] * m_scale[i] for i, j in T.Parallel(BLOCK_M, BLOCK_N): acc_s[i, j] T.exp2(acc_s[i, j] * scale - m_cur[i] * scale) # apply mask here T.gemm(acc_s, V_tile, acc_o) for i in T.Parallel(BLOCK_M): logsum[i] logsum[i] * m_scale[i] T.sum(acc_s[i, :]) # update m_prev T.copy(m_cur, m_prev) # final scale by logsum for i, j in T.Parallel(BLOCK_M, HEAD_DIM): acc_o[i, j] / logsum[i] T.copy(acc_o, O[bx * BLOCK_M, 0])这段代码跑通前向在 A100 上大概是 86% 的 MFU相对于理论峰值 312 TFLOPS 的 FP16 计算量对于第一版迁移来说已经算不错了。需要说明的是上面的代码里我保留了 mask 的占位注释实际应用中变长 mask 的处理方式会稍微复杂一点但骨架是对的。4.3 为什么T.gemm能自动适配转置transpose_BTrue这个参数背后是有硬件逻辑的。Hopper 和 Ampere 的 tensor core 都支持从 shared memory 读取转置后的 B 矩阵本质上是通过ldmatrix指令配合共享内存的地址映射实现的。TileLang 的编译器会生成对应版本的布局推导你不用手动做重排。如果你从 CUDA 转过来第一次看到这种代码很容易担心transpose_B是否意味着多一次显存拷贝实际上没有它是在 MMA 指令级别解决了数据排布问题代价只是共享内存地址模式略有不同。这种看起来像 numpy跑起来是 CUDA的体验是我用下来最大的感受。4.4 反向传播迁移中遇到的数值问题反向传播的迁移没有前向那么顺利。问题出在exp2数值稳定性上。我在前向用scale和 running max 的方式做了数值补偿之后结果和 PyTorch 的差异大约在 1e-5 级别这是正常的。但反向里当我直接用前向保存的logsum反向计算概率矩阵时出现了个别元素误差到 1e-2 的情况。查了半天发现问题根源是反向里对exp2的中间结果先转成了 float16 再参与后续计算而这个中间结果的范围跨了多个数量级。解决方式是所有 softmax 相关的中间量强制保持 float32只在T.gemm的输入处再转回 float16。改完之后误差立刻回到 1e-5 级别。5. 和 Triton 的实测对比不是谁取代谁是分工不同5.1 相同 GEMM 下的数字对比为了公平对比我分别用 Triton 和 TileLang 写了同一版本的 GEMMtile 大小统一为 128x128x32线程数保持 256A100 上跑 FP16结果如下方案GEMM TFLOPS达到峰值比例代码行数cuBLAS (torch)105.784.8%-Triton (手动调优)91.373.1%~45TileLang (默认)108.286.7%~20TileLang (swizzle 优化)111.489.3%~25单独看这个表Triton 在默认配置下确实打不过 TileLang但它并不会因此失去意义。Triton 对新手极其友好pytorch 的生态里很多扩展都是基于它写的。TileLang 更像是一个高天花板、多一层控制权的进阶工具。5.2 Attention 场景下的行为差异在 Flash Attention 这种 S 矩阵和 O 矩阵形状不对称的场景里差异更为显著。Triton 由于对tl.dot的 block 形状有较强的限制你在调BLOCK_M和BLOCK_N的时候经常要迁就硬件。TileLang 则因为 fragment 和 layout 是分离的理论上你可以用更大的BLOCK_M而不影响 inner loop 的访存效率。我做了一个变量实验固定BLOCK_N64改变BLOCK_MBLOCK_MTriton 耗时(ms)TileLang 耗时(ms)641.821.711281.641.522561.931.44Triton 在BLOCK_M到 256 时性能回退原因是它的线程布局没法在块内灵活重排。TileLang 则表现稳定得益于布局推导时的自动调整。这个实验让我更坚定了一个判断TileLang 在处理不规则 tile 形状的算子时优化空间明显更大。5.3 什么时候我还是会选 Triton说句公道话如果你的目标是快速验证想法、写一次性性能还行的 kernelTriton 依然是最好的选择。它的报错信息友好度目前比 TileLang 好不少而且社区样本大搜一下基本能找到现成方案。但如果是生产环境的算子库要反复研究不同硬件上的调度行为需要精确控制 TMA、warp specialize 这类特性我现在会更倾向 TileLang。本质上不是比谁强而是看你需不需要那一层控制权。如果需要Triton 给不了你CUTLASS 又太费人TileLang 正好卡在中间。6. 踩坑记录布局推导失败和 TMA 相关的问题6.1 你以为的 shape 合法编译器说不行我第一次写一个BLOCK_M64, BLOCK_N64, HEAD_DIM128的 Flash Attention 前向时在T.gemm(acc_s, V_tile, acc_o)这行直接报了一个布局推导的错误大意是 fragment shape 与 MMA 指令的布局不匹配。排查链路是这样的先检查acc_s的 shape 是(64, 64)看起来完全正常。再检查V_tile的 shape 是(64, 128)也正常。最后发现问题是acc_o是(64, 128)但T.gemm的累加维度要求它和V_tile的列维度对齐的布局规则在 MMA 的 warp-level 切分上产生了冲突。解决方式看起来有点反直觉把acc_o声明 shape 改成(64, 128)不变但把V_tile的加载方式改成按(128, 64)加载再转置。也就是让T.gemm自动推导时能匹配到ldmatrix支持的转置模式。改完之后一切正常。这类问题的本质是TileLang 会自动推导 layout但推导的前提是它与硬件指令的输入约束一致。大多数时候编译器会帮你做好但极端 shape 下你就得反过来配合它。6.2 Hopper TMA 的 swizzle 模式坑如果你在 H100 上使用 TMA共享内存的 swizzle 模式不是随便选的。TMA 要求shared memory的地址映射使用特定的 128 字节 swizzle 模式否则cp.async.bulk.tensor的行为就会出错。TileLang 的T.copy默认情况下会为 TMA 自动选择正确的 swizzle但在某些手写 layout 介入的场景下这个自动选择可能和你预设的 layout 冲突。我踩的一个实际案例是对一个(128, 64)的共享内存块手动做annotate_layout之后TMA 的拷贝结果变成了每行错位 16 字节。排查方法比较原始——把共享内存的值 dump 出来和 PyTorch 的 reference 对比然后逐步缩小到是拷贝阶段的问题。最后解决方式是让 swizzle 跟随 TMA 的默认设置而不是自己指定。后来的经验是能不动 swizzle 就不动让编译器做它擅长的地址映射你只需要关注 tile 形状是否正确。6.3 编译缓存导致的改代码无效假象还有一个容易误导新手的点TileLang 的 JIT 缓存是以源码哈希为准的但这里的源码包括整个 Python 函数的源码而不是只有 Kernel 循环体。所以如果你在外部修改了一个被闭包捕获的变量的默认值比如把BLOCK_M从 128 改成 64缓存会失效没问题但如果你改的是 kernel 内部逻辑而这段逻辑被 Python 的某种方式字符串化了可能不会触发重新编译。我的建议是在调试阶段给tilelang.jit加一个force_compileTrue参数源码里有这个选项等稳定之后再去掉。不然你很容易陷入改了代码但性能没变化的迷茫中。7. 我对 TileLang 后续路线的个人判断TileLang 目前还处在快速迭代期主分支几乎每周都有更新。从我看到的项目动向来说有几个方向我是比较看好的对 Hopper 架构的完整支持正在逐步成熟包括 TMA 和 warp specialization 的更多自动化组合。对 FP8 的支持在推进中包括 CUTLASS 那样的按 block 缩放策略。自动调优器和布局搜索框架如果做完善会大幅降低使用门槛到时候它的竞争力会明显上升。不过目前它还没到开箱即用、全程自动的程度。你在享受高控制权的同时也要承担对应的调试成本。个人建议的路线是先用 Triton 把正确性跑通再用 TileLang 重写关键算子做性能压榨而不是一上来就用它从零写所有东西。最后分享一个实用技巧如果你要对比不同 kernel 的性能记得把 L2 cache 的持久化行为考虑进去——单独跑一遍 warmup 再计时否则反复跑同一个 kernel 时L2 命中率会虚高导致数据不真实。这个细节在 TileLang 和 Triton 的 benchmark 脚本里都存在特别是小规模矩阵时影响显著。
前段时间做高性能算子优化的时候一直在跟 Triton 和 TVM 较劲卡在 Flash Attention 的调度细节上整整一周。后来在一个交流群里看到有人提了一嘴 tilelang说它在 GEMM 上的表现比手写 CUDA 还稳我当时第一反应是不太信直到自己跑完一轮 benchmark 才被折服。这篇文章就把我调研、上手、移植算子、最终落地在 A100 上的完整过程写出来涵盖了 TileLang 的编译链路、线程抽象模型、自动布局推导机制、与 Triton 的实测差异以及我踩过的几个比较隐蔽的坑希望对正在研究 GPU Kernel 编译方案的朋友有点参考价值。1. 从 Flash Attention 的调度之痛说起先说背景。我当时在做一个长序列场景的 Attention 算子优化针对的是 A100 40GB 的算力环境目标很直接在前向和反向都跑出接近 cuDNN 或者 FlashAttention-2 的水平同时还要处理好 mask 和变长序列。一开始用 Triton 写确实很爽。指针语义、自动向量化、块级别的调度上手成本低到一个下午就能把 naive 的 attention 跑通。但随着优化深入问题逐渐暴露出来主要是三件事Triton 的线程块布局大多数时候是你给它一个BLOCK_M和BLOCK_N它自己去决定线程怎么排布、访存怎么向量化你很难精确控制 SM 内部 warp 的物理排布。对 shared memory 的使用Triton 抽象得比较黑盒像 swizzle、padding 来消 bank conflict 这类操作大多数时候依赖编译器自动做结果不可预期。想接入 Hopper 的 TMATensor Memory Accelerator或者做 warp specializationTriton 目前的支持还是偏试探性的。这种情况下最朴素的想法是回去写 CUDA。但手写 kernel 的维护成本实在太高尤其是反向传播里那堆 tile 索引计算稍不留神就是一个越界 bug。正当我在犹豫要不要用 CUTLASS 的时候TileLang 进入了视野。TileLang 这个项目简单来说是一个 tile 级 DSL领域特定语言但它不像 Triton 那样把线程抽象几乎完全隐藏掉。它直接构建在 TVM TIR 和 MLIR 的基础上让你可以用 Python 风格写出接近 CUDA 的调度逻辑同时又把跨线程协作、共享内存分配、寄存器切分这些繁琐的事留给编译器处理。它解决的核心问题就是 Triton 太隐式、CUDA 太费手之间的那条鸿沟。尤其是它对线程布局的显式建模让我这种习惯看 CUDA 代码的人有一种莫名的亲切感。项目虽然年轻但背后是 TVM 社区那套经过大规模验证的编译基础设施所以不是从零造轮子而是把轮子重新设计了一遍形状。关于适合谁来学我觉得优先级是这样的已经被 Triton 的隐式调度折磨过、想精确控制硬件行为的进阶用户排第一有 CUDA 经验但不想再手写模板元编程的人排第二做推理引擎、算子编译器方向的工程同学也可以重点关注。如果你只是业务侧调模型那 TileLang 目前对你的帮助不大它更偏底层编译器基础设施。2. 装环境和最小复现场景编译链路整体走一遍TileLang 的环境安装不算复杂但有几个细节容易坑人我按自己的实操路径逐步说。2.1 源码编译的完整步骤我的环境是 CUDA 12.1、PyTorch 2.1、Python 3.10系统是 Ubuntu 22.04。TileLang 目前建议直接从源码编译因为 PyPI 上的版本滞后且经常没跟上主分支的更新。git clone --recursive https://github.com/tile-ai/tilelang.git cd tilelang mkdir build cd build cmake .. -DCMAKE_BUILD_TYPERelease -DPython3_EXECUTABLE$(which python) make -j$(nproc) cd .. python setup.py install这里有个关键点--recursive一定不能漏因为 TileLang 依赖了部分子模块的固定版本特别是它对 TVM 的 fork 做了 pin。如果漏了这一步后面跑import tilelang就会出现一堆类似 TIR node 定义找不到的诡异报错。编译大概需要十几分钟机器差一点的话可能半小时左右期间会编译 LLVM 相关组件建议提前确认磁盘剩余空间在 10GB 以上。2.2 一个可以直接复现的 GEMM 脚本装完之后我习惯性先用一个最小 GEMM 验证环境是否正常而不是直接上手大算子。import tilelang import tilelang.language as T import torch M, N, K 4096, 4096, 4096 tilelang.jit(out_idx[2]) def gemm(A: T.Tensor((M, K), dtypefloat16), B: T.Tensor((K, N), dtypefloat16), C: T.Tensor((M, N), dtypefloat16)): with T.Kernel(T.grid(16, 16), threads256) as (bx, by): A_tile T.alloc_shared((128, 32), float16) B_tile T.alloc_shared((32, 128), float16) C_accum T.alloc_fragment((128, 128), float16) T.clear(C_accum) for ko in T.serial(T.ceildiv(K, 32)): T.copy(A[bx * 128, ko * 32], A_tile) T.copy(B[ko * 32, by * 128], B_tile) T.gemm(A_tile, B_tile, C_accum) T.copy(C_accum, C[bx * 128, by * 128])这个脚本就是标准的 GEMM 模板分配共享内存的 A 和 B tile分配寄存器上的累加 fragment循环搬运数据到 shared memory然后调T.gemm做矩阵乘法。跑一遍对比 PyTorch 的 baselinepython gemm_benchmark.py我这边第一次跑出来的结果是 TileLang 的 float16 GEMM 在 4096 立方规模下能达到约 111 TFLOPS 左右而 PyTorch 的 cuBLAS 大概在 105 TFLOPS 上下等于说编译生成的代码已经能跟 cuBLAS 的主流实现打成平手甚至隐隐超过一点这还是在没有做任何手动调参的情况下。对于第一次接触这个框架的人来说这个结果会非常直观地建立信心。2.3 JIT 缓存行为的观察TileLang 的 JIT 机制很有意思。每次以相同的参数和代码结构运行 kernel 时它会在本地缓存编译产物后续跑就不再触发完整的编译流程而是直接加载二进制。缓存目录通常在~/.cache/tilelang里源码里以 kernel 源码的哈希值作为文件命名。我第一次跑 GEMM 大约等了十几秒因为要把整个 TIR 构建、MLIR 生成、CUDA 代码生成、nvcc编译几步全走完第二次跑就降到毫秒级的加载时间了。日常调参过程中改动任何一个 tile 大小都会触发重新编译所以建议把TILE_M、TILE_N、TILE_K作为 Python 变量传入编译缓存会基于最终生成的源码哈希进行区分不至于每次只改一个数字就全部重来。3. 线程布局的手感TileLang 和 Triton 的根本差异3.1 为什么说 TileLang 是一种显式隐式编程模型Triton 给你的编程接口是num_warps和num_stages至于这些 warps 怎么映射到 tensor 的维度上编译器内部有一套规则。多数情况下它做得很好但当你需要的不是合理而是最优时问题就会出现——尤其是在 Attention 这种每个 SM 承担的 tile 形状不规则的场景里。TileLang 的抽象层级很有意思它把 fragment 的概念保留为逻辑上的寄存器布局但你在代码里写的T.gemm实际会展开成对不同线程的 MMA 指令调用。线程布局并不是完全由你手动指定而是由编译器基于 tile 形状和硬件能力自动推断但这套推断是透明的你可以通过T.annotate_layout主动干预。这就有一种显式的隐式的手感不写死 threadIdx.x 怎么映射但你清楚地知道当前这个 fragment 里的每一个元素对应哪个线程的哪个寄存器。这种透明性对理解性能瓶颈至关重要。3.2 Warp 级编程与 fragment 的心智对齐熟悉 CUDA 的同学都知道经典的 warp-level GEMM 写法你把一个 thread block 按照 8x4 的方式布置成 warp tile 网格然后每个 warp 处理多个 16x16 的 fragment。这在 CUTLASS 里是模板参数的排列组合游戏而在 TileLang 里你只需要声明一个(128, 128)的alloc_fragment编译器会基于 MMA 指令的形状自动切分到 warp 和 lane。我一开始最大的困惑是我到底有没有完全控制权答案是你可以在 macro 层面控制但 lane-level 的 swizzle 它帮你做了。比如你写一个 128x128 的累加 fragmentT.gemm内部会使用 HMMA 16816 指令线程映射遵循硬件规定好的规则这层你不用管也最好不要管。从工程实践角度看这种设计是很聪明的。它把开发者的心智从我的 lane 0 应该存哪个元素中解放出来但又比 Triton 的黑盒多给了一层可见性——当你需要特殊 layout 时T.annotate_layout可以介入。3.3 一个直观的 layout 对比实验为了验证布局系统的灵活性我做了个简单测试分别用 Triton 和 TileLang 写一段相同的 GEMM然后对 tile 的共享内存访问模式做 profiling看 bank conflict 的情况。Triton 在默认情况下对 A 的共享内存访问做了一定的 padding但维度一超过某个阈值冲突率会出现跳变。TileLang 这边我可以在拷贝那一步手动做一个 swizzle写法如下T.copy(A[bx * 128, ko * 32], A_tile, swizzlebank)这个swizzle参数支持几种内置策略bank就是针对 32-bank 结构做的地址交错。实测显示在 K32 的 GEMM 里bank conflict 从大约 12% 降到了 3% 以内。虽然整体 TFLOPS 提升只有几个点但在更复杂的 Attention 算子中这种手动干预能力往往就是最后那 5%~10% 性能的关键。4. 把 Flash Attention 前向迁到 TileLang一次实操复盘4.1 从 naive 到可用的三段式结构先说我迁移 Flash Attention 的整体思路。标准的 forward kernel 大体分三步加载 query 的一个 block 到寄存器然后循环加载 key 的 block。计算S Q K^T做 online softmax更新统计量。计算P V累加输出。这三步在 TileLang 里的写法和 CUDA 几乎是线性对应的但它帮你处理了跨线程的 reduce 和 broadcast。我用tilelang.language里的T.reduce_max和T.reduce_sum实现了 online softmax 里的 running max 和 running sum。4.2 实际代码片段和解说下面这个是 head_dim128 时TileLang 版的 Attention 核心循环片段关键是T.gemm和 reducer 的配合。tilelang.jit(out_idx[2]) def attn_fwd(Q, K, V, O, scale): with T.Kernel(T.grid(num_tiles_m, num_tiles_n), threads128) as (bx, by): Q_tile T.alloc_fragment((BLOCK_M, HEAD_DIM), float16) K_tile T.alloc_fragment((BLOCK_N, HEAD_DIM), float16) V_tile T.alloc_fragment((BLOCK_N, HEAD_DIM), float16) acc_s T.alloc_fragment((BLOCK_M, BLOCK_N), float16) acc_o T.alloc_fragment((BLOCK_M, HEAD_DIM), float16) m_prev T.alloc_fragment((BLOCK_M,), float32) m_cur T.alloc_fragment((BLOCK_M,), float32) m_scale T.alloc_fragment((BLOCK_M,), float32) logsum T.alloc_fragment((BLOCK_M,), float32) T.clear(acc_o) T.fill(logsum, 0.0) T.fill(m_prev, -1e30) T.copy(Q[bx * BLOCK_M, 0], Q_tile) for ko in T.serial(T.ceildiv(seq_len, BLOCK_N)): T.copy(K[ko * BLOCK_N, 0], K_tile) T.gemm(Q_tile, K_tile, acc_s, transpose_BTrue) T.copy(V[ko * BLOCK_N, 0], V_tile) T.reduce_max(acc_s, m_cur, dim1, clearTrue) # online softmax update for i in T.Parallel(BLOCK_M): m_scale[i] T.exp2(m_prev[i] * scale - m_cur[i] * scale) # rescale acc_o for i, j in T.Parallel(BLOCK_M, HEAD_DIM): acc_o[i, j] * m_scale[i] for i, j in T.Parallel(BLOCK_M, BLOCK_N): acc_s[i, j] T.exp2(acc_s[i, j] * scale - m_cur[i] * scale) # apply mask here T.gemm(acc_s, V_tile, acc_o) for i in T.Parallel(BLOCK_M): logsum[i] logsum[i] * m_scale[i] T.sum(acc_s[i, :]) # update m_prev T.copy(m_cur, m_prev) # final scale by logsum for i, j in T.Parallel(BLOCK_M, HEAD_DIM): acc_o[i, j] / logsum[i] T.copy(acc_o, O[bx * BLOCK_M, 0])这段代码跑通前向在 A100 上大概是 86% 的 MFU相对于理论峰值 312 TFLOPS 的 FP16 计算量对于第一版迁移来说已经算不错了。需要说明的是上面的代码里我保留了 mask 的占位注释实际应用中变长 mask 的处理方式会稍微复杂一点但骨架是对的。4.3 为什么T.gemm能自动适配转置transpose_BTrue这个参数背后是有硬件逻辑的。Hopper 和 Ampere 的 tensor core 都支持从 shared memory 读取转置后的 B 矩阵本质上是通过ldmatrix指令配合共享内存的地址映射实现的。TileLang 的编译器会生成对应版本的布局推导你不用手动做重排。如果你从 CUDA 转过来第一次看到这种代码很容易担心transpose_B是否意味着多一次显存拷贝实际上没有它是在 MMA 指令级别解决了数据排布问题代价只是共享内存地址模式略有不同。这种看起来像 numpy跑起来是 CUDA的体验是我用下来最大的感受。4.4 反向传播迁移中遇到的数值问题反向传播的迁移没有前向那么顺利。问题出在exp2数值稳定性上。我在前向用scale和 running max 的方式做了数值补偿之后结果和 PyTorch 的差异大约在 1e-5 级别这是正常的。但反向里当我直接用前向保存的logsum反向计算概率矩阵时出现了个别元素误差到 1e-2 的情况。查了半天发现问题根源是反向里对exp2的中间结果先转成了 float16 再参与后续计算而这个中间结果的范围跨了多个数量级。解决方式是所有 softmax 相关的中间量强制保持 float32只在T.gemm的输入处再转回 float16。改完之后误差立刻回到 1e-5 级别。5. 和 Triton 的实测对比不是谁取代谁是分工不同5.1 相同 GEMM 下的数字对比为了公平对比我分别用 Triton 和 TileLang 写了同一版本的 GEMMtile 大小统一为 128x128x32线程数保持 256A100 上跑 FP16结果如下方案GEMM TFLOPS达到峰值比例代码行数cuBLAS (torch)105.784.8%-Triton (手动调优)91.373.1%~45TileLang (默认)108.286.7%~20TileLang (swizzle 优化)111.489.3%~25单独看这个表Triton 在默认配置下确实打不过 TileLang但它并不会因此失去意义。Triton 对新手极其友好pytorch 的生态里很多扩展都是基于它写的。TileLang 更像是一个高天花板、多一层控制权的进阶工具。5.2 Attention 场景下的行为差异在 Flash Attention 这种 S 矩阵和 O 矩阵形状不对称的场景里差异更为显著。Triton 由于对tl.dot的 block 形状有较强的限制你在调BLOCK_M和BLOCK_N的时候经常要迁就硬件。TileLang 则因为 fragment 和 layout 是分离的理论上你可以用更大的BLOCK_M而不影响 inner loop 的访存效率。我做了一个变量实验固定BLOCK_N64改变BLOCK_MBLOCK_MTriton 耗时(ms)TileLang 耗时(ms)641.821.711281.641.522561.931.44Triton 在BLOCK_M到 256 时性能回退原因是它的线程布局没法在块内灵活重排。TileLang 则表现稳定得益于布局推导时的自动调整。这个实验让我更坚定了一个判断TileLang 在处理不规则 tile 形状的算子时优化空间明显更大。5.3 什么时候我还是会选 Triton说句公道话如果你的目标是快速验证想法、写一次性性能还行的 kernelTriton 依然是最好的选择。它的报错信息友好度目前比 TileLang 好不少而且社区样本大搜一下基本能找到现成方案。但如果是生产环境的算子库要反复研究不同硬件上的调度行为需要精确控制 TMA、warp specialize 这类特性我现在会更倾向 TileLang。本质上不是比谁强而是看你需不需要那一层控制权。如果需要Triton 给不了你CUTLASS 又太费人TileLang 正好卡在中间。6. 踩坑记录布局推导失败和 TMA 相关的问题6.1 你以为的 shape 合法编译器说不行我第一次写一个BLOCK_M64, BLOCK_N64, HEAD_DIM128的 Flash Attention 前向时在T.gemm(acc_s, V_tile, acc_o)这行直接报了一个布局推导的错误大意是 fragment shape 与 MMA 指令的布局不匹配。排查链路是这样的先检查acc_s的 shape 是(64, 64)看起来完全正常。再检查V_tile的 shape 是(64, 128)也正常。最后发现问题是acc_o是(64, 128)但T.gemm的累加维度要求它和V_tile的列维度对齐的布局规则在 MMA 的 warp-level 切分上产生了冲突。解决方式看起来有点反直觉把acc_o声明 shape 改成(64, 128)不变但把V_tile的加载方式改成按(128, 64)加载再转置。也就是让T.gemm自动推导时能匹配到ldmatrix支持的转置模式。改完之后一切正常。这类问题的本质是TileLang 会自动推导 layout但推导的前提是它与硬件指令的输入约束一致。大多数时候编译器会帮你做好但极端 shape 下你就得反过来配合它。6.2 Hopper TMA 的 swizzle 模式坑如果你在 H100 上使用 TMA共享内存的 swizzle 模式不是随便选的。TMA 要求shared memory的地址映射使用特定的 128 字节 swizzle 模式否则cp.async.bulk.tensor的行为就会出错。TileLang 的T.copy默认情况下会为 TMA 自动选择正确的 swizzle但在某些手写 layout 介入的场景下这个自动选择可能和你预设的 layout 冲突。我踩的一个实际案例是对一个(128, 64)的共享内存块手动做annotate_layout之后TMA 的拷贝结果变成了每行错位 16 字节。排查方法比较原始——把共享内存的值 dump 出来和 PyTorch 的 reference 对比然后逐步缩小到是拷贝阶段的问题。最后解决方式是让 swizzle 跟随 TMA 的默认设置而不是自己指定。后来的经验是能不动 swizzle 就不动让编译器做它擅长的地址映射你只需要关注 tile 形状是否正确。6.3 编译缓存导致的改代码无效假象还有一个容易误导新手的点TileLang 的 JIT 缓存是以源码哈希为准的但这里的源码包括整个 Python 函数的源码而不是只有 Kernel 循环体。所以如果你在外部修改了一个被闭包捕获的变量的默认值比如把BLOCK_M从 128 改成 64缓存会失效没问题但如果你改的是 kernel 内部逻辑而这段逻辑被 Python 的某种方式字符串化了可能不会触发重新编译。我的建议是在调试阶段给tilelang.jit加一个force_compileTrue参数源码里有这个选项等稳定之后再去掉。不然你很容易陷入改了代码但性能没变化的迷茫中。7. 我对 TileLang 后续路线的个人判断TileLang 目前还处在快速迭代期主分支几乎每周都有更新。从我看到的项目动向来说有几个方向我是比较看好的对 Hopper 架构的完整支持正在逐步成熟包括 TMA 和 warp specialization 的更多自动化组合。对 FP8 的支持在推进中包括 CUTLASS 那样的按 block 缩放策略。自动调优器和布局搜索框架如果做完善会大幅降低使用门槛到时候它的竞争力会明显上升。不过目前它还没到开箱即用、全程自动的程度。你在享受高控制权的同时也要承担对应的调试成本。个人建议的路线是先用 Triton 把正确性跑通再用 TileLang 重写关键算子做性能压榨而不是一上来就用它从零写所有东西。最后分享一个实用技巧如果你要对比不同 kernel 的性能记得把 L2 cache 的持久化行为考虑进去——单独跑一遍 warmup 再计时否则反复跑同一个 kernel 时L2 命中率会虚高导致数据不真实。这个细节在 TileLang 和 Triton 的 benchmark 脚本里都存在特别是小规模矩阵时影响显著。
