TileLang：写 GPU kernel 的专用语言

它站在哪一层

先用一张图定位关系，否则容易搞错用法：

• PyTorch 负责搭模型、跑训练、给张量
• TileLang 负责其中少数几个 kernel 写得更准更快
• TVM 底层的编译基础设施，TileLang 是架在它上面的 Python 前端

所以它的场景是「PyTorch 已经能跑，但某个算子成了瓶颈」。整个模型都是你自己的话，大概不需要它。

三种安装

稳定版（推荐）：

bashpip install tilelang
python -c "import tilelang; print(tilelang.__version__)"

Nightly（要新特性时用）：

bashpip install tilelang --find-links https://tile-ai.github.io/whl/nightly

AMD 卡（ROCm）需要两步，顺序不能反：

bashpip install torch --index-url https://download.pytorch.org/whl/rocm7.0
pip install tilelang

运行时主机侧必须已装 ROCm。同样的 Linux wheel 可以直接用。

如果你是源码构建、editable install、Docker、想让 pip 装 CUDA 工具链、或者要用自定义的 TVM checkout —— 这五种才需要读完整安装指南（在新标签页打开），只从 PyPI 装的人跳过。

跑通第一个 kernel

这是官方 Quick Start 的完整代码，一个 FP16 GEMM 带 FP32 累加和融合的 ReLU 尾处理。读它比读文档快：

pythonimport torch
import tilelang
import tilelang.language as T


@tilelang.jit
def matmul_relu(A, B, block_M: int = 128, block_N: int = 128, block_K: int = 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)

        T.clear(C_local)
        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)

        for i, j in T.Parallel(block_M, block_N):
            C_local[i, j] = T.max(C_local[i, j], 0)

        T.copy(C_local, C[by * block_M, bx * block_N])

    return C


M = N = K = 1024
a = torch.randn((M, K), device="cuda", dtype=torch.float16)
b = torch.randn((K, N), device="cuda", dtype=torch.float16)
c = matmul_relu(a, b)
torch.testing.assert_close(c, torch.relu(a @ b), rtol=1e-2, atol=1e-2)
print("GEMM + ReLU passed.")

跑起来会看到 GEMM + ReLU passed.。注意 PyTorch 在 ROCm 系统上用的设备名也叫 cuda，TileLang 会从当前环境自动选目标后端。

这段代码在说什么

四个关键构造，恰好是 TileLang 的全部心智负担：

• T.Kernel(..., threads=128) 声明网格与线程数，得到 (bx, by) 两个 block 索引
• T.alloc_shared / T.alloc_fragment 分别申请共享内存和寄存器片段——这两个名字对应了 CUDA 里手动写 __shared__ 和局部数组
• T.Pipelined(..., num_stages=3) 是流水：三次数据搬运与计算重叠，CUDA 里要手写双缓冲或多缓冲，编译器帮你做
• T.gemm 把 tile 级的矩阵乘映射到目标后端；T.Parallel 表达逐元素的 ReLU 尾处理

@tilelang.jit 值得单独说。 它会按输入形状和编译期参数做特化，第一次调用触发的是编译不是执行。所以切换一个形状等于换一次编译——调试时别把第一次的慢当成性能问题。

从哪个例子开始

官方给的路线是：先跑 examples/quickstart.py 和 elementwise kernels，再进 gemm 目录看布局和自动调优。目录里值得知道的几类：

注意 GEMM 那几个是不同类别的量化/缩放路径，不是同一件事的五个版本——dequantize、FP8、block-scaled 各对应一种数值格式。

后端覆盖

CUDA、ROCm、Metal、LLVM（CPU），2026-09-30 起官方支持华为昇腾 950 NPU。目标平台从环境自动判断，一般不用手选。

两个版本相关的提醒

一是它还在 0.1.x。2026-08-03 的 v0.1.13 移除了若干旧 API，跨版本升级前读兼容性说明。

二是nightly 别当稳定版用。官方原话是 nightly 可能不如正式发布稳。

调试工具

仓库里 examples/analyze、examples/plot_layout、examples/iket 分别对应分析、布局可视化、时间线插桩。社区另有独立的 LSP 实现（tile-ai/tilelang-lsp），提供 buffer 形状、dtype、scope 和推断布局的 inlay hint 以及精确诊断——写 TileLang 时把编辑器支持打开会省不少事。

上手顺序

1. pip install tilelang
2. 跑上面那段 GEMM + ReLU，确认输出 passed
3. 读语言基础（language_basics）
4. 复制 examples/quickstart.py 到本地改参数，先改 blockM/blockN 看性能怎么变
5. 打开 LSP，理解自己写的那几行分别对应什么硬件操作