作者：GlowLED（zpeiyu11@gamil.com）

本文将对torch compile技术进行总览性介绍，虽然不全也不精，但可以快速建立起一个对于torch compile的整体认知，在进行相关工程时不会变成完全黑盒。

故命名为二十一世纪火柴编译目录。

Background

PyTorch默认为eager模式，即“立即计算”，在计算过程中动态构建计算图。这使得研究人员可以方便地调试，可以在任何地方插入print，可以直接使用Python的if/for来控制计算流。 但是这种灵活性是有性能代价的：

• 增大了启动开销：每个PyTorch算子都要被独立调用，有很繁复的Kernel Launch开销。
• 增大优化难度：这也使得跨算子优化（算子融合）比较困难。

在PyTorch 1.x的时代，对于这个问题也有解决方案，就是TorchScript。这是一种DSL（类似Triton），使用torch.jit装饰器来对被编译的区域进行warp。TorchScript存在这样一些问题：

• 有自己独立于Python的类型系统和定义控制流的语法，例如不能用NumPy，不能用Python的字典推导式、很多标准库函数不能用等
• 图捕获能力有点残疾。TorchScript使用torch.jit.trace进行图捕获，它从执行时调用的算子Op获取图，但是对于数据依赖的情况（比较典型的就是if-else依赖数据值时），torch.jit.trace不能感知到语法层面的控制流，只能感知到实际执行时调用的算子Op，因此会固化第一次运行时所走的分支。

所以，综上，TorchScript在Pythonic和无痛使用方面做得比较一般，大家还是更愿意在Triton里痛苦一会拿到更好的性能。

为了彻底解决这个问题，PyTorch 2.x大版本引入了一整套JIT编译系统。为了实现：

• 完整的图捕获能力
• 为了易用性牺牲一些强迫性：接受代码部分回退到eager模式
• 性能提升应当显著和可预测

Overview

torch的编译系统体系主要包含三大核心组件：

• TorchDynamo：用于捕获运行时生成的计算图，转化为一种叫FX Graph的计算图。
• AOTAutograd/AOTDispatcher：规范化FX Graph，将节点内容dispatch到aten ops上，生成反向计算图；前向后向中间激活值最小化（训练模式下）。
• TorchInductor：对FX Graph进行算子融合优化，然后生成Triton代码，调用Triton编译器进行编译并缓存kernel。

torch.compile的实际性能收益主要来源于：

• 算子融合，减小memory bound和kernel launch开销。
• 在训练时，AOTAutograd的激活值联合优化，减少显存开销。
• cuda graph（可选），降低CPU启动开销。

  flowchart LR
    A[代码] -->|Python代码| B[TorchDynamo]
    B -->|FX Graph| C[AOTAutograd
AOTDispatcher]

    C -->|Joint Graph| D[TorchInductor]

    D -->|Triton代码| E[Triton编译器]
    E --> F[可运行的kernel]

TorchDynamo

TorchDynamo 是整个编译系统的入口，负责从 Python 代码中提取计算图。

可以用自定义后端的方式来拦截TorchDynamo -> AOTAutograd/TorchInductor的过程，查看捕获的FX Graph：

import torch
from typing import Sequence, Callable

def custom_backend(
    gm: torch.fx.GraphModule,
    example_inputs: Sequence[torch.Tensor],
):
    print("FX Graph:")
    
    gm.graph.print_tabular()
    
    print("\nGenerated Python code:")
    print(gm.code)
    
    return gm.forward
    
@torch.compile(backend=custom_backend)
def func(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    z = x + y
    return torch.relu(z) * 2

x = torch.randn((4,), dtype=torch.float32, device="cuda:1")
y = torch.randn((4,), dtype=torch.float32, device="cuda:1")

o = func(x, y)

x = torch.randn((4,), dtype=torch.float32, device="cuda:1")
y = torch.randn((4,), dtype=torch.float32, device="cuda:1")

o = func(x, y)

输出为：

本文2026-07-20首发于GlowLED的后花园，最后修改于2026-07-20