PyTorch 训练框架内核深度解析

第15章 CUDA Graph 与 torch.compile 端到端

作者 杨艺韬 · 13,068 字 · 发布于 · 更新于

"CUDA Graphs will free tensors of a prior iteration. A new iteration is started on each invocation of torch.compile, so long as there is not a pending backward that has not been called."

—— torch.compiler.cudagraph_mark_step_begin 的 docstring（torch/compiler/__init__.py:379）

本章要点

  • CUDA Graph 把多次 kernel launch 录制成一个 "graph"：replay 时一次 launch 完成几百个 kernel，CPU 开销几乎归零
  • Inductor 的 reduce-overhead 模式自动用 CUDA Graph 包裹生成的 kernel：用户加 mode='reduce-overhead' 就行
  • cudagraph_trees.py（~2700 行）解决"多 graph 共存"的内存协作：同一设备上由它管理的 graph 共用一个 PrivatePool（第 4 章 §4.9），保证 replay 时显存地址稳定
  • 完整端到端：@torch.compile → Dynamo trace → AOTAutograd functionalize+partition → Inductor lowering+codegen → Triton kernel → CUDA Graph 包装 → replay
  • 小模型 / 小 batch 推理收益最大：单 token decode 每步要 launch 几百个小 kernel，CPU launch 开销占主导，包成一次 replay 后这部分几乎归零
  • 大模型训练收益小：单步 GPU 计算本身就很久，CPU launch 占比低，CUDA Graph 能省的那部分不显著

15.1 CUDA Graph 是什么

CUDA 10.0 引入的 API：把一段 GPU 操作"录制"成 cudaGraph_t 结构，replay 时用一次 cudaGraphLaunch 提交所有操作。

不用 CUDA Graph 的传统流程：

CPU: launch kernel 1 → launch kernel 2 → launch kernel 3 → ...
              ↓ 每次几微秒   ↓            ↓
              GPU 接收并排队

每次 kernel launch 的 CPU 侧开销是微秒量级。一个模型 forward 如果调上千个 kernel，光 launch 就要几毫秒；当每个 kernel 本身只跑几微秒（小算子密集的 decode 场景）时，CPU 就成了瓶颈 —— GPU 在等 CPU 喂。

用 CUDA Graph：

录制阶段（只做一次）:  capture_begin → 跑一遍 → capture_end → 得到 graph
replay 阶段（每次跑）:  cudaGraphLaunch(graph)   ← 一次调用！

一次 launch 完成几百个 kernel，CPU 开销几乎归零。这是 CUDA Graph 在 小算子密集 + 重复执行 场景的杀手级优化。

15.2 CUDA Graph 的硬约束

CUDA Graph 不是免费午餐，有几条硬约束让用它不那么直接：

1. 内存地址必须固定：录制时 kernel 看到的 input / output 张量地址在 replay 时也必须一样。如果 caching allocator 在两次 replay 之间分配了别的张量到那个地址 —— 数据被覆盖。第 4 章 §4.9 的 PrivatePool 就是为这个设计的。

2. shape 必须固定：录制时 input shape = [4, 4]，replay 时如果是 [8, 8]，graph 不能动态调整。

3. 控制流必须一致：如果 forward 里有 if x.sum() > 0，录制时进入 if 分支，replay 时不会重新判断。

4. 不能有 CPU-GPU 同步：录制中调 .item() 这类让 CPU 等 GPU 的操作会让 capture 报错。张量分配本身是允许的 —— capture 期间 caching allocator 把分配导向这张 graph 的私有池（第 4 章 §4.9）。

这四条约束让"任意 PyTorch 代码自动包 CUDA Graph"几乎不可能。所以 CUDA Graph 必须配合编译器使用 —— Inductor 把"已编译、shape 固定、无 host 控制流"的代码段交给 CUDA Graph 包装。

15.3 Inductor 的 reduce-overhead 模式

torch.compile(model, mode='reduce-overhead') 自动用 CUDA Graph 包 Inductor 生成的 kernel。流程：

flowchart TB
    User["@torch.compile(mode='reduce-overhead')"]
    User --> Cap[Inductor 生成 Triton kernel]
    Cap --> Pool[登记到本设备 cudagraph 树<br/>共用树管理器的 PrivatePool]
    Pool --> Warm[Warmup: 首次调用在 cudagraph 池里 eager 跑 1 次]
    Warm --> Rec[Capture: 再次调用时 cudaStreamBeginCapture → 调 kernel → end]
    Rec --> Cache[录好的 graph 挂成 cudagraph_trees 的 CUDAGraphNode]
    Run[每次 forward 调用]
    Run --> Check{shape / device 与录制时一致?}
    Check -->|是| Replay[cudaGraphLaunch]
    Check -->|否| Recap[重新 capture 一份 graph]
    style Replay fill:#dcfce7,stroke:#22c55e,stroke-width:2px
    style Recap fill:#fef3c7,stroke:#f59e0b

第一次 forward 时 Inductor 生成完 kernel 后并不立刻 capture，而是先在 cudagraph 内存池里 eager 跑一次做 warmup（torch/_inductor/cudagraph_trees.py 的 CUDAGraphTreeManager._run → run_eager）；下一次调用才进入录制：用 cudaStreamBeginCapture 录制这次调用的所有 GPU 操作，cudaStreamEndCapture 得到一张 graph。之后相同 shape 的 forward 直接 replay。

收益结构（本机无 GPU，不给具体毫秒数；自己机器上量的正确方式是 torch.profiler 看 CPU 侧 launch 时间占比）：

  • 不用 torch.compile：每个 ATen 算子各走一遍 Python + dispatcher + launch
  • mode='default'（仅 Inductor）：算子被 fuse，kernel 数量大幅下降，但每个 kernel 仍是一次 launch
  • mode='reduce-overhead'（+ CUDA Graph）：剩下的这些 launch 再被压成一次 cudaGraphLaunch

三档的收益是叠加的，而且越是"kernel 多、每个 kernel 又短"的场景，第三档的增量越大 —— 这是 vLLM 等推理引擎大量用 CUDA Graph 的原因。

15.4 cudagraph_trees：多 graph 协作

实际工程里一次训练有几张 graph：forward 一张、backward 一张、可能还有 optimizer step 一张。它们之间的内存怎么协调？

torch/_inductor/cudagraph_trees.py（2711 行）实现了一套"graph 树"管理器：

  • 每个 graph 节点对应一段编译产物
  • 节点之间通过 input/output tensor 连接
  • 一个设备上的 CUDAGraphTreeManager 只开一个 PrivatePool（cuda_graphs_thread_pool），它管理的所有 graph 都录在这个池里（在 caching allocator 里与普通分配隔离）
  • 靠跟踪哪些输出张量仍然存活，保证后录的 graph 不会覆盖前一张 graph 仍被引用的输出

这套树形结构让"forward → backward → optimizer step"三个 graph 能流畅协作 —— forward 输出是 backward 的输入，backward 输出是 optimizer 的输入，全程地址稳定。

如果模型有 graph break（第 12 章），每段 graph 成为树上的一个节点 —— 该文件的模块文档串说跨 graph break 使用正是它的主要场景。中间 eager 执行的代码段不被 CUDA Graph 包裹，开销比纯 graph 情况高。这是为什么 graph break 越多、reduce-overhead 模式收益越小。

15.5 完整端到端：一行代码到 GPU 二进制

把 12-14 章串起来，看 @torch.compile 全旅程：

flowchart TB
    U["@torch.compile(model, mode='reduce-overhead')<br/>用户代码"]
    U --> D[Dynamo<br/>PEP 523 拦截 forward<br/>字节码符号执行<br/>输出 FX Graph + Guards]
    D --> A[AOTAutograd<br/>fake tensor 跑反向<br/>functionalize<br/>min-cut partition<br/>输出 fw + bw 子图]
    A --> I[Inductor<br/>Lowering: ATen→IR<br/>Scheduling: fusion 决策<br/>Codegen: 生成 Triton]
    I --> T[Triton 编译<br/>Triton DSL → PTX → CUBIN]
    T --> CG[CUDA Graph 录制<br/>warmup 1 次 + capture]
    CG --> Run[Runtime<br/>guards check → graph replay]
    Run -.guards 失败.-> D
    style D fill:#fef3c7
    style A fill:#dbeafe
    style I fill:#dcfce7,stroke:#22c55e,stroke-width:2px
    style CG fill:#fce7f3

每一步对应的源码主目录：

阶段 关键代码
Dynamo torch/_dynamo/ (v2.11 实测 ~100000 行 Python)
AOTAutograd torch/_functorch/ (~32000 行)
Inductor torch/_inductor/ (~213000 行 —— 单 namespace 最大)
Triton triton/ (独立项目，几十万行)
CUDA Graph cudagraph_trees.py (~2700 行)

