Ch 08 · TileLang
封面/第二部分 · 基础设施/Ch 08
第二部分 · 基础设施 · 08

TileLang — 写给 GPU 的方言

V4 的细粒度架构会产生数百个 ATen operators。TileLang 通过轻量 host code、编译期验证和显式布局控制降低调度开销，并在需要时支持与 CUDA 对齐的位级输出。

名词速通 · 一分钟看懂 TileLang

TileLang = 一种把 GPU kernel 写成 tile-级 Python DSL 的语言，自带 SMT 求解器、host codegen、位级可复现的 lowering 钉死机制

一句话：把"写 CUDA"变成"写 Python 描述 tile 怎么搬"，编译器负责把 host 代码静态生成、用 Z3 解决索引非线性、用 layout 标注钉死代码生成顺序。 报告给出的结果是：CPU 侧参数验证从几十到数百微秒降至低于 1 微秒。通过让简化规则、layout 与 lowering 对齐 CUDA 实现，TileLang 可以在需要时生成位级一致的输出；这不等于 SASS 指令字节完全相同。

ATen（被替代的对象）
PyTorch 的 ATen 算子层会进行参数检查、调度和 kernel launch。V4 报告称细粒度算子使 CPU 侧验证开销不可忽略，但没有给出“每次 30–100µs”或“占 wall-time 30%”的统一测量。
Host 开销 vs Device 开销
Device 是 GPU 执行 kernel 的时间；Host 是 CPU 侧验证、参数封送和发起调用的开销。kernel 越短，固定 host 开销占比通常越高。报告量化的是 CPU-side validation 低于 1µs，不是整个 CUDA launch 的物理下限。
DSL（Domain-Specific Language）
领域专用语言。TileLang 提供面向 tile、shared memory、barrier 与 reduce 等 GPU kernel 抽象的编程接口，让编译器能够生成和优化设备代码。
Z3 / QF_NIA（编译器里的 SMT 求解器）
Z3 是微软的 SMT 求解器（Satisfiability Modulo Theories）。QF_NIA = "Quantifier-Free Nonlinear Integer Arithmetic" —— 处理无量词、整数、含乘除的算术约束。TileLang 用它判断"这个张量索引 `i*M + j*K + k` 在所有合法 `(i,j,k)` 下是否唯一"等问题，从而决定是否能向量化、能否合并访存。把传统编译器需要"启发式 + fallback"的部分，换成可证明的精确判断。
T.annotate_layout（钉死 lowering）
TileLang 允许开发者显式标注布局，并让编译器的简化与 lowering 规则贴近 CUDA 实现，从而在需要时获得位级一致的输出。报告没有声称每次编译都会生成与手写 CUDA 完全相同的 SASS。
IEEE-754 严格 vs fast-math
TileLang 默认关闭编译器级 fast-math，把影响精度的近似操作设为显式 opt-in；需要严格语义时，还提供带明确舍入模式的 IEEE-754 intrinsics。fast-math 可能改变运算与舍入结果，但不等于结果随机。
位级可复现（bit-identical reproducibility）
在指定输入、实现与执行条件下得到逐 bit 一致的输出，比数值近似一致更严格。它便于核对 TileLang 与 CUDA baseline，也有助于复查训练异常；并非所有调试都必须以逐位一致为前提。
一句话定位：TileLang 用 DSL、轻量 host code、SMT 辅助验证和显式 layout 支撑 V4 的融合 kernel 开发。论文明确量化的是 CPU 参数验证开销降至 1µs 以下，没有给出全模型节省 30% wall-time 的结果。

1. 为什么要减少细粒度 ATen 调用

通过 Python / ATen 调用细粒度 kernel 时，运行期检查和参数封送会产生固定成本。对短 kernel，这类 CPU 开销更容易限制利用率。V4 的新模块又会展开成大量细粒度操作，因此作者引入 TileLang host codegen：

  • 架构侧：mHC、hybrid attention 与低精度路径都需要定制 kernel；
  • 调用侧：模型结构会展开成数百个 ATen operators，放大逐次 validation 的固定成本；
  • 编译侧：复杂 layout、索引和数值语义需要比普通算子组合更强的分析能力。

示意：当一层由大量细粒度 operators 组成时，每次参数验证与 launch 的固定成本会累积。V4 报告只说架构会产生数百个 ATen operators，没有公布每层实际 launch 数、device 时间或 wall-time 占比。

怎样理解固定 host 开销 若一段计算由 $N$ 次细粒度调用组成，CPU 侧固定开销可近似写成 $N\cdot T_{\text{validation}}$。kernel 越短，这一固定项在 wall-time 中越显眼。报告称生成的轻量 host code 将参数验证开销从每次调用几十或数百微秒降到 1µs 以下；这个数字只针对 CPU-side validation，不是整个 kernel launch 的耗时，也不能单独推出固定的端到端加速倍数。

2. TileLang 怎么做到 < 1 µs

