torch.compile 2.x 学习笔记

torch.compile粒度

torch.compile 的编译入口通常是一个 Python function/callable，但它的实际 JIT 粒度并不等于“整个函数大小”。更准确地说，Dynamo 会在函数执行过程中捕获能够连续 trace 的 tensor computation，并把这段连续区域表示成 FX graph；因此 torch.compile 的实际编译粒度是 FX graph region。一个函数可能被完整捕获成一张 FX graph，也可能因为 graph break 被切成多个 FX graph region；同时，同一个 region 还可能因为不同的 shape、dtype、device、Python flag 或控制流路径生成多个 specialized graph。

例如：

@torch.compile
def f(x, use_moe: bool):
    y = x + 1
    if use_moe:
        y = y * 2
    else:
        y = y - 2
    return y.relu()

这里是以整个f作为一整个compiled region如果第一次调用是 f(x, True)，Dynamo 实际 trace 到的是 x -> add -> mul -> relu 这条路径，并生成一张带有 guard 的 specialized graph：

graph_true:
    y = x + 1
    y = y * 2
    out = relu(y)

guard:
    use_moe == True

如果之后又第一次调用 f(x, False)，原来的 graph_true guard 不满足，Dynamo 不会错误地继续执行包含 mul 的 fused graph，而是重新 trace 当前实际路径，生成另一张 graph：

graph_false:
    y = x + 1
    y = y - 2
    out = relu(y)

guard:
    use_moe == False

因此，一个函数内部可以有多份 specialized graph。它们不是同时表示完整的 if/else 语义，而是分别对应某一次实际执行时走到的控制流路径。后续如果再次遇到 use_moe=True，就复用 graph_true；如果遇到 use_moe=False，就复用 graph_false。这也是为什么如果 if 条件在进入 compiled region 之前就已经由 Python flag、配置项或 shape 决定，那么 if 外部的 op 和当前分支内部的 op 有机会被放进同一张 FX graph，甚至被 Inductor fuse 成同一个 kernel。因为这张 graph 的入口 guard 已经保证了“这次一定会走这条路径”。

但是，如果控制流条件必须运行到中间才能知道，例如条件依赖 tensor value：

@torch.compile
def f(x):
    a = x + 1
    if a.sum() > 0:
        b = a * 2
    else:
        b = a - 2
    return b.relu()

它有可能被这样切：

FXGraph 0:
    a = x + 1
    s = a.sum()

Python/runtime:
    判断 s > 0

FXGraph 1 true:
    b = a * 2
    out = relu(b)

FXGraph 1 false:
    b = a - 2
    out = relu(b)

这种情况下，在进入 graph 之前，compiler 无法知道 a.sum() > 0 是否成立，所以我们不能和之前一样以整个f作为compiled region。这里a.sum() > 0造成了graph break，因此将两个分支内部的语句作为compiled region了。

因此，torch.compile 的逻辑可以总结为：函数只是编译入口（你把 torch.compile 套在哪个函数上，Dynamo 就从哪个函数开始进入 tracing/JIT），不是严格的编译边界；真正的 JIT 粒度是 Dynamo 能连续捕获出来的 FX graph region。控制流如果能在 graph 入口处通过 guard 确定，就可以为当前路径生成 specialized graph，并允许跨当前分支边界的优化；如果控制流必须依赖运行时 tensor value 才能决定，则会且成subgraph，subgraph内部在入口处能确定整个控制流。

Fusion

Pointwise fusion 容易，Reduction fusion 难

对于 pointwise：

y[i] = relu((x[i] + 1) * 2)

每个 i 独立，一个 thread 处理一个元素：

load x[i] -> register
add -> mul -> relu
最后 store y[i]

中间结果一直在同一个 thread 的 register 里，所以可以很自然地 fuse，避免中间 tensor 写回 global memory。

但 reduction 不一样：

s = sum(x[0], x[1], ..., x[1023])

s 由很多元素共同决定，通常需要多个 thread 各算一部分 partial sum，然后再合并。问题是：register 是 thread 私有的，thread_0 的 register 不能直接被 thread_1 使用，所以必须通过 warp shuffle、shared memory、atomic 或 global memory 来通信和同步。

所以 reduction 的本质难点是：

many input elements -> one output element

它改变了并行模式，不再是简单的“一元素一个 thread”。

但 reduction 不是不能 fuse。比如：

sum(x[0:128])

如果 128 个元素可以在一个 warp 或一个 block 内规约，就可以用 shuffle/shared memory 在一个 kernel 内完成，不一定要把中间结果写回 global memory。

又比如：

y = sum(x * scale + bias)

x * scale + bias 可以融合进 reduction 前面，叫 map-reduce fusion。

再比如：

y = relu(sum(x))