仅 PyTorch 主仓这三段加起来就是 34 万行以上的 Python（还不算 torch/csrc/dynamo/ 的 C/C++ 与 Triton 本身），协同工作让用户写一个装饰器就完成。每个 minor version 这套代码量还在持续增长（同口径下 v2.4.0 的 torch/_inductor/ 约 8.9 万行，到 v2.11.0 已是约 21.3 万行）。

15.6 模式选择决策

torch.compile(model, mode=...) 的 mode 取值见 torch/__init__.py:2625 的文档串与 torch/_inductor/__init__.py:336 的 list_mode_options()。常用的四档：

Mode 行为 适用
default Inductor 编译，无 CUDA Graph 大多数训练
reduce-overhead + CUDA Graph 包 推理 / 小 batch decode
max-autotune + 更激进的 autotune（max_autotune 与 coordinate_descent_tuning），并开 CUDA Graph 性能最敏感场景
max-autotune-no-cudagraphs max-autotune 但不开 CUDA Graph dynamic shape 多但还想 autotune

（list_mode_options() 里另有一档 lite，面向按需开启优化的场景，不在上面四档之列。）

15.6.1 怎么选

  • 大模型训练（70B+）：用 default。CUDA Graph 收益小（每步本身慢、launch 占比低），且 graph 容量限制可能导致显存爆
  • 大模型推理（prefill）：用 default。prefill 一次处理整 batch、kernel 大，graph 收益小
  • 大模型推理（decode）：用 reduce-overhead。decode 是单 token、kernel 小，CPU launch 占主导，graph 收益巨大
  • 小模型（如 ResNet）训练：用 reduce-overhead。每步 kernel 多但小，类似 decode 场景
  • 极致性能调优：用 max-autotune，编译时间成倍上涨，换取更好的 kernel 配置

vLLM 并不走这几个 mode 开关：它有自己的编译后端，按 cudagraph_capture_sizes 列出的若干档 token 数分段录 CUDA Graph（《vLLM 推理内核深度解析》专栏第 8 章 §8.4）。

15.6.2 什么时候 CUDA Graph 才划算

不给倍数，给判据。CUDA Graph 消掉的是 CPU 侧的 launch 时间，所以收益上限就是"CPU launch 在 step 墙钟时间里的占比"：

场景 launch 占比 CUDA Graph 收益
小模型 / batch=1 推理（kernel 多且短） 高 大
大模型 decode（每 token 几百个小 kernel） 高 大
大模型 prefill（kernel 少而大，compute-bound） 低 小
大模型训练（每步 GPU 计算几十毫秒起） 低 小

判断方法很直接：用 torch.profiler 或 Nsight Systems 看 GPU 是否有明显空档在等 CPU。有空档 → launch 是瓶颈 → 值得开；GPU 时间线密不透风 → 开了也白开，还要多付 graph capture 的时间与 PrivatePool 的显存。

先 profile、看 launch 开销占比、再决定，不要盲目开 —— 盲目开会导致显存涨、编译变慢，收益却不大。

15.6.3 编译时间的工程权衡

torch.compile 第一次跑慢，到底慢在哪？按阶段拆（本机无 GPU，不给秒数；torch._dynamo.utils.compile_times() 会按阶段给出你机器上的真实耗时表）：

阶段 相对量级 备注
Dynamo trace 小 inline 整个 forward
AOTAutograd functionalize + partition 中 trace 反向 + min-cut
Inductor lowering（ATen → IR） 中
Inductor scheduling + fusion 中
Inductor codegen（生成 Triton 源码） 中 几百个 kernel
Triton 编译（Triton DSL → PTX → cubin） 最大头 每个 kernel 一次完整的编译器流程
autotune（max-autotune 额外） 大 每个候选配置都要真跑一遍 benchmark

Triton 编译通常是压倒性的大头。优化思路：

  • compile-cache 预热：第一次跑完把 cubin 落盘，下次进程直接 load
  • AsyncCompile（torch/_inductor/async_compile.py:232，§14.11.1）：worker 进程池并发编译不同 kernel
  • autotune 结果复用：max_autotune 选出的配置会写进本地 autotune 缓存（torch/_inductor/config.py:134 的 autotune_local_cache，默认 True），同一 mode 下次编译直接复用。注意别指望"先用 max-autotune 编、再切 default 吃缓存"——FX graph cache 的 key 含 Inductor config（torch/_inductor/codecache.py 的 inductor_config），换 mode 就换了 key

并发编译把墙钟时间从"所有 kernel 编译时间之和"压到"最慢那几个 kernel"；TORCHINDUCTOR_CACHE_DIR 命中后连编译都不用跑，只从盘上 load。

工程决策：

  • 训练任务：编译时间一次性、可接受
  • 推理服务冷启动：编译时间 = 服务上线延迟。预编译 + 落盘 + 镜像层包含 cache 是标准做法
  • debug 调代码：每次改代码会让 cache 失效、再编一次、几秒到几十秒。工作流痛点

理解编译时间分布让你针对性优化。如果你的痛点是"频繁改代码 + 编译变慢"，看 cache 命中率 + 是否可以 partial cache（改某段代码不让全部失效）；如果痛点是"冷启动慢"，看 cache 落盘 + 镜像化。

15.6.4 graph instance 的显存成本从哪来

每张录制好的 graph 都要占显存，来源有两块，量级完全不同：

  1. cudaGraphExec_t 本身：graph 结构 + 每个节点的参数，正比于 kernel 数量，通常远小于下面的第二块
  2. PrivatePool（第 4 章 §4.9）：这张 graph 录制期间分配的所有 input / output / 中间张量，必须一直保活，否则 replay 时地址失效。这才是大头，正比于模型 activation 规模

所以"录 N 个 batch size bucket"时，如果每张 graph 各用独立的池（手工 torch.cuda.graph 不传 pool 时就是如此），显存开销近似 N × 单张 graph 的 PrivatePool。当 bucket 划得很细时很容易把余量吃光 —— 这是生产上开 reduce-overhead 最常见的翻车点。

缓解手段：

  1. 共享 pool：cudagraph_trees 让同一设备上的多张 graph 共用一个 PrivatePool（§15.4），这是它存在的主要理由；手工路径可以给 torch.cuda.graph(..., pool=...) 传同一个池，但官方文档（docs/source/notes/cuda.rst 的 "Sharing memory across captures"）要求这些 graph 总按录制顺序 replay、且不并发 replay 才安全
  2. 少切 bucket：torch._inductor.config.triton.cudagraph_dynamic_shape_warn_limit（torch/_inductor/config.py:1545，默认 8）就是在 graph 数超过阈值时提醒你 bucket 切太细了
  3. 该关就关：确认 launch 不是瓶颈时（§15.6.2），直接用 default 模式，一分显存不用付

评估方法是实测而不是估算：开与不开各跑一次，用 torch.cuda.max_memory_allocated() / memory_reserved() 对比峰值。

15.7 CUDA Graph 的捕获与运行

15.7.1 CUDAGraph capture 的具体 API 流程

CUDA 10 起就提供 stream-based graph capture API。PyTorch 包装它的关键函数在 aten/src/ATen/cuda/CUDAGraph.cpp：

// 简化版 capture 流程（据 aten/src/ATen/cuda/CUDAGraph.cpp，省略 RNG、host 分配器等细节）
void CUDAGraph::capture_begin(MempoolId_t pool, cudaStreamCaptureMode capture_mode) {
    // 1. 切换 caching allocator 到 PrivatePool (第 4 章 §4.9)
    CUDACachingAllocator::beginAllocateToPool(capture_dev_, mempool_id_, filter);
    // 2. 让当前 stream 开始 capture（capture_mode 由 Python 侧传入，默认 global）
    cudaStreamBeginCapture(capture_stream_, capture_mode);
}

void CUDAGraph::capture_end() {
    cudaStreamEndCapture(capture_stream_, &graph_);   // 拿到 cudaGraph_t
    CUDACachingAllocator::endAllocateToPool(capture_dev_, mempool_id_);
    instantiate();   // 内部 cudaGraphInstantiateWithFlags → cudaGraphExec_t
}

void CUDAGraph::replay() {
    cudaGraphLaunch(graph_exec_, at::cuda::getCurrentCUDAStream());   // 一次 launch 跑完整段
}

三个关键 CUDA API：

  • cudaStreamBeginCapture：标记 stream 进入 capture 模式。之后所有 launch 进 stream 都被记录、不真执行
  • cudaStreamEndCapture：结束 capture，拿到 cudaGraph_t（graph 描述，未编译）
  • cudaGraphInstantiate：把 graph 编译成 cudaGraphExec_t（可执行版本）。这一步不便宜，所以后面才有 cudaGraphExecUpdate 这种"改参数不重建"的路径（§15.7.5）

