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

TileLang — 写给 GPU 的方言

为什么 V4 这种"几百个细粒度算子"的架构在 PyTorch + ATen 下会被 host 开销吃掉性能，TileLang 是怎么用 SMT 求解器 + host codegen 把 per-invocation 开销从几十微秒压到 1µs 以内，以及为什么"位级可复现"是这次替换最值钱的副产品。

名词速通 · 一分钟看懂 TileLang

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

一句话：把"写 CUDA"变成"写 Python 描述 tile 怎么搬"，编译器负责把 host 代码静态生成、用 Z3 解决索引非线性、用 layout 标注钉死代码生成顺序。 净效果：每次 kernel 调用的 host 端开销从 PyTorch + ATen 的 30–100 µs 压到 < 1 µs，并且生成的 SASS 与手写 CUDA bit-identical。

ATen（被替代的对象）
PyTorch 的算子库（"A Tensor library"）。每个 op 走"Python → C++ dispatcher → 算子选择 → CUDA launch"四层。per-call 开销 30–100 µs，对粗粒度算子（大 matmul）无所谓，但 V4 一层有 几百 个细算子（CSA 的 indexer / mHC 的 sum reduce / FP4 dequant…），累计后 host 开销吃掉 30%+ wall-time。
Host 开销 vs Device 开销
Device：GPU 真正算的时间；Host：CPU 把 kernel 排上队的时间（参数打包、stream 选择、launch 系统调用）。kernel 越小，host 占比越高。1µs 这个数字是 PCIe + ioctl 的物理下限附近 —— TileLang 几乎做到了硬件能给的极限。
DSL（Domain-Specific Language）
领域专用语言。TileLang 是嵌入 Python 的 DSL：你写 Python，但只能用 TileLang 提供的"tile / shared memory / barrier / reduce"等抽象。没有 numpy / for-loop 自由 —— 因为编译器要能静态推理。
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 提供的标注：告诉编译器"这块 shared memory 必须按 `swizzle(8, 64)` 排列、这个 thread 必须负责 idx ∈ [0, 32) 的元素"。把"编译器自由发挥"的部分变成显式约束。代价是写 kernel 时多花心思；收益是同一段 TileLang 代码每次编译都生成同一份 SASS → 与手写 CUDA bit-identical。
IEEE-754 严格 vs fast-math
IEEE-754 严格：浮点遵守标准，结合律不成立但每步可预测；fast-math：编译器可任意重排（如 `(a*b)+c` → `fma(a,b,c)`，把 a+b+c 拆成树形归约），速度↑但结果不可预测。TileLang 默认严格，opt-in 才用 fast-math（如 T.__exp 走快速指数）。这是"位级可复现"的根基。
位级可复现（bit-identical reproducibility）
同一段代码、同一份输入，每次跑都逐 bit 相同。比"数值精度一致"更严格。是 V4 工程哲学的核心：没有 bit-identity，loss spike 无法复现，调试就是赌博。Ch 09 会展开。
一句话定位：当一层网络从 5 个粗算子变成 200 个细算子时，"PyTorch 调度开销" 会突然变成主要瓶颈。TileLang 是 V4 拿来止血的工具：用 DSL 替手写 CUDA，用 SMT 替启发式优化，用 layout 标注替"编译器自由"，以"少一点写代码自由 + 多一点编译期可证明"换 30% wall-time 与位级可调试性。

1. 为什么 V4 不能再用 ATen

PyTorch 的 ATen 算子库是给"中等粒度算子"设计的：你写 linear(x, W)，它走 Python → C++ dispatcher → 选择最优 kernel → cudaLaunchKernel。调度路径长，但因为 kernel 本身要算 100 µs+，host 的 30 µs 占比可接受。 V4 改变了这个前提：

  • CSA 的 Lightning Indexer：每层一次 FP4 QK + top-k，单次 kernel ~50 µs，host 开销和它一个量级；
  • mHC 残差：每层 8 次小 reduce + 8 次 sum，每次 kernel 1–5 µs，host 开销远大于 device；
  • FP4 dequant：每个 GEMM 前一次小 kernel，单次 ~3 µs；
  • Sparse attention 反向：上百个 SM-level 的 reduce kernel。

一层 V4 有 200+ 个 kernel launch，按 ATen 的 30 µs 算 host 开销 = 6 ms，而 device 真正算的时间 ~10 ms。host 占了 wall-time 的 30%+。

数值演练 · host 开销吃掉多少 wall-time 假设 V4-Pro 一层 attention：
  • SWA 主干 GEMM × 4：device 1.2 ms，host 4×30 µs = 0.12 ms（10%）；
  • CSA Indexer + topk + gather：device 0.3 ms，host 6×30 µs = 0.18 ms（60%）；
  • mHC 残差：device 0.05 ms，host 16×30 µs = 0.48 ms（~10×）；
