TileLang：用 Python 写出接近手写的 GPU 内核 | 三步快速上手完整指南

TileLang：用 Python 写出接近手写的 GPU 内核 | 三步快速上手完整指南

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

TileLang 是一个开源的 GPU 编程 DSL：用类 Python 的语法写内核，编译器为 NVIDIA CUDA、AMD ROCm、CPU 等后端生成高性能代码。不需要 CUDA 基础，它就是 TileLang GPU 编程的实战入口，适合做深度学习算子开发、大模型推理优化，以及想免 CUDA 自定义算子的开发者。

项目速览

• 🚀Python 写内核：线程与内存交给编译器处理
• 多后端：NVIDIA、AMD、CPU、Metal 都能跑
• 性能接近手写：FP16 GEMM 与 cuBLAS 打平
• 40+ 示例：GEMM、注意力开箱可运行

它是如何工作的：分块 + 内存层次

TileLang 的核心概念是"tile"（瓦片）。以矩阵乘法为例：把大矩阵切成block_M × block_K的小块，每块走一轮"全局内存 → 共享内存 → 寄存器"的流程。你只描述取哪一块、算哪一步，哪个线程负责哪个元素由编译器推导。

内存层次被压缩成三个 API：T.alloc_shared申请片上共享内存，T.alloc_fragment用寄存器保存中间结果，T.Pipelined(num_stages=3)自动把"搬数据"和"做计算"错开重叠，隐藏访存延迟。并行细节（线程划分、向量化）同样不用手写，T.Parallel一行即可表达二维并行。

@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_shared = T.alloc_shared((block_M, block_K), T.float16) B_shared = T.alloc_shared((block_K, block_N), T.float16) C_local = T.alloc_fragment((block_M, block_N), T.float32) for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3): T.copy(A[by * block_M, k * block_K], A_shared) T.copy(B[k * block_K, bx * block_N], B_shared) T.gemm(A_shared, B_shared, C_local) T.copy(C_local, C[by * block_M, bx * block_N])

这就是 examples/quickstart.py 里 GEMM 的核心骨架（完整示例还会先用T.clear清零C_local再累加，并带 ReLU epilogue）。@tilelang.jit在首次调用时按输入形状特化编译内核；T.gemm会自动分发到目标硬件的对应矩阵指令上。

典型场景：哪些场景下它更划算

• 矩阵乘法（GEMM）：FP16/FP8/INT4、grouped GEMM、split-K 都有现成内核，见 examples/gemm/ 与 examples/gemm_fp8/。当 cuBLAS 无法满足融合需求（如额外挂 ReLU、缩放）时，自己写一个往往更划算。
• 注意力机制：examples/flash_attention/ 覆盖 FlashAttention 前向、反向与变长输入；examples/deepseek_mla/ 用约 80 行 Python 实现 DeepSeek MLA 解码，在 H100 上性能接近手写汇编水平。这是 LLM 推理优化的主战场。
• 量化推理：examples/dequantize_gemm/ 提供 W4A8、W8A16 等反量化矩阵乘内核，可直接用于大模型权重量化部署。

快速上手：三步跑通第一个内核

① 安装：

git clone https://gitcode.com/GitHub_Trending/ti/tilelang cd tilelang && pip install -e . python examples/quickstart.py

三条命令完成克隆、源码安装并跑通示例；如果只想试用，也可以直接pip install tilelang装预编译轮子。

② 跑 examples/ 下的示例：quickstart.py自带正确性校验；想改参数就打开 examples/gemm/example_gemm.py，它是另一个最小可跑版本。

③ 测量性能：kernel.get_profiler().do_bench(backend="cupti")返回延迟（毫秒）；再调kernel.get_kernel_source()直接查看生成的 CUDA 代码，方便定位瓶颈、调整 block 大小与流水级数。

性能：和 cuBLAS 有多接近

⚡ 两组公开的基准结论：

• FP16 GEMM：在 RTX 4090、A100、H100、MI300X 四款卡上，多数尺寸与 cuBLAS/rocBLAS 打平（加速比 ≈ 1.0），部分尺寸 MI300X 上接近 1.8 倍；Triton 整体略低。
• 量化 GEMV：A100 混合精度 GEMV 测试中，BitBLAS-TileLang 的 WINT2/AFPP16 组合峰值约 8 倍于 cuBLAS 的 FP16 路径，是推理场景里优势最大的一块。

进阶能力

• 自动调优：@tilelang.autotune自动搜索 block 大小、线程数、流水级数的最优组合，省去手动试错，参考 examples/gemm/example_gemm_autotune.py。
• 稀疏计算：T.gemm_sp支持 2:4 结构化稀疏张量核心，示例在 examples/gemm_sp/。
• 多后端编译：同一份内核代码可编译到cuda/hip/cpu/metal等目标，也支持接入新后端，设计见 tilelang/backend/。

适合谁用

它适合做推理算子、大模型推理链路和自定义融合算子、又不想陷进 CUDA 细节的工程师，尤其是需要同一份代码同时跑 NVIDIA 与 AMD 的团队。暂时不适合想开箱即用、一行代码都不改的人——TileLang 是编程语言而不是算子库，内核要你自己写（好消息是示例库很全）；如果还没建立 GPU 内存层次的基本概念，前期调性能会吃力。

学习路径与资源

• docs/：官方文档，含安装、语言基础、软件流水线、类型系统等章节
• examples/：示例库，从 examples/quickstart.py 起步，再按 examples/gemm/ → examples/flash_attention/ → examples/deepseek_mla/ 递进
• benchmark/：基准脚本与对比配置
• testing/：各功能的 Python 测试与回归用例

🚀 现在就克隆仓库跑一遍 quickstart，今晚你就能拥有第一个 GPU 内核。

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