capture 模式由 Python 侧 capture_error_mode 参数决定（torch/cuda/graphs.py，默认 "global"，可选 "thread_local" / "relaxed"）：global 最保守，capture 期间其他线程调了不安全的 CUDA API（如 cudaMalloc）也会报错；thread_local 只管本线程。

PrivatePool 与 capture 的协作（第 4 章 §4.9）很关键：capture 期间所有 alloc 走专属 pool，replay 时这些地址保持有效；capture 之外其他张量分配不进 pool，避免污染 graph 录制的地址。

15.7.2 capture 中的 PyTorch wrapper API

PyTorch 暴露原始的 torch.cuda.CUDAGraph 类，以及上下文管理器 torch.cuda.graph 作为 capture 的 high-level wrapper：

# 按官方文档 docs/source/notes/cuda.rst「Whole-network capture」的写法
# model / loss_fn / optimizer / loader 与 N、D_in、D_out 自备；
# optimizer 用 SGD 这类无 step 计数的；Adam 等需构造时传 capturable=True（第 10 章 §10.6.3）
static_input = torch.randn(N, D_in, device="cuda")     # 固定地址的输入占位
static_target = torch.randn(N, D_out, device="cuda")

# warmup：必须在 side stream 上跑
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
    for _ in range(3):
        optimizer.zero_grad(set_to_none=True)
        loss = loss_fn(model(static_input), static_target)
        loss.backward()
        optimizer.step()
torch.cuda.current_stream().wait_stream(s)

# capture：先把 grad 置 None，让 backward 在 graph 私有池里新建 .grad
g = torch.cuda.CUDAGraph()
optimizer.zero_grad(set_to_none=True)
with torch.cuda.graph(g):
    static_loss = loss_fn(model(static_input), static_target)
    static_loss.backward()
    optimizer.step()

# replay
for data, target in loader:
    static_input.copy_(data)        # 用同一个 tensor 装数据
    static_target.copy_(target)
    g.replay()                      # forward + backward + step；无需再 zero_grad
    print(static_loss.item())       # static_loss 也是同一个 tensor

关键约束：

  • 输入 / 输出是固定 tensor：每次 replay 前用 .copy_() 写入新数据，graph 内部 kernel 看到的还是原 tensor
  • 整个 step 在 capture 内：forward + backward + optimizer 全部一起录，否则 grad 状态在 graph 内 vs 外不一致；capture 前的 zero_grad(set_to_none=True) 让 .grad 在私有池里分配，之后每次 replay 原地重写，不会跨步累加
  • warmup 必须，且在 side stream 上做：让 cuDNN / cuBLAS 选算法、各种懒初始化等一次性工作在 capture 之前完成（原因见 §15.7.3）

这套 API 比 torch.compile(mode='reduce-overhead') 更底层，让用户精确控制哪段进 graph。Megatron-LM、DeepSpeed 等训练框架内部都直接用这层 API（不用 torch.compile，自己手写 graph capture）。

15.7.3 warmup 必须的内部原因

torch.cuda.graph(...) 之前必须 warmup 几次（make_graphed_callables 的 num_warmup_iters 默认 3），原因有四：

1. 一次性分配放在 capture 之外：第 4 章 §4.4 的 caching allocator 第一次见到一个新 shape 时会 cudaMalloc；optimizer state 这类首次 step 才创建的张量也在这时建好。warmup 让这些一次性分配发生在 capture 之前，不落进 graph 的私有池。

2. cuDNN / cuBLAS algorithm 选择：开了 torch.backends.cudnn.benchmark 时，第一次调用 at::cudnn_convolution（aten/src/ATen/native/cudnn/ConvShared.cpp）等会触发 algo selection（试几种 algo、选最快的、cache 起来）。warmup 让 algo 选择完成，capture 时不再有这个一次性开销。

3. NCCL communicator init：分布式训练里 NCCL collective 第一次调用做 communicator setup（占几十 ms）。warmup 把它跑完。

4. kernel 模块的懒加载 / PTX JIT：CUDA 的 lazy loading 让 kernel 所在模块在第一次 launch 时才加载；二进制里若没有当前架构的 SASS，驱动还要在首次加载时把 PTX JIT 成 SASS。warmup 让这步完成。

为什么默认是 3 次？3 是经验默认值，并非对所有模型都够：make_graphed_callables 的文档串自己就注明 DDP 需要 11 次 warmup（torch/cuda/graphs.py 中 num_warmup_iters 参数说明）。

实战：如果你的模型某些 op 第一次跑特别慢（如 cudnn algo selection 很慢），可以多跑几次 warmup 再 capture。手工路径上这个次数是 torch.cuda.make_graphed_callables(..., num_warmup_iters=3) 的参数（torch/cuda/graphs.py:277，默认 3）。

理解 warmup 不是迷信，是有具体技术原因的。生产代码"3 次 warmup 后再 capture"是把这些一次性开销摊销到 warmup 阶段，让 capture 录到的是"稳态"。

15.7.4 cudagraph_mark_step_begin：分隔训练步的 hint

torch._dynamo.mark_static_address（torch/_dynamo/decorators.py:1136）与 torch.compiler.cudagraph_mark_step_begin（torch/compiler/__init__.py:379，内部转到 torch/_inductor/cudagraph_trees.py:324 的 mark_step_begin）是 v2.x 给用户提供的 hint API。前者声明某个输入 tensor 的地址跨调用不变，cudagraphs 就不必为它在池里另做一份拷贝（地址真变了：guard=False 时重录 graph，guard=True 时整体重编），后者告诉 cudagraph_trees "新的训练步开始了"。

# 训练循环
for batch in dataloader:
    torch.compiler.cudagraph_mark_step_begin()    # hint!
    loss = model(batch).sum()
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

为什么需要这个？因为 cudagraph_trees 维护一棵"graph 树"（§15.4），需要知道两个 invocation 之间有没有跨越"step 边界"。step 边界处会做：

  • 把上一代生成的 outputs 视为可覆盖（PrivatePool 里那块内存可被本代复用；之后再访问旧 output 会报错）
  • 把当前所在的树节点复位，下一次调用从树根重新匹配已录制的 graph（CUDAGraphTreeManager.can_start_new_generation / try_end_curr_execution）

不调用它时，cudagraph_trees 的启发式是"每次进入 torch.compile 编译产物就算新一轮迭代，除非还有未调用的 pending backward"（见该函数 docstring）。这个启发式在某些写法下会判错 —— docstring 自己给的反例就是 rand_foo() + rand_foo() 这种一轮里调两次编译产物的代码 —— 这时就要手工插这个调用。

理解这个 hint 让你看到"compile 与 PyTorch 训练循环之间的协议"—— 不是黑盒，有可干预的接口。生产代码里如果遇到 reduce-overhead 模式报 "Error: accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run"（cudagraph_trees.py 的 format_dealloc_msg），错误信息给的两招就是：在 torch.compile 外 clone 该输出，或每次调用前加 mark_step_begin。

15.7.5 cudaGraphExecUpdate：动态更新 graph

CUDA 提供 cudaGraphExecUpdate API：让一个已 instantiate 的 graph 在不重新 capture 的情况下更新少量参数（如 kernel 参数、memory copy 的 size）。

例子：推理时 batch_size 不变但 input pointer 变化（不同 request 的 input 在不同地址）。传统做法是重新 capture + instantiate，新做法是 cudaGraphExecUpdate(exec, new_graph) 把新 graph 的 metadata patch 到老 exec 上，比重新 instantiate 便宜得多。

PyTorch 与 vLLM 都没走这条路：v2.11 源码里 cudaGraphExecUpdate 只出现在 hipify 的名字映射表（torch/utils/hipify/cuda_to_hip_mappings.py），vLLM v0.8.5 源码里也搜不到。它们用的是"static input tensor + copy_ 写入新数据"（见 §15.9.2）。

限制：update 不能改 graph 结构（节点数、依赖边），只能改节点的参数（kernel arg、memory address）。结构变化仍要重新 capture。

所以 vLLM 能让不同 request 复用同一份 graph，靠的不是 update 参数，而是把每拍的输入 copy 进固定地址的 buffer、再 pad 到预录的尺寸（《vLLM 推理内核深度解析》专栏第 8 章 §8.5）。

15.7.6 多 stream 编排：graph 内的并行

CUDA Graph 不只是"一串顺序 kernel"，它是有向无环图（DAG），节点间可以并行：capture 时分到不同 stream、彼此无依赖的 kernel，会成为 graph 里可并发执行的节点。

