Overview 首先介绍一下使用 tilelang 写 kernel 的两个视角： Block 视角 从一个 BLock 的视角写 Kernel, 操作的对象是整个 tile（比如一个 (BLOCK_M, BLOCK_N) 的矩阵片段），并行的粒度和分工方式通过 T.Parallel/T.copy 等声明式原语来表达，不关心具体哪个 threadIdx 做哪个元素。 换句话说，我们只需要声明，具体的线程分配由 tilelang 框架来完成 Thread 视角 从单个 Thread 的视角写 Kernel，block 内的并发是隐含的——因为 BLOCK_N 个 thread 各自都在执行同一份代码，天然并行，但代码里不体现这个并行结构。 Note Thread 视角可以让我们更贴近 CUDA 代码，但好像 tilelang example 中的用法基本都是 Block 视角，所以我们的重心应该放在 Block 视角上 ...
在生产环境下使用 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 ...