torch.compile、Inductor 与 XLA

手写内核和 Triton 都是"逐个算子"的工作。torch.compile 把这件事自动化：捕获整个模型的前向图，编译器自动做算子融合、内存规划、代码生成——一行装饰器，典型提速 30%–2 倍。XLA 是 Google 的同位对应物（JAX/PyTorch-XLA 的后端）。本章讲清楚图编译器的三步工作、以及"什么时候它快、什么时候它没用了"。

三步工作：捕获 → 优化 → 生成

所有训练/推理编译器都是同一套流程：

一、图捕获：把 PyTorch 的命令式执行捕获成静态计算图。torch.compile 的捕获机制（dynamo）追踪 Python 执行、把张量操作记录成图；动态控制流（if/循环随数据变化）会被切成多个子图（graph break）——子图越少，优化空间越大。

二、图优化：在图上做全局优化：

优化	做什么	收益
算子融合	相邻的逐元素/reduction 算子合成一个内核	最大头：显存往返按融合数减少
死代码消除	没用到的分支/节点删掉	小
常量折叠	编译期算完常量	小
内存规划	中间张量复用同一块显存	显存峰值降低

融合的极限是"整个模型一个内核"——实际做不到（矩阵乘要库、注意力要专门内核），所以现代编译器的标准做法是模板匹配：识别出 attention → 用 FlashAttention 内核；识别出 GEMM → 用 cuBLAS/CUTLASS；剩下的逐元素算子 → 自动生成融合内核。

三、代码生成：Inductor 是 torch.compile 的默认后端，把优化后的图生成 Triton 代码（GPU）或 C++（CPU）——它生成的是 Triton，所以 Inductor 的质量直接继承 Triton 的能力边界。

model = torch.compile(model)          # 一行
# 模式可选：mode="max-autotune"（更激进，编译更慢）
# 或只编译推理/训练的一部分（dynamo 的 graph break 日志值得看）

XLA：另一条路线

XLA（Accelerated Linear Algebra）是 Google 的编译器，服务 JAX 与 PyTorch-XLA：

** philosophy 相同**：捕获（jax.jit / XLA graph）→ 优化（融合、内存规划）→ 生成（TPU/GPU 机器码）
更强的地方：与 TPU 的脉动阵列共同设计（见加速器格局），全局形状推断与整图优化更彻底；JAX 的函数式编程模型让捕获无 graph break
更弱的地方：PyTorch 生态的算子覆盖与调试体验；动态 shape 处理要显式标注

TPU 上 XLA 是必经之路（没有别的选择）；GPU 上 torch.compile 是主路线。两者的工程经验互相通用。

什么时候它快，什么时候没用

编译器的收益来源决定了适用场景：

场景	收益	原因
小/中模型，算子密集	30%–2 倍	逐元素算子多，融合收益大（Roofline：访存受限算子融合翻倍）
大模型训练（TP/PP 混合并行）	5–15%	计算已被矩阵乘主导，融合空间小；通信优化是另一回事
推理（batch 小）	显著	kernel launch 开销与小算子主导，编译 + CUDA Graph 直接受益
动态控制流重的模型	可能负收益	graph break 频繁，编译缓存失效
极致性能的内核	没用	编译器生成的代码不如手写/库（FlashAttention）

判断方法：看编译日志里的内核数与融合情况——"200 个内核 → 60 个"说明融合生效；"graph break × 20"说明捕获碎了，收益有限。

训练侧的注意点

一、编译时间与 warmup：首次编译分钟级（max-autotune 更久），自动调优（autotune）对每个 shape 逐个试——动态 shape 频繁变化的负载（变长序列）会让编译缓存不断失效，用 dynamic=True 或 shape bucketing。

二、与分布式并行的交互：TP/PP 的集合通信插入在图里，编译器不优化通信本身（NCCL 调优另算，见NCCL）；FSDP 与 torch.compile 可叠加但 wrap 粒度影响图边界。

三、数值等价性：融合改变浮点求和顺序，结果与 eager 模式有微小差异（正常）；torch.compiler.assert_equal 可验证。loss 突然跳变才是 bug，微小差异是特性。

与其他优化手段的关系

编译器是"免费的优化"，它不与其他手段冲突，反而互补：

编译器做融合（省显存往返）→ FlashAttention 做算法级重写（省注意力访存）→ 量化做精度压缩（省字节）→ 三者叠加
推理侧再叠加 CUDA Graph（消 launch 开销）与连续批处理（消空闲）

优化顺序：先 compile（免费）、再专用内核（FlashAttention 等库已内置）、最后手写——大多数负载在第一步就吃到了大部分收益。

本章要点

编译器三步：捕获（graph break 是天敌）→ 优化（融合是最大头）→ 生成（Inductor 生成 Triton）。
模板匹配是现代编译器的标准做法：attention → FlashAttention，GEMM → cuBLAS，剩下融合。
收益场景：小/中模型、算子密集、推理 batch 小（30%–2 倍）；大模型训练收益 5–15%。
XLA 与 TPU 共同设计，TPU 上必经；JAX 函数式让捕获无 break。
编译时间与动态 shape 是工程成本；融合改变浮点顺序的微小差异是特性不是 bug。
优化顺序：先 compile（免费），再专用内核，最后手写。