5分钟上手TileLang：GPU内核开发指南

5分钟上手TileLang：GPU内核开发指南

【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang

手写一个打平 cuBLAS 的 FP16 GEMM 内核，通常是编译器团队的活。TileLang 把 GPU 内核开发拉回应用开发者的桌面：它是构建在 TVM 之上的 Python 风格领域特定语言，你用 tile（数据块）为单位描述数据搬移和计算，编译器自动处理内存层级映射、软件流水线与张量核心调度，同一份源码可编译到 CUDA、ROCm、Metal 甚至 CPU。

核心机制：Tile 抽象如何替掉手写 CUDA

传统 CUDA 里，哪个线程读哪个地址、在哪个 barrier 上等待都要你自己定；TileLang 只向你收"什么数据、做什么运算"，效率细节交给编译器。真正起作用的是下面两处机制。

内存层级就是三次分配

GPU 三级存储对应三次分配调用：从全局内存用T.copy切出需要的 tile，T.alloc_shared把它放进共享内存，T.alloc_fragment把累加器放进寄存器文件。好处是数据重用变得显式——一个 tile 只从 HBM 读一次，在多次T.gemm里反复使用，带宽问题基本归结为"块大小怎么选"。

一行 T.Pipelined 换自动软件流水线

手写流水线意味着把共享内存切多份、重排 copy 与 gemm 的顺序、插入异步 barrier，是最繁琐的部分。TileLang 里写T.Pipelined(n, num_stages=3)就够了：编译器自行推断运算顺序、各操作落在哪一级 stage、哪些 copy 该走异步，还能在新架构上进一步降低为 TMA 批量搬运和 warp 特化流水线。你只写朴素循环，流水线由编译器生成。

一次注解，多后端生成

@tilelang.jit首次调用时按具体形状特化内核并缓存，目标默认从环境自动探测。同一份源码经过不同后端落到 CUDA、ROCm 和 Metal 上，不用为不同显卡维护多套代码。

左半部分是三级存储上的 tile 数据流，右半部分是逐行对应的 TileLang 内核代码，这张图基本概括了"tile 抽象"的完整含义。

安装与第一个内核：一条命令跑通

稳定版在 PyPI 上：

pip install tilelang

文件头import tilelang和import tilelang.language as T之后，一个完整可跑的 FP16 GEMM 长这样（官方 quickstart 还带融合 ReLU 的 epilogue，这里取核心 15 行）：

@tilelang.jit def matmul(A, B, block_M=128, block_N=128, block_K=32): M, N, K = T.const("M, N, K") A: T.Tensor((M, K), T.float16); B: T.Tensor((K, N), T.float16) C = T.empty((M, N), T.float16) with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by): A_s, B_s = T.alloc_shared((block_M, block_K), T.float16), T.alloc_shared((block_K, block_N), T.float16) C_l = T.alloc_fragment((block_M, block_N), T.float32) T.clear(C_l) for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3): T.copy(A[by*block_M, k*block_K], A_s) T.copy(B[k*block_K, bx*block_N], B_s) T.gemm(A_s, B_s, C_l) T.copy(C_l, C[by*block_M, bx*block_N]) return C

首次调用matmul(a, b)会触发编译，之后直接运行；结果可用 PyTorch 的a @ b交叉验证，示例里还会顺手打印生成的 CUDA 源码和 CUPTI 实测延迟。

H100 实测：GEMM 与 Attention 的基准数据

📊 仓库 README 里公开了一组基准图，挑两组最代表性的。

这张图以 TileLang 为 1.0 基线（越低越快）：FP16 GEMM 与 cuBLAS 基本持平；WFP4 混合精度 GEMM 优势最大，PyTorch（bitsandbytes）参考实现慢约 1.5-1.9 倍；注意力一组里 FA3 与 TileLang 互有胜负，PyTorch 参考在长序列下最慢接近 2 倍。

换成跨平台视角：RTX 4090、A100、H100、MI300X 四张卡上，TileLang 的 FP16 GEMM 大体贴着或高于 1.0 基线（cuBLAS/rocBLAS），MI300X 的大尺寸用例接近 2 倍。单一 DSL 在 NVIDIA 和 AMD 上都有竞争力，这点在同类工具里不多见。

进阶方向：从 GEMM 到真实模型算子

仓库的 examples/ 目录相当于一本内核 cookbook，按算子分了几个值得翻的子目录：

• examples/gemm/：自动调优、persistent 调度、intrinsics 等 GEMM 变体，适合把同一次矩阵乘调到极限
• examples/deepseek_mla/：FlashMLA 解码实现，单文件约 270 行，附带 H100 基准图和优化笔记，适合 decode/注意力场景
• dequantize_gemm 与 flash_attention 目录：低比特量化 GEMM（W4A8、MXFP4 等）和完整 MHA 前向/反向，适合推理与训练算子

避坑与调参：五个实用建议

⏱ 几条能省时间的经验：

• 首次调用慢是正常的：@tilelang.jit会触发完整 TVM 编译流水线；服务化集成时用.compile(...)预编译并复用缓存
• 最有效的调优参数是块大小和num_stages：从 128/128/32 + stages=3 起步，再交给 AutoTuner 搜索，别手动猜
• 性能不对就先kernel.get_kernel_source()看生成的 CUDA，确认落到了哪条硬件路径；要看内核中间值用T.print
• 累加器 dtype 保持 float32，写回 fp16 放在最后一次 copy，两者混用结果会漂移
• 源码安装：git clone https://gitcode.com/GitHub_Trending/ti/tilelang之后pip install -e .

TileLang 的价值一句话：用你思考问题的粒度（tile、流水线）描述高性能内核，把繁琐留给编译器。系统资料看 docs/ 目录，examples/ 有逐算子的可运行示例，benchmark/ 有可复现的测试脚本。

【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang