Production Ready Tilelang Ops Framework

Mon, 17 Aug 2026 00:00:00 +0000

在生产环境下使用 Tilelang，需要考虑几个问题：

算子缓存
与 torch.compile 兼容

算子缓存

算子定义

在了解算子缓存之前，我们需要了解一个 GPU 算子应该如何定义，通常由三部分组成：

参数类型	RMSNorm 中的实例	作用
问题规格（Workload / specialization parameters）	`M, N, eps, dtype`	确定一类具体计算问题
调度参数（Schedule / config / meta-parameters）	`block_m, threads`	确定该 workload 的一种实现方案
运行时参数（Runtime arguments / operands）	`x, weight`	每次调用时传入的实际数据

在代码里对应得非常直接：

_rms_norm_kernel(M, N, eps, dtype)   # workload parameters

    def _func(block_m, threads):     # schedule/meta-parameters

        def main(x, weight, y):      # runtime arguments

这三类参数不是一次性传给同一个函数，而是在 kernel 从定义到执行的过程中逐层绑定：

jit_impl = _rms_norm_kernel(M, N, eps, dtype)
jit_kernel = jit_impl(block_m=4, threads=128)
y = jit_kernel(x, weight)

第一步绑定 workload 参数，得到 TileLang JITImpl
- 描述一个确定的计算问题，例如处理形状为 (M, N)、数据类型为 dtype 的 RMSNorm，但尚未确定具体的调度方案。
第二步绑定 schedule 参数，得到可以执行的 JITKernel
- block_m/threads 可以来自默认配置、手动配置或 autotune
- 不同 schedule 对应同一个 workload 的不同实现版本
最后传入 x/weight 等 runtime arguments，启动已经选定的 JITKernel
- 运行数据只参与本次调用，不改变前面已经确定的 workload 和 schedule

算子调用流程

一次真实的算子调用从业务 Op 开始，依次确定 workload、选择 schedule、取得编译结果，最后才把 Tensor 作为 runtime arguments 启动 kernel

Tilelang on KKKZOZ's Blog

Production Ready Tilelang Ops Framework

算子缓存

算子定义

算子调用流程