graph TB
    Start[graph 开始]
    Start --> K1[kernel 1: layer norm]
    Start --> K2[kernel 2: bias load<br/>无依赖]
    K1 --> K3[kernel 3: matmul]
    K2 --> K3
    K3 --> K4[kernel 4: activation]
    K4 --> End[graph 结束]
    style K1 fill:#dcfce7
    style K2 fill:#fef3c7
    style K3 fill:#dbeafe
    style K4 fill:#fce7f3

实现方式：record 时在多个 stream 各自跑 kernel、用 events 同步 → graph capture 把 stream 关系记录成 graph node 的 dependency edges。replay 时 GPU 调度器按 DAG 并发执行无依赖节点。

实战收益：单 stream record 出来的 graph 是一条链，节点间没有并行余地；多 stream record + 合理切分能让 GPU 在无依赖的节点之间并发。但多 stream graph 调试难度大（错误信号难定位）。Inductor 生成的计算 kernel 都发在同一条 stream 上（torch/_inductor/ 里新建 torch.cuda.Stream() 的只有 compile_fx.py 与 cudagraph_trees.py 各一处，都不是给 kernel 分流用的），所以 reduce-overhead 录出来的是一条链。

想要 graph 内并行，只能在手工 capture 里自己分流：在 capture 上下文里让其他 stream wait_stream capture stream、干完活再汇合回来，这些 stream 上的 kernel 会录进同一张 graph（官方文档 docs/source/notes/cuda.rst 的 "Usage with multiple streams" 一节；stream 与分配器的配合见第 4 章 §4.6）。注意 torch.cuda.graph(stream=...) 这个参数指定的是 capture 用哪条 stream，不是开多 stream。

15.8 CUDA Graph 的边界与兼容性

15.8.1 dynamic shape × CUDA Graph：根本冲突的工程妥协

CUDA Graph 要求 shape 固定，但生产推理需要变长输入（如 LLM decode 的 KV cache 长度持续增长）。两个解决方案都被业界采用：

方案 A：Padding + Mask

把变长输入 pad 到几个固定 shape（如 [128, 256, 512, 1024]），graph 录制时用最大 shape：

@torch.compile(mode='reduce-overhead')
def model_forward(input_ids, attention_mask):
    return model(input_ids, attention_mask)

# 用户层 pad
input_padded = pad_to_bucket(input_ids, [128, 256, 512, 1024])
mask = make_attention_mask(input_ids, input_padded)
out = model_forward(input_padded, mask)
out = unpad(out, original_lengths)

代价：padding 出来的那部分算力是白烧的，浪费比例 =（bucket 大小 − 真实长度）/ bucket 大小，bucket 划得越粗浪费越多。

方案 B：Split Graph (Bucketed Capture)

为每个常用 shape 单独录一张 graph，运行时按 input shape lookup：

graphs, static_in, static_out = {}, {}, {}   # bucket → CUDAGraph / 固定输入 / 固定输出
for b in [128, 256, 512, 1024]:               # 录制前各 bucket 已在 side stream 上 warmup（§15.7.2）
    static_in[b] = make_dummy_input(b)
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g):                # torch.cuda.graph 必须传入 CUDAGraph 对象
        static_out[b] = model(static_in[b])
    graphs[b] = g

def forward(input):
    b = next_bucket(input.shape[0])          # 向上取整到最近的 bucket
    static_in[b][: input.shape[0]].copy_(input)
    graphs[b].replay()
    return static_out[b][: input.shape[0]]

vLLM 的 V1 引擎用这套：按 cudagraph_capture_sizes 列出的若干档 token 数分别录 graph，运行时把本拍 token 数向上 pad 到最近的档位（《vLLM 推理内核深度解析》专栏第 8 章 §8.5）。bucket 数越多、覆盖越好，显存开销也线性上涨（§15.6.4）。

方案选择：

  • batch 内 sequence 变长（如 prefill）→ 用 Padding + Mask（FlashAttention 帮助处理变长）
  • batch 数变化（如不同请求数）→ 用 Bucketed Capture
  • 彻底没办法 graph 的代码段（如 dynamic 控制流）→ graph break，那段走 eager

torch._inductor.config.triton.cudagraph_dynamic_shape_warn_limit = N 让 Inductor 在 graph 数超过 N 时警告 —— 提醒你 bucket 是否划分太细。

15.8.2 CUDA Graph × NCCL 通信的兼容性

DDP / FSDP 训练里有 collective 操作（AllReduce、AllGather）。CUDA Graph 能否包住这些 collective？答案是 能但有限制。

支持路径：NCCL 从 2.9 起支持被 CUDA stream capture 录入 graph（collective 本来就是 stream-based 的异步调用）。PyTorch 在 ProcessGroupNCCL 里配合 capture 模式走这条路径（第 16 章会展开 ProcessGroupNCCL 的实现）。

// 录制阶段
cudaStreamBeginCapture(stream, ...);
ncclAllReduce(input, output, count, ncclFloat, ncclSum, comm, stream);  // 能被 capture
cudaStreamEndCapture(stream, &graph);

// replay
cudaGraphLaunch(exec, stream);   // collective 也被 replay

限制条件：

  • comm 必须固定：record 时用 commA、replay 时不能换 commB（comm 切换需要重新 capture）
  • input/output 地址必须固定：与普通 CUDA Graph 一致
  • 不能跨 graph 用同一 comm：每个 graph 应该有独占的 comm（避免 NCCL 内部状态混乱）
  • collective 大小必须固定：count 在 record 时固化、replay 不能改

实战 corner case：

  • DDP gradient bucketing：第 17 章 DDP 章会讲，AllReduce 由 autograd hook 在梯度 ready 时触发、bucket 还会在首轮后按 ready 顺序重建 → 与 CUDA Graph 的静态录制天然冲突。Dynamo 对 DDP 另有处理：torch/_dynamo/backends/distributed.py 的 DDPOptimizer 按 DDP 的 bucket 边界把图切成多段分别编译，让 AllReduce 能与反向计算重叠
  • FSDP-2 all-gather：分片 shape 固定时 all-gather 可以进 graph，但 FSDP 自身的 prefetch / reshard 调度带来的动态性往往迫使这段留在 graph 外
  • EP / TP 训练：experts collective 在 graph 内复杂度高，多数实现走 graph break 把 collective 留在 eager

调优建议：分布式训练默认 mode='default'（不开 graph）；推理（无 collective）大胆开 reduce-overhead；分布式推理（如 vLLM TP=8）需要逐 op 验证 graph capture 不出错。

15.8.3 多 GPU 上的 CUDA Graph

多卡训练 / 推理时每张 GPU 独立录制 graph。关键点：

  • 每张 GPU 一份 PrivatePool：互不干扰
  • 每张 GPU 一份 graph instance：rank 0 与 rank 1 各自独立 capture
  • collective 协调：record 时所有 rank 同时调 NCCL collective，capture 把 collective 录入；replay 时所有 rank 同时 launch graph，collective 自然同步

启动顺序很关键：所有 rank 必须同时进入 capture，否则有 rank 在等 collective 时另一些 rank 在 capture，会 hang。常见做法是在 capture 前用 dist.barrier() 同步所有 rank。

@torch.compile(mode='reduce-overhead')
def step(model, batch):
    out = model(batch)
    return out.sum().backward()

# 4 卡训练
dist.barrier()    # 必须！同步所有 rank
for i in range(num_warmup):
    step(model, batch)
dist.barrier()    # warmup 完成后再次同步

故障模式：rank 间 capture 时机不一致 → "watchdog timeout"（NCCL 内部 watchdog 检测到部分 rank 不响应）。第 16 章 process group 章会展开 watchdog 机制。

实战经验：分布式 + reduce-overhead 上手前，先在单卡跑通 + profile 收益、确认值得切；分布式部署再加 barrier 同步逻辑。

15.8.4 reduce-overhead × autocast 的兼容性

第 20 章会讲 mixed precision (autocast)。autocast 与 CUDA Graph 的兼容性如何？

手工 capture 路径：官方文档（docs/source/notes/cuda.rst 的 "Usage with torch.cuda.amp"）指出，GradScaler.step 会做 CPU-GPU 同步，不能录进 graph；做法是只录 forward、loss 与 backward，optimizer step 留在 graph 外（或改用 partial-network capture）。

torch.compile 路径：autocast 上下文在 Dynamo trace 阶段就被捕获（第 13 章 §13.17），AOTAutograd 拿到的 graph 里每个 op 的 dtype 已经定死 —— graph 内部不再有运行时的类型分支，capture 出来的结构自然是稳定的。

两个优化的收益来源不同（autocast 省的是显存带宽与 Tensor Core 吞吐，CUDA Graph 省的是 CPU launch），因此是可以叠加的 —— 只开一个都拿不到峰值。

