上篇文章把我和两位同事的对话内容po出来之后,他们很震惊很宠幸,今天吃饭我赶紧把他们拉着别跑了,这次决定在饭前就进行一次首席科学家级别的高质量圆桌对话,以凸显我们的内涵。
同事1:上次聊完triton、cuTile和Tilelang三者的定位和技术差异,我下来后去浅显地学习了一下Tilelang,有两个疑问:一是 TileLang 基于 TVM 这个点,具体能带来什么优势?二是有没有实际的性能对比数据,能直观看出它们在不同场景下的表现?
我:这两个问题不错,先说说 TileLang 基于 TVM 的优势 ——TVM 本身就是成熟的跨平台编译器框架,TileLang 直接把这个 “底座” 用起来了,最核心的好处就是跨硬件适配能力。比如 TVM 已经做了 NVIDIA CUDA、AMD ROCm、x86 CPU 及各国产芯片的编译后端,TileLang 不用从头开发,直接继承这些能力,所以它能同时在 H100、MI300X、昇腾 910B 上跑,而且性能不会打折扣。
同事 2:对,还有编译优化的复用性。TVM 有一套完整的multi stage lowering编译优化流程,比如自动地内存布局调整、软件流水线调优、tensorcore映射这些,TileLang 不用自己再做一遍,只需要专注在 “Tile 级编程抽象” 上。比如你写一个 GEMM 算子,TVM 的scheduler会帮 TileLang 算出最优的 tile 大小和thread binding策略,比 Triton 自己做的heuristic search更灵活,也比 cuTile 只能依赖 NVIDIA 显卡优化的范围更广泛。
我:还有一点很重要 ——TVM 的 “算子融合” 和 auto tune 能力,TileLang 也能直接用。比如你写一个包含卷积 + 激活的计算,TVM 会自动把它们融合成一个 kernel;而且 TileLang 还能调用 TVM 的 AutoTVM/Ansor,自动搜索最优的调度方案,不用像 Triton 那样手动调 num_warps、num_stages 这些参数,虽然我觉得二者对新手门槛都不算高
同事 1:TileLang看来确实踩在了TVM的肩膀上啊,那性能对比呢?比如在常用的 GEMM、FlashAttention 这些算子上,三者谁更快?
同事 2:我特意查了最近的测试数据,分场景说会更清楚。先看 NVIDIA GPU 场景,比如 H100 上的 GEMM(1024x1024x1024,float16):cuTile 因为是 NVIDIA 原生极致优化, latency 大概在 0.012ms 左右,接近手写 CUTLASS 的性能;Triton 的官方测试数据印象中是 0.016ms,差了一点(记错勿喷);TileLang latency 能做到 0.013ms左右,比 Triton 好,略逊于 cuTile,但胜在不用绑死 NVIDIA 硬件。(绝对数据仅供参考,不一定来自权威榜单,但相对数据在Tilelang repo里面有)
同事 2:再看 FlashAttention 这个高频场景,DeepSeek 团队做过测试 —— 他们用 TileLang 写的 FlashAttention 算子,代码量从原版 500 多行减到 几十行,在 H100 上处理 4,48,4096,4096,64 的输入时, latency 相比FA3和Triton都更快(来自Tilelang repo)而且 TileLang 能在 AMD MI300X GPU上跑同样的算子,cuTile 不支持 AMD GPU。
同事 1:那国产硬件场景呢?比如昇腾 910B,三者表现怎么样?
我:这方面那 TileLang 支持一些,昇腾 910B 上,TileLang 因为 TVM 已经做了深度适配,而 Triton 目前只有社区的非官方适配,cuTile 完全不支持。华为在某个大会上还展示过,用 TileLang 写的反量化 GEMM 算子,在昇腾上比手写的 AscendCL 代码快 12%,而且代码量少了 3/4。
同事 1:总结一下性能:在 NVIDIA 专属硬件上,cuTile 性能略优;跨硬件场景(尤其是国产芯片、AMD),TileLang 靠 TVM 底座占优;Triton 则看起来在生态上相对cuTile和Tilelang本身走在相对靠前
同事 2:对的。
同事 1:让我们借着Tilelang来聊聊TVM,大模型出来后这个项目没有以前火了,都快被很多人遗忘了,甚至23年入行AI infra的萌新根本么听过TVM。
我:是啊,TVM作为一个非常有野心的想要统一所有硬件后端的开源深度学习编译器,先不看商业和战略层面,只看技术层面就有很多硬核的玩意儿,TVM 而是 “图级 + 算子级 + runtime级” 的三层体系,外加Auto-tune。用TVM写一个矩阵乘法可能需要手动管理循环拆分、并行化,对新手太不友好。我先放个 基于TVM Tensor Expression写的的矩阵乘法代码,你们看看:
import tvmfrom tvm import teN = 1024k = te.reduce_axis((0, N), "k")A = te.placeholder((N, N), name="A")B = te.placeholder((N, N), name="B")# 定义计算逻辑C = te.compute((N, N), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k))# 手动写调度s = te.create_schedule(C.op)i, j = C.op.axisio, ii = s[C].split(i, 32) # 手动拆分维度jo, ji = s[C].split(j, 32)s[C].reorder(io, jo, ii, ji, k) # 手动调整计算顺序s[C].parallel(io)s[C].unroll(ji)
调度部分要写一堆代码,需要调用 te.schedule、split 这些 schedule原语,如果用TensorIR来写,那又会是另一套API,学习成本相对比较高。
同事 2:这其实就是 TileLang 想要解决的问题,它在 TVM 基础上做了语法简化和统一,以及dataflow和schedule的分离,同样的矩阵乘法,TileLang如下,不过说老实话,也有一些API和schedule原语要去熟悉,但是好在个人感觉比较清爽
import tilelangimport tilelang.language as Tdef matmul(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float"):@T.prim_funcdef main(A: T.Buffer((M, K), dtype),B: T.Buffer((K, N), dtype),C: T.Buffer((M, N), dtype),):# Initialize Kernel Contextwith T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) 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_stages=3):T.copy(A[by * block_M, ko * block_K], A_shared)for k, j in T.Parallel(block_K, block_N):B_shared[k, j] = B[ko * block_K + k, bx * block_N + j]T.gemm(A_shared, B_shared, C_local)T.copy(C_local, C[by * block_M, bx * block_N])return mainfunc = matmul(1024, 1024, 1024, 128, 128, 32)jit_kernel = tilelang.JITKernel(func, out_idx=[2], target="cuda")
我:那 cuTile 和 Triton 的矩阵乘法代码是什么样的?正好对比下四个工具的差异。
同事 1:cuTile 作为 NVIDIA 原生工具,老实说,代码最简洁,深度绑定 CUDA 生态,不用额外配置各种schedule原语:
def matmul_kernel(A, B, C,tm: ConstInt, # Tile size along M dimension (rows of C)tn: ConstInt, # Tile size along N dimension (columns of C)tk: ConstInt): # Tile size along K dimension (inner product dimension)GROUP_SIZE_M = 8M = A.shape[0]N = B.shape[1]bidx, bidy = swizzle_2d(M, N, tm, tn, GROUP_SIZE_M)num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk))accumulator = ct.full((tm, tn), 0, dtype=ct.float32)zero_pad = ct.PaddingMode.ZEROdtype = ct.tfloat32 if A.dtype == ct.float32 else A.dtypefor k in range(num_tiles_k):a = ct.load(A, index=(bidx, k), shape=(tm, tk), padding_mode=zero_pad).astype(dtype)b = ct.load(B, index=(k, bidy), shape=(tk, tn), padding_mode=zero_pad).astype(dtype)accumulator = ct.mma(a, b, accumulator)accumulator = ct.astype(accumulator, C.dtype)ct.store(C, index=(bidx, bidy), tile=accumulator)
你看,这不干净么?各种优化帮你做了,不需要各种schedule balabala的,帮你写第5代Tensor Core MMA的代码,不用像 TVM 那样手动配置,在 NVIDIA GPU 上跑,性能直接拉满。
import tritonimport triton.language as tl@triton.jitdef matmul_kernel(A, B, C, M, N, K,BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr):# 手动定义块ID和线程ID,但逻辑比TVM简单pid = tl.program_id(0)i = pid / (N / BLOCK_SIZE_N)j = pid % (N / BLOCK_SIZE_N)for k ....a = tl.load(A[i*BLOCK_SIZE_M : (i+1)*BLOCK_SIZE_M, k*BLOCK_SIZE_K : (k+1)*BLOCK_SIZE_K])b = tl.load(B[k*BLOCK_SIZE_K : (k+1)*BLOCK_SIZE_K, j*BLOCK_SIZE_N : (j+1)*BLOCK_SIZE_N])c += tl.dot(a, b)tl.store(C[i*BLOCK_SIZE_M : (i+1)*BLOCK_SIZE_M, j*BLOCK_SIZE_N : (j+1)*BLOCK_SIZE_N], c)
同事 1:这四个代码一对比,风格还是差异比较较明显,TVM最底层,代码最繁琐,但灵活性最高;TileLang 基于 TVM 做了语法简化, 且听说面向三种不同level的算子编写人员都有相应的接口,野心和TVM一样都不小,想要通吃;cuTile 最省心,NVIDIA 全自动化优化;Triton 则是跨平台场景的平衡之选。
我:而且性能表现也能看出来 cuTile 是NVIDIA亲儿子,多数workload在NVIDIA GPU上表现最优;TileLang凭借TVM 的多个auto tuner,比如 AutoTVM/Ansor/meta schedule做 自动调优,略逊于 cuTile;Triton相对二者都要略逊但胜在均衡,支持硬件广;不过TVM的 auto tune一直以后诟病的一点是时间过长,因此涌现了很多减少tune时间的工作,诞生了很多OSDI级别的paper。
同事 1:现在真是一目了然了,选工具的时候,看代码写法和性能需求就能定:要 NVIDIA 专属极致性能,选 cuTile,最省心;要跨平台 + 中等学习成本,选 Triton,代码灵活;要跨平台 + 自主可控schedule + 自动调优,选 TileLang,借力 TVM 生态。
同事 2:对的,这几个代码示例正好印证了咱们之前的结论 —— 这些工具的核心差异,本质是 “抽象层次” 和 “生态绑定” 的差异。cuTile 抽象层次最高,适合rookie,绑定 NVIDIA 生态;TileLang 抽象层次不高不低,适合各级水平的人,不管你是rookie还是expert都可以,绑定 TVM 跨平台生态;Triton 抽象层次中等,适合rookie,自建跨平台生态;TVM 抽象层次最低,适合expert,也是跨平台生态。
我:哎,和两位首席科学家交流真是醍醐灌顶啊,好了,我得回去benchmark了[大哭]。

