第14章 TorchInductor：从 ATen IR 到 Triton kernel

"A Scheduler is a graph of BaseSchedulerNodes. It is responsible for optimizations such as fusion, reorder, and graph partition."

—— torch/_inductor/scheduler.py 中 Scheduler 类的 docstring

14.1 Inductor 在编译栈的位置

回顾整个 torch.compile 链路：

graph LR
    User["@torch.compile<br/>用户代码"]
    User --> Dy[Dynamo<br/>第 12 章<br/>FX Graph]
    Dy --> Aot[AOTAutograd<br/>第 13 章<br/>fw + bw 子图]
    Aot --> Ind[Inductor<br/>本章]
    Ind --> Tri[Triton kernel<br/>GPU]
    Ind --> Cpp[C++ + OpenMP<br/>CPU]
    style Ind fill:#fef3c7,stroke:#f59e0b,stroke-width:2px

Inductor 拿到的是一张 纯函数式的 ATen FX Graph（AOTAutograd 已经 functionalize 过），它要把这张图变成实际可执行的代码。源码在 torch/_inductor/，其中几个核心文件（行数按 v2.11.0）：

• compile_fx.py (3047 行) — 编译入口
• lowering.py (7935 行) — ATen → Inductor IR
• ir.py (9966 行) — Inductor IR 节点定义
• scheduler.py (7158 行) — fusion 调度
• torch/_inductor/codegen/triton.py (6536 行) — Triton 代码生成

torch/_inductor/ 在 v2.11.0 共约 213,000 行 Python（同样口径下 v2.4.0 是约 89,000 行；v2.4.0 的 tag 提交于 2024-07，不到两年涨了一倍多），是 PyTorch 编译器栈里体量最大、增长很快的一块。

14.2 Lowering：ATen → Inductor IR

打开 torch/_inductor/lowering.py，会看到一两百处 @register_lowering / register_pointwise 注册（按 overload 展开后，运行时 lowerings 表在 v2.11.0 有约 1800 项；PyTorch 2 论文写作时是 433 个算子、含 overload 1605 个）。以 aten.add 为例，它走的是逐元素算子的通用注册器：

# torch/_inductor/lowering.py
add = register_pointwise(
    aten.add, allow_alpha=True, override_fn_when_input_bool="logical_or"
)

register_pointwise（:903）内部用 make_pointwise（:630）生成 lowering：alpha != 1 时先把最后一个输入乘上 alpha，再构造一个 inner_fn 调 ops.add 的 Pointwise 节点，最后交给 register_lowering（:529）登记。每个 ATen 算子被注册一个 lowering 函数，把它翻译成 Inductor IR ops。Inductor IR 比 ATen 更底层 —— 它用循环 + 索引表达式描述计算，而不是 high-level 算子。

核心 IR 类（ir.py）：

类	含义
Pointwise (line 1071)	逐元素操作，output[i] = f(input[i])
Reduction (line 1221)	归约操作，output[i] = sum(input[i, j] for j)
Scan (line 2367)	扫描，如 cumsum
Sort (line 2575)	排序
FixedLayout (line 4045)	内存布局描述（strides + offset）
ExternKernel	外部 kernel（如 cuBLAS gemm，不自己生成）

关键设计：每个 IR 节点都用 indexing_fn(index) → 标量计算表达式 描述。比如 relu(x) 是一个 Pointwise，它的 inner_fn 大致是 lambda idx: ops.relu(loader(idx))（Triton 后端再把 ops.relu 展开成 ops.maximum(0, x)，见 torch/_inductor/codegen/triton.py:1269）。这种"循环 body 用纯函数描述"是 Inductor 实现 fusion 的基础 —— 两个 Pointwise 的 inner_fn 可以直接拼起来。

14.3 Scheduling：fusion 决策的核心算法

Lowering 结果是一组 IR 节点，每个节点是一段循环。Scheduler 决定哪些节点合并成一个 kernel：

节点 A: out_a[i] = relu(x[i])         # Pointwise
节点 B: out_b[i] = out_a[i] * 2       # Pointwise

→ 合并成: out_b[i] = relu(x[i]) * 2   # 一个 kernel, 一次内存读, 一次写

fusion 是 GPU 性能最主要的杠杆之一。HBM 访问延迟在数百纳秒量级、带宽是逐元素算子的硬瓶颈；融合两个 pointwise 把 HBM 往返从 4 次降到 2 次，理论上限就是接近 2x。

graph LR
    subgraph NoFuse[不 fuse: 4 次 HBM 读写]
        X1[load x] --> A1[relu]
        A1 --> S1[store out_a]
        S1 -.HBM.-> L2[load out_a]
        L2 --> M1[mul 2]
        M1 --> S2[store out_b]
    end
    subgraph Fuse[fuse: 2 次 HBM 读写]
        X2[load x] --> A2[relu]
        A2 --> M2[mul 2]
        M2 --> S3[store out_b]
    end
    style NoFuse fill:#fee2e2
    style Fuse fill:#dcfce7

scheduler.py 实现的 fusion 决策考虑：

• 依赖关系：B 依赖 A（纵向），或二者互不依赖但读同一个 buffer（横向）→ 可以考虑 fuse
• 循环可对齐：A 是 [N] Pointwise、B 是 [N] Pointwise → 循环维度一致 → 可 fuse
• 规模与峰值内存：合并后的节点数超过 config.max_fusion_size（默认 64）、或融合会抬高峰值内存，就不 fuse
• 内存读写量：两节点共享的读写字节数（score_fusion_memory）为 0 就不 fuse
• 共享中间值：A 输出被 B 和 C 共用 → fuse 哪个更优

最终算法是 贪心 + 启发式：先枚举所有合法的候选对，按分数从高到低排序，依次尝试融合（细节见 §14.9.1）。

FusedSchedulerNode（line 1866）表示融合后的节点，ForeachKernelSchedulerNode（line 2216）是 multi-tensor 优化的 _foreach_* 类操作专门融合节点。

14.4 Triton Codegen：把 IR 转成 Triton DSL

Triton 是 OpenAI 开发的 GPU DSL，比 CUDA C++ 更高级 —— 用户写"逻辑上的 block-level 循环"，Triton 编译器自动决定 thread mapping、内存合并、shared memory 等。

Inductor 生成的 Triton 代码大致长这样（简化版）：