工程提示：autocast 决策在 Dynamo trace 期间确定。外层的 autocast 开关与 dtype 属于 Dynamo 的全局状态 guard（torch/csrc/dynamo/guards.cpp 的 AutocastState），trace 时开、调用时不开，会触发重编译而不是悄悄用错 dtype（本机 CPU 实测：同一函数在 torch.autocast("cpu", dtype=torch.bfloat16) 内外各调一次，输出分别为 bfloat16 / float32，共编译 2 帧）；但来回切换就会反复重编。建议把 autocast 包在 compile 函数内部，让 Dynamo 能完整 trace 进去：

@torch.compile(mode='reduce-overhead')
def step(model, batch):
    with torch.autocast('cuda', dtype=torch.float16):
        return model(batch)

这样 autocast 与 compile 都进 graph，没有外层 / 内层不一致问题。

15.8.5 capture 失败的常见原因与诊断

CUDA Graph capture 失败通常报 cudaErrorStreamCaptureUnsupported 或 cudaErrorStreamCaptureInvalidated。失败场景：

错误 触发原因
StreamCaptureUnsupported 调了不能被 capture 的 API（如 cudaMemcpy 同步版本）
StreamCaptureInvalidated capture 过程中其他线程在同 stream 做了 GPU 操作
IllegalAddress (replay 时) replay 时 input 地址与 record 时不一致（caching allocator 误分配）
LaunchTimeout graph 内某 kernel 运行时间超过 driver 的 watchdog 阈值（该阈值由驱动/平台决定）
OOM during instantiate cudaGraphInstantiate 需要把整张 graph 拷贝到显存，超出余量

诊断命令：

# 开详细 CUDA log（inductor 日志组已包含 torch._inductor.cudagraph_trees 模块）
CUDA_LAUNCH_BLOCKING=1 \
TORCH_LOGS="+inductor,cudagraphs,perf_hints" \
python train.py

具体看 cudagraph_trees.py 输出的 trace：warmup / 录制的时机、因何跳过 cudagraph（perf_hints）、重录的原因。多数 capture 问题从这里就能定位。

实战经验：第一次开 reduce-overhead 必碰几个 corner case。常见解法：

  1. 把不能 capture 的 op（如 print / .item()）移到 compile 包外
  2. 用 torch.cuda.is_current_stream_capturing() 在代码里判断当前是否在 capture，做差异化处理
  3. 关闭 triton.cudagraphs flag 让 Inductor 走非 graph 路径，先排查是否 graph 本身有问题

15.9 Inductor 的配置与产物

15.9.1 _inductor.config 关键 flags

torch/_inductor/config.py（2526 行）的 flag 控制编译行为。最常调的几个：

Flag 默认 用途
max_autotune False 编译时 autotune 试更多 block size 配置（编译时间成倍上涨，换更优 kernel 配置）
max_autotune_gemm False 专门给 GEMM autotune（让 Inductor 试 cuBLAS / Triton GEMM templates / handwritten）
triton.cudagraphs False 配 reduce-overhead 模式时启用，自动给 Triton kernel 包 CUDA Graph
aggressive_fusion False 更激进 fuse（可能产生超大 kernel，占太多 register）
coordinate_descent_tuning False 用坐标下降算法做 autotune（找比 grid search 更优配置）
cpp.threads -1 CPU 后端 OpenMP 线程数，-1 表示取 torch.get_num_threads()
freezing False 推理时把常量参数 freeze 进 kernel（更激进 fuse，但模型变 inference-only）

环境变量等价物：TORCHINDUCTOR_MAX_AUTOTUNE=1 等。

调优顺序：先开 max_autotune 看效果；推理服务再开 freezing；如果某 kernel 用 aggressive_fusion 后崩了说明 register 压力过大，关掉。

理解这些 flag 让你在编译性能不达预期时有"杠杆"可调，而不是只盯着 mode='reduce-overhead' 一个粗粒度选择。

15.9.2 cudagraphify：Inductor 与 graph 的接口

torch/_inductor/cudagraph_trees.py 的 cudagraphify 函数（v2.11 实测 line 472）是 Inductor 调用 CUDA Graph 的入口：

# torch/_inductor/cudagraph_trees.py（省略部分关键字参数）
def cudagraphify(
    model: ModelType,
    inputs: list[InputType],
    static_input_idxs: Sequence[int] = (),
    *,
    device_index: int,
    is_backward: bool,
    is_inference: bool,
    ...
) -> tuple[ModelType, OutputType]:
    ...
    manager = get_container(device_index).get_tree_manager()
    return manager.add_function(
        model,
        inputs,
        static_input_idxs,
        ...
    )

它本身不做录制，而是把函数登记到本设备的 CUDAGraphTreeManager，返回（包装后的函数，首次调用的输出）。真正的"拷输入 → replay → 返回固定输出"发生在树节点 CUDAGraphNode 里。按 symint 取值分别录图的那层缓存在同文件的 cudagraphify_impl（:393）。

关键概念：

  • static_input_idxs：标记哪些 input "地址固定"（如 model 权重 —— 永远是同一张 tensor）。这些不需要每次 replay 前 copy
  • dynamic input：用户每次传新数据的 tensor（如 input batch）—— graph 录制时用占位 tensor、replay 时把新数据 copy 到占位 tensor

具体逻辑（示意，非源码原文）：

# 录制时
static_inputs = [t.clone() if i not in static_input_idxs else t for i, t in enumerate(inputs)]
with torch.cuda.graph(graph):
    output = model(*static_inputs)
static_output = output

# replay 时
def wrapped(*new_inputs):
    for i, new_t in enumerate(new_inputs):
        if i not in static_input_idxs:
            static_inputs[i].copy_(new_t)
    graph.replay()
    return static_output

这种"用 copy 替换 dynamic input"的设计让 graph 既能处理"每次输入数据不同"又能保持地址固定。代价是每次 replay 前几 us 的 copy 开销（远小于 launch 开销，仍划算）。

注意这些拷贝不在 Inductor 生成的 wrapper code 里 —— output_code.py 里看不到它们，它们发生在 cudagraph_trees 运行时包的那一层。

15.9.3 编译产物的磁盘缓存

TORCHINDUCTOR_CACHE_DIR 让编译产物落盘，下次进程启动跳过编译。默认位置不是 ~/.cache，而是 $TMPDIR/torchinductor_<用户名>（torch/_inductor/runtime/cache_dir_utils.py 的 default_cache_dir()）。本机 CPU 实跑一个只含 a @ b 的函数（matmul 走 extern kernel，不需要编 C++），目录是这样的：

$TMPDIR/torchinductor_<user>/
├── fxgraph/                 # FxGraphCache：<key 第 2-3 位>/<key>/<条目>，条目是 pickle 后的 CompiledFxGraph
│   └── oe/foekkccj.../54xcf3pp...
├── aotautograd/             # AOTAutograd cache：<key>/<条目>
│   └── as6ldhva.../szudgkh2...
└── vw/                      # PyCodeCache：Inductor 生成的 wrapper 源码（output code）
    └── cvw5zwpd....py

GPU 上还会多出 triton/<设备号>/（triton_cache_dir()，除非设了 TRITON_CACHE_DIR），里面是 Triton 自己的编译缓存（kernel 源码、cubin 与 metadata）。

工作流：编译时对"FX graph + 输入的 shape/dtype/stride + Inductor config + PyTorch/Triton 版本与 GPU 型号"算 hash 作为 key（torch/_inductor/codecache.py 的 FxGraphHashDetails、CacheBase.get_system()）。下次进程启动时 lookup key → 命中直接反序列化编译产物、不命中才走完整编译。

命中率取决于：模型固定（同一 nn.Module 类）+ 输入 shape 固定（或同样的 dynamic shape 假设）+ Inductor config 与 PyTorch 版本不变。只要这些都稳定，冷启动就从"跑完整条编译链"变成"从盘上 load"。

注意事项：

  • PyTorch 版本升级后旧条目不会再命中：key 里含 torch_key()（PyTorch 源码的 hash），不会误用旧产物；删掉 $TMPDIR/torchinductor_<user>（或你自己设的 TORCHINDUCTOR_CACHE_DIR）只是为了回收空间
  • 缓存目录大小要监控：长时间训练 + 多模型可能占几 GB；定期清理 find -mtime +30 -delete
  • 多种 GPU 共享 cache 目录：不同 SM 架构（A100 vs H100）的 cubin 不通用，但 FX graph cache 的 key 已含 GPU 型号与 CUDA 版本，跨型号不会误命中，只是彼此命中不了

torch._inductor.config.fx_graph_cache（torch/_inductor/config.py:104）默认 True，可用 TORCHINDUCTOR_FX_GRAPH_CACHE 环境变量强制开关。另有 fx_graph_remote_cache（:117）走远端缓存后端，让多机训练共享编译产物；它在开源版里默认不开（值为 None，源码注释写明 "Off for OSS"），需要自己配后端。

