Sealessland logo Sealessland
LLM Inference

TVM 与 TileLang：深度学习编译器

整理深度学习编译器和高性能 kernel DSL 的实际作用，避免和推理引擎搞混。

发布时间
标签
TVMTileLangCompilerKernel

TVM 与 TileLang：深度学习编译器

我之前容易把 TVM 当成”另一个 ONNX”，或者把 TileLang 当成”Python 写 CUDA 的语法糖”。聊完之后才发现，编译器这条线和推理引擎（vLLM、TensorRT-LLM）解决的问题不一样，而且 TileLang 的低延迟也不是靠”语法糖”实现的。

这篇把编译器和推理引擎的分工理清楚，再讲 TileLang 到底在做什么。

编译器 vs 推理引擎

TVM / TileLangvLLM / TensorRT-LLM
输入计算图（Relay/Relax/TIR）模型权重 + 配置
输出硬件指令（CUDA/Metal/OpenCL/LLVM）推理服务（请求调度 + kernel 执行）
优化层级图级 + 算子级 + 自动调优手工优化 kernel + 系统调度
硬件覆盖CPU/GPU/FPGA/自定义加速器主要是 NVIDIA GPU

推理引擎管的是”怎么把请求排好队、喂给 GPU、把结果送回去”。编译器管的是”这个算子在这个硬件上怎么跑最快”。

两者可以组合：MLC-LLM 就是基于 TVM Unity（Relax），把 Llama、Mistral 编译到手机、浏览器 WebGPU、苹果 Metal 上运行。编译器生成 kernel，推理引擎做调度。

TVM 的工作流程

TVM 的核心是分层 lowering：

PyTorch/TensorFlow/ONNX 模型
  │
  ▼
Relay / Relax（高层计算图表示）
  │── 算子融合、常量折叠、死代码消除
  ▼
TensorIR / TIR（底层张量运算表示）
  │── loop split、reorder、vectorize、bind to thread
  ▼
AutoTVM / MetaSchedule（自动搜索最优调度）
  │── 搜索 tile size、loop order、unroll factor
  ▼
Codegen（生成目标硬件代码）
  │── CUDA、OpenCL、Vulkan、LLVM
  ▼
Runtime（在目标硬件上加载执行）

Relay / Relax 是图级别的 IR，类似 ONNX，但可编程、可优化。可以做算子融合、布局转换这些图级优化。

TensorIR / TIR 是循环级别的 IR，描述具体的 for-loop 嵌套。这里的关键抽象是 Block（带边界谓词的计算单元）和 Schedule（一系列变换原语）。优化和正确性分离，调度变换不会改变计算结果。

AutoTVM / MetaSchedule 是自动调优。以矩阵乘法为例，TVM 会搜索：

  • tile size（分块大小）
  • loop order（循环顺序）
  • vectorize width（向量化宽度）
  • thread binding（绑到 GPU 几维 block/thread）

AutoTVM 需要手写搜索模板，Ansor / MetaSchedule 是无模板搜索，用进化算法 + 成本模型探索调度空间。

TileLang 的定位

TileLang 是 MLC 团队（TVM 生态）推出的 kernel DSL，目标是：用接近 Python 的语法，写出媲美手写 CUDA（cuBLAS / Cutlass）性能的算子，同时保持低发射延迟。

低延迟从哪来

TileLang 的 Python 代码只是前端 DSL，实际运行时：

  1. 编译时全特化：shape、dtype、tile size 这些在编译时确定，运行时直接调用，没有动态 dispatch
  2. 零开销抽象：tl.copy、tl.load 这些不是函数调用，是编译期内联原语，没有临时对象分配
  3. 显式内存层级控制：强制程序员管理 Global Memory → Shared Memory → Register 的数据流动，减少冗余搬运
  4. TVM 后端：lower 到 TIR 后，继承所有循环优化和代码生成能力
# TileLang 风格（示意）
@tl.jit
def matmul(A: tl.Buffer, B: tl.Buffer, C: tl.Buffer):
    with tl.Tile([128, 128]):
        shared_A = tl.copy(A)      # Global → Shared
        shared_B = tl.copy(B)
        acc = tl.zeros([128, 128])
        for k in tl.range(0, K, 32):
            local_A = tl.load(shared_A)   # Shared → Register
            local_B = tl.load(shared_B)
            acc += local_A * local_B
        tl.store(C, acc)

编译后，tl.copy 变成直接的 ld.global → st.shared PTX，没有 Python runtime 介入。

TileLang vs Triton vs Cutlass

维度TileLangTritonCutlass
抽象层级显式 tile + 内存层级隐式 tile（block 级）C++ template，最底层
编译方式TVM 编译栈MLIR/Triton IR → LLVMC++ 编译器
延迟优化编译时全特化 + 零开销抽象编译优化好，少量 runtime dispatch极致，但开发成本高
易用性Python DSL，中等Python DSL，较易C++，难
硬件覆盖TVM 支持的后端主要是 NVIDIA GPUNVIDIA only

Triton 的 tile 管理是半自动的，编译器帮你决定一些细节。TileLang 是显式的，控制精度更高，代价是程序员需要理解内存层级。

我之前搞混的地方

“编译优化”和”推理框架优化”的区别。编译器优化的是单次 kernel 的执行效率，推理框架优化的是多个请求在 GPU 上的排布效率。两者正交：

  • 编译器把 MLP 融合成一个 kernel，减少 HBM 访问
  • 推理框架把多个请求的 MLP 打包成一个 batch，提高 GPU 利用率

AutoTVM 的搜索成本。我原来以为自动调优是”免费”的，但其实：

  • 搜索一个算子的最优调度可能需要几百次 trial
  • 每次 trial 要在目标硬件上实际运行测量
  • 搜索时间可能几十分钟到几小时

所以 AutoTVM 的结果通常要缓存（tuning log），不是每次编译都重新搜。

实际应用场景

  • TVM：适合需要跨平台部署的场景（手机、边缘设备、国产芯片），或者需要自定义算子和硬件后端
  • TileLang：适合写高性能 GPU kernel，特别是融合算子、MoE、量化 kernel 这些推理框架里没有现成实现的
  • MLC-LLM：基于 TVM Unity，把大模型编译到手机、浏览器，做端侧推理

写到这才把”编译器管什么、推理引擎管什么”分清楚。之前看 MLC-LLM 的代码，总觉得和 vLLM 重叠，现在明白一个是生成 kernel 的，一个是调度请求的，分工很清晰。