relu 可以融合到 reduction 后面，作为 epilogue。

真正容易打断 fusion 的情况是：reduction 太大，一个 block 算不完，必须：

block0 -> partial sum
block1 -> partial sum
...
partial sums 写回 global memory
第二个 kernel 再合并

这时 global memory materialization 才会成为 fusion boundary。

一句话：pointwise fusion 受寄存器压力限制；reduction fusion 受跨 thread 通信、同步范围和 partial result 是否必须落 global memory 限制。

AOT Autograd

ReLU 是：

h2 = max(h1, 0)

它的导数是：

relu'(h1) = 1 if h1 > 0 else 0

所以：

d loss / d h1 = d loss / d h2 * (h1 > 0)

也就是说，ReLU 后向不需要完整的 h1 数值，它只需要知道每个位置是不是大于 0。

CUDA graph

CUDA Graph 的本质是：把一段固定的 CUDA kernel launch 序列录下来，之后用一次 cudaGraphLaunch 重放。它合并的是 CPU 发射过程，不是把多个 kernel 融合成一个 kernel。

普通执行：

CPU: launch K1 → launch K2 → launch K3
GPU:      K1   →      K2   →      K3

CUDA Graph replay：

CPU: cudaGraphLaunch(graph)
GPU: K1 → K2 → K3

所以 GPU 上还是多个 kernel，只是 CPU 不再逐个 launch，减少 launch overhead。

它和 FX Graph 的关系是上下游关系：

Python code
  ↓ Dynamo trace
FX Graph
  ↓ Inductor / backend lowering
wrapper code + kernel launches
  ↓ CUDA Graph capture
CUDA Graph
  ↓ replay
cudaGraphLaunch

FX Graph 是编译器语义图，描述“要算什么”：matmul、relu、attention、all_reduce 等。CUDA Graph 是 runtime launch 图，描述“实际发射哪些 kernel、什么顺序、什么参数地址”。一个 FX node 可能对应多个 CUDA kernel，多个 FX node 也可能被 Inductor fuse 成一个 kernel；一个 CUDA Graph 也可能覆盖多个 FX node，甚至包含 FX Graph 里没有显式出现的 runtime helper kernel。

CUDA Graph 的粒度不是固定的，取决于开发者/框架选的 captured region（在调用后面加.replay()）。它可以包：

一个 custom op 内部
一个 Inductor compiled wrapper
一个 transformer layer
一个完整 decode step
一个 piecewise decode segment

所以你说得对：如果 K1/K2/K3 要被 replay，那么一定要在发射 K1 之前，在它们上层的某个 region/OP/runner 入口处决定：

key 命中 → cudaGraphLaunch(graph)
key miss → 普通 launch / capture 新 graph

不是执行到 K1 后再决定 K2/K3。

replay 时，framework 会先选中一个 graph runner：

key = batch bucket + seq bucket + dtype + device
    + backend + KV layout + execution path
    + shape/stride + TP/rank config ...

这个 key 是 runtime 算出来的，但字段范围是框架/开发者设计的。凡是会影响 kernel topology、grid/block、shape、stride、dtype、device、buffer layout、执行路径的参数，都必须进入 key/guard。否则就是框架 bug。CUDA runtime 本身不会判断语义正确性。

选中 runner 后：

runner = graph_cache[key]
copy real_input → runner.static_input_buffer（但是在replay前要验证这个static input buffer仍然存在）
cudaGraphLaunch(runner.graph)
read runner.static_output_buffer

关键是：不是要求当前 tensor 地址等于 capture 时地址，而是把当前数据 copy 到 graph 绑定的 static buffer 地址里。 CUDA Graph 记录的是 K1/K2/K3 的 kernel function、grid/block、参数地址和依赖顺序。它不关心 static buffer 里的具体数据值。

例如可以 replay：

K1(input_static → tmp_static)
K2(tmp_static → output_static)

每次 token 值不同没关系，因为只是写进同一个 input_static 地址。

不能直接 replay 的情况：

K1 计算 accept_num
CPU 读取 accept_num
if accept_num == 1: launch K2_small
if accept_num == 4: launch K2_large

因为 K2 的 launch 参数依赖 K1 的运行结果，路径不固定。除非 bucket 化、固定最大 shape、把动态信息 tensor 化，或者在这里 graph break。

注意，虽然我们的CUDA graph作用在inductor compile之后，但是Inductor compiled region 也不等于一个 fused kernel。它通常是一个 wrapper，里面可能 launch GEMM、Triton fused kernel、FlashAttention、NCCL 等多个 kernel。CUDA Graph 可以 capture 这个 wrapper 实际发射出来的一串 kernel。

torch.compile 2.x 学习笔记

torch.compile粒度

Fusion

AOT Autograd

CUDA graph

Comments