15.9.4 一个 Inductor output_code 的真实样子

为了具体感受"compile 产出什么"，看一份 output_code.py 的形态（按 v2.11 的 Triton codegen 与 wrapper 模板整理的精简示意，对应 y = relu(x + bias)；本机无 GPU，未实跑）：

# output_code.py (Inductor 生成)
import triton
import triton.language as tl
from torch._inductor.runtime import triton_helpers

@triton.jit
def triton_poi_fused_add_relu_0(in_ptr0, in_ptr1, 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)        # x
    tmp1 = tl.load(in_ptr1 + xindex, xmask)        # bias
    tmp2 = tmp0 + tmp1                              # add
    tmp3 = tl.full([1], 0, tl.float32)
    tmp4 = triton_helpers.maximum(tmp3, tmp2)      # relu
    tl.store(out_ptr0 + xindex, tmp4, xmask)


def call(args):
    arg0_x, arg1_bias = args
    args.clear()
    out = empty_strided_cuda((1024,), (1,), torch.float32)
    triton_poi_fused_add_relu_0.run(arg0_x, arg1_bias, out, 1024, stream=stream0)
    return (out,)

可以看到：

  • Inductor fuse 了 add + relu 进同一个 Triton kernel（triton_poi_fused_add_relu_0，名字里是参与 fuse 的原始 ATen op 名按字母序排列，见 torch/_inductor/utils.py 的 get_fused_kernel_name）
  • kernel 直接读 in_ptr0 / in_ptr1，加完立刻 relu 立刻写回 —— 没有中间 tensor，最优内存模式
  • wrapper call 负责：创建 output buffer、调 kernel、返回。这是被 cudagraphify 包装的"static callable"

如果不开 compile，eager 执行同样代码：1 次 add（读 2 个 tensor、写 1 个 tensor）+ 1 次 relu（读 1 个 tensor、写 1 个 tensor）= 5 次 memory access。compile 后：2 次 read + 1 次 write = 3 次。对纯 memory-bound 的算子，访存量降到 3/5，理论加速上限约 5/3 倍。

理解 output_code 让你看到 compile 的实质收益不是"魔法"，是"把 ATen op 的中间 tensor 全部消除"。这是 Inductor 给 PyTorch 加速的根本机制。

把这段产物与 §12.31 的 Dynamo OutputGraph 对照看：Dynamo 那层产出的 fx Graph 是"逻辑算子序列"，Inductor 这层把它编译成"具体 GPU kernel + memory plan"。两层职责明确分开，让每层都能独立优化、独立替换。这是为什么 PyTorch 团队能在 v2.x 的几年里持续给 Inductor 加新优化（fusion patterns、autotune 改进、cudagraph_trees 演进），而不影响上层 Dynamo 接口稳定性 —— 接口稳、内部演进是这套编译栈最大的工程价值，也是看完 12-15 章源码后最值得记住的一条架构原则。

15.10 AOT 编译与部署终点

15.10.1 AOT Inductor (AOTI)：ahead-of-time 编译

@torch.compile 是 JIT —— 第一次跑时编译，依赖 Python 解释器 + Dynamo + Inductor 全部活着。AOTI 是另一条路（torch._inductor.aot_compile 在 v2.1.0 的 torch/_inductor/__init__.py 里就已存在，打包成 .pt2 的 aoti_compile_and_package 到 v2.6.0 才出现）：编译前提前做、产物落盘、运行时不需要 Python。

入口（torch/_inductor/__init__.py:55 的 aoti_compile_and_package，旧接口 aot_compile 在 :272）：

import torch
import torch._inductor

# 1. export 模型成 ExportedProgram (静态图)；example_inputs 是 tuple
exported = torch.export.export(model, example_inputs)

# 2. AOT 编译, 产出打包好的 .pt2（里面含编好的 .so）
output_path = torch._inductor.aoti_compile_and_package(
    exported,
    package_path="model_aoti.pt2",
)

# 3. Python 侧可用 torch._inductor.aoti_load_package(output_path) 加载验证；
#    C++ 端加载则不需要 Python

输出的 .pt2 是个 zip 包，内含：

  • 编译好的 .so（C++ wrapper + kernel；GPU 上另有各 Triton kernel 的 cubin）
  • 模型权重：默认直接链进 .so（torch/_inductor/config.py 的 aot_inductor.package_constants_in_so，默认 True，源码标注为 Experimental），也可配置成单独打包
  • 生成的 C++ 源码、元数据 JSON 等辅助文件

C++ 端用 AOTIModelPackageLoader 打开 .pt2，内部创建 AOTIModelContainerRunner 加载 .so，再用 run() 接口直接推理。整套部署不依赖 Python 解释器，对只有 C++ 运行时的服务端部署很关键（移动端 / 嵌入式另走 ExecuTorch，见 §15.10.4）。

AOTI 与普通 torch.compile(mode='reduce-overhead') 的区别：

维度 torch.compile (JIT) AOTI
编译时机 第一次调用 离线提前
部署依赖 Python + PyTorch 只需 libtorch (C++)
启动时间 首调用要跑完整条编译链 只需 dlopen 已编好的 .so
灵活性 支持 dynamic shape / graph break 不允许 graph break（要整图可 export）；只有 export 时用 dynamic_shapes 声明的维度可变
目标场景 训练 + JIT 推理 生产推理服务 / 边缘部署

AOTI 内部仍走 Inductor 整套 pipeline（lowering / scheduling / Triton codegen），但产物从"运行时 callable"变成"持久化 .so"。Wrapper codegen 也从 Python wrapper 切换成 C++ wrapper（codegen/cpp_wrapper_*.py，§14.10.2 提过），让 .so 内部完全 C++、零 Python overhead。

实战决策：训练用 torch.compile、需要脱离 Python 的推理服务用 AOTI —— 模型分发时直接给 .pt2 包，部署时不用 Python 环境。这是 PyTorch 走向 "一次训练、多端部署" 的关键工程基础。

15.10.2 .pt2 包：AOTI 的产物格式

§15.10.1 提到 AOTI 输出 .pt2 文件。它的布局由 torch/csrc/export/pt2_archive_constants.h 定死（Python 侧在 torch/export/pt2_archive/constants.py 转出），unzip -l 大致能看到（示意：本机编不出 AOTI 的 .so，下面的顶层几项是用 PT2ArchiveWriter 在 CPU 上实写一个归档确认的，data/aotinductor/ 下的文件名按 torch/_inductor/codecache.py 的产物列表整理）：

model_aoti/archive_format                  # 内容为 "pt2"
model_aoti/archive_version                 # PT2 归档规范版本，v2.11 为 "0"
model_aoti/byteorder
model_aoti/.data/version                   # zip 容器（PyTorchFileWriter）自己的版本
model_aoti/.data/serialization_id
model_aoti/data/aotinductor/model/...      # 生成的 wrapper / kernel C++ 源码、元数据 JSON、
                                           # 编好的 .so（GPU 上还有各 kernel 的 cubin）
model_aoti/data/constants/...              # 仅当有 custom object 等常量需单独存放时

顶层目录名取自归档文件名；model 是默认的模型名（torch/_inductor/package/package.py 的 model_name: str = "model"），一个 .pt2 里可以放多个模型。

加载流程（C++ 端）：

  1. AOTIModelPackageLoader（torch/csrc/inductor/aoti_package/model_package_loader.cpp）打开 .pt2、解压到临时目录
  2. 按设备创建 AOTIModelContainerRunner（CPU / CUDA 等子类），用 at::DynamicLibrary（即 dlopen）加载 .so
  3. 用 dlsym 拿到 AOTInductorModelContainerCreateWithDevice 等函数指针（torch/csrc/inductor/aoti_runner/model_container_runner.cpp）
  4. 调用它创建 model container，常量（权重）随之就位
  5. 用户每次 runner.run(inputs) 时直接进 .so 内部的 forward 函数

.pt2 格式很可能是 PyTorch 模型分发的"未来"（这是预测）：它要承接的是 v1 的 .pt（pickle）+ torchscript .pt（C++ 推理）这两类用途。归档里的目录名不是随手起的，都定义在 torch/export/pt2_archive/constants.py（WEIGHTS_DIR / CONSTANTS_DIR / AOTINDUCTOR_DIR / ARCHIVE_VERSION_PATH 等）。注意 v2.11 的 .pt2 不用 safetensors 存权重 —— 全仓里提到 safetensors 的几个文件（torch/distributed/checkpoint/hf_storage.py、_hf_utils.py、_consolidate_hf_safetensors.py、quantized_hf_storage.py、filesystem.py）全在 DCP ↔ HuggingFace 那条通路上（第 19 章会讲），两件事别混为一谈。