三件事：

  1. 共同生成 device kernel 与轻量 host launcher：把 dtype、rank / shape 约束和 stride / layout 假设写入生成代码；
  2. TVM-FFI 调用约定：通过紧凑接口与 zero-copy tensor interop 降低参数封送成本；
  3. 移出 Python 热路径：运行时验证和 marshaling 在生成的 host code 中完成。
📖 公式白话翻译 · launch 时间预算

把一次 kernel 调用的 host 时间拆开：

$$ T_{\text{host}} \;=\; \underbrace{T_{\text{py}}}_{\text{Python 字节码}} + \underbrace{T_{\text{disp}}}_{\text{ATen 选 kernel}} + \underbrace{T_{\text{check}}}_{\text{shape/dtype 校验}} + \underbrace{T_{\text{launch}}}_{\text{cudaLaunchKernel}} $$

这只是概念拆分。V4 报告没有分别给出 $T_{\text{py}}$、$T_{\text{disp}}$ 与 $T_{\text{launch}}$；它测得的是 CPU-side validation overhead 从几十或数百微秒降至每次调用低于 1 微秒。

3. 把 Z3 SMT 搬进编译器：QF_NIA 解决索引非线性

传统 GPU 编译器（NVCC、TVM）面对张量索引时，用启发式判断"两次访存能不能合并"、"循环能不能向量化"。启发式的问题：能处理简单情况（i, i+1 相邻），但碰到 i*M + j*K 这种带乘法的就退化为保守（不敢合并），白白浪费带宽。

TileLang 把这步换成可证明：用 Z3 求解器对索引表达式建立 QF_NIA 约束，问"是否存在两个不同 thread 写入同一地址"。

📖 公式白话翻译 · 索引唯一性的 SMT 编码

例子：要判断 idx = i*M + j（$0 \le i < N$, $0 \le j < M$）能否安全向量化（即不同 thread 写不同地址）。SMT 编码：

$$ \exists\, i_1, j_1, i_2, j_2 :\; (i_1, j_1) \ne (i_2, j_2) \;\land\; i_1 M + j_1 = i_2 M + j_2 $$

翻译成大白话："存在两组下标，下标本身不同但映射到的地址相同" —— 如果存在解 → 有冲突，不能向量化；不存在解 → 安全。

SMT solver 可以形式化判断这类整数表达式是否满足所需性质，并把结果用于 layout inference、memory hazard detection、bound analysis、vectorization 与 barrier insertion。报告只说在合理资源限制下，整体编译开销控制在数秒；没有给出单次查询 10 ms，也没有承诺自动找到全局最优 layout，开发者仍可用 T.annotate_layout 固定布局决策。

Demo · 单次 kernel launch 的时间分布对比（拖动 kernel device 时长看 host 占比）
交互
PyTorch + ATen TileLang 总 wall-time = 0 µs · host 占 0% 总 wall-time = 0 µs · host 占 0% Python 字节码 dispatcher shape check launch device

读图法：每条横向条形是一组教学时间预算。红、黄、橙、紫段的数值是为了演示固定 host 成本，不是报告实测，也不是硬件下限；蓝段表示可调的 device 时间。
把 device 时间调短，可以观察固定 host 成本占比上升。可确认的论文结果只有 validation overhead 低于 1µs；端到端差异仍取决于 kernel、硬件、调用路径和 workload。

4. 数值精度与按需位级复现

浮点加法不满足结合律：$(a+b)+c \ne a+(b+c)$，因此代数化简、lowering 与布局选择都可能改变最后几位。若要与手写 CUDA baseline 逐位比较，编译器需要尽量保持相同的求值与累加顺序。

TileLang 默认关闭编译器级 fast-math，把影响精度的近似作为显式 opt-in 操作，例如 T.__exp、T.__log 和 T.__sin。需要严格 IEEE-754 语义时，还可以使用带明确舍入模式的 T.ieee_fsqrt、T.ieee_fdiv 和 T.ieee_add。配合 T.annotate_layout 和对齐 CUDA toolchain 的化简、lowering 规则，可以在需要时得到与参考 CUDA 实现位级一致的输出；这不等于生成完全相同的 SASS。

位级复现为什么有用

与参考 kernel 位级对齐，可以把差异定位到具体算子，而不必先判断偏差来自算法、布局还是浮点重排。报告没有给出某个训练 step 的调试案例，也没有声称 TileLang 单独保证整次训练从 checkpoint 开始逐字节复现。

本章小结

  • V4 的细粒度 operators 让 host 侧验证和调度值得专门优化；报告未给出其占 wall-time 的百分比。
  • TileLang 的 host codegen 将运行时检查和参数封装移出 Python 路径，使 CPU 侧 validation overhead 从几十或数百微秒降到每次调用低于 1 微秒。
  • Z3 SMT 进编译器解决张量索引的非线性约束，把启发式的保守换成可证明的精确。
  • 通过显式 layout 与对齐 CUDA 的简化/lowering 规则，TileLang 可以在需要时获得位级一致的输出；不应扩写成“同一 SASS”。