一层合计：device 1.55 ms，host 0.78 ms → wall-time 2.33 ms，其中 host 占 33%。一个 step 61 层 × 2 forward+backward = ~95 ms 单步浪费在 launch。换 TileLang 后 host 单次 0.8 µs：61×2×200×0.8µs = ~20 ms → 每步省 75 ms，相当于 train throughput +10–15%。

这纯粹是把 launch 的系统调用成本压下去带来的 free lunch —— 不改算法、不改硬件、不改精度，只换语言。

2. TileLang 怎么做到 < 1 µs

三件事：

  1. Host codegen：把 shape 校验、stride 计算、grid/block size 选择全部静态生成成 C++，编译时就钉好。运行时 host 路径只剩"打包参数 + ioctl"。
  2. 无 dispatcher：每个 TileLang 函数直接编译为一个 CUDA kernel + 一段 host stub，没有 PyTorch 的 op 注册 / dtype 分发链路。
  3. 编译产物长存：编译结果按签名 cache，第二次调用直接走 hot path，不再编译。
📖 公式白话翻译 · 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}} $$
  • $T_{\text{py}}$ ~10 µs：Python 解释器栈帧、参数对象创建。TileLang 把它变成一次 C++ stub 调用 ≈ 0.2 µs。
  • $T_{\text{disp}}$ ~5 µs：ATen 按 dtype/device 在哈希表里找实现。TileLang 没 dispatcher = 0。
  • $T_{\text{check}}$ ~5–10 µs：检查 contiguous、stride、broadcast。TileLang 在编译期就把这些条件钉成静态约束 ≈ 0。
  • $T_{\text{launch}}$ ~5 µs：CUDA 驱动走 ioctl 进 kernel 队列。无法消除，是物理下限。

所以 ATen 的 30 µs 里有 25 µs 是语言/框架开销，可以全部砍掉，剩下 5 µs 就是 cudaLaunchKernel 本身。1 µs 的承诺 = 0.2 µs Python stub + 0.6 µs launch + 0.2 µs 杂项。

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 $$

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

对启发式来说这条要靠"看 stride 是不是常数"猜，错了直接 segfault。对 SMT 来说这是10ms 内可证明的判定问题。整个编译过程多花几秒，换来运行时每个 kernel 自动选最优 layout，不需要手写 hint。

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

读图法：每条横向条形是一次 kernel launch 的"时间预算"。ATen 的红/黄/橙段是语言开销（python 10µs + dispatcher 5µs + check 8µs），紫段是 cudaLaunchKernel（5µs，物理下限），蓝段是真正在 GPU 上算的时间。
把 device 滑到 10 µs 以下（V4 里 mHC 小算子的真实区间）你会看到 ATen 条形里蓝色不到 1/3 —— 这就是为什么细粒度算子下 host 是主瓶颈。TileLang 把红/黄/橙合到 0.2 µs 以内，紫的 launch 是省不掉的 5 µs，但相对占比一下子降到 50% 以下，device 越大 TileLang 优势越不明显，反过来device 越小 TileLang 越是救命稻草。

4. 位级可复现：默认严格、显式 opt-in fast-math

浮点加法不满足结合律：$(a+b)+c \ne a+(b+c)$ 在 FP32 下也常出现 1ULP 的差异。一旦编译器自由地把 $a+b+c+d$ 重排成树形归约或SIMD 横向求和，同一段代码每次运行结果就不一样。这对调试是灾难。

TileLang 默认走 IEEE-754 严格：所有归约保持源码顺序，FMA 不自动融合（除非显式写 T.fma）。需要 fast-math 时必须显式调用 T.__exp、T.__div 等。配合 T.annotate_layout 钉死 shared memory 排布，同一份 TileLang 源码 → 同一份 SASS → 同一字节输出。这就是 Ch 09 整章的"批不变 + 确定性"赖以成立的语言层基础。

"语言级 bit-identical" 在调试里的实战价值

假设训练在第 13417 步出现 loss spike。如果代码不是 bit-identical：你重启复现 → 在第 13420 才 spike，参数已经偏了几步，找不到根因。
如果是 bit-identical：你保存 13416 步 checkpoint，重新跑一遍 → 每个 logits 字节相同，可以在出现 NaN 的那一层、那一个 token 上接 debugger，直接看是哪个 attention head 在哪几步爆掉。这就是 V4 把 1.6T 参数训练做工程化的根本姿态：所有"难以复现的 bug"先用工具变成"必然复现"，再去解。

本章小结

  • V4 的细粒度算子化让 host 开销从"可忽略"变成"占 30% wall-time"，必须换语言止血。
  • TileLang 用 host codegen + 无 dispatcher + 编译 cache 把 per-launch 开销从 30 µs 压到 < 1 µs。
  • Z3 SMT 进编译器解决张量索引的非线性约束，把启发式的保守换成可证明的精确。
  • 默认 IEEE-754 严格 + 显式 layout 钉死 → 同源同 SASS 同输出 → Ch09 的 bit-identical 调试由此成立。