打开 .pt2 看里面是 zip + 标准结构，让你做 deployment debug 时（如检查 weight 是否对、cubin 是否匹配 SM 架构）能直接 unzip 看。这种"开放格式"是 PyTorch 团队工程哲学的体现。

15.10.3 编译产物的 cross-platform 兼容性

.pt2 编译产物在不同环境间的兼容性：

维度 兼容性
CUDA driver 版本 向后兼容（编译时用 12.4，运行时 12.5+ OK）
CUDA toolkit 版本 不向前兼容（编译时用 12.4，运行时若 12.0 会报"PTX version too high"）
SM 架构 完全不兼容（A100 sm_80 编译产物不能在 H100 sm_90 跑，反之亦然）
PyTorch 版本 无明确的跨版本承诺，升级后应重编验证
操作系统 Linux x86_64 与 ARM 不兼容（.so 是平台特定的）

工程后果：

  • 多 GPU 类型部署：要为每种 SM 架构编一份 .pt2（Triton 产出的 cubin 只针对单一 SM 架构）
  • PyTorch 升级：旧 .pt2 在新 PyTorch 上不一定能跑，建议每次升级重编
  • Docker 镜像分层：base image 锁 CUDA + PyTorch 版本，减少兼容性问题

正确部署流程：

  1. 训练机 export 出 .pt2（含元信息）
  2. 部署 pipeline 检测目标机 SM 架构、CUDA / PyTorch 版本
  3. 不匹配 → 在目标机上 re-export（用同一份 PyTorch model + 重新 AOT 编）
  4. 匹配 → 直接 deploy

vLLM / TensorRT-LLM 等引擎都内置这套"按目标硬件编译"流程。理解兼容性矩阵让你看到 .pt2 不是"训一次到处跑"的银弹，而是"为目标硬件编译"的中间制品。

15.10.4 ExecuTorch / Mobile 部署：编译栈的另一终点

ExecuTorch 是 PyTorch 官方的"on edge devices"解决方案（独立仓库 pytorch/executorch，与主仓分开发布）。它与 AOTI 的关系：

  • AOTI：服务器端 ahead-of-time 编译，输出 .pt2 含 cubin（GPU 二进制）
  • ExecuTorch：移动端 / 嵌入式编译，输出 .pte 文件（不含 cubin，含 PyTorch IR + delegate 注册）
graph TB
    Model[PyTorch nn.Module]
    Model --> Exp[torch.export]
    Exp --> EP[ExportedProgram]
    EP -->|服务器 GPU| AOTI[AOTI compile]
    AOTI --> PT2[.pt2<br/>含 cubin]
    EP -->|移动端| ET[ExecuTorch backend]
    ET --> PTE[.pte<br/>不含 GPU 二进制]
    PTE --> XNN[XNNPACK]
    PTE --> COR[Core ML]
    PTE --> NN[Vulkan]
    PTE --> QC[Qualcomm QNN]
    style AOTI fill:#dcfce7
    style ET fill:#fef3c7

ExecuTorch 的核心机制是 delegate：把 PyTorch IR 交给目标平台的 native engine（XNNPACK on CPU、Core ML on Apple、Vulkan on GPU、QNN on Qualcomm 等，以 pytorch/executorch 仓库 backends/ 目录为准）跑。.pte 内部存的是"分割好的 IR + 委托表"，runtime 把每段交给对应 engine 执行。

工程实践：

  • 服务器推理：AOTI（.pt2 + libtorch）
  • iOS 应用：ExecuTorch + Core ML delegate（由 Core ML 决定是否调度到 Apple Neural Engine）
  • Android 应用：ExecuTorch + XNNPACK / Vulkan，或 Qualcomm、MediaTek 等厂商 delegate
  • 嵌入式 / IoT：ExecuTorch + 自定义 delegate（需要硬件厂商支持）

PyTorch 团队在 v2.x 把 export 路径设计为"统一前端"，AOTI 与 ExecuTorch 共用 torch.export 这一入口、产出一致的 ExportedProgram、再分流到不同 backend。这是"一次 trace、多端部署"的架构基础。

理解这个全景让你看到 PyTorch 的野心 —— 不只是"训练框架"，是"从训练到部署的端到端 stack"。第 12-15 章讲的编译栈是这套战略的核心引擎。

15.11 排错与观测

15.11.1 编译栈错误的诊断流程

报错时怎么定位是哪一层的问题？标准流程：

flowchart TD
    Err[编译报错 / 性能差]
    Err --> Q1{有错误信息?}
    Q1 -->|是| L1[看 traceback 第一行<br/>哪个文件]
    Q1 -->|否, 性能问题| Profile[torch.profiler / Nsight]
    L1 --> CD[在 torch/_dynamo/]
    L1 --> CA[在 torch/_functorch/]
    L1 --> CI[在 torch/_inductor/]
    L1 --> CT[在 triton/]
    CD --> SD["TORCH_LOGS=dynamo,recompiles<br/>看 trace 卡在哪里"]
    CA --> SA["TORCH_LOGS=aot_graphs<br/>看反向 graph"]
    CI --> SI["TORCH_LOGS=output_code<br/>看 Inductor 生成什么"]
    CT --> ST["手动打 print 测 kernel 输入"]
    Profile --> P1{CPU launch 开销占比高?}
    P1 -->|是| RO[改用 reduce-overhead]
    P1 -->|否| P2{kernel 内 SM 利用率偏低?}
    P2 -->|是| MA[max_autotune + aggressive_fusion]
    P2 -->|否| P3{有 graph break?}
    P3 -->|是| FG[fullgraph=True 重构代码]
    style L1 fill:#fef3c7
    style Profile fill:#dbeafe

每层的常见错误：

层 典型错误 解决思路
Dynamo "Unsupported: ..." graph break 该 op，或 disable trace
AOTAutograd UnsupportedOperatorException（fake tensor 跑不了该 op） op 没注册 fake / meta 实现，补上或加 decomp
Inductor lowering "missing lowering"（MissingOperatorWithoutDecomp） op 没注册 lowering，要写 custom
Inductor scheduling 不报错，只是没 fuse 不影响正确性，对比 ir_pre_fusion.txt / ir_post_fusion.txt
Triton compile OutOfResources（如 shared memory 不够） 关 aggressive_fusion、缩小 block
CUDA Graph "accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run" clone 输出，或加 cudagraph_mark_step_begin

按这个流程能快速定位多数 compile 问题。盲改 mode flag 是浪费时间。

15.11.2 TORCH_COMPILE_DEBUG：编译流程全景透视

TORCH_COMPILE_DEBUG=1 让 PyTorch 把整个编译过程的中间产物全部 dump 到磁盘。是调试 compile 问题的"核武器"：

TORCH_COMPILE_DEBUG=1 python train.py

产物目录（本机 CPU 实跑一个只含 a @ b 的函数得到的真实结构）：

torch_compile_debug/run_2026_09_25_08_25_14_870062-pid_3545545/
├── torchdynamo/
│   └── debug.log                          # Dynamo trace 详细日志
└── torchinductor/
    ├── aot_model___0_debug.log            # AOTAutograd 日志
    └── model__0_inference_0.0/
        ├── fx_graph_readable.py           # 进入 Inductor 的 fx graph（可读版）
        ├── fx_graph_runnable.py           # 同一张图的可独立运行复现脚本
        ├── fx_graph_transformed.py        # Inductor 自己的 fx pass 之后
        ├── ir_pre_fusion.txt              # Inductor IR fusion 前
        ├── ir_post_fusion.txt             # fusion 后
        ├── inductor_provenance_tracking_node_mappings.json
        └── output_code.py                 # 最终生成的代码（GPU 上 Triton kernel 源码也内嵌在这里）

训练图会按 forward / backward 各出一个子目录（名字里的 inference 换成相应阶段）。每个文件都是可读的 Python / 文本。能看到：

  • output_code.py 是 Inductor 生成的最终代码，含每个 kernel 的源码与调用顺序、shape、stride；它可以直接 python output_code.py 跑（文件末尾自带 benchmark 入口）
  • ir_post_fusion.txt 让你看每个 Inductor IR node 是怎么 fuse 的
  • fx_graph_runnable.py 是最小复现：报 bug 时附上它最省事

实战调试场景：

  • 某 op 报错：看 fx_graph_readable.py 找该 op 在 fx graph 的位置
  • fusion 不生效：看 ir_pre_fusion.txt 与 ir_post_fusion.txt 对比，找为什么没 fuse
  • kernel 性能差：看 output_code.py 里对应 kernel，对比 expected shape vs actual

更细的：TORCH_LOGS="+inductor" 输出更详细日志（含每个 fusion 决策的原因）。组合 TORCH_COMPILE_DEBUG=1 TORCH_LOGS="+inductor,+aot,+dynamo" 让你看到编译栈每一层在做什么。