@triton.jit
def fused_relu_mul_kernel(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel

    tmp0 = tl.load(in_ptr0 + xindex, xmask)
    tmp1 = triton_helpers.maximum(0, tmp0)   # relu
    tmp2 = tmp1 * 2                    # mul
    tl.store(out_ptr0 + xindex, tmp2, xmask)

codegen/triton.py:TritonKernel（line 2503）是核心生成类。它遍历 fused IR 节点的 inner_fn，把每条 Inductor IR op 翻译成 Triton 操作：

Inductor op	Triton 翻译
ops.load(buf, idx)	tl.load(buf + idx)
ops.add(a, b)	a + b
ops.maximum(a, b)	triton_helpers.maximum(a, b)（CUDA 路径；ROCm 才直接发 tl.maximum(a, b, tl.PropagateNan.ALL)）
ops.reduction(..., "sum")	tl.sum(...)
ops.reduction(..., "max")	triton_helpers.max2(...)（any / max / min / prod 都走 helper，且 max/min 带 2 后缀）
ops.store(buf, idx, val)	tl.store(buf + idx, val)

注意 maximum 与 max 归约这两行——并不是所有 op 都直接翻译成 tl.*。triton_helpers 是 PyTorch 自带的一层包装（torch/_inductor/runtime/triton_helpers.py），用来补 NaN 传播语义与归约的边界处理；选哪一条路由由 get_triton_reduction_function（torch/_inductor/codegen/triton.py:134）和 TritonOverrides.maximum（同文件 :1294）决定。读者拿生成出来的 kernel 对照本表时，看到的会是 triton_helpers.max2 而不是 tl.max。

整个 codegen 是 string templating：拼接 Python 字符串组装 Triton 源码，再交给 Triton 的 triton.jit 编译成 PTX → CUBIN。

14.4.1 为什么用 Triton 而不是手写 CUDA

写一个 element-wise CUDA kernel 要决定：

• block size（一般 128 / 256 / 512，但最优值依 dtype / shape 变化）
• thread coalescing（让相邻 thread 访问相邻地址）
• shared memory 用法
• vectorized load（float4 / __half2）

每个算子都人工调这些参数太累。Triton 把其中大部分决策交给编译器 —— 用户只写"block-level 循环"，block 内的线程映射、向量化、合并访存由 Triton 编译器处理；block size 本身是 tl.constexpr 参数，由调用方选（Inductor 用启发式 + autotune 选，见 §14.4.2）。

这让 Inductor codegen 大大简化：只生成 logical Triton 代码，底层优化交给 Triton 编译器。PyTorch 2 论文（ASPLOS'24，§4.1 与 §4.8）给出的选型理由是：Triton 已被证明能用简单的输入代码生成快过手写库的 kernel，而且 PyTorch 用户本来就越来越多地用 Triton 写高性能 kernel。

14.4.2 配 autotune 找最优 block size

Triton 自己提供 @triton.autotune，但 Inductor 没用它，而是自带一套：生成的 kernel 挂的是 Inductor 自己的装饰器（示意）：

@triton_heuristics.pointwise(
    size_hints={'x': 1048576},
    filename=__file__,
    triton_meta={...},
    inductor_meta={...},
    min_elem_per_thread=0
)
@triton.jit
def triton_poi_fused_mul_relu_0(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    ...

装饰器这几行由 Triton codegen 在生成 kernel 源码时拼出（codegen_kernel 里的 heuristics_line）。triton_heuristics.pointwise（torch/_inductor/runtime/triton_heuristics.py:2818）按 size_hints 生成候选 triton.Config：一维 pointwise 默认只有两三个候选（XBLOCK = bs 与 bs // 2，bs = max(256, min(numel // 128, 1024))，外加 autotune hint 给出的配置）；二维 tiling 的 pointwise 默认有六个左右的 (XBLOCK, YBLOCK) 候选（方形 tile 且非 ROCm 时只留一个 32×32，开 max_autotune / max_autotune_pointwise 才恢复全部候选）。候选交给 CachingAutotuner（同文件 :313），第一次跑时逐个 benchmark 选最快的（autotune_to_one_config），结果写进本地 autotune 缓存，下次编译出同一份 kernel 时直接取用。这套机制让 Inductor 在没人手动调参的情况下也能选到合适的 block size。

14.5 CPU 后端：C++ + OpenMP

GPU 上生成 Triton，CPU 上 Inductor 生成 C++ + OpenMP 代码。torch/_inductor/codegen/cpp.py 生成形如：

#pragma omp parallel for
for (int i = 0; i < N; ++i) {
    auto tmp0 = in_ptr0[i];
    auto tmp1 = std::max(tmp0, 0.f);
    out_ptr0[i] = tmp1 * 2;
}

生成代码开头是 #include <torch/csrc/inductor/cpp_prefix.h>（它带进 omp.h 与 ATen 的 at::vec::Vectorized）；SIMD 不是靠编译器自动向量化，而是 Inductor 直接生成 Vectorized<float> 调用（§14.10.3）。整段 C++ 由 torch/_inductor/cpp_builder.py 的 CppBuilder 调系统 C++ 编译器编成 .so（只借 torch.utils.cpp_extension 取头文件与库路径）；默认的 CPU kernel 走 async_compile.cpp_pybinding，编成 Python 扩展模块后用 importlib 加载（CppPythonBindingsCodeCache，torch/_inductor/codecache.py:3055）。

CPU 上的加速比通常比 GPU 小，原因见 §14.18 第 6 条：ATen 的 CPU 算子本来就在调 oneDNN / MKL 这类高度优化的库，留给 Inductor 优化的胶水更少。

14.6 一个完整 trace：a + b * c

@torch.compile
def f(a, b, c):
    return a + b * c

Inductor 收到 ATen graph：mul(b, c) → add(a, _)。流程：

flowchart LR
    FX["fx Graph<br/>mul b c → add a _"]
    FX --> LO["Lowering<br/>每个 op 转 Pointwise IR"]
    LO --> IR["IR Nodes<br/>node mul: lambda i  load b i * load c i<br/>node add: lambda i  load a i + intermediate"]
    IR --> SC["Scheduling<br/>检查 add 依赖 mul  循环维度对齐<br/>决定 fuse 成一个 kernel"]
    SC --> CG["Codegen<br/>生成 Triton kernel<br/>load b/c → mul → load a → add → store"]
    CG --> KER["fused_mul_add 一个 GPU kernel<br/>4 次 HBM 访问替代 6 次"]
    style FX fill:#fef3c7
    style SC fill:#dbeafe
    style CG fill:#dcfce7
    style KER fill:#fce7f3

各步细节：

1. Lowering：mul → Pointwise(inner_fn=lambda i: load(b,i) * load(c,i))，add → Pointwise(inner_fn=lambda i: load(a,i) + load(intermediate,i))
2. Scheduling：发现 add 依赖 mul，循环维度都是 [N]，可 fuse
3. Codegen：生成一个 Triton kernel（简化示意：实际 kernel 名形如 triton_poi_fused_add_mul_0，参数叫 in_ptr0 / out_ptr0；本机 CPU 上 TORCH_LOGS=output_code 实跑得到的对应 C++ kernel 是 cpp_fused_add_mul_0，同样只有一个 kernel）：

@triton.jit
def fused_mul_add(a_ptr, b_ptr, c_ptr, out_ptr, xnumel, XBLOCK: tl.constexpr):
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)
    xmask = xindex < xnumel

    tmp0 = tl.load(b_ptr + xindex, xmask)
    tmp1 = tl.load(c_ptr + xindex, xmask)
    tmp2 = tmp0 * tmp1
    tmp3 = tl.load(a_ptr + xindex, xmask)
    tmp4 = tmp3 + tmp2
    tl.store(out_ptr + xindex, tmp4, xmask)

——两个原始 ATen 算子（mul / add）变成一个 GPU kernel：不融合时是"读 b、读 c、写 tmp"加"读 a、读 tmp、写 out"共 6 次 HBM 访问，融合后只剩"读 a、读 b、读 c、写 out"4 次。对这种纯 memory-bound 的图，省下的访存就是省下的时间。

14.7 Lowering 之前：FX 图上的 Pass

14.7.1 fx_passes：Lowering 之前的图变换

torch/_inductor/fx_passes/（含子目录约 37700 行、七十多个文件）有一批 Pass 在 Lowering 之前对 FX Graph 做改写：

Pass	文件	职责
reinplace	reinplace.py (936 行)	把 functional op 重新转成 inplace（撤销 AOTAutograd 的 functionalize）省一次 alloc
split_cat	split_cat.py (3041 行)	识别 split → 多个 op → cat 这种"切了再合"模式，整段消除
replace_random	replace_random.py	把 aten.rand / aten.randn 替换成 inductor_prims.random(size, seed, mode)，把 seed 显式化
reduced_atomic_contention	reduced_atomic_contention.py	把高争用的 index_put 累加写分散到多个分区 buffer 再 sum 合并，降低 atomic 争用（config.partitioned_scatter_enabled，默认关）
mkldnn_fusion	mkldnn_fusion.py	CPU 上 conv + bn + relu 等的 oneDNN 融合
joint_graph	joint_graph.py	把 fw + bw 合在一起的图做联合优化

每个 Pass 是一个独立的图变换，按 compile_fx.py 里固定顺序跑。这些 Pass 让 Inductor 拿到的不是"AOTAutograd 原始输出"，而是已经做过几轮高级优化的图 —— 后面的 Lowering / Scheduling 在更优的 IR 上工作。

reinplace 特别有意思：第 13 章 §13.4 提过 AOTAutograd 把 inplace 全部 functionalize 掉，让中端图分析能干净做。但最终 codegen 出 Triton 时，inplace 才能省内存。reinplace Pass 在 Lowering 前把"逻辑上独立但实际可复用"的张量重新 inplace。这是"先 functionalize 简化分析、再 reinplace 恢复性能"的工程闭环。

14.7.2 Pattern Matcher：自动重写为 fused 算子

pattern_matcher.py（2376 行）实现了 Inductor 的"模式匹配重写"。它能识别 FX Graph 里特定子图，整段替换成更高效的算子调用。

经典例子：SDPA pattern matching。用户写：

# 朴素 attention 实现
scores = q @ k.transpose(-2, -1) / sqrt(d)
attn = F.softmax(scores, dim=-1)
out = attn @ v

Pattern matcher（torch/_inductor/fx_passes/fuse_attention.py）识别这一段，自动重写成 aten.scaled_dot_product_attention，再由 SDPA 按设备与输入条件挑 FlashAttention / memory-efficient / math 等后端。用户什么都不用改，条件满足时就能用上 FlashAttention。

类似地有：

• 多个互相独立、形状相同的 layer_norm / linear 批量合并（fx_passes/group_batch_fusion.py，默认不开，要在 config.pre_grad_fusion_options 里点名）
• BN 处于 eval 模式时 conv + bn 合并（fx_passes/efficient_conv_bn_eval.py，config.efficient_conv_bn_eval_fx_passes 默认 False）；CPU 上 fx_passes/mkldnn_fusion.py 把 conv 与后续 relu 等融进 oneDNN 算子
• matmul 维度补齐到对齐尺寸（fx_passes/pad_mm.py）

每个 pattern 是一段"匹配子图模板 + 替换函数"（入口是 register_replacement()）。Pattern matcher 把模式表示成 PatternExpr 节点组成的 DAG，从图中的锚点节点出发逐个递归 _match（见 torch/_inductor/pattern_matcher.py 文件头的说明）。

14.8 IR 节点层次与内存规划

14.8.1 IR 节点的完整层次：Buffer / ComputedBuffer / TensorBox / StorageBox

§14.2 提到 Inductor IR 的核心是 Loops（Pointwise / Reduction），但实际 IR 是个有层次的对象图。打开 ir.py（9966 行），会看到几个不同抽象层级的类：

层	类	职责
计算层	Pointwise / Reduction / Scan (Loops 子类)	描述"循环 body" + indexing function
缓冲层	Buffer (基类) / ComputedBuffer / ConcatKernel	把计算结果存成具体内存 buffer
包装层	TensorBox / StorageBox	给上层一个"像 Tensor 一样的句柄"
布局层	FixedLayout / FlexibleLayout / MultiOutputLayout	描述 strides / 内存布局

调用关系：TensorBox 包 StorageBox，StorageBox 包 Buffer，ComputedBuffer 持有 Loops（具体计算）。这种"四层包装"是 PyTorch 内部 Tensor 的镜像（第 2 章 §2.2 的 Tensor / TensorBase / TensorImpl / Storage 三件套思想）。

graph TB
    TB[TensorBox<br/>用户视角的 tensor]
    SB[StorageBox<br/>共享存储]
    Buf[Buffer<br/>具体内存 buffer]
    CB[ComputedBuffer<br/>需要计算才能填的 buffer]
    Loops[Pointwise / Reduction / Scan<br/>计算定义]
    Layout[FixedLayout / FlexibleLayout<br/>strides + offset]
    TB --> SB
    SB --> Buf
    Buf -.子类.-> CB
    CB --> Loops
    Buf --> Layout
    style TB fill:#fef3c7
    style SB fill:#dbeafe
    style Buf fill:#dcfce7
    style Loops fill:#fce7f3

为什么要这么多层？因为 Inductor 要在不同阶段做不同变换：

• StorageBox 让 view 操作（reshape / view）零成本 —— 多个 TensorBox 共享同一个 StorageBox，类似 ATen 的 view
• FlexibleLayout 表示"layout 还没决定，编译器可以选最优 strides"，到后期 decide_layout() 才固化
• MultiOutputLayout 用于返回多个张量的外部 kernel（如回退到多输出 ATen 算子的 FallbackKernel，以及 Conditional / WhileLoop 等控制流节点），各个输出再由 MultiOutput 节点取出

理解这套层级，再读 ir.py 就不会被几千行类定义绕晕 —— 每层的职责清晰、互不重叠。

14.8.2 Memory Planning：buffer 复用与生命周期

Inductor 的 buffer 复用发生在 wrapper codegen 阶段。基本原理：

Buffer A: alive [step 1, step 5]
Buffer B: alive [step 6, step 9]  ← B 完全在 A 死亡之后

→ B 可以复用 A 的内存地址 (assign 同一段)

具体实现是 torch/_inductor/codegen/wrapper.py:1949 的 memory_plan_reuse，它按顺序扫一遍 wrapper 里的内存规划行：

1. 某个 buffer 最后一次被使用后生成一条"若未被复用就释放"的行（FreeIfNotReusedLine），把它放进复用池，键是 buffer_reuse_key（device、dtype、符号化的 storage 大小、对齐与否）
2. 后面遇到分配行（AllocateLine.plan）时，若池里有同键的已死 buffer，就改写成复用行（ReuseLine），不再新分配
3. 被复用的旧 buffer 不再单独释放

torch/_inductor/memory.py（1108 行）则负责另一件事：估算每个调度顺序下的峰值内存，并据此重排节点（下文的 reorder_for_peak_memory）；其中 assign_memory_planning_info_for_scheduler_buffers（torch/_inductor/memory.py:226）只是给每个 SchedulerBuffer 标上大小与后继节点，供峰值估算用。

需要区分两个开关：torch._inductor.config.allow_buffer_reuse（torch/_inductor/config.py:242）默认 True，这是"buffer 死了就让下一个 buffer 顶上"的基础复用；config.memory_planning（:245，由 TORCHINDUCTOR_MEMORY_PLANNING 控制）默认关闭，它是更激进的"把 buffer 池化到一块大 storage 上"。另有 reorder_for_peak_memory（:415，默认 True）在调度前先按峰值内存重排节点，实现是 torch/_inductor/memory.py:913，内部的拓扑序策略叫 topological_sort_lpmf（Least Peak Memory First，:540）。

最终 codegen 出的 wrapper 代码里会看到：

buf3 = buf0; del buf0  # reuse

（codegen/wrapper.py:3313 的 codegen_exact_buffer_reuse；size/stride 不完全一致时走 make_buffer_reuse 里的 reinterpret_tensor 分支。）

这种"显式 reuse"让 Triton kernel 不需要新 alloc，直接写入复用 buffer。不过 wrapper 里剩下的分配（empty_strided_cuda）仍然走第 4 章的 caching allocator；编译期复用只是减少了分配与释放的次数，并不替代 allocator。

14.9 Scheduler 的融合细节

14.9.1 Scheduler 的 fusion 算法细节

§14.3 给了 Scheduler 的高层视角，这里展开 scheduler.py（7158 行）的具体 fusion 算法。核心两个函数：

can_fuse(node1, node2) -> bool（torch/_inductor/scheduler.py:5412）：判断两个节点能否融合。主要关卡：

1. 节点类型：extern / nop 节点（如 cuBLAS 调用）不参与融合；GEMM 等 template 后面接 epilogue 需要 config.epilogue_fusion（默认 True）且被接的节点不是 reduction
2. 依赖与设备：node1 不能依赖 node2，两者必须在同一 device；node2 依赖 node1 时走纵向检查 can_fuse_vertical，互不依赖时走横向检查
3. 启发式（torch/_inductor/choices.py:513）：两节点共享的读写量 shared_data_score 为 0 不融；合并后节点数超过 config.max_fusion_size（默认 64）不融；会抬高峰值内存不融
4. 循环 / reduction 兼容：交给后端判断（Triton 路径是 torch/_inductor/codegen/simd.py:1305 的 can_fuse），两节点的 numel 与 reduction 长度（rnumel）要能对上

任意一关不过就拒绝 fuse。第 4 条是 reduction 类算子（如 sum / mean）与形状不同的逐元素算子经常无法 fuse 的主要原因。

score_fusion_memory(node1, node2)（torch/_inductor/scheduler.py:5741，签名里还带 count_bytes 开关）估算两节点共同读写的字节数，即融合能省下的访存量。排序用的完整分数由 score_fusion_key（:5826）转调 torch/_inductor/choices.py:615 的 score_fusion，是一个按顺序比较的元组：

• template 项：prologue 融合（把节点融进后面的 template）永远排最后
• 类型项：两节点同为 reduction 或同为非 reduction，且有共享数据，优先
• 内存项：score_fusion_memory 算出的共享字节数
• 距离项：两节点在原图顺序里越近越优先

Scheduler 的做法是 一次性排序 + 贪心：fuse_nodes_once（:4592）先用 get_possible_fusions（:4715）在读同一 buffer 的节点之间枚举所有合法候选对，按上面的分数从高到低排序，再依次尝试融合（每次还要检查融合不会在图里造成环）。这种局部贪心不保证全局最优，但工程上够用。fuse_nodes（torch/_inductor/scheduler.py:3697）会反复迭代最多 10 轮，每轮把还能合的再合一次，直到节点数不再减少。

torch/_inductor/scheduler.py:1866 的 FusedSchedulerNode 是融合后的产物。一个 FusedSchedulerNode 可能包几个原始 SchedulerNode 的 inner_fn，运行时它们一起进同一个 Triton kernel。

14.9.2 Foreach Fusion

torch/_inductor/scheduler.py:2216 的 ForeachKernelSchedulerNode 是特殊融合：把多个独立但 shape 相似的算子（如 _foreach_add_(params, grads) 处理几百个张量）一次 launch。第 10 章 §10.6.1 提过 torch._foreach_* 算子，Inductor 在编译路径上对它们做单独优化 —— 多张量并行处理在 Inductor codegen 出来的是一个 combo kernel（参数过多时拆成少数几个），kernel 内按 program id 所在区间分派到各个子张量的计算（torch/_inductor/codegen/triton_combo_kernel.py）。

14.9.3 CSE：消除冗余索引计算

Inductor 的 IR 大量使用索引表达式（如 idx0 = pid * BLOCK + offset、mask = idx0 < n）。同一个表达式可能在 kernel 不同位置重复出现。CSE（Common Subexpression Elimination）让相同表达式只算一次。

具体在 torch/_inductor/codegen/common.py:1927 的 CSE 类。每次生成代码时：

# 朴素 codegen 可能产出
tmp0 = tl.load(x + idx)      # idx = pid*BLOCK + arange
tmp1 = tl.load(y + idx)      # 又算一遍 idx

# CSE 后
idx_cse = pid*BLOCK + arange
tmp0 = tl.load(x + idx_cse)
tmp1 = tl.load(y + idx_cse)

CSE 直接减少 kernel 里的重复算术，对寄存器受限的复杂 kernel 尤其有意义。

与它配套的是 torch/_inductor/index_propagation.py:201 的 IndexPropagation —— 它在 codegen 过程中做有限的常量传播，并把 ops.index_expr 下游的计算保留为 sympy 表达式（例如把由 arange 算出的间接索引还原成静态索引 x * 2），能做比通用 CSE 更激进的化简。这套优化是"Triton 编译器拿不到、必须 Inductor 做"的工程价值 —— Triton 看到的是 IR 里已经化简过的表达式。

14.10 Codegen 的专用路径：GEMM、wrapper 与 CPU 向量化

14.10.1 GEMM Template：matmul 的专用 codegen

§14.13 会讲 ExternKernel（cuBLAS / cuDNN 直接调用），但 Inductor 还有第三条 GEMM 路径：自家 GEMM template。torch/_inductor/kernel/mm.py:305 tuned_mm 是入口（下面是简化示意，源码里模板先收进 templates_to_use，再由 V.choices.get_template_configs 展开成候选）：

def tuned_mm(mat1, mat2, out_dtype=None, *, layout=None):
    # 收集所有候选 GEMM 实现
    choices = []
    if use_aten_gemm_kernels():
        choices.append(aten_mm.bind(...))                # cuBLAS / cuDNN
    if use_triton_template(layout):
        for config in mm_configs(...):
            choices.append(mm_template.maybe_append_choice(config))   # Triton GEMM
    if use_cutlass_template(layout):
        choices.append(cutlass_template.maybe_append_choice(...))     # CUTLASS

    # autotune 在所有候选里选最快的
    return autotune_select_algorithm("mm", choices, ...)

这里有个默认值必须先说清楚：use_triton_template() 与 use_cutlass_template() 内部都带一句 config.max_autotune or config.max_autotune_gemm 的闸门（torch/_inductor/utils.py:1836 与 :2128），而这两个 flag 默认都是 False。所以默认的 torch.compile() 下 tuned_mm 的 choices 里只有 aten 那一项，等于直接调 cuBLAS，根本没有 autotune 可言。只有开了 mode="max-autotune" 或显式 max_autotune_gemm=True，Triton template 与 CUTLASS 候选才会被加进来、才真的去实测比较。这与普通 Pointwise 算子（无条件直接 codegen Triton）路径完全不同。

为什么 GEMM 要这样？因为 matmul 性能差异巨大：

• cuBLAS 在大 matmul（如 4096×4096）上接近峰值，但作为外部 kernel 调用，没法与前后的 epilogue 融合
• Triton GEMM template 可以与 epilogue（如 add bias、relu）fuse，端到端更优；但小 matmul 上不如 cuBLAS
• CUTLASS 在部分 dtype / 形状（如 fp8）上提供额外候选

mm.py 还有几个特殊 template（:216 DecomposeKSugraphTemplate / :254 ContiguousTemplate）：

• DecomposeKSugraphTemplate：把大 K 维 matmul 拆成多个小 K + accumulate（避免 K 维太大让单 kernel reduction 慢）
• ContiguousTemplate：只处理 contiguous input 的快路径

所以"让 LLM 训练 / 推理里的 attention（QKV 投影 + softmax 后 mm）自动选最优 GEMM 实现 + fuse epilogue"这件事，是 max-autotune 才解锁的能力，不是默认行为；max_autotune_gemm_search_space="EXHAUSTIVE"（取值是 "DEFAULT" / "EXHAUSTIVE"，torch/_inductor/config.py:609）再往上穷举更多 config，代价是编译时间明显变长（§15.9.1）。

cpp_gemm_template.py 是 CPU 端 GEMM template，1828 行实现 oneDNN / mkldnn 集成 + 自家 fp32 / bf16 GEMM 微 kernel 的 C++ codegen。CPU 上 LLM 推理的核心性能优化集中在这里。

14.10.2 Wrapper Codegen：粘合 Python 与 Triton

Inductor 生成的不只是 Triton kernel，还有包装这些 kernel 的 Python 代码（torch/_inductor/codegen/wrapper.py）。简化后产物形如：

# Inductor 生成的 wrapper 函数
def call(args):
    arg0_1, arg1_1 = args
    args.clear()
    buf0 = empty_strided_cuda((1024, 1024), (1024, 1), torch.float32)
    triton_poi_fused_mul_relu_0.run(arg0_1, arg1_1, buf0, 1048576, stream=stream0)
    del arg0_1, arg1_1
    return (buf0,)

（v2.11.0 的真实产物里 call 是 class Runner 的方法，kernel launch 形如 .run(..., stream=stream0)，见 torch/_inductor/codegen/wrapper.py:3181；本机 CPU 上 TORCH_LOGS=output_code 实跑可以看到同样结构的 call。）

这段 Python 代码包了 kernel launch 与中间 buffer 分配。它本身由 Python 解释器逐行执行，每次 launch 都有一点 Python 开销；对大 kernel 可以忽略，对小 kernel 密集的图就明显了 —— 这正是 CUDA Graph（mode="reduce-overhead"）与 cpp_wrapper 要解决的。

torch/_inductor/codegen/cpp_wrapper_cpu.py 与 cpp_wrapper_gpu.py 是更激进的版本：把 wrapper 也编译成 C++（避免 Python 解释器开销）。AOTInductor（§15.10.1）生成不依赖 Python 的部署产物，用的就是这种 C++ wrapper。

14.10.3 CPU 后端的 Vectorized 抽象

§14.5 提过 CPU 用 C++ + OpenMP，这里展开 SIMD 细节。torch/_inductor/codegen/cpp.py（CPU codegen 主文件）大量用 PyTorch 自家的 Vectorized<T> 抽象（aten/src/ATen/cpu/vec/）：

// Inductor 生成的 (简化) CPU kernel
#pragma omp parallel for
for (int64_t i = 0; i < N; i += Vectorized<float>::size()) {
    auto a = Vectorized<float>::loadu(in_a + i);
    auto b = Vectorized<float>::loadu(in_b + i);
    auto c = a * b + a;       // 自动向量化
    c.store(out + i);
}

Vectorized<float> 在 x86 是 AVX2/AVX512 寄存器（8/16 元素）、ARM（含 Apple Silicon）是 NEON（4 元素）或 SVE（见 aten/src/ATen/cpu/vec/ 下的 vec256 / vec512 / vec128 / sve）。同一份 codegen 代码在不同平台编译出最优 SIMD，无需 Inductor 写多套。

特别值得提的是 bf16 / fp16 在 x86 上的处理：x86 没有原生 fp16 指令（除了最新 Sapphire Rapids），Vectorized 通过"fp16 → fp32 → 计算 → fp16"的 cast 链实现。Inductor codegen 出来的 cast 会被 C++ 编译器进一步优化成最少的 cast 指令。

CPU codegen 还有专门的 cpp_gemm_template.py（§14.10.1）：CPU GEMM 用 oneDNN 或者自家 micro-kernel（cpp_micro_gemm.py）。LLM 在 CPU 推理时这条路径承担核心性能。

14.11 编译期工程：并行编译、产物缓存与 autotune

14.11.1 AsyncCompile：多进程编译加速

Triton 编译单个 kernel 本身就有可观开销（要走 MLIR → PTX → CUBIN）。一个大模型 forward + backward 可能产生几百个 kernel —— 串行编译会很慢。

async_compile.py（755 行）实现了多进程并行编译：

• 主进程生成 Triton 源码 + 计算 hash
• 派给一个 worker 进程池（worker 数由 config.compile_threads 决定，默认取 min(32, 可用 CPU 核数)，torch/_inductor/config.py:1083 的 decide_compile_threads）
• 每个 worker 调 Triton compiler 编译
• 主进程并发等待所有 future
• 所有 kernel 编译完才返回最终的 compiled callable

几十上百个 kernel 从串行编译变成 worker 池并发编译，墙钟时间基本被最慢的那几个 kernel 决定。这是缩短 torch.compile 首次编译墙钟时间的主要手段之一。

worker 进程间通过 hash 共享缓存：同一个 Triton 源码（不同 input shape 但生成的代码一样）被认成同一份编译产物，避免重复编译。

14.11.2 FxGraphCache：跨进程编译产物缓存

torch/_inductor/codecache.py:1206 的 FxGraphCache 让 Inductor 编译产物持久化到磁盘：

# 第一次跑: 编译 + 写盘
torch.compile(model)(x)    # Inductor 编译, 写到 $TMPDIR/torchinductor_<user>/ (见 runtime/cache_dir_utils.py: default_cache_dir)

# 重启进程后跑同样代码
torch.compile(model)(x)    # cache hit, 跳过编译, 直接 load .so

cache key 怎么算？FxGraphCachePickler (:485) 把 FX Graph + example_inputs shape/dtype + Inductor config + PyTorch version 等一起 hash（参与 hash 的字段见 FxGraphHashDetails，:784）。任何一项变化就 cache miss，避免拿到错误的旧产物。

对生产推理服务，FxGraphCache 把"重启后首次请求要重跑整条编译链"变成"Dynamo 仍要重新 trace，但 Inductor 编译产物直接从盘上 load"；上游还有一层 AOTAutogradCache（torch/_functorch/_aot_autograd/autograd_cache.py:699，torch._functorch.config.enable_autograd_cache 默认开），命中时连 AOTAutograd 这一步也一起省掉。设 TORCHINDUCTOR_CACHE_DIR=/path/to/persistent 让 cache 跨容器共享 —— 同一个模型多个 pod 可以共用编译产物。

cache 失效的常见原因：

• 改了模型架构（FX Graph 变化）
• 升级了 PyTorch 版本
• 改了 _inductor.config 任何 flag
• 输入 shape 偏离了 trace 时的假设

第一次部署时预编译 + warmup 能避免线上首请求慢，是大模型推理服务的标准操作。

14.11.3 ChoiceCaller / autotune select_algorithm

ChoiceCaller（定义在 torch/_inductor/ir.py:5277）是 GEMM template autotune（§14.10.1）的核心，入口函数 autotune_select_algorithm 在 torch/_inductor/select_algorithm.py:4488。每个 ChoiceCaller 是"一种算法实现 + 它的 benchmark 接口"（简化示意）：

class ChoiceCaller:
    def benchmark(self, *args, out) -> float:
        # 取 to_callable() 得到的实现, 交给 benchmarker 计时
        ...

    def to_callable(self):
        # 返回可直接调用的实现 (由子类实现)
        ...

autotune_select_algorithm("mm", choices, ...) 内部：

1. 对每个 choice 跑几次 benchmark（用真实 input shape，不是假数据）
2. 选最快那个
3. 把"shape -> 最优 choice"缓存到 disk
4. 后续相同 shape 的 mm 直接命中缓存

ExternKernelCaller（torch/_inductor/select_algorithm.py:2438）是接 cuBLAS / cuDNN 等外部 kernel 的 ChoiceCaller 实现，TritonTemplateCaller（同文件 :2342）对应 Triton 模板。用户也能把自己的 matmul 实现挂进候选 —— tuned_mm 末尾会把 config.external_matmul 列表里的实现逐个包成外部候选加进 choices。

这套机制让 Inductor 在 GEMM 等关键算子上性能自适应 —— 用户不调参，编译器自己选最优。

14.11.4 Subgraph 编译与 HigherOrderOperator

某些用户代码包含"子图"结构 —— 如 torch.utils.checkpoint（第 7 章 §7.5.4）、while_loop / cond 控制流、flex_attention 自定义 attention。它们在 Dynamo 里都被包成 HOP，但到 Inductor 的处理方式不一样：

def forward(pred, x):
    # cond 的两个分支各是一段 subgraph
    return torch.cond(pred, true_fn, false_fn, (x,))

# Inductor 里对应 ir.Conditional: 两个分支子图各自 lower,
# 运行时按 pred 选一个执行

• 控制流：torch.cond / torch.while_loop 在 Inductor 里对应 ir.Conditional / ir.WhileLoop（torch/_inductor/ir.py:8977 / :9186），子图单独 lower
• checkpoint：Dynamo 把它包成 tag_activation_checkpoint（torch/_higher_order_ops/wrap.py:226），它只给子图节点打上"可重算"标记；AOTAutograd 展开后由 partitioner 负责复制重算节点，到 Inductor 时已是普通节点，不是独立编译单元
• flex_attention：lower 成模板 kernel（GPU 上默认是 Triton 模板），用户写的 score_mod / mask 子图经 build_subgraph_buffer 内联进模板；lowering 入口见 torch/_inductor/kernel/flex/flex_attention.py:107 的 flex_attention

HigherOrderOperator (HOP) 是编译栈里的概念：算子的"参数"可以是另一段计算图（v2.1 起叫这个名字，v1.13 / v2.0 时类名是 PyOperator）。torch.cond、torch.while_loop、torch._higher_order_ops.flex_attention 等都是 HOP。

这套机制让 FlexAttention 等高级用法在 Inductor 路径上能完整编译，而不是 fall back 到 eager。

14.12 配置项与一段完整的编译产物

14.12.1 inductor.config 的 200+ flags 概览

torch/_inductor/config.py 有 2526 行，注册了几百个 config flag。按功能分组：

分组	关键 flag
autotune	max_autotune / max_autotune_gemm / coordinate_descent_tuning
fusion	epilogue_fusion / aggressive_fusion / benchmark_epilogue_fusion / loop_ordering_after_fusion
memory	memory_planning / memory_pool / realize_acc_reads_threshold
codegen	cpp_wrapper / triton.cudagraphs / assert_indirect_indexing
compile_cache	force_disable_caches / fx_graph_cache
debug	debug / verbose_progress / trace.output_code
optimization	freezing / pattern_matcher / reorder_for_locality / reorder_for_peak_memory
CPU specific	cpp.threads / cpp.simdlen / cpp.no_redundant_loops

每个 flag 都对应一个具体的 codegen / 优化 pass。生产级调优要熟悉其中几十个才能在每个场景找最优组合。但大多数用户用默认值即可。

config.fallback_random（torch/_inductor/config.py:760）是一个有意思的 flag：默认 False，也就是默认让 random 算子走 Inductor 自己的 functional random（§13.25 的 inductor_prims.random），从而能参与 fusion。想让 compiled 的随机数与 eager 逐位对齐（调试数值差异时常用）就把它设成 True，代价是这些算子退回 ATen、fusion 被打断。

14.12.2 一段 LayerNorm 的完整 Inductor 编译产物

把整章串起来，看 nn.LayerNorm 在 Inductor 下编译成什么。源码：

@torch.compile
def f(x, weight, bias):
    return F.layer_norm(x, [768], weight, bias, 1e-5)

经过完整链路：

1. Dynamo trace：捕获 fx graph，含一个 torch.nn.functional.layer_norm 调用节点（torch 层，还不是 ATen）
2. AOTAutograd functionalize：layer_norm 没 inplace，跳过
3. AOTAutograd decomposition（§13.10）：拆成 var_mean → add(eps) → rsqrt → sub → mul → mul(weight) → add(bias) 共 7 个计算节点（另有 2 个 getitem 取出 var / mean）
4. Inductor Lowering：每个基础节点变成 IR 节点（Pointwise / Reduction）；var_mean 变成一个 Welford reduction，一趟同时求出 mean 与方差
5. Scheduler fusion：后面的 sub / mul / add 串成 pointwise，作为 reduction 的 epilogue 一起 fuse
6. Triton codegen：生成单个 Triton kernel，含 Welford reduction 部分 + epilogue（normalize 部分）
7. Triton 编译：Triton DSL → MLIR → PTX → CUBIN
8. 运行时：Wrapper 调用此 kernel + autotune 选最优 BLOCK 配置

第 3 步的图与"只有一个 kernel"这两点在本机 CPU 上实跑确认过（TORCH_LOGS=aot_graphs,output_code，输入 [32, 768]）：生成的是单个 cpp_fused_native_layer_norm_0，先用 Welford<Vectorized<float>> 求 mean / m2，再在同一个 kernel 里做归一化与仿射；GPU 上的 Triton 版按 codegen 源码判断结构相同，未实跑。

最终一行 F.layer_norm(...) 调用变成一次 kernel launch。与 ATen 自带的 layer_norm CUDA kernel 相比，它的额外好处是能继续与前后的逐元素 op（残差 add、dropout 等）fuse，独立的 ATen kernel 做不到这一点。

理解这条全旅程，你就理解了 torch.compile 在 Transformer 训练上的真实加速来源 —— 不是某个魔法算法，是端到端编译链的累积优化。

14.13 Inductor 的"非 Pointwise"路径

复杂算子（matmul、conv、SDPA）不容易自动 fuse 出最优代码，Inductor 不自己生成，而是调用 现成的 cuBLAS / cuDNN / FlashAttention kernel。这些走 ExternKernel（IR 类型之一）路径，Inductor 只决定 layout 与 stride 转换，不生成 GEMM 代码。

在 attention 这类场景，SDPA 本身作为外部 kernel 不参与融合，但它前后的 layernorm / dropout / 残差 add 等会被 Inductor 彼此融合，整体优于 eager。开了 max-autotune 后，GEMM 还能在 cuBLAS 与 Triton 模板之间挑选并带上 epilogue 融合（§14.10.1）。

14.14 编译错误与 fallback 机制

Inductor 不是万能的。某些算子组合编译会失败 —— 复杂 indirect indexing、动态 shape 边界、外部库依赖等。torch._inductor 有完善的 fallback 体系：

flowchart TB
    Try[尝试 Inductor 编译]
    Try --> OK{成功?}
    OK -->|是| Triton[生成 Triton kernel]
    OK -->|否| Cause{失败原因?}
    Cause -->|某 op 无 lowering| Per[per-op fallback<br/>这个 op 用 ATen, 其他用 Triton]
    Cause -->|编译整体出错| Graph[默认抛 InductorError<br/>suppress_errors=True 才退回 eager]
    Cause -->|某个 BLOCK 配置超出硬件资源| Smaller[丢弃该配置<br/>用其余候选]
    style Per fill:#fef3c7
    style Graph fill:#fee2e2

具体机制：

• lowerings 里没有这个算子 → 默认 config.implicit_fallbacks=True，Inductor 只在 info 日志里记一条 "Creating implicit fallback"，自动 fallback：把这个 op 留给 ATen 跑，前后的 Triton kernel 把数据传过来（torch/_inductor/graph.py:1304 起）；关掉 implicit_fallbacks 才会抛 MissingOperatorWithoutDecomp / MissingOperatorWithDecomp
• 编译过程本身出错（lowering 异常、C++ 编译失败等）→ 默认直接抛 InductorError（本机缺 C++ 编译器时实跑就是这个结果）；只有设了 torch._dynamo.config.suppress_errors=True（或环境变量 TORCHDYNAMO_SUPPRESS_ERRORS）才吞掉错误、退回 eager
• 某个 Triton 配置超出硬件资源（OutOfResources / ptxas 报错）→ 丢弃这个配置，用其余候选；全部失败才报错（torch/_inductor/runtime/triton_heuristics.py:505 的 _precompile_worker）

也就是说，per-op fallback 是默认行为，整图退回 eager 却不是：想在生产上"编译失败也不崩"，要显式打开 suppress_errors。

14.15 自定义算子的 Inductor lowering 注册

§22.2 讲过 torch.library.custom_op 写自定义算子。如果想让自定义算子也走 Inductor 编译（不只是当不透明 op），需要注册 lowering：

import torch
from torch._inductor.lowering import register_lowering, lowerings

@register_lowering(torch.ops.my_lib.my_mul.default)
def my_mul_lowering(x, y):
    # lowering 收到的是 TensorBox, 要返回 IR 节点;
    # 这里复用 aten.mul 的 lowering, 产出一个 Pointwise
    return lowerings[torch.ops.aten.mul.Tensor](x, y)

注意不能直接 return ops.mul(x, y)：ops.* 只能在 inner_fn 里对标量值使用，放在 lowering 顶层会得到 LoweringException（"not a supported top level IR node"）。上面的写法在本机实跑，my_mul(a, b) + c 生成的是一个融合 kernel cpp_fused_add_my_mul_0；想手写 Pointwise.create(inner_fn=...) 的版本见 §22.18。

这种"用 Inductor IR 重新表达"让自定义算子能与周围算子 fuse，而不是当不透明 op 阻断 fusion。代价是要写两套实现（一份 eager kernel + 一份 Inductor lowering），但对性能敏感的核心算子值得。

第三方硬件接 Inductor 的官方入口是 register_backend_for_device（torch/_inductor/codegen/common.py:408）：为自家 device 注册 scheduling 与 wrapper codegen，复用 Inductor 的 lowering 与 scheduler。这是除了第 5 章 dispatcher 之外的另一条接入路径，让国产芯片也能用上 Inductor 的 fusion / autotune 优化。

14.16 Triton kernel 的完整生命周期

讲了 Inductor 这么多，最后看一下产物（Triton kernel）从生成到运行的完整旅程：

1. Inductor 生成 Triton DSL (Python 代码)
   ↓
2. Triton compiler: Triton DSL → Triton IR (MLIR-based)
   ↓
3. Triton lowering: Triton IR → MLIR LLVM Dialect
   ↓
4. LLVM: → PTX (NVIDIA) / AMDGPU IR (AMD)
   ↓
5. NVCC ptxas / AMD lld: → CUBIN (NVIDIA) / HSA code object (AMD)
   ↓
6. CUDA driver / HIP runtime: load → 运行时 launch

整套链路在 torch/_inductor/async_compile.py（§14.11.1，class AsyncCompile 在 line 232）的 worker 进程里跑。每一步都是开源工具：

• Triton compiler 是 OpenAI 开源（独立项目 github.com/triton-lang/triton）
• MLIR 是 LLVM 子项目
• PTX/CUBIN 是 NVIDIA 闭源但有公开规范

CUBIN 二进制最终被 CUDA driver 加载到 GPU：用的是 driver API 的 cuModuleLoadData 取得 module handle、cuLaunchKernel 发起执行（cpp_wrapper 路径的生成代码见 torch/_inductor/codegen/cuda/device_op_overrides.py；Python wrapper 路径由 Triton 的 launcher 或 torch/csrc/inductor/static_launcher/cuda.cpp 发起）。Triton kernel 的 launch 不经过 PyTorch 的 dispatcher（第 5 章）—— 生成 kernel 这部分运行时没有 dispatcher 开销；wrapper 里剩下的 extern kernel 调用（如 aten.mm）与张量分配仍走 ATen。

理解这条链路你就明白：torch.compile 的"加速"本质之一是把"用 Python 解释 + dispatcher 调度"的运行时开销，提前到编译期消化掉。

14.17 Inductor 与 Pattern Matcher 在 attention 路径上的协奏

LLM 训练 / 推理最热的是 attention。Inductor 在 attention 上做的事最丰富，把多个机制综合起来：

1. Pattern matcher 识别：用户写的"naive attention"代码（QK / sqrt / softmax / V）被识别成 SDPA 模式
2. 替换为 aten.scaled_dot_product_attention：再由 SDPA 按条件选 FlashAttention 等后端
3. 周边融合：SDPA 作为外部 kernel 不接受 epilogue 融合，但它前后的 dropout、残差 add、norm 等逐元素 / 归约算子彼此融合
4. GEMM template 选择（仅 max-autotune 下）：QKV 投影 qkv = x @ Wqkv 在 cuBLAS / Triton / CUTLASS 候选里 autotune，Triton 模板还能带 epilogue 融合
5. Buffer 复用：QKV 中间张量、attention scores 等的内存复用（§14.8.2）

整套综合优化让 Inductor 在 transformer 训练上拿到可观加速（相对 eager），而且几乎不需要用户手写优化 —— 写朴素的 attention 代码，条件满足时就能用上 FlashAttention 与融合后的周边算子。

这就是 Inductor 真正的工程价值：让用户专注算法、编译器负责性能。这与传统 CUDA C++ "性能必须靠手写优化" 的模式形成鲜明对比。

14.18 几条工程经验

实战 Inductor：

1. TORCH_LOGS=inductor,output_code 看生成的 Triton 源码：能看到 fusion 决策、kernel 数量、autotune 结果

2. fusion 失败常见原因：循环维度不一致（如 [N] 和 [N, M]）、有 reduction 隔断、有 stride 不匹配的 view

3. 第一次跑会编译几秒到几分钟：Triton 编译 + autotune。设 TORCHINDUCTOR_CACHE_DIR 落盘，下次进程直接用

4. shape 频繁变化会触发重编译：默认第一次按具体 shape 特化编译；遇到不同尺寸时 automatic_dynamic_shapes（默认开）把该维改成符号维再编一次，之后一份 kernel 处理多种尺寸。设 dynamic=True 则一开始就按动态编

5. 自定义 op 想参与融合要注册 register_lowering：否则它在 Inductor 里是一个 fallback（外部 kernel）节点，照常执行、不会 graph break，但会打断前后的融合

6. CPU 上加速比通常比 GPU 小：Inductor 在 CPU 上生成 C++ + OpenMP，而 ATen 本来就在调 oneDNN / MKL 这类高度优化的库，可优化的胶水更少

7. mode='reduce-overhead' 启用 CUDA Graph：把多个 kernel launch 合成一次 graph replay，对 launch 开销占比高的小 batch 推理收益明显

14.19 历史决策：为什么不直接生成 CUDA C++

torch/_inductor/ 在 v1.13.0 里还不存在，首个含它的 release 是 v2.0.0。它的 GPU 输出语言选了 Triton 而不是 CUDA C++：PyTorch 2 论文（ASPLOS'24，§4.8）给的主要理由是 Triton 已被证明能用简单的输入代码生成快过手写库的 kernel。从工程角度对比两条路：

CUDA C++ 路径的问题：

• 每个 kernel 要手写 launch config（block / grid / shared memory size）
• 跨 GPU 架构（A100 / H100 / L4）调优要分别写
• 编译路径长（NVCC 慢、cubin 不可移植）
• 与 PyTorch C++ 端集成要 build .so，部署复杂

Triton 路径的优势：

• block-level 抽象让 codegen 不用管 thread mapping
• 跨架构统一编译（Triton 自动选最优）
• Python 端直接 jit，无 build 步骤
• autotune 框架成熟

代价是 Triton 对某些"激进 GPU 编程模式"（如 warp specialization、H100 的 TMA）的支持来得比手写 CUDA 晚；v2.11.0 的 torch/_inductor/kernel/mm.py 已有 persistent_tma_mm_template 与 blackwell_ws_persistent_device_tma_mm_template 这类 TMA / warp-specialized 模板，但它们只在 max-autotune 候选里出现。Inductor 在专用算子上仍走 ExternKernel（如 SDPA 调 FlashAttention 等专门的 CUDA kernel，不走 Triton codegen）。这是"通用 codegen + 关键 op 走专用实现"的混合策略。

理解这条决策，你就明白为什么 Inductor 不试图与 cuBLAS 在 GEMM 上完全替代 —— 通用 codegen 路径在 GEMM 这种"几十年优化的算子"上很难超越专用库。Inductor 的价值在 fusion + 通用算子的 codegen，不在 reinvent GEMM。

14.20 Halide 思想的影响

Inductor IR 的"循环 body 用纯函数描述、融合与调度另行决定"，和 Halide（image processing DSL，2012 年 MIT）开创的"compute / schedule 分离"思路相通。不过 PyTorch 2 论文（§4.8）点名受 Halide 启发的是 TVM、nvFuser、NNC，对 Inductor 自己的 IR 则称为"define-by-run 的循环级 IR"（论文 §4.3），下面的对应是本专栏的类比。Halide 的核心范式：

• Compute：定义"算什么"（数学公式）
• Schedule：定义"怎么算"（loop ordering / vectorize / parallelize）

Inductor 对应到：

• Loops.inner_fn：compute 部分 —— 描述循环 body 怎么算
• Scheduler 的 fusion 决策：schedule 部分 —— 决定哪些 loop 合并、什么顺序

这种分离让"数学语义"与"性能调优"解耦。Halide 用户可以写一份 compute、试多种 schedule 找最优。Inductor 把"试 schedule"自动化（autotune），用户连 schedule 都不用写。

这条思想也影响了 TVM、nvFuser、NNC 等编译器（论文 §4.8 的列举）。Inductor 不是凭空发明，而是站在 Halide / TVM / Triton 等前人工作的肩膀上，把这些思想集成到 PyTorch 生态。

14.21 dynamic shape 在 Inductor 的处理

§12.5 提过 Dynamo 的 guards 处理 dynamic shape。Inductor 这一层也有相应支持：

# 用户指定 batch dim 是动态的
torch._dynamo.mark_dynamic(input, 0)
output = compiled_model(input)

Inductor 收到的 graph 里 batch dim 是 SymInt（符号整数），不是具体值。codegen 出来的 Triton kernel 把与它相关的尺寸作为普通整型参数，运行时传入（简化示意）：

@triton.jit
def triton_poi_fused_0(in_ptr0, out_ptr0, ks0, xnumel, XBLOCK: tl.constexpr):
    # ks0 / xnumel 由 batch dim 算出, 是普通运行时参数
    ...

这些尺寸参数不是 tl.constexpr：torch/_inductor/codegen/triton.py:5423 起，xnumel 之类以 SizeArg 普通参数传入，旁边的注释写明 constexpr 版本会出问题，已被注释掉；只有 XBLOCK 这类块大小才是 tl.constexpr。所以同一份 kernel 服务所有 B 值，B 变化时既不用重新 trace，也不用重编 kernel。

这套机制让 LLM 推理（不同 batch size、不同 sequence length）的 dynamic shape 路径在 Inductor 下能正确编译，而不是退回 eager。代价是 kernel 里要带符号化的索引运算，少了针对具体尺寸的特化；另外 Dynamo 层若因 guard 失败反复重编译，受 torch._dynamo.config.recompile_limit（旧名 cache_size_limit，默认 8，见第 12 章）限制。

14.22 Inductor 之外的两条编译路径

虽然 Inductor 是 torch.compile 默认 backend，PyTorch 主仓还有两条编译实验路径：

• aot_eager：只跑 AOTAutograd 不上 Inductor，主要用于调试。Inductor 编译出问题时切到这个 backend 看是不是 AOT 阶段错的
• TorchScript (torch.jit.script / torch.jit.trace)：v1.x 时代的编译路径，自 v2.5 起在 docstring 中标为 deprecated（"TorchScript is deprecated, please use torch.compile instead"，见 torch/jit/__init__.py、torch/jit/_script.py）。某些 mobile 部署场景（不能跑 Triton）仍用 TorchScript

新代码不要用 TorchScript —— 它已被标为 deprecated，也不走 Inductor 路径。torch.export + AOTI（§15.10.1）是 v2.x 的统一部署方向。

14.23 跨专栏关联

• 《Rust 编译器与运行时揭秘》专栏第 16 章 LLVM codegen：Inductor 的 IR-Lowering-Codegen 三段式可以和 LLVM 的 IR-MachineIR-Asm 分层对照着看
• 《vLLM 推理内核深度解析》专栏第 8 章 model runner：vLLM 也用 torch.compile + Inductor 编译模型 forward。理解本章能帮你解释 vLLM 启动慢的原因（首次 Inductor 编译）
• Triton 论文（Tillet, Kung, Cox，Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations，MAPL 2019）：理解 Inductor 必读 Triton 原理。Triton 把"GPU 编程"抽象到 block 级别让 Inductor 这种自动 codegen 成为可能

14.24 Inductor 加速比的结构：什么场景收益大

本机无 GPU，本节不给具体倍数（要看数就跑 PyTorch 主仓的 benchmarks/dynamo/，它对 TorchBench / HuggingFace / TIMM 三套模型集有公开的 CI 结果面板）。但加速比的结构是可以讲清楚的：

场景特征	Inductor 收益	原因
小算子密集、kernel launch 多（如 batch=1 的自回归 decode）	最大	eager 里 Python + dispatcher + launch 开销占比高，编译后几乎全部消化掉
memory-bound 的逐元素 / norm / activation 链	大	fusion 直接砍掉中间张量的 HBM 往返
大 GEMM / conv 主导的训练（compute-bound）	小	时间本来就花在已经接近峰值的 cuBLAS / cuDNN 上，Inductor 只能优化胶水
图里 graph break 多	小甚至负	每段图各编各的，跨段没法 fuse，还多出编译成本

Inductor 不能 100% 替代手写优化：FlashAttention v3、cuDNN convolution、cuBLAS GEMM 等长期专门优化的 kernel 仍然手写更优。Inductor 的价值在 算子之间的胶水（fusion）和 没人专门优化的小算子，这两块它能在不手写代码的前提下补上。

14.25 Inductor 与 profiler 的协作

第 21 章会讲 PyTorch profiler。Inductor 编译产物在 profile chrome trace 里有特殊命名规则：

triton_poi_fused_mul_relu_0    ← Inductor 生成的 fused kernel
triton_<类别>_fused_<算子名...>_<序号>

类别是 poi（pointwise）/ red（reduction）/ per（persistent reduction）等，取 kernel 所用 triton_heuristics 装饰器名的前三个字母；中间的算子名默认取被融合节点对应的原始 ATen 算子名（config.triton.descriptive_names="original_aten"），去重后按字母序排列，不是执行顺序（torch/_inductor/codegen/triton.py:6152 起，get_fused_kernel_name 在 torch/_inductor/utils.py:826）。本机 CPU 实跑时同样规则的名字是 cpp_fused_mul_relu_0、cpp_fused_native_layer_norm_0。

这个命名让你看 trace 时能立刻识别"这个 kernel fuse 了哪些 op"。如果你看到 triton_per_fused_add_native_layer_norm_3 这种名字，说明残差 add 与 layer_norm 被 fuse 成了一个 kernel。如果看到很多独立的 triton_poi_fused_relu_0 / triton_poi_fused_mul_1，说明 fusion 没起作用 —— 该排查为什么 scheduler 拒绝了这些 fuse（往往是循环维度不对齐或 reduction 隔断）。

TORCH_LOGS=output_code 让 Inductor 把生成的 Triton 源码 + wrapper Python 都打印到日志，结合 profiler 的 kernel 时间，能精确定位"哪段编译产物慢"。这套调试链路是 Inductor 性能调优的标配，比单纯看 chrome trace 信息丰富得多。

14.26 Inductor 与 torch.export 协作

torch.export(model) 把模型 trace 成 ExportedProgram（静态图，更严格）。AOTI（§15.10.1）就是基于 ExportedProgram 调 Inductor 编译。但有些算子在 torch.export 路径下需要特殊处理：

• HigherOrderOperator (§14.11.4)：control flow / checkpoint / flex_attention 等
• Custom op 必须有 register_fake（§22.7）：否则 export 时 shape 推导失败
• Effect token op：print、torchbind 方法调用等带副作用的算子（torch/_higher_order_ops/effects.py 里登记为 ORDERED）在 ExportedProgram 里要保留 side effect ordering

torch._inductor 暴露的 aoti_compile_and_package(exported, ...) 内部就是把 ExportedProgram 喂给 Inductor 后端 + 生成 .so + 打包。整套链路与 JIT torch.compile 共享 lowering / scheduler / kernel codegen，差异主要在最后的 wrapper codegen（C++ vs Python）。

理解这条链路你就明白：无论是训练时的 JIT 加速、还是 AOTI 的部署二进制，底层都是 Inductor 在支撑；它是 PyTorch 通向"非 Python runtime 部署"的关键基础设施之一（另一条终点是 §15.10.4 的 ExecuTorch）。

14.27 Inductor 在 LLM 推理上的关键优化路径

LLM 推理引擎里，vLLM V1 默认启用 torch.compile（《vLLM 推理内核深度解析》专栏第 8 章 §8.6）。按本章的机制，Inductor 在推理路径上能确定的收益集中在 attention 之外：

• Norm / 残差 / 激活链的融合：RMSNorm、残差 add、SiLU × mul、RoPE 这类逐元素与归约算子被融合成少数几个 kernel
• attention 本身不进 Inductor 融合：在 vLLM 里 attention 是图的切分点，走自家 kernel（该专栏 §8.4）；SDPA 在 Inductor 里也只是外部 kernel
• 线性层默认走 cuBLAS：Triton GEMM 模板与 epilogue 融合要开 max-autotune（§14.10.1）

这些都是"通用 fusion + pattern matcher 自动识别"的产物，LLM 框架开发者不必为 Norm / RoPE 这类胶水手写 fused kernel；而 attention、KV cache 读写这类核心 kernel 仍由框架自己提供。理解 Inductor 让你能在调优 LLM 推理时知道"哪些性能问题该让 Inductor 自动解决、哪些需要手写 CUDA"。

14.28 数学函数精度对齐

sin / cos / exp / log / sqrt / pow / erf 等数学函数有个常被忽略的工程问题：PyTorch 内置 ATen 实现与 Triton 库的数值精度可能不同。

举例：Triton 的 tl.exp 可能走快速近似实现，精度比 libdevice 的 exp 低一些 —— 所以 Inductor 默认生成 libdevice.exp，只有打开 fast math 才改用 tl_math.exp（见下）。打开这类开关后，如果用户代码对精度敏感（如 numerical analysis 类训练），eager 与 compiled 的 loss 曲线可能分叉。

Inductor 的解法：

• Triton codegen 侧的开关是 torch._inductor.config.use_fast_math（torch/_inductor/config.py:248，由环境变量 TORCHINDUCTOR_USE_FAST_MATH 打开，默认关），打开后 exp 生成 ftz（flush-to-zero）变体 tl_math.exp，默认则生成 libdevice.exp（见 torch/_inductor/codegen/triton.py:1240 起的 exp）
• CUTLASS / ROCm 编译侧另有各自的 use_fast_math：config.cutlass.use_fast_math 默认关，config.rocm.use_fast_math（torch/_inductor/config.py:2223）默认开、给编译器加 -ffast-math；CPU C++ 路径没有单独开关，torch/_inductor/cpp_builder.py:833 的 _get_ffast_math_flags 总会加上 -fno-trapping-math、-funsafe-math-optimizations、-ffinite-math-only 等一组等价于 fast-math 的选项
• 关键数学函数（如 exp 在 softmax 里）有专门的 codegen 路径与 ATen 对齐

这种"精度 vs 性能"取舍在 ML 训练里通常无影响（梯度有噪声，几个 ULP 差异不可见），但在 numerical 任务（如物理模拟、金融计算）必须知道这个开关存在。

14.29 几个用户视角的"为什么"

整章信息密集，最后用 Q&A 形式回答几个常被问的问题：

• 为什么 torch.compile 第一次特别慢？ —— 因为整条 Dynamo + AOTAutograd + Inductor + Triton 编译链路要跑一遍，还要编译并 autotune 大量 Triton kernel。同一进程里第二次调用直接用 Dynamo 缓存的编译产物；换进程重启后，命中 AOTAutogradCache / FxGraphCache 可省掉大部分编译
• 为什么我 fuse 不了？ —— 常见原因是循环维度不对齐（一个 [N] 一个 [N, M]）或 reduction 隔断。开 TORCH_LOGS=fusion 看 scheduler 拒绝原因
• 为什么 Inductor 编出来的代码比我手写的 CUDA 快？ —— 不是单 kernel 比手写快，是 fusion 让多个 kernel 合并、整体内存读写减少
• 可以用 Inductor 编译我的 ResNet 部署吗？ —— 用 AOTI（§15.10.1）→ 输出 .pt2 → C++ 加载，整套不依赖 Python
• Inductor 与 TensorRT 比怎么样？ —— 两者侧重不同：TensorRT 专做推理，Inductor 同时覆盖训练与推理、对动态 shape 与任意 PyTorch 代码的支持更直接。本专栏不给跨环境的性能对比，生产推理服务取决于 latency / 部署灵活性的权衡，需在目标环境实测

这些 Q&A 把整章串起来。每个问题的答案都对应前面某节的具体技术点，能反过来回查源码。

14.30 Scheduler 的 reorder passes

Inductor 还有一组 reorder pass，分布在几个文件里：

• reorder_for_locality（fx_passes/post_grad.py:751，config reorder_for_locality 默认 True）：这是 post-grad FX 层的 pass，把"读同一 buffer 的多个 op"挪到一起，让 cache 命中率高
• reorder_for_peak_memory（torch/_inductor/memory.py:913，config 同名，默认 True）：在调度前调整 op 顺序让 peak memory 最小
• reorder_compute_for_overlap（torch/_inductor/comms.py:108，由 reorder_for_compute_comm_overlap 打开，默认 False）：分布式场景下让 collective 与 compute 充分 overlap；具体跑哪几个 pass 由 reorder_for_compute_comm_overlap_passes 列表指定

这些 reorder 分别发生在 FX 层与 scheduler 层：reorder_for_locality 在 lowering 之前的 post-grad 图上跑，另外两个在 SchedulerNode 层面跑 —— fusion 决定了哪些 op 已经合体（合体后内部顺序固定），剩下的"独立 SchedulerNode 之间"还能调顺序。

reorder_for_peak_memory 与 §14.8.2 的 buffer 复用互补：buffer 复用是"已知 schedule 后在 wrapper 里复用内存"，reorder_for_peak_memory 是"先调 schedule 让内存峰值更小"。两者一前一后压低最终的内存占用。

14.31 Inductor 调试武器库

Inductor 编译失败 / 性能不达预期 / 数值不对时的诊断工具：

• TORCH_LOGS=inductor：开 Inductor 全套日志，看每段编译做了什么
• TORCH_LOGS=output_code：dump 生成的 Triton + wrapper Python 源码
• TORCH_LOGS=fusion：看 Scheduler 的 fusion 决策（哪些 op 合并、哪些拒绝、原因）
• TORCH_LOGS=schedule：看 scheduler 的调度信息（每个 kernel 包含哪些节点），比 fusion 粗
• TORCH_LOGS=aot_graphs：看 AOTAutograd 输出给 Inductor 的 graph
• TORCH_COMPILE_DEBUG=1：把每次编译的中间产物（fx graph / lowered IR / 生成代码）保存到 torch_compile_debug/ 目录（位置由 torch._dynamo.config.debug_dir_root 决定，默认是当前工作目录下的 torch_compile_debug/，可用环境变量 TORCH_COMPILE_DEBUG_DIR 改），方便事后分析
• torch._dynamo.config.verbose=True：Dynamo 的详细输出

这套日志组合让你能从"编译失败的 stack trace"一路追到"哪个算子的 lowering 拒绝了什么 fusion"。生产级 Inductor 调优必须熟悉这些工具。

特别有用的是 TORCH_COMPILE_DEBUG=1 产生的目录结构：

torch_compile_debug/
  run_<timestamp>-pid_<pid>/
    torchdynamo/debug.log    ← Dynamo 日志
    torchinductor/model__0_forward_1.0/   (推理图叫 model__0_inference_0.0)
      fx_graph_runnable.py     ← 可以独立运行的 fx graph repro
      fx_graph_readable.py     ← 人类友好的 fx graph
      fx_graph_transformed.py  ← 经过 fx_passes 后的 graph
      ir_pre_fusion.txt        ← 融合前的 Inductor IR
      ir_post_fusion.txt       ← 融合后的 Inductor IR
      output_code.py           ← 最终生成的 wrapper + kernel 源码

（本机 CPU 实跑的目录结构；每个图目录里的 debug.log 要另开 torch._inductor.config.trace.debug_log=True 才会生成。）

把 fx_graph_runnable.py 单独跑能复现编译失败 —— 这是给 PyTorch 提 Inductor bug issue 时常用的 repro 格式。

14.32 设计启示

Inductor 的几条核心思想：

第一：用更高级的 DSL 做 codegen 后端：选 Triton 而非 CUDA C++ 让 codegen 复杂度大幅降低。这套思想可以借鉴到任何"自动生成 GPU 代码"的场景

第二：用 inner_fn 描述循环 body：把"逻辑循环"和"实际生成"解耦，让 fusion 决策能在纯函数式抽象上做，避免循环结构的干扰

第三：贪心 fusion + autotune 双保险：编译期用启发式决定融合范围（贪心可能不是全局最优），运行期 autotune 选 block size。两层一起把性能榨到接近最优

第四：ExternKernel 接现成库：不是所有算子都要自己生成。matmul / conv / attention 这些有几十年优化的算子用 cuBLAS / cuDNN / FlashAttention 调用即可，Inductor 只优化"算子之间的胶水"

下一章把 Dynamo / AOTAutograd / Inductor 串起来，加上 CUDA Graph，看 torch.compile 完整端到端流程。