15.11.3 用 Nsight 看 CUDA Graph 内部

NVIDIA Nsight Systems 是看 CUDA Graph 性能的标准工具：

nsys profile --trace=cuda,nvtx --output=report python train.py
nsys-ui report.nsys-rep

打开 timeline 能看到：

  • CUDA Graph Launch：一次 cudaGraphLaunch 出现在 CPU timeline 上，仅几个 us 宽
  • Graph 内部 kernel：在 GPU timeline 上展开，每个 kernel 各自显示
  • stream 之间依赖：箭头连接表明 graph 内部的并行结构
  • Memory copy：static_input copy 显示为 H2D / D2D 短矩形

对比方式很直观：不开 graph 时 CPU 时间线上是一长串密集的 launch 方块，开 graph 后塌缩成一个 cudaGraphLaunch；GPU 时间线上的 kernel 序列不变，但 kernel 之间等 CPU 的空档消失了。

进阶分析：

  • kernel 之间的 gap：暴露调度延迟，graph 也无法消除（这部分要靠 fusion 解决）
  • 不必要的 stream sync：发现某些 event wait 是多余的，可优化 graph 拓扑
  • memory allocation 异常：graph 录制时若有意外 alloc，profile 会显示 cudaMalloc 出现在 graph 区域（不应该出现）

实战：调优 reduce-overhead 模式时，先开 Nsight 看 timeline，再决定调哪个 flag。盲调 max_autotune / cudagraph_dynamic_shape_warn_limit 是浪费时间。

15.12 生产实践与横向对照

15.12.1 vLLM piecewise CUDA Graph：生产推理引擎的实战形态

vLLM 的 V1 引擎是 CUDA Graph 在生产里用得最复杂的例子之一。它面对的难题是：

  • 每拍的 token 数动态变化（每个推理请求独立调度，prefill 与 decode 混在同一批）
  • attention 依赖每个请求的 seq len 与 block table，形状组合多到没法逐一录图

piecewise CUDA Graph 是 vLLM 的解（以 vLLM v0.8.5 为准）：借 torch.compile 拿到整图后，在 attention 算子处把图切开（切分点是 vllm.unified_attention / vllm.unified_attention_with_output），attention 之外的各段按 cudagraph_capture_sizes 列出的若干档 token 数分别录 graph，attention 本身每拍以 eager 方式调用。

graph LR
    Input[input tensors<br/>pad 到最近的捕获档位] --> S1[Segment 1: embedding + 第 1 层非 attention 部分<br/>CUDA Graph]
    S1 --> A1[Attention 1<br/>eager]
    A1 --> S2[Segment 2: 第 1 层 MLP + 第 2 层非 attention 部分<br/>CUDA Graph]
    S2 --> A2[Attention 2<br/>eager]
    A2 --> More[...]
    More --> Out[norm + lm_head → output]
    style A1 fill:#fee2e2,stroke:#ef4444
    style A2 fill:#fee2e2,stroke:#ef4444
    style S1 fill:#dcfce7
    style S2 fill:#dcfce7

这样切的好处是：进 graph 的段只依赖总 token 数这一维，档位有限（v0.8.5 V1 默认 [1, 2, 4] + 8 的倍数直到 512，再按 max_num_batched_tokens 截断）；真正形状敏感的 attention 留在图外，不必为每种 seq len / block table 组合各录一份。具体省了多少启动时间和显存，以 vLLM 自己 capture_model() 收尾时打的日志为准（《vLLM 推理内核深度解析》专栏第 8 章 §8.4）。

vLLM 的 V1 引擎把 CUDA Graph 的 capture / replay 直接建在 PyTorch 这层 API 上（具体模块随版本变动，参考《vLLM 推理内核深度解析》专栏第 8 章）。理解 PyTorch 这层 API 让你看明白 vLLM 怎么把这套打造出来。这是"compile 栈 + 生产引擎"配合的最佳样本。

15.12.2 PyTorch compile vs JAX vs TF graph：哲学对比

把 PyTorch 编译栈与同类系统对比，能看清各自的设计选择：

维度 PyTorch (torch.compile) JAX (jit) TensorFlow Graph
程序模型 命令式 + 编译 函数式（纯函数 + jit） 声明式（先建图再跑）
副作用支持 是（SideEffects 追踪） 否（要求纯函数） 否（变量是图节点）
控制流 graph break + 重 trace jit 后 jax.lax.cond / jax.lax.scan tf.cond / tf.while_loop
trace 粒度 per-frame（PEP 523） per-function（@jit 装饰） 整图
反向计算 AOTAutograd 自动 grad / vjp / jvp tf.GradientTape
编译产物 Triton kernel + Python wrapper XLA HLO → cubin XLA HLO → cubin
部署路径 AOTI + ExecuTorch jax.export + IREE SavedModel + TFLite

核心差异：

  • PyTorch：把"动态图"作为顶层语义，编译是"事后包装"。代价是 trace 复杂、graph break 处理；好处是用户完全不用改代码就能享受加速
  • JAX：把"纯函数"作为顶层语义，trace 简单（pytree-in / pytree-out）。代价是用户要按函数式风格写代码、副作用要显式表达
  • TF Graph (1.x)：先 build graph 再 sess.run，最静态、最易优化、但最难写。TF 2.x 转向 eager + tf.function 抄 JAX 路径

为什么 PyTorch 赢了：用户用 Python 命令式代码已经几年，改代码成本巨高。PyTorch v2.0 选择"trace 复杂度由内部承担、用户接口零改动"。@torch.compile 一行装饰器背后的工程复杂度，是为了换"用户不改代码"这个核心目标。

理解这个哲学差异让你看到：好的工程不只是技术实现，更是"为谁优化哪个维度"的产品决策。PyTorch 选择"为研究者优化迭代速度"，TF 选择"为部署优化静态图"，结果在 LLM 时代 PyTorch 全面胜出。

15.13 几条工程经验

实战 torch.compile 端到端：

1. warmup 是必须的：第一次 forward 编译几秒到几分钟。生产服务前必须 warmup 几个 batch，否则首请求超时

2. dynamic shape 用 mark_dynamic：torch._dynamo.mark_dynamic(input, 0) 让 Inductor 一开始就生成处理多个 batch size 的 kernel。不标的话（dynamic=None 默认），第一次遇到新 shape 会重编一次并自动转成 dynamic（本机 CPU 实测：同一函数喂 5 种长度，共编译 2 帧）

3. 用 TORCHINDUCTOR_CACHE_DIR 持久化编译产物：进程重启不用重新编译。生产服务节省冷启动时间

4. graph break 会让 reduce-overhead 失效那一段：fullgraph=True 强制不允许 break，能逼自己写 trace-friendly 代码

5. 不要在 compiled 函数内做 .item() / .tolist() 这类把数据拉回 Python 的操作：默认会触发 graph break（本机 torch._dynamo.explain 实测 .item() 产生 1 次 break；.cpu() / .to(device) 本身能进图）。在外面做完再传进来

6. CUDA Graph 显存占用：每个 graph 至少占一份 PrivatePool。多 shape 训练 / 推理时显存可能暴增

7. torch._dynamo.reset() 清空进程内的编译缓存：换模型架构 / debug 时常用；它不动磁盘上的 Inductor 缓存（见其文档串）

8. distributed 训练慎用 reduce-overhead：CUDA Graph 不与 NCCL collective 完美兼容（v2.x 在改善但仍有 corner case）。DDP 与 torch.compile 的配合见第 17 章 §17.19

15.14 跨专栏关联

  • 《vLLM 推理内核深度解析》专栏第 8 章 model runner：vLLM 的 V1 引擎大量用分段 CUDA Graph 处理变长 prefill / decode
  • 第 4 章 §4.9 PrivatePool：CUDA Graph 与 caching allocator 的协作核心
  • 第 12-14 章：本章是这三章的串联，理解前置链路是理解本章的基础

15.15 整个编译栈的"一句话总结"

把 12-15 章的核心思想浓缩成一句话：

torch.compile = "Python 帧拦截 + 反向自动捕获 + 函数化 IR + 贪心 fusion + Triton codegen + CUDA Graph 包装"，每一段都解决一个具体问题，串起来让 PyTorch 在保留动态图体验的同时拿到接近静态图的性能。

理解了这条链路，你就明白为什么 v2.x 的编译栈如此庞大 —— 每一段都是几年研究 + 几万行工程，叠在一起形成 @torch.compile 这个一行装饰器的入口。

下一章进分布式训练第七篇 —— 拆 ProcessGroup 与 NCCL，看多卡训练的通信底座，理解 collective 是怎么走到 GPU、watchdog 是怎么发现卡死、communicator 又是怎么